123456789101112131415161718192021222324252627 |
- """
- Helper code to change YOLO's 0-79 class IDs to 0-90 SSD class IDs.
- """
- COCO_CLASSES_LIST = [
- 'person', 'bicycle','car','motorbike','aeroplane','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','sofa','pottedplant','bed','diningtable','toilet','tvmonitor','laptop','mouse','remote','keyboard',
- 'cell phone','microwave','oven','toaster','sink','refrigerator','book','clock','vase','scissors','teddy bear','hair drier','toothbrush',
- ]
- yolo_to_ssd_classes = [
- 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20,
- 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40,
- 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58,
- 59, 60, 61, 62, 63, 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79,
- 80, 81, 82, 84, 85, 86, 87, 88, 89, 90,
- ]
- def get_cls_dict(category_num):
- if category_num == 80:
- return {i: n for i, n in enumerate(COCO_CLASSES_LIST)}
- else:
- return {i: 'CLS%d' % i for i in range(category_num)}
|