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
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))
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'}
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With