295 lines
14 KiB
Python
295 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
|
# @Author : LG
|
|
|
|
from json import load, dump
|
|
import os
|
|
from pycocotools import mask as coco_mask
|
|
import cv2
|
|
import imgviz
|
|
import yaml
|
|
import numpy as np
|
|
|
|
|
|
class COCOConverter:
|
|
def convert_to_coco(self, isat_json_root:str, to_path:str):
|
|
coco_anno = {}
|
|
# info
|
|
coco_anno['info'] = {}
|
|
coco_anno['info']['description'] = 'coco from ISAT'
|
|
coco_anno['info']['version'] = None
|
|
coco_anno['info']['year'] = None
|
|
coco_anno['info']['contributor'] = None
|
|
coco_anno['info']['date_created'] = None
|
|
|
|
# licenses
|
|
coco_anno['licenses'] = []
|
|
license1 = {}
|
|
license1['url'] = None
|
|
license1['id'] = 0
|
|
license1['name'] = None
|
|
coco_anno['licenses'].append(license1)
|
|
|
|
# images and annotations
|
|
coco_anno['images'] = []
|
|
coco_anno['annotations'] = []
|
|
coco_anno['categories'] = []
|
|
|
|
categories_dict = {}
|
|
|
|
jsons = [f for f in os.listdir(isat_json_root) if f.endswith('.json')]
|
|
for file_index, json in enumerate(jsons):
|
|
print('Load ISAT: {}'.format(json))
|
|
try:
|
|
with open(os.path.join(isat_json_root, json), 'r') as f:
|
|
dataset = load(f)
|
|
info = dataset.get('info', {})
|
|
description = info.get('description', '')
|
|
if not description.startswith('ISAT'):
|
|
# 不是ISAT格式json
|
|
continue
|
|
|
|
img_name = info.get('name', '')
|
|
width = info.get('width', None)
|
|
height = info.get('height', None)
|
|
depth = info.get('depth', None)
|
|
note = info.get('note', '')
|
|
objects = dataset.get('objects', [])
|
|
|
|
# image
|
|
coco_image_info = {}
|
|
coco_image_info['license'] = None
|
|
coco_image_info['url'] = None
|
|
coco_image_info['file_name'] = img_name
|
|
coco_image_info['height'] = height
|
|
coco_image_info['width'] = width
|
|
coco_image_info['date_captured'] = None
|
|
coco_image_info['id'] = file_index
|
|
coco_anno['images'].append(coco_image_info)
|
|
|
|
objects_groups = [obj.get('group', 0) for obj in objects]
|
|
objects_groups.sort()
|
|
objects_groups = set(objects_groups)
|
|
# 同group
|
|
for group_index, group in enumerate(objects_groups):
|
|
objs_with_group = [obj for obj in objects if obj.get('group', 0) == group]
|
|
cats = [obj.get('category', 'unknow') for obj in objs_with_group]
|
|
cats = set(cats)
|
|
# 同category
|
|
for cat in cats:
|
|
if cat not in categories_dict:
|
|
categories_dict[cat] = len(categories_dict)
|
|
category_index = categories_dict.get(cat)
|
|
|
|
objs_with_cat = [obj for obj in objs_with_group if obj.get('category', 0) == cat]
|
|
crowds = [obj.get('iscrowd', 'unknow') for obj in objs_with_group]
|
|
crowds = set(crowds)
|
|
# 同iscrowd
|
|
for crowd in crowds:
|
|
objs_with_crowd = [obj for obj in objs_with_cat if obj.get('iscrowd', 0) == crowd]
|
|
# anno
|
|
coco_anno_info = {}
|
|
coco_anno_info['iscrowd'] = crowd
|
|
coco_anno_info['image_id'] = file_index
|
|
coco_anno_info['image_name'] = img_name
|
|
coco_anno_info['category_id'] = category_index
|
|
coco_anno_info['id'] = len(coco_anno['annotations'])
|
|
coco_anno_info['segmentation'] = []
|
|
coco_anno_info['area'] = 0.
|
|
coco_anno_info['bbox'] = []
|
|
|
|
for obj in objs_with_crowd:
|
|
|
|
segmentation = obj.get('segmentation', [])
|
|
area = obj.get('area', 0)
|
|
bbox = obj.get('bbox', [])
|
|
if bbox is None:
|
|
segmentation_nd = np.array(segmentation)
|
|
bbox = [min(segmentation_nd[:, 0]), min(segmentation_nd[:, 1]),
|
|
max(segmentation_nd[:, 0]), max(segmentation_nd[:, 1])]
|
|
del segmentation_nd
|
|
segmentation = [e for p in segmentation for e in p]
|
|
|
|
|
|
|
|
if bbox != []:
|
|
if coco_anno_info['bbox'] == []:
|
|
coco_anno_info['bbox'] = bbox
|
|
else:
|
|
bbox_tmp = coco_anno_info['bbox']
|
|
bbox_tmp = [min(bbox_tmp[0], bbox[0]), min(bbox_tmp[1], bbox[1]),
|
|
max(bbox_tmp[2], bbox[2]), max(bbox_tmp[3], bbox[3])]
|
|
coco_anno_info['bbox'] = bbox_tmp
|
|
coco_anno_info['segmentation'].append(segmentation)
|
|
if area is not None:
|
|
coco_anno_info['area'] += float(area)
|
|
|
|
# (xmin, ymin, xmax, ymax) 2 (xmin, ymin, w, h)
|
|
bbox_tmp = coco_anno_info['bbox']
|
|
coco_anno_info['bbox'] = [bbox_tmp[0], bbox_tmp[1],
|
|
bbox_tmp[2] - bbox_tmp[0], bbox_tmp[3] - bbox_tmp[1]]
|
|
|
|
coco_anno['annotations'].append(coco_anno_info)
|
|
except Exception as e:
|
|
print('Load ISAT: {}, error: {}'.format(json, e))
|
|
|
|
categories_dict = sorted(categories_dict.items(), key=lambda x:x[1])
|
|
coco_anno['categories'] = [{'name': name, 'id': id, 'supercategory': None} for name, id in categories_dict]
|
|
|
|
with open(to_path, 'w') as f:
|
|
try:
|
|
dump(coco_anno, f)
|
|
print('Save coco json to {}'.format(to_path))
|
|
except Exception as e:
|
|
print('Save {} error :{}'.format(to_path, e))
|
|
|
|
|
|
def convert_from_coco(self, coco_json_path:str, to_root:str, keep_crowd:bool=False):
|
|
assert coco_json_path.endswith('.json')
|
|
annos = {}
|
|
if os.path.exists(coco_json_path):
|
|
with open(coco_json_path, 'r') as f:
|
|
dataset = load(f)
|
|
images = {image.get('id', None):{
|
|
'file_name': image.get('file_name', ''),
|
|
'height': image.get('height', ''),
|
|
'width': image.get('width', ''),
|
|
} for image in dataset.get('images', [])}
|
|
annotations = dataset.get('annotations', [])
|
|
categories = {categorie.get('id', None): {'name': categorie.get('name', '')} for categorie in dataset.get('categories', [])}
|
|
for index, annotation in enumerate(annotations):
|
|
|
|
annotation_index = annotation.get('id')
|
|
annotation_image_id = annotation.get('image_id')
|
|
annotation_category_id = annotation.get('category_id')
|
|
|
|
file_name = images[annotation_image_id].get('file_name')
|
|
height = images[annotation_image_id].get('height')
|
|
width = images[annotation_image_id].get('width')
|
|
iscrowd = annotation["iscrowd"]
|
|
|
|
if file_name == '000000279278.jpg':
|
|
continue
|
|
if annotation_image_id not in annos:
|
|
annos[annotation_image_id] = {}
|
|
|
|
objects = annos[annotation_image_id].get('objects', [])
|
|
|
|
if iscrowd == 0:
|
|
# polygon
|
|
segmentations = annotation.get('segmentation')
|
|
for segmentation in segmentations:
|
|
xs = segmentation[::2]
|
|
ys = segmentation[1::2]
|
|
points = [[x, y] for x ,y in zip(xs, ys)]
|
|
obj = {
|
|
'category': categories.get(annotation_category_id).get('name'),
|
|
'group': annotation_index,
|
|
'area': None,
|
|
'segmentation': points,
|
|
'layer': 1,
|
|
'bbox': None,
|
|
'iscrowd': iscrowd,
|
|
'note': ''
|
|
}
|
|
objects.append(obj)
|
|
elif iscrowd == 1 and keep_crowd:
|
|
segmentations = annotation.get('segmentation', {})
|
|
if isinstance(segmentations, dict) and 'counts' in segmentations:
|
|
# RLE
|
|
rles = coco_mask.frPyObjects(segmentations, height, width)
|
|
masks = coco_mask.decode(rles)
|
|
contours, _ = cv2.findContours(masks, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_KCOS)
|
|
for contour in contours:
|
|
points = []
|
|
for point in contour:
|
|
x, y = point[0]
|
|
points.append([float(x), float(y)])
|
|
obj = {
|
|
'category': categories.get(annotation_category_id).get('name'),
|
|
'group': annotation_index,
|
|
'area': None,
|
|
'segmentation': points,
|
|
'layer': 1,
|
|
'bbox': None,
|
|
'iscrowd': iscrowd,
|
|
'note': ''
|
|
}
|
|
objects.append(obj)
|
|
else:
|
|
# polygon
|
|
for segmentation in segmentations:
|
|
xs = segmentation[::2]
|
|
ys = segmentation[1::2]
|
|
points = [[x, y] for x, y in zip(xs, ys)]
|
|
obj = {
|
|
'category': categories.get(annotation_category_id).get('name'),
|
|
'group': annotation_index,
|
|
'area': None,
|
|
'segmentation': points,
|
|
'layer': 1,
|
|
'bbox': None,
|
|
'iscrowd': iscrowd,
|
|
'note': ''
|
|
}
|
|
objects.append(obj)
|
|
else:
|
|
pass
|
|
annos[annotation_image_id]['objects'] = objects
|
|
|
|
|
|
for image_id, values in annos.items():
|
|
image_path = images[image_id].get('file_name')
|
|
folder, name = os.path.split(image_path)
|
|
height = images[image_id].get('height')
|
|
width = images[image_id].get('width')
|
|
objects = values.get('objects', [])
|
|
|
|
isat_anno = {}
|
|
isat_anno['info'] = {}
|
|
isat_anno['info']['description'] = 'ISAT'
|
|
isat_anno['info']['folder'] = folder
|
|
isat_anno['info']['name'] = name
|
|
isat_anno['info']['width'] = width
|
|
isat_anno['info']['height'] = height
|
|
isat_anno['info']['depth'] = None
|
|
isat_anno['info']['note'] = ''
|
|
isat_anno['objects'] = []
|
|
# coco annotation的id 太大了，这里缩一下，每张图片重新开始计数
|
|
groups_dict = {}
|
|
for obj in objects:
|
|
group = obj.get('group', 0)
|
|
if group not in groups_dict:
|
|
groups_dict[group] = len(groups_dict)+1
|
|
for obj in objects:
|
|
object = {}
|
|
object['category'] = obj.get('category', '')
|
|
object['group'] = groups_dict.get(obj.get('group', 0))
|
|
object['segmentation'] = obj.get('segmentation', [])
|
|
object['area'] = obj.get('area', None)
|
|
object['layer'] = obj.get('layer', None)
|
|
object['bbox'] = obj.get('bbox', None)
|
|
object['iscrowd'] = obj.get('iscrowd', 0)
|
|
object['note'] = obj.get('note', '')
|
|
isat_anno['objects'].append(object)
|
|
json_name = '.'.join(name.split('.')[:-1]) + '.json'
|
|
save_json = os.path.join(to_root, json_name)
|
|
with open(save_json, 'w') as f:
|
|
try:
|
|
dump(isat_anno, f)
|
|
print('Converted coco to ISAT: {}'.format(json_name))
|
|
|
|
except Exception as e:
|
|
print('Convert coco to ISAT {} ,error: {}'.format(json_name, e))
|
|
|
|
### 类别文件
|
|
cmap = imgviz.label_colormap()
|
|
sorted(categories)
|
|
for index, (k, categorie_dict) in enumerate(categories.items()):
|
|
r, g, b = cmap[index+1]
|
|
categorie_dict['color'] = "#{:02x}{:02x}{:02x}".format(r, g, b)
|
|
print(categories)
|
|
|
|
s = yaml.dump({'label': list(categories.values())})
|
|
with open(os.path.join(to_root, 'categorys.yaml'), 'w') as f:
|
|
f.write(s)
|