from detectron2.data.datasets.register_coco import register_coco_instances
import os

categories = [
{'id': 1, 'name': 'person'},
 {'id': 2, 'name': 'sneakers'},
 {'id': 3, 'name': 'chair'},
 {'id': 4, 'name': 'hat'},
 {'id': 5, 'name': 'lamp'},
 {'id': 6, 'name': 'bottle'},
 {'id': 7, 'name': 'cabinet/shelf'},
 {'id': 8, 'name': 'cup'},
 {'id': 9, 'name': 'car'},
 {'id': 10, 'name': 'glasses'},
 {'id': 11, 'name': 'picture/frame'},
 {'id': 12, 'name': 'desk'},
 {'id': 13, 'name': 'handbag'},
 {'id': 14, 'name': 'street lights'},
 {'id': 15, 'name': 'book'},
 {'id': 16, 'name': 'plate'},
 {'id': 17, 'name': 'helmet'},
 {'id': 18, 'name': 'leather shoes'},
 {'id': 19, 'name': 'pillow'},
 {'id': 20, 'name': 'glove'},
 {'id': 21, 'name': 'potted plant'},
 {'id': 22, 'name': 'bracelet'},
 {'id': 23, 'name': 'flower'},
 {'id': 24, 'name': 'tv'},
 {'id': 25, 'name': 'storage box'},
 {'id': 26, 'name': 'vase'},
 {'id': 27, 'name': 'bench'},
 {'id': 28, 'name': 'wine glass'},
 {'id': 29, 'name': 'boots'},
 {'id': 30, 'name': 'bowl'},
 {'id': 31, 'name': 'dining table'},
 {'id': 32, 'name': 'umbrella'},
 {'id': 33, 'name': 'boat'},
 {'id': 34, 'name': 'flag'},
 {'id': 35, 'name': 'speaker'},
 {'id': 36, 'name': 'trash bin/can'},
 {'id': 37, 'name': 'stool'},
 {'id': 38, 'name': 'backpack'},
 {'id': 39, 'name': 'couch'},
 {'id': 40, 'name': 'belt'},
 {'id': 41, 'name': 'carpet'},
 {'id': 42, 'name': 'basket'},
 {'id': 43, 'name': 'towel/napkin'},
 {'id': 44, 'name': 'slippers'},
 {'id': 45, 'name': 'barrel/bucket'},
 {'id': 46, 'name': 'coffee table'},
 {'id': 47, 'name': 'suv'},
 {'id': 48, 'name': 'toy'},
 {'id': 49, 'name': 'tie'},
 {'id': 50, 'name': 'bed'},
 {'id': 51, 'name': 'traffic light'},
 {'id': 52, 'name': 'pen/pencil'},
 {'id': 53, 'name': 'microphone'},
 {'id': 54, 'name': 'sandals'},
 {'id': 55, 'name': 'canned'},
 {'id': 56, 'name': 'necklace'},
 {'id': 57, 'name': 'mirror'},
 {'id': 58, 'name': 'faucet'},
 {'id': 59, 'name': 'bicycle'},
 {'id': 60, 'name': 'bread'},
 {'id': 61, 'name': 'high heels'},
 {'id': 62, 'name': 'ring'},
 {'id': 63, 'name': 'van'},
 {'id': 64, 'name': 'watch'},
 {'id': 65, 'name': 'sink'},
 {'id': 66, 'name': 'horse'},
 {'id': 67, 'name': 'fish'},
 {'id': 68, 'name': 'apple'},
 {'id': 69, 'name': 'camera'},
 {'id': 70, 'name': 'candle'},
 {'id': 71, 'name': 'teddy bear'},
 {'id': 72, 'name': 'cake'},
 {'id': 73, 'name': 'motorcycle'},
 {'id': 74, 'name': 'wild bird'},
 {'id': 75, 'name': 'laptop'},
 {'id': 76, 'name': 'knife'},
 {'id': 77, 'name': 'traffic sign'},
 {'id': 78, 'name': 'cell phone'},
 {'id': 79, 'name': 'paddle'},
 {'id': 80, 'name': 'truck'},
 {'id': 81, 'name': 'cow'},
 {'id': 82, 'name': 'power outlet'},
 {'id': 83, 'name': 'clock'},
 {'id': 84, 'name': 'drum'},
 {'id': 85, 'name': 'fork'},
 {'id': 86, 'name': 'bus'},
 {'id': 87, 'name': 'hanger'},
 {'id': 88, 'name': 'nightstand'},
 {'id': 89, 'name': 'pot/pan'},
 {'id': 90, 'name': 'sheep'},
 {'id': 91, 'name': 'guitar'},
 {'id': 92, 'name': 'traffic cone'},
 {'id': 93, 'name': 'tea pot'},
 {'id': 94, 'name': 'keyboard'},
 {'id': 95, 'name': 'tripod'},
 {'id': 96, 'name': 'hockey'},
 {'id': 97, 'name': 'fan'},
 {'id': 98, 'name': 'dog'},
 {'id': 99, 'name': 'spoon'},
 {'id': 100, 'name': 'blackboard/whiteboard'},
 {'id': 101, 'name': 'balloon'},
 {'id': 102, 'name': 'air conditioner'},
 {'id': 103, 'name': 'cymbal'},
 {'id': 104, 'name': 'mouse'},
 {'id': 105, 'name': 'telephone'},
 {'id': 106, 'name': 'pickup truck'},
 {'id': 107, 'name': 'orange'},
 {'id': 108, 'name': 'banana'},
 {'id': 109, 'name': 'airplane'},
 {'id': 110, 'name': 'luggage'},
 {'id': 111, 'name': 'skis'},
 {'id': 112, 'name': 'soccer'},
 {'id': 113, 'name': 'trolley'},
 {'id': 114, 'name': 'oven'},
 {'id': 115, 'name': 'remote'},
 {'id': 116, 'name': 'baseball glove'},
 {'id': 117, 'name': 'paper towel'},
 {'id': 118, 'name': 'refrigerator'},
 {'id': 119, 'name': 'train'},
 {'id': 120, 'name': 'tomato'},
 {'id': 121, 'name': 'machinery vehicle'},
 {'id': 122, 'name': 'tent'},
 {'id': 123, 'name': 'shampoo/shower gel'},
 {'id': 124, 'name': 'head phone'},
 {'id': 125, 'name': 'lantern'},
 {'id': 126, 'name': 'donut'},
 {'id': 127, 'name': 'cleaning products'},
 {'id': 128, 'name': 'sailboat'},
 {'id': 129, 'name': 'tangerine'},
 {'id': 130, 'name': 'pizza'},
 {'id': 131, 'name': 'kite'},
 {'id': 132, 'name': 'computer box'},
 {'id': 133, 'name': 'elephant'},
 {'id': 134, 'name': 'toiletries'},
 {'id': 135, 'name': 'gas stove'},
 {'id': 136, 'name': 'broccoli'},
 {'id': 137, 'name': 'toilet'},
 {'id': 138, 'name': 'stroller'},
 {'id': 139, 'name': 'shovel'},
 {'id': 140, 'name': 'baseball bat'},
 {'id': 141, 'name': 'microwave'},
 {'id': 142, 'name': 'skateboard'},
 {'id': 143, 'name': 'surfboard'},
 {'id': 144, 'name': 'surveillance camera'},
 {'id': 145, 'name': 'gun'},
 {'id': 146, 'name': 'life saver'},
 {'id': 147, 'name': 'cat'},
 {'id': 148, 'name': 'lemon'},
 {'id': 149, 'name': 'liquid soap'},
 {'id': 150, 'name': 'zebra'},
 {'id': 151, 'name': 'duck'},
 {'id': 152, 'name': 'sports car'},
 {'id': 153, 'name': 'giraffe'},
 {'id': 154, 'name': 'pumpkin'},
 {'id': 155, 'name': 'piano'},
 {'id': 156, 'name': 'stop sign'},
 {'id': 157, 'name': 'radiator'},
 {'id': 158, 'name': 'converter'},
 {'id': 159, 'name': 'tissue '},
 {'id': 160, 'name': 'carrot'},
 {'id': 161, 'name': 'washing machine'},
 {'id': 162, 'name': 'vent'},
 {'id': 163, 'name': 'cookies'},
 {'id': 164, 'name': 'cutting/chopping board'},
 {'id': 165, 'name': 'tennis racket'},
 {'id': 166, 'name': 'candy'},
 {'id': 167, 'name': 'skating and skiing shoes'},
 {'id': 168, 'name': 'scissors'},
 {'id': 169, 'name': 'folder'},
 {'id': 170, 'name': 'baseball'},
 {'id': 171, 'name': 'strawberry'},
 {'id': 172, 'name': 'bow tie'},
 {'id': 173, 'name': 'pigeon'},
 {'id': 174, 'name': 'pepper'},
 {'id': 175, 'name': 'coffee machine'},
 {'id': 176, 'name': 'bathtub'},
 {'id': 177, 'name': 'snowboard'},
 {'id': 178, 'name': 'suitcase'},
 {'id': 179, 'name': 'grapes'},
 {'id': 180, 'name': 'ladder'},
 {'id': 181, 'name': 'pear'},
 {'id': 182, 'name': 'american football'},
 {'id': 183, 'name': 'basketball'},
 {'id': 184, 'name': 'potato'},
 {'id': 185, 'name': 'paint brush'},
 {'id': 186, 'name': 'printer'},
 {'id': 187, 'name': 'billiards'},
 {'id': 188, 'name': 'fire hydrant'},
 {'id': 189, 'name': 'goose'},
 {'id': 190, 'name': 'projector'},
 {'id': 191, 'name': 'sausage'},
 {'id': 192, 'name': 'fire extinguisher'},
 {'id': 193, 'name': 'extension cord'},
 {'id': 194, 'name': 'facial mask'},
 {'id': 195, 'name': 'tennis ball'},
 {'id': 196, 'name': 'chopsticks'},
 {'id': 197, 'name': 'electronic stove and gas stove'},
 {'id': 198, 'name': 'pie'},
 {'id': 199, 'name': 'frisbee'},
 {'id': 200, 'name': 'kettle'},
 {'id': 201, 'name': 'hamburger'},
 {'id': 202, 'name': 'golf club'},
 {'id': 203, 'name': 'cucumber'},
 {'id': 204, 'name': 'clutch'},
 {'id': 205, 'name': 'blender'},
 {'id': 206, 'name': 'tong'},
 {'id': 207, 'name': 'slide'},
 {'id': 208, 'name': 'hot dog'},
 {'id': 209, 'name': 'toothbrush'},
 {'id': 210, 'name': 'facial cleanser'},
 {'id': 211, 'name': 'mango'},
 {'id': 212, 'name': 'deer'},
 {'id': 213, 'name': 'egg'},
 {'id': 214, 'name': 'violin'},
 {'id': 215, 'name': 'marker'},
 {'id': 216, 'name': 'ship'},
 {'id': 217, 'name': 'chicken'},
 {'id': 218, 'name': 'onion'},
 {'id': 219, 'name': 'ice cream'},
 {'id': 220, 'name': 'tape'},
 {'id': 221, 'name': 'wheelchair'},
 {'id': 222, 'name': 'plum'},
 {'id': 223, 'name': 'bar soap'},
 {'id': 224, 'name': 'scale'},
 {'id': 225, 'name': 'watermelon'},
 {'id': 226, 'name': 'cabbage'},
 {'id': 227, 'name': 'router/modem'},
 {'id': 228, 'name': 'golf ball'},
 {'id': 229, 'name': 'pine apple'},
 {'id': 230, 'name': 'crane'},
 {'id': 231, 'name': 'fire truck'},
 {'id': 232, 'name': 'peach'},
 {'id': 233, 'name': 'cello'},
 {'id': 234, 'name': 'notepaper'},
 {'id': 235, 'name': 'tricycle'},
 {'id': 236, 'name': 'toaster'},
 {'id': 237, 'name': 'helicopter'},
 {'id': 238, 'name': 'green beans'},
 {'id': 239, 'name': 'brush'},
 {'id': 240, 'name': 'carriage'},
 {'id': 241, 'name': 'cigar'},
 {'id': 242, 'name': 'earphone'},
 {'id': 243, 'name': 'penguin'},
 {'id': 244, 'name': 'hurdle'},
 {'id': 245, 'name': 'swing'},
 {'id': 246, 'name': 'radio'},
 {'id': 247, 'name': 'CD'},
 {'id': 248, 'name': 'parking meter'},
 {'id': 249, 'name': 'swan'},
 {'id': 250, 'name': 'garlic'},
 {'id': 251, 'name': 'french fries'},
 {'id': 252, 'name': 'horn'},
 {'id': 253, 'name': 'avocado'},
 {'id': 254, 'name': 'saxophone'},
 {'id': 255, 'name': 'trumpet'},
 {'id': 256, 'name': 'sandwich'},
 {'id': 257, 'name': 'cue'},
 {'id': 258, 'name': 'kiwi fruit'},
 {'id': 259, 'name': 'bear'},
 {'id': 260, 'name': 'fishing rod'},
 {'id': 261, 'name': 'cherry'},
 {'id': 262, 'name': 'tablet'},
 {'id': 263, 'name': 'green vegetables'},
 {'id': 264, 'name': 'nuts'},
 {'id': 265, 'name': 'corn'},
 {'id': 266, 'name': 'key'},
 {'id': 267, 'name': 'screwdriver'},
 {'id': 268, 'name': 'globe'},
 {'id': 269, 'name': 'broom'},
 {'id': 270, 'name': 'pliers'},
 {'id': 271, 'name': 'volleyball'},
 {'id': 272, 'name': 'hammer'},
 {'id': 273, 'name': 'eggplant'},
 {'id': 274, 'name': 'trophy'},
 {'id': 275, 'name': 'dates'},
 {'id': 276, 'name': 'board eraser'},
 {'id': 277, 'name': 'rice'},
 {'id': 278, 'name': 'tape measure/ruler'},
 {'id': 279, 'name': 'dumbbell'},
 {'id': 280, 'name': 'hamimelon'},
 {'id': 281, 'name': 'stapler'},
 {'id': 282, 'name': 'camel'},
 {'id': 283, 'name': 'lettuce'},
 {'id': 284, 'name': 'goldfish'},
 {'id': 285, 'name': 'meat balls'},
 {'id': 286, 'name': 'medal'},
 {'id': 287, 'name': 'toothpaste'},
 {'id': 288, 'name': 'antelope'},
 {'id': 289, 'name': 'shrimp'},
 {'id': 290, 'name': 'rickshaw'},
 {'id': 291, 'name': 'trombone'},
 {'id': 292, 'name': 'pomegranate'},
 {'id': 293, 'name': 'coconut'},
 {'id': 294, 'name': 'jellyfish'},
 {'id': 295, 'name': 'mushroom'},
 {'id': 296, 'name': 'calculator'},
 {'id': 297, 'name': 'treadmill'},
 {'id': 298, 'name': 'butterfly'},
 {'id': 299, 'name': 'egg tart'},
 {'id': 300, 'name': 'cheese'},
 {'id': 301, 'name': 'pig'},
 {'id': 302, 'name': 'pomelo'},
 {'id': 303, 'name': 'race car'},
 {'id': 304, 'name': 'rice cooker'},
 {'id': 305, 'name': 'tuba'},
 {'id': 306, 'name': 'crosswalk sign'},
 {'id': 307, 'name': 'papaya'},
 {'id': 308, 'name': 'hair drier'},
 {'id': 309, 'name': 'green onion'},
 {'id': 310, 'name': 'chips'},
 {'id': 311, 'name': 'dolphin'},
 {'id': 312, 'name': 'sushi'},
 {'id': 313, 'name': 'urinal'},
 {'id': 314, 'name': 'donkey'},
 {'id': 315, 'name': 'electric drill'},
 {'id': 316, 'name': 'spring rolls'},
 {'id': 317, 'name': 'tortoise/turtle'},
 {'id': 318, 'name': 'parrot'},
 {'id': 319, 'name': 'flute'},
 {'id': 320, 'name': 'measuring cup'},
 {'id': 321, 'name': 'shark'},
 {'id': 322, 'name': 'steak'},
 {'id': 323, 'name': 'poker card'},
 {'id': 324, 'name': 'binoculars'},
 {'id': 325, 'name': 'llama'},
 {'id': 326, 'name': 'radish'},
 {'id': 327, 'name': 'noodles'},
 {'id': 328, 'name': 'yak'},
 {'id': 329, 'name': 'mop'},
 {'id': 330, 'name': 'crab'},
 {'id': 331, 'name': 'microscope'},
 {'id': 332, 'name': 'barbell'},
 {'id': 333, 'name': 'bread/bun'},
 {'id': 334, 'name': 'baozi'},
 {'id': 335, 'name': 'lion'},
 {'id': 336, 'name': 'red cabbage'},
 {'id': 337, 'name': 'polar bear'},
 {'id': 338, 'name': 'lighter'},
 {'id': 339, 'name': 'seal'},
 {'id': 340, 'name': 'mangosteen'},
 {'id': 341, 'name': 'comb'},
 {'id': 342, 'name': 'eraser'},
 {'id': 343, 'name': 'pitaya'},
 {'id': 344, 'name': 'scallop'},
 {'id': 345, 'name': 'pencil case'},
 {'id': 346, 'name': 'saw'},
 {'id': 347, 'name': 'table tennis paddle'},
 {'id': 348, 'name': 'okra'},
 {'id': 349, 'name': 'starfish'},
 {'id': 350, 'name': 'eagle'},
 {'id': 351, 'name': 'monkey'},
 {'id': 352, 'name': 'durian'},
 {'id': 353, 'name': 'game board'},
 {'id': 354, 'name': 'rabbit'},
 {'id': 355, 'name': 'french horn'},
 {'id': 356, 'name': 'ambulance'},
 {'id': 357, 'name': 'asparagus'},
 {'id': 358, 'name': 'hoverboard'},
 {'id': 359, 'name': 'pasta'},
 {'id': 360, 'name': 'target'},
 {'id': 361, 'name': 'hotair balloon'},
 {'id': 362, 'name': 'chainsaw'},
 {'id': 363, 'name': 'lobster'},
 {'id': 364, 'name': 'iron'},
 {'id': 365, 'name': 'flashlight'}]

