|
|
@@ -0,0 +1,226 @@
|
|
|
+#!/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()
|
LaraStuStu