PaddleX/DataAnnotation/labelme2coco.py

#!/usr/bin/env python
# coding: utf-8
import argparse
import glob
import json
import os
import os.path as osp
import sys

import numpy as np
import PIL.ImageDraw


class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(MyEncoder, self).default(obj)



def images(data, num):
    image = {}
    image['height'] = data['imageHeight']
    image['width'] = data['imageWidth']
    image['id'] = num + 1
    image['file_name'] = data['imagePath'].split('/')[-1]
    return image


def categories(label, labels_list):
    category = {}
    category['supercategory'] = 'component'
    category['id'] = len(labels_list) + 1
    category['name'] = label
    return category


def annotations_rectangle(iscrowd, points, label, num, label_to_num, count):
    annotation = {}
    seg_points = np.asarray(points).copy()
    seg_points[1, :] = np.asarray(points)[2, :]
    seg_points[2, :] = np.asarray(points)[1, :]
    annotation['segmentation'] = [list(seg_points.flatten())]
    annotation['iscrowd'] = iscrowd
    annotation['image_id'] = num + 1
    annotation['bbox'] = list(
        map(
            float,
            [
                points[0][0],
                points[0][1],
                points[1][0] - points[0][0],
                points[1][1] - points[0][1],
            ], ), )
    annotation['area'] = annotation['bbox'][2] * annotation['bbox'][3]
    annotation['category_id'] = label_to_num[label]
    annotation['id'] = count
    return annotation


def annotations_polygon(annotation, iscrowd, height, width, points, label, num,
                        label_to_num, count):
    
    if len(annotation) == 0:
        annotation['segmentation'] = [list(np.asarray(points).flatten())]
        annotation['iscrowd'] = iscrowd
        annotation['image_id'] = num + 1
        annotation['bbox'] = list(map(float, get_bbox(height, width, points)))
        annotation['area'] = annotation['bbox'][2] * annotation['bbox'][3]
        annotation['category_id'] = label_to_num[label]
        annotation['id'] = count
    else:
        annotation['segmentation'].append(list(np.asarray(points).flatten()))
        box1 = annotation['bbox']
        box2 = list(map(float, get_bbox(height, width, points)))
        x11, y11, x12, y12 = box1[0], box1[1], box1[0] + box1[2], box1[1] + box1[3]
        x21, y21, x22, y22 = box2[0], box2[1], box2[0] + box2[2], box2[1] + box2[3]
        x1 = x21 if x11 > x21 else x11
        y1 = y21 if y11 > y21 else y11
        x2 = x22 if x12 < x22 else x12
        y2 = y22 if y12 < y22 else y12
        annotation['bbox'] = [x1, y1, x2 - x1, y2 - y1]



def get_bbox(height, width, points):
    polygons = points
    mask = np.zeros([height, width], dtype=np.uint8)
    mask = PIL.Image.fromarray(mask)
    xy = list(map(tuple, polygons))
    PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
    mask = np.array(mask, dtype=bool)
    index = np.argwhere(mask == 1)
    rows = index[:, 0]
    clos = index[:, 1]
    left_top_r = np.min(rows)
    left_top_c = np.min(clos)
    right_bottom_r = np.max(rows)
    right_bottom_c = np.max(clos)
    return [
        left_top_c,
        left_top_r,
        right_bottom_c - left_top_c,
        right_bottom_r - left_top_r,
    ]


def deal_json(img_path, json_path):
    data_coco = {}
    label_to_num = {}
    images_list = []
    categories_list = []
    annotations_list = []
    labels_list = []
    num = -1
    for img_file in os.listdir(img_path):
        img_label = img_file.split('.')[0]
        if img_label == '':
            continue
        label_file = osp.join(json_path, img_label + '.json')
        assert os.path.exists(label_file), \
            'The .json file of {} is not exists!'.format(img_file)
        print('Generating dataset from:', label_file)
        num = num + 1
        with open(label_file) as f:
            data = json.load(f)
            images_list.append(images(data, num))
            count = 0
            lmid_count = {}
            for shapes in data['shapes']:
                count += 1
                label = shapes['label']
                part = label.split('_')
                iscrowd = int(part[-1][0])
                label = label.split('_' + part[-1])[0]
                if label not in labels_list:
                    categories_list.append(categories(label, labels_list))
                    labels_list.append(label)
                    label_to_num[label] = len(labels_list)
                points = shapes['points']
                p_type = shapes['shape_type']
                if p_type == 'polygon':
                    if len(part[-1]) > 1:
                        lmid = part[-1][1:]
                        if lmid in lmid_count:
                            real_count = lmid_count[lmid]
                            real_anno = None
                            for anno in annotations_list:
                                if anno['id'] == real_count:
                                    real_anno = anno
                                    break
                            annotations_polygon(anno, iscrowd, data['imageHeight'], data[
                                'imageWidth'], points, label, num, label_to_num,
                                                real_count)
                            count -= 1
                        else:
                            lmid_count[lmid] = count
                            anno = {}
                            annotations_polygon(anno, iscrowd, data['imageHeight'], data[
                                'imageWidth'], points, label, num, label_to_num,
                                                count)
                            annotations_list.append(anno)
                    else:
                        anno = {}
                        annotations_polygon(anno, iscrowd, data['imageHeight'], data[
                            'imageWidth'], points, label, num, label_to_num,
                                            count)
                        annotations_list.append(anno)
                if p_type == 'rectangle':
                    points.append([points[0][0], points[1][1]])
                    points.append([points[1][0], points[0][1]])
                    annotations_list.append(
                        annotations_rectangle(iscrowd, points, label, num,
                                              label_to_num, count))
    data_coco['images'] = images_list
    data_coco['categories'] = categories_list
    data_coco['annotations'] = annotations_list
    return data_coco


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter, )
    parser.add_argument('--json_input_dir', help='input annotated directory')
    parser.add_argument('--image_input_dir', help='image directory')
    args = parser.parse_args()
    try:
        assert os.path.exists(args.json_input_dir)
    except AssertionError as e:
        print('The json folder does not exist!')
        os._exit(0)
    try:
        assert os.path.exists(args.image_input_dir)
    except AssertionError as e:
        print('The image folder does not exist!')
        os._exit(0)

    # Allocate the dataset.
    total_num = len(glob.glob(osp.join(args.json_input_dir, '*.json')))

    # Deal with the json files.
    res_dir = os.path.abspath(os.path.join(args.image_input_dir, '..'))
    if not os.path.exists(res_dir + '/annotations'):
        os.makedirs(res_dir + '/annotations')
    train_data_coco = deal_json(args.image_input_dir, args.json_input_dir)
    train_json_path = osp.join(
        res_dir + '/annotations',
        'instance_{}.json'.format(
            os.path.basename(os.path.abspath(args.image_input_dir))))
    json.dump(
        train_data_coco,
        open(
            train_json_path,
            'w'),
        indent=4,
        cls=MyEncoder)


if __name__ == '__main__':
    main()