4

In the documentation for detectron2, it states that class labels are located in output_dict['Instances'].pred_classes. This is all fine, and I can access this easily, but at no point in the documentation (or the output dictionary, as far as I can tell) does it specify which integer label refers to which class. I am assuming that somewhere there is a dictionary that contains something like {0: 'Person', 1: 'Bicycle', 2: 'Car',... }, but I can't seem to find it. Does anyone know where to find this?

Here are the specs for the output dictionary: https://detectron2.readthedocs.io/tutorials/models.html

user8577930
  • 193
  • 2
  • 9

3 Answers3

4

The dataset metadata class has a property thing_classes which has a list of class names of that dataset. Simply by passing the class id from pred_classes one may get the class names as seen here.

pred_classes = output_dict['instances'].pred_classes.cpu().tolist()
class_names = MetadataCatalog.get("mydataset").thing_classes
pred_class_names = list(map(lambda x: class_names[x], pred_classes))
gmuraleekrishna
  • 3,375
  • 1
  • 27
  • 45
2

you can find it here:

https://gist.github.com/AruniRC/7b3dadd004da04c80198557db5da4bda

consider that the value the model output is the key from this dict + 1.

{0: u'__background__',
 1: u'person',
 2: u'bicycle',
 3: u'car',
 4: u'motorcycle',
 5: u'airplane',
 6: u'bus',
 7: u'train',
 8: u'truck',
 9: u'boat',
 10: u'traffic light',
 11: u'fire hydrant',
 12: u'stop sign',
 13: u'parking meter',
 14: u'bench',
 15: u'bird',
 16: u'cat',
 17: u'dog',
 18: u'horse',
 19: u'sheep',
 20: u'cow',
 21: u'elephant',
 22: u'bear',
 23: u'zebra',
 24: u'giraffe',
 25: u'backpack',
 26: u'umbrella',
 27: u'handbag',
 28: u'tie',
 29: u'suitcase',
 30: u'frisbee',
 31: u'skis',
 32: u'snowboard',
 33: u'sports ball',
 34: u'kite',
 35: u'baseball bat',
 36: u'baseball glove',
 37: u'skateboard',
 38: u'surfboard',
 39: u'tennis racket',
 40: u'bottle',
 41: u'wine glass',
 42: u'cup',
 43: u'fork',
 44: u'knife',
 45: u'spoon',
 46: u'bowl',
 47: u'banana',
 48: u'apple',
 49: u'sandwich',
 50: u'orange',
 51: u'broccoli',
 52: u'carrot',
 53: u'hot dog',
 54: u'pizza',
 55: u'donut',
 56: u'cake',
 57: u'chair',
 58: u'couch',
 59: u'potted plant',
 60: u'bed',
 61: u'dining table',
 62: u'toilet',
 63: u'tv',
 64: u'laptop',
 65: u'mouse',
 66: u'remote',
 67: u'keyboard',
 68: u'cell phone',
 69: u'microwave',
 70: u'oven',
 71: u'toaster',
 72: u'sink',
 73: u'refrigerator',
 74: u'book',
 75: u'clock',
 76: u'vase',
 77: u'scissors',
 78: u'teddy bear',
 79: u'hair drier',
 80: u'toothbrush'}
Guinther Kovalski
  • 1,629
  • 1
  • 7
  • 15
0

For Detectron2 use this link to copy & store the classes as a dict (className in this case)

This is the code if you want to get the classes for default Detectron2 model

instances = outputs["instances"]
detected_class_indexes = instances.pred_classes.tolist()
print(detected_class_indexes)
pred_class_names = list(map(lambda x: className[x+1], detected_class_indexes))
print(pred_class_names)
gamingflexer
  • 250
  • 3
  • 5