def _get_builtin_metadata_obj365v1():
    id_to_name = {x['id']: x['name'] for x in categories}
    thing_dataset_id_to_contiguous_id = {i + 1: i for i in range(365)}
    thing_classes = [id_to_name[k] for k in sorted(id_to_name)]
    return {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes}

_PREDEFINED_SPLITS_OBJECTS365 = {
    "objects365_train": ("objects365/train", "objects365/annotations/objects365_train.json"),
    "objects365_val": ("objects365/val", "objects365/annotations/objects365_val.json"),
}

# for key, (image_root, json_file) in _PREDEFINED_SPLITS_OBJECTS365.items():
#     register_coco_instances(
#         key,
#         _get_builtin_metadata(),
#         os.path.join("datasets", json_file) if "://" not in json_file else json_file,
#         os.path.join("datasets", image_root),
#     )

_PREDEFINED_SPLITS_OBJECTS365V1 = {
    "objects365_v1_train": ("Objects365v1/train", "Objects365v1/objects365_train.json"),
    "objects365_v1_masktrain": ("Objects365v1/train", "Objects365v1/annotations/filtered_objects365_v1_train_with_mask.json"),
    "objects365_v1_val": ("Objects365v1/val/val", "Objects365v1/objects365_val.json"),
    "objects365_v1_val_mini": ("Objects365v1/val/val", "Objects365v1/objects365_val_mini.json"),
}

def register_all_obj365v1(root):
    for key, (image_root, json_file) in _PREDEFINED_SPLITS_OBJECTS365V1.items():
        register_coco_instances(
            key,
            _get_builtin_metadata_obj365v1(),
            os.path.join(root, json_file) if "://" not in json_file else json_file,
            os.path.join(root, image_root),
        )

_root = os.getenv("DETECTRON2_DATASETS", "datasets")
register_all_obj365v1(_root)