딥러닝

convert_to_coco

승무_ 2022. 9. 14. 19:51

def convert_to_coco(
    json_paths,
    save_path,
):
    """
        only for train dataset
    """
    res = defaultdict(list)
    # json_paths = glob(os.path.join(root_path, 'train', '*.json'))
    categories = {
        '01_ulcer': 1,
        '02_mass': 2,
        '04_lymph': 3,
        '05_bleeding': 4
    }
    
    n_id = 0
    for json_path in tqdm(json_paths):
        with open(json_path, 'r') as f:
            tmp = json.load(f)
            
        image = BytesIO(base64.b64decode(tmp['imageData']))
        image = Image.open(image).convert('RGB')
        
        image.save(os.path.join(base_dir, "train_img", tmp['file_name'].split(".")[0]+".jpg"))
        
        image_id = int(tmp['file_name'].split('_')[-1][:6])
        res['images'].append({
            'id': image_id,
            'width': tmp['imageWidth'],
            'height': tmp['imageHeight'],
            'file_name': tmp['file_name'].split(".")[0]+".jpg", 
        })
        
        for shape in tmp['shapes']:
            box = np.array(shape['points'])
            x1, y1, x2, y2 = \
                    min(box[:, 0]), min(box[:, 1]), max(box[:, 0]), max(box[:, 1])
            
            w, h = x2 - x1, y2 - y1
            
            res['annotations'].append({
                'id': n_id,
                'image_id': image_id,
                'category_id': categories[shape['label']],
                'area': w * h,
                'bbox': [x1, y1, w, h],
                'iscrowd': 0,
            })
            n_id += 1
    
    for name, id in categories.items():
        res['categories'].append({
            'id': id,
            'name': name,
        })
        
    with open(save_path, 'w') as f:
        json.dump(res, f)

'딥러닝' 카테고리의 다른 글

[MMdetection] AP for each class  (0) 2022.09.21
[MMdetection] samples_per_gpu  (0) 2022.09.16
[Colab] unzip  (0) 2022.09.14
[Colab] MMdetection Custom Data 학습  (1) 2022.09.10
COCO json file 병합  (0) 2022.09.08