#!/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()