datasets/split_handal_byobj.py

import json
import os
from PIL import Image
import numpy as np
from pycocotools.mask import encode, decode, frPyObjects
from tqdm import tqdm
import copy
from natsort import natsorted

if __name__ == '__main__':
    json_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/handal_test_all.json"
    root_path = "/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL"
    with open(json_path, "r") as fp:
        datas = json.load(fp)
    objs = os.listdir(root_path)
    objs.remove("vis_rondom_check")
    objs = [file for file in objs if not file.endswith('.json')]
    for obj in tqdm(objs):
        data_save = []
        for data in datas:
            if data['image'].split('/')[0] == obj:
                data_save.append(data)
        print(f"num of {obj}", len(data_save))
        save_path = f"/data/work-gcp-europe-west4-a/yuqian_fu/datasets/HANDAL/{obj}_test.json"
        with open(save_path, "w") as fp:
            json.dump(data_save, fp)