How to display subset of classes during inference?

❓ Questions and Help

I want to output masks of a particular class in instance segmentation say person (id: 0 according to COCO). In order to do so (on the same image in colab notebook), I first find the indexes of non-zero classes then remove the corresponding tensors from pred_classses, scores, etc. as follows:

cls = outputs['instances'].pred_classes
scores = outputs["instances"].scores
masks = outputs['instances'].pred_masks
print(cls)
# tensor([17,  0,  0,  0,  0,  0,  0,  0, 25,  0, 25, 25,  0,  0, 24], device='cuda:0')

# non-zero elements (non-person class categories)
indx_to_remove = cls.nonzero().flatten().tolist()
print(indx_to_remove)
# [0, 8, 10, 11, 14]

# delete corresponding arrays
cls = np.delete(cls.cpu().numpy(), indx_to_remove)
scores = np.delete(scores.cpu().numpy(), indx_to_remove)
masks = np.delete(masks.cpu().numpy(), indx_to_remove, axis=0)

# convert back to tensor and move to cuda
cls = torch.tensor(cls).to('cuda:0')
scores = torch.tensor(scores).to('cuda:0')
masks = torch.tensor(masks).to('cuda:0')
print(cls)  # similar to as it was before except unwanted classes removed
# tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0], device='cuda:0')

# not interested in boxes
outputs['instances'].remove('pred_boxes')

outputs['instances'].pred_classes = cls
outputs["instances"].scores = scores
outputs['instances'].pred_masks = masks

It gives me error:

AssertionError                            Traceback (most recent call last)

<ipython-input-47-695b068ba737> in <module>()
     15 masks = torch.tensor(masks).to('cuda:0')
     16 
---> 17 outputs['instances'].pred_classes = cls
     18 outputs["instances"].pred_boxes = None
     19 outputs["instances"].scores = scores

1 frames

/content/detectron2_repo/detectron2/structures/instances.py in set(self, name, value)
     69             assert (
     70                 len(self) == data_len
---> 71             ), "Adding a field of length {} to a Instances of length {}".format(data_len, len(self))
     72         self._fields[name] = value
     73 

AssertionError: Adding a field of length 10 to a Instances of length 15

Is it because the instances structure has a fixed label 15 in it {'instances': Instances(num_instances=15, image_height=480, image_width=640, fields=[... ?

How do I solve this problem, and is there any better method to do so?

[kHarshit]

Thanks, solved it by creating new Instances object as follows:

# create new instance obj and set its fields
obj = detectron2.structures.Instances(image_size=(480, 640))
obj.set('pred_classes', cls)
obj.set('scores', scores)
obj.set('pred_masks', masks)

# now, pass the obj to visualize fn

[SharifElfouly]

Thanks, solved it by creating new Instances object as follows:

# create new instance obj and set its fields
obj = detectron2.structures.Instances(image_size=(480, 640))
obj.set('pred_classes', cls)
obj.set('scores', scores)
obj.set('pred_masks', masks)

# now, pass the obj to visualize fn

I get the following error. It seems that the visualization needs the bboxes.

_DetectedInstance(classes[i], boxes[i], mask_rle=None, color=None, ttl=8)
TypeError: 'NoneType' object is not subscriptable