Spaces:
Runtime error
Runtime error
| import os | |
| import os.path as osp | |
| from pathlib import Path | |
| import tqdm | |
| from glob import glob | |
| import numpy as np | |
| from PIL import Image | |
| COCO_CATEGORIES = [{'color': [220, 20, 60], 'isthing': 1, 'id': 0, 'name': 'person', 'trainId': 0}, | |
| {'color': [119, 11, 32], 'isthing': 1, 'id': 1, 'name': 'bicycle', 'trainId': 1}, | |
| {'color': [0, 0, 142], 'isthing': 1, 'id': 2, 'name': 'car', 'trainId': 2}, | |
| {'color': [0, 0, 230], 'isthing': 1, 'id': 3, 'name': 'motorcycle', 'trainId': 3}, | |
| {'color': [106, 0, 228], 'isthing': 1, 'id': 4, 'name': 'airplane', 'trainId': 4}, | |
| {'color': [0, 60, 100], 'isthing': 1, 'id': 5, 'name': 'bus', 'trainId': 5}, | |
| {'color': [0, 80, 100], 'isthing': 1, 'id': 6, 'name': 'train', 'trainId': 6}, | |
| {'color': [0, 0, 70], 'isthing': 1, 'id': 7, 'name': 'truck', 'trainId': 7}, | |
| {'color': [0, 0, 192], 'isthing': 1, 'id': 8, 'name': 'boat', 'trainId': 8}, | |
| {'color': [250, 170, 30], 'isthing': 1, 'id': 9, 'name': 'traffic light', 'trainId': 9}, | |
| {'color': [100, 170, 30], 'isthing': 1, 'id': 10, 'name': 'fire hydrant', 'trainId': 10}, | |
| {'color': [220, 220, 0], 'isthing': 1, 'id': 12, 'name': 'stop sign', 'trainId': 11}, | |
| {'color': [175, 116, 175], 'isthing': 1, 'id': 13, 'name': 'parking meter', 'trainId': 12}, | |
| {'color': [250, 0, 30], 'isthing': 1, 'id': 14, 'name': 'bench', 'trainId': 13}, | |
| {'color': [165, 42, 42], 'isthing': 1, 'id': 15, 'name': 'bird', 'trainId': 14}, | |
| {'color': [255, 77, 255], 'isthing': 1, 'id': 16, 'name': 'cat', 'trainId': 15}, | |
| {'color': [0, 226, 252], 'isthing': 1, 'id': 17, 'name': 'dog', 'trainId': 16}, | |
| {'color': [182, 182, 255], 'isthing': 1, 'id': 18, 'name': 'horse', 'trainId': 17}, | |
| {'color': [0, 82, 0], 'isthing': 1, 'id': 19, 'name': 'sheep', 'trainId': 18}, | |
| {'color': [120, 166, 157], 'isthing': 1, 'id': 20, 'name': 'cow', 'trainId': 19}, | |
| {'color': [110, 76, 0], 'isthing': 1, 'id': 21, 'name': 'elephant', 'trainId': 20}, | |
| {'color': [174, 57, 255], 'isthing': 1, 'id': 22, 'name': 'bear', 'trainId': 21}, | |
| {'color': [199, 100, 0], 'isthing': 1, 'id': 23, 'name': 'zebra', 'trainId': 22}, | |
| {'color': [72, 0, 118], 'isthing': 1, 'id': 24, 'name': 'giraffe', 'trainId': 23}, | |
| {'color': [255, 179, 240], 'isthing': 1, 'id': 26, 'name': 'backpack', 'trainId': 24}, | |
| {'color': [0, 125, 92], 'isthing': 1, 'id': 27, 'name': 'umbrella', 'trainId': 25}, | |
| {'color': [209, 0, 151], 'isthing': 1, 'id': 30, 'name': 'handbag', 'trainId': 26}, | |
| {'color': [188, 208, 182], 'isthing': 1, 'id': 31, 'name': 'tie', 'trainId': 27}, | |
| {'color': [0, 220, 176], 'isthing': 1, 'id': 32, 'name': 'suitcase', 'trainId': 28}, | |
| {'color': [255, 99, 164], 'isthing': 1, 'id': 33, 'name': 'frisbee', 'trainId': 29}, | |
| {'color': [92, 0, 73], 'isthing': 1, 'id': 34, 'name': 'skis', 'trainId': 30}, | |
| {'color': [133, 129, 255], 'isthing': 1, 'id': 35, 'name': 'snowboard', 'trainId': 31}, | |
| {'color': [78, 180, 255], 'isthing': 1, 'id': 36, 'name': 'sports ball', 'trainId': 32}, | |
| {'color': [0, 228, 0], 'isthing': 1, 'id': 37, 'name': 'kite', 'trainId': 33}, | |
| {'color': [174, 255, 243], 'isthing': 1, 'id': 38, 'name': 'baseball bat', 'trainId': 34}, | |
| {'color': [45, 89, 255], 'isthing': 1, 'id': 39, 'name': 'baseball glove', 'trainId': 35}, | |
| {'color': [134, 134, 103], 'isthing': 1, 'id': 40, 'name': 'skateboard', 'trainId': 36}, | |
| {'color': [145, 148, 174], 'isthing': 1, 'id': 41, 'name': 'surfboard', 'trainId': 37}, | |
| {'color': [255, 208, 186], 'isthing': 1, 'id': 42, 'name': 'tennis racket', 'trainId': 38}, | |
| {'color': [197, 226, 255], 'isthing': 1, 'id': 43, 'name': 'bottle', 'trainId': 39}, | |
| {'color': [171, 134, 1], 'isthing': 1, 'id': 45, 'name': 'wine glass', 'trainId': 40}, | |
| {'color': [109, 63, 54], 'isthing': 1, 'id': 46, 'name': 'cup', 'trainId': 41}, | |
| {'color': [207, 138, 255], 'isthing': 1, 'id': 47, 'name': 'fork', 'trainId': 42}, | |
| {'color': [151, 0, 95], 'isthing': 1, 'id': 48, 'name': 'knife', 'trainId': 43}, | |
| {'color': [9, 80, 61], 'isthing': 1, 'id': 49, 'name': 'spoon', 'trainId': 44}, | |
| {'color': [84, 105, 51], 'isthing': 1, 'id': 50, 'name': 'bowl', 'trainId': 45}, | |
| {'color': [74, 65, 105], 'isthing': 1, 'id': 51, 'name': 'banana', 'trainId': 46}, | |
| {'color': [166, 196, 102], 'isthing': 1, 'id': 52, 'name': 'apple', 'trainId': 47}, | |
| {'color': [208, 195, 210], 'isthing': 1, 'id': 53, 'name': 'sandwich', 'trainId': 48}, | |
| {'color': [255, 109, 65], 'isthing': 1, 'id': 54, 'name': 'orange', 'trainId': 49}, | |
| {'color': [0, 143, 149], 'isthing': 1, 'id': 55, 'name': 'broccoli', 'trainId': 50}, | |
| {'color': [179, 0, 194], 'isthing': 1, 'id': 56, 'name': 'carrot', 'trainId': 51}, | |
| {'color': [209, 99, 106], 'isthing': 1, 'id': 57, 'name': 'hot dog', 'trainId': 52}, | |
| {'color': [5, 121, 0], 'isthing': 1, 'id': 58, 'name': 'pizza', 'trainId': 53}, | |
| {'color': [227, 255, 205], 'isthing': 1, 'id': 59, 'name': 'donut', 'trainId': 54}, | |
| {'color': [147, 186, 208], 'isthing': 1, 'id': 60, 'name': 'cake', 'trainId': 55}, | |
| {'color': [153, 69, 1], 'isthing': 1, 'id': 61, 'name': 'chair', 'trainId': 56}, | |
| {'color': [3, 95, 161], 'isthing': 1, 'id': 62, 'name': 'couch', 'trainId': 57}, | |
| {'color': [163, 255, 0], 'isthing': 1, 'id': 63, 'name': 'potted plant', 'trainId': 58}, | |
| {'color': [119, 0, 170], 'isthing': 1, 'id': 64, 'name': 'bed', 'trainId': 59}, | |
| {'color': [0, 182, 199], 'isthing': 1, 'id': 66, 'name': 'dining table', 'trainId': 60}, | |
| {'color': [0, 165, 120], 'isthing': 1, 'id': 69, 'name': 'toilet', 'trainId': 61}, | |
| {'color': [183, 130, 88], 'isthing': 1, 'id': 71, 'name': 'tv', 'trainId': 62}, | |
| {'color': [95, 32, 0], 'isthing': 1, 'id': 72, 'name': 'laptop', 'trainId': 63}, | |
| {'color': [130, 114, 135], 'isthing': 1, 'id': 73, 'name': 'mouse', 'trainId': 64}, | |
| {'color': [110, 129, 133], 'isthing': 1, 'id': 74, 'name': 'remote', 'trainId': 65}, | |
| {'color': [166, 74, 118], 'isthing': 1, 'id': 75, 'name': 'keyboard', 'trainId': 66}, | |
| {'color': [219, 142, 185], 'isthing': 1, 'id': 76, 'name': 'cell phone', 'trainId': 67}, | |
| {'color': [79, 210, 114], 'isthing': 1, 'id': 77, 'name': 'microwave', 'trainId': 68}, | |
| {'color': [178, 90, 62], 'isthing': 1, 'id': 78, 'name': 'oven', 'trainId': 69}, | |
| {'color': [65, 70, 15], 'isthing': 1, 'id': 79, 'name': 'toaster', 'trainId': 70}, | |
| {'color': [127, 167, 115], 'isthing': 1, 'id': 80, 'name': 'sink', 'trainId': 71}, | |
| {'color': [59, 105, 106], 'isthing': 1, 'id': 81, 'name': 'refrigerator', 'trainId': 72}, | |
| {'color': [142, 108, 45], 'isthing': 1, 'id': 83, 'name': 'book', 'trainId': 73}, | |
| {'color': [196, 172, 0], 'isthing': 1, 'id': 84, 'name': 'clock', 'trainId': 74}, | |
| {'color': [95, 54, 80], 'isthing': 1, 'id': 85, 'name': 'vase', 'trainId': 75}, | |
| {'color': [128, 76, 255], 'isthing': 1, 'id': 86, 'name': 'scissors', 'trainId': 76}, | |
| {'color': [201, 57, 1], 'isthing': 1, 'id': 87, 'name': 'teddy bear', 'trainId': 77}, | |
| {'color': [246, 0, 122], 'isthing': 1, 'id': 88, 'name': 'hair drier', 'trainId': 78}, | |
| {'color': [191, 162, 208], 'isthing': 1, 'id': 89, 'name': 'toothbrush', 'trainId': 79}, | |
| {'id': 91, 'name': 'banner', 'supercategory': 'textile', 'trainId': 80}, | |
| {'id': 92, 'name': 'blanket', 'supercategory': 'textile', 'trainId': 81}, | |
| {'id': 93, 'name': 'branch', 'supercategory': 'plant', 'trainId': 82}, | |
| {'id': 94, 'name': 'bridge', 'supercategory': 'building', 'trainId': 83}, | |
| {'id': 95, 'name': 'building-other', 'supercategory': 'building', 'trainId': 84}, | |
| {'id': 96, 'name': 'bush', 'supercategory': 'plant', 'trainId': 85}, | |
| {'id': 97, 'name': 'cabinet', 'supercategory': 'furniture-stuff', 'trainId': 86}, | |
| {'id': 98, 'name': 'cage', 'supercategory': 'structural', 'trainId': 87}, | |
| {'id': 99, 'name': 'cardboard', 'supercategory': 'raw-material', 'trainId': 88}, | |
| {'id': 100, 'name': 'carpet', 'supercategory': 'floor', 'trainId': 89}, | |
| {'id': 101, 'name': 'ceiling-other', 'supercategory': 'ceiling', 'trainId': 90}, | |
| {'id': 102, 'name': 'ceiling-tile', 'supercategory': 'ceiling', 'trainId': 91}, | |
| {'id': 103, 'name': 'cloth', 'supercategory': 'textile', 'trainId': 92}, | |
| {'id': 104, 'name': 'clothes', 'supercategory': 'textile', 'trainId': 93}, | |
| {'id': 105, 'name': 'clouds', 'supercategory': 'sky', 'trainId': 94}, | |
| {'id': 106, 'name': 'counter', 'supercategory': 'furniture-stuff', 'trainId': 95}, | |
| {'id': 107, 'name': 'cupboard', 'supercategory': 'furniture-stuff', 'trainId': 96}, | |
| {'id': 108, 'name': 'curtain', 'supercategory': 'textile', 'trainId': 97}, | |
| {'id': 109, 'name': 'desk-stuff', 'supercategory': 'furniture-stuff', 'trainId': 98}, | |
| {'id': 110, 'name': 'dirt', 'supercategory': 'ground', 'trainId': 99}, | |
| {'id': 111, 'name': 'door-stuff', 'supercategory': 'furniture-stuff', 'trainId': 100}, | |
| {'id': 112, 'name': 'fence', 'supercategory': 'structural', 'trainId': 101}, | |
| {'id': 113, 'name': 'floor-marble', 'supercategory': 'floor', 'trainId': 102}, | |
| {'id': 114, 'name': 'floor-other', 'supercategory': 'floor', 'trainId': 103}, | |
| {'id': 115, 'name': 'floor-stone', 'supercategory': 'floor', 'trainId': 104}, | |
| {'id': 116, 'name': 'floor-tile', 'supercategory': 'floor', 'trainId': 105}, | |
| {'id': 117, 'name': 'floor-wood', 'supercategory': 'floor', 'trainId': 106}, | |
| {'id': 118, 'name': 'flower', 'supercategory': 'plant', 'trainId': 107}, | |
| {'id': 119, 'name': 'fog', 'supercategory': 'water', 'trainId': 108}, | |
| {'id': 120, 'name': 'food-other', 'supercategory': 'food-stuff', 'trainId': 109}, | |
| {'id': 121, 'name': 'fruit', 'supercategory': 'food-stuff', 'trainId': 110}, | |
| {'id': 122, 'name': 'furniture-other', 'supercategory': 'furniture-stuff', 'trainId': 111}, | |
| {'id': 123, 'name': 'grass', 'supercategory': 'plant', 'trainId': 112}, | |
| {'id': 124, 'name': 'gravel', 'supercategory': 'ground', 'trainId': 113}, | |
| {'id': 125, 'name': 'ground-other', 'supercategory': 'ground', 'trainId': 114}, | |
| {'id': 126, 'name': 'hill', 'supercategory': 'solid', 'trainId': 115}, | |
| {'id': 127, 'name': 'house', 'supercategory': 'building', 'trainId': 116}, | |
| {'id': 128, 'name': 'leaves', 'supercategory': 'plant', 'trainId': 117}, | |
| {'id': 129, 'name': 'light', 'supercategory': 'furniture-stuff', 'trainId': 118}, | |
| {'id': 130, 'name': 'mat', 'supercategory': 'textile', 'trainId': 119}, | |
| {'id': 131, 'name': 'metal', 'supercategory': 'raw-material', 'trainId': 120}, | |
| {'id': 132, 'name': 'mirror-stuff', 'supercategory': 'furniture-stuff', 'trainId': 121}, | |
| {'id': 133, 'name': 'moss', 'supercategory': 'plant', 'trainId': 122}, | |
| {'id': 134, 'name': 'mountain', 'supercategory': 'solid', 'trainId': 123}, | |
| {'id': 135, 'name': 'mud', 'supercategory': 'ground', 'trainId': 124}, | |
| {'id': 136, 'name': 'napkin', 'supercategory': 'textile', 'trainId': 125}, | |
| {'id': 137, 'name': 'net', 'supercategory': 'structural', 'trainId': 126}, | |
| {'id': 138, 'name': 'paper', 'supercategory': 'raw-material', 'trainId': 127}, | |
| {'id': 139, 'name': 'pavement', 'supercategory': 'ground', 'trainId': 128}, | |
| {'id': 140, 'name': 'pillow', 'supercategory': 'textile', 'trainId': 129}, | |
| {'id': 141, 'name': 'plant-other', 'supercategory': 'plant', 'trainId': 130}, | |
| {'id': 142, 'name': 'plastic', 'supercategory': 'raw-material', 'trainId': 131}, | |
| {'id': 143, 'name': 'platform', 'supercategory': 'ground', 'trainId': 132}, | |
| {'id': 144, 'name': 'playingfield', 'supercategory': 'ground', 'trainId': 133}, | |
| {'id': 145, 'name': 'railing', 'supercategory': 'structural', 'trainId': 134}, | |
| {'id': 146, 'name': 'railroad', 'supercategory': 'ground', 'trainId': 135}, | |
| {'id': 147, 'name': 'river', 'supercategory': 'water', 'trainId': 136}, | |
| {'id': 148, 'name': 'road', 'supercategory': 'ground', 'trainId': 137}, | |
| {'id': 149, 'name': 'rock', 'supercategory': 'solid', 'trainId': 138}, | |
| {'id': 150, 'name': 'roof', 'supercategory': 'building', 'trainId': 139}, | |
| {'id': 151, 'name': 'rug', 'supercategory': 'textile', 'trainId': 140}, | |
| {'id': 152, 'name': 'salad', 'supercategory': 'food-stuff', 'trainId': 141}, | |
| {'id': 153, 'name': 'sand', 'supercategory': 'ground', 'trainId': 142}, | |
| {'id': 154, 'name': 'sea', 'supercategory': 'water', 'trainId': 143}, | |
| {'id': 155, 'name': 'shelf', 'supercategory': 'furniture-stuff', 'trainId': 144}, | |
| {'id': 156, 'name': 'sky-other', 'supercategory': 'sky', 'trainId': 145}, | |
| {'id': 157, 'name': 'skyscraper', 'supercategory': 'building', 'trainId': 146}, | |
| {'id': 158, 'name': 'snow', 'supercategory': 'ground', 'trainId': 147}, | |
| {'id': 159, 'name': 'solid-other', 'supercategory': 'solid', 'trainId': 148}, | |
| {'id': 160, 'name': 'stairs', 'supercategory': 'furniture-stuff', 'trainId': 149}, | |
| {'id': 161, 'name': 'stone', 'supercategory': 'solid', 'trainId': 150}, | |
| {'id': 162, 'name': 'straw', 'supercategory': 'plant', 'trainId': 151}, | |
| {'id': 163, 'name': 'structural-other', 'supercategory': 'structural', 'trainId': 152}, | |
| {'id': 164, 'name': 'table', 'supercategory': 'furniture-stuff', 'trainId': 153}, | |
| {'id': 165, 'name': 'tent', 'supercategory': 'building', 'trainId': 154}, | |
| {'id': 166, 'name': 'textile-other', 'supercategory': 'textile', 'trainId': 155}, | |
| {'id': 167, 'name': 'towel', 'supercategory': 'textile', 'trainId': 156}, | |
| {'id': 168, 'name': 'tree', 'supercategory': 'plant', 'trainId': 157}, | |
| {'id': 169, 'name': 'vegetable', 'supercategory': 'food-stuff', 'trainId': 158}, | |
| {'id': 170, 'name': 'wall-brick', 'supercategory': 'wall', 'trainId': 159}, | |
| {'id': 171, 'name': 'wall-concrete', 'supercategory': 'wall', 'trainId': 160}, | |
| {'id': 172, 'name': 'wall-other', 'supercategory': 'wall', 'trainId': 161}, | |
| {'id': 173, 'name': 'wall-panel', 'supercategory': 'wall', 'trainId': 162}, | |
| {'id': 174, 'name': 'wall-stone', 'supercategory': 'wall', 'trainId': 163}, | |
| {'id': 175, 'name': 'wall-tile', 'supercategory': 'wall', 'trainId': 164}, | |
| {'id': 176, 'name': 'wall-wood', 'supercategory': 'wall', 'trainId': 165}, | |
| {'id': 177, 'name': 'water-other', 'supercategory': 'water', 'trainId': 166}, | |
| {'id': 178, 'name': 'waterdrops', 'supercategory': 'water', 'trainId': 167}, | |
| {'id': 179, 'name': 'window-blind', 'supercategory': 'window', 'trainId': 168}, | |
| {'id': 180, 'name': 'window-other', 'supercategory': 'window', 'trainId': 169}, | |
| {'id': 181, 'name': 'wood', 'supercategory': 'solid', 'trainId': 170}] | |
| if __name__ == "__main__": | |
| dataset_dir = Path(os.getenv("DETECTRON2_DATASETS", "datasets")) / "coco-stuff" | |
| id_map = {} | |
| for cat in COCO_CATEGORIES: | |
| id_map[cat["id"]] = cat["trainId"] | |
| for name in ["train2017", "val2017"]: | |
| annotation_dir = dataset_dir / "annotations" / name | |
| output_dir = dataset_dir / "annotations_detectron2" / name | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| for file in tqdm.tqdm(list(annotation_dir.iterdir())): | |
| output_file = output_dir / file.name | |
| lab = np.asarray(Image.open(file)) | |
| assert lab.dtype == np.uint8 | |
| output = np.zeros_like(lab, dtype=np.uint8) + 255 | |
| for obj_id in np.unique(lab): | |
| if obj_id in id_map: | |
| output[lab == obj_id] = id_map[obj_id] | |
| Image.fromarray(output).save(output_file) |