<filename>create_panoptic_dataset.py
import imagesize
from avcv.utils import *
from avcv.vision import *
import pprint


import argparse
parser = argparse.ArgumentParser()
parser.add_argument('subset', choices=['training', 'validation'])
# parser.add_argument('--arguments', '-a', nargs='+', default=None)
args = parser.parse_args()


pp = pprint.PrettyPrinter(4)

# sample = read_json('/root/detr/data/coco/coco-panoptic/annotations/panoptic_val2017.json')



#---- Create categories
categories = []
config_data = read_json("/mnt/vinai/haianh/mappilari-v2/config_v2.0.json")


name_to_cat_id = dict()
for cat_id, label in enumerate(config_data['labels']):
    cat = {'supercategory': 'unknown', 'isthing': label['instances'], 'id': cat_id, 'name': label['name'], 'readable':label['readable']}
    name_to_cat_id[label['name']] = cat_id
    categories += [cat]


# Create images list

VIS = None # './cache/vis'
segment_id = 0
def get_segment_id():
    global segment_id
    segment_id += 1
    return segment_id

resize = 0.25

images = []
annotations = []
data_root = "/mnt/vinai/haianh/mappilari-v2/"

subset = args.subset
# subset = "validation" # comment this if you want to generate training
image_dir = osp.join(data_root, subset, "images")

image_paths = get_paths(image_dir, 'jpg')
assert len(image_paths), image_dir

resize_rate = int(resize*100)
out_root = f"./data/mapillary_{resize_rate}"
image_dir_out = osp.join(out_root, subset, "images")
annotation_dir_out = osp.join(out_root, 'annotations')
panoptic_dir_out = osp.join(annotation_dir_out, "panoptic_val2017" if "val" in subset else "panoptic_train2017")

def get_polygon(filename):
    json_path = osp.join(data_root, subset, "v2.0/polygons", filename.split('.')[0]+'.json')
    return read_json(json_path)

def get_random_color(id):
    np.random.seed(id)
    return tuple(np.random.choice(256, 3).tolist())

def f(inp):
    image_id, path = inp
    file_name = osp.basename(path)


    polygons = get_polygon(file_name)
    objects, height, width = [polygons[k] for k in ["objects","height","width"]]
    
    if resize != 1:
        img = mmcv.imread(path)
        assert resize >0.2 and resize <1, resize
        img = cv2.resize(img, None, fx=resize, fy=resize)
        path = osp.join(image_dir_out, file_name)
        mmcv.imwrite(img, path)
        height = int(height*resize)
        width = int(width*resize)

        panoptic = np.zeros_like(img)

    image = {'file_name': file_name,
        'height': height,
        'width': width,
        'id': image_id}

    # pp.pprint(image)

    annotation = dict(
        image_id=image_id,
        file_name=file_name,
        segments_info=[]
    )
    if VIS:
        # mask = np.zeros([height, width, 3], 'uint8')
        
        img = mmcv.imread(path)
        mask = img.copy()

    for object in objects:
        segment_id = get_segment_id()
        cat_id = name_to_cat_id[object['label']]
        
        cnt = np.array(object['polygon'])
        if resize != 1:
            cnt = (cnt*resize).astype('int')
            cv2.drawContours(panoptic, [cnt], -1, get_random_color(segment_id), -1)


        cnt = cnt.astype('int')
        x,y,w,h = cv2.boundingRect(cnt)
        if VIS:
            color = tuple(np.random.choice(256, 3).tolist())
            cv2.drawContours(mask, [cnt], -1, color, -1)

        segment_info = {'id': segment_id,
                    'category_id': cat_id,
                    'iscrowd': 0,
                    'bbox': [x,y,w,h],
                    'area': h*w}
        annotation['segments_info'].append(segment_info)
    
    if resize != 1:
        out_panoptic_path = osp.join(panoptic_dir_out, file_name).replace('.jpg', '.png')
        mmcv.imwrite(panoptic, out_panoptic_path)

    if VIS:
        vis = (img*0.5+mask*0.5).astype('uint8')
        mmcv.imwrite(vis, osp.join(VIS, file_name))
    annotations.append(annotation)
    images.append(image)

multi_thread(f, enumerate(image_paths), verbose=True, max_workers=16)


ouput = dict(
    images=images,
    annotations=annotations,
    categories=categories,
)
mmcv.dump(ouput, osp.join(annotation_dir_out, subset+".json"))