partialweights.py

#source https://github.com/matterport/Mask_RCNN/issues/1361

#You can 'filter out' classes you don't need after you run the detection but it doesn't prune the weights though. So, I'm not sure it would speed things up. Hope it helps.

# COCO class names
class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
               'bus', 'train', 'truck', 'boat', 'traffic light',
               'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
               'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
               'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
               'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
               'kite', 'baseball bat', 'baseball glove', 'skateboard',
               'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
               'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
               'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
               'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
               'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
               'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
               'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
               'teddy bear', 'hair drier', 'toothbrush']

# Class Index
id_car = class_names.index('car')
id_truck = class_names.index('truck')
id_bus = class_names.index('bus')

# Read Image
image = cv2.imread(os.path.join(path_input)) # BGR
image = image[..., ::-1] # BGR --> RGB

# Run detection
r = model.detect([image], verbose=0)[0]

# Exclude other classes except car,bus, & truck
r['class_ids_new'], r['rois_new']= [],[]
r['masks_new'], r['scores_new']= [],[]
r['masks'] = np.transpose(r['masks'], (2, 1, 0))
for idx in range(len(r['class_ids'])):
     if r['class_ids'][idx] == id_car or r['class_ids'][idx] == id_truck or r['class_ids'][idx] == id_bus:
          r['class_ids_new'].append(r['class_ids'][idx])
          r['rois_new'].append(r['rois'][idx])
          r['masks_new'].append(r['masks'][idx])
          r['scores_new'].append(r['scores'][idx])  
del(r['class_ids'], r['rois'], r['masks'], r['scores'])

source