Az/coco format support (#623)
* added coco dumper * added license note * revert changes of launch.json * fixed codacy issues * fixed codacy issuesmain
Andrey Zhavoronkov
7 years ago
committed by
Nikita Manovich
parent
098827568e
commit
7467e298fd
@ -0,0 +1,293 @@
|
|||||||
|
# Copyright (C) 2018 Intel Corporation
|
||||||
|
#
|
||||||
|
# SPDX-License-Identifier: MIT
|
||||||
|
|
||||||
|
format_spec = {
|
||||||
|
"name": "COCO",
|
||||||
|
"dumpers": [
|
||||||
|
{
|
||||||
|
"display_name": "{name} {format} {version}",
|
||||||
|
"format": "JSON",
|
||||||
|
"version": "1.0",
|
||||||
|
"handler": "dump"
|
||||||
|
},
|
||||||
|
],
|
||||||
|
"loaders": [
|
||||||
|
],
|
||||||
|
}
|
||||||
|
|
||||||
|
def dump(file_object, annotations):
|
||||||
|
import numpy as np
|
||||||
|
import json
|
||||||
|
from skimage import measure
|
||||||
|
from pycocotools import mask as mask_util
|
||||||
|
from pycocotools import coco as coco_loader
|
||||||
|
|
||||||
|
def mask_to_polygon(mask, tolerance=1.0, area_threshold=1):
|
||||||
|
"""Convert object's mask to polygon [[x1,y1, x2,y2 ...], [...]]
|
||||||
|
Args:
|
||||||
|
mask: object's mask presented as 2D array of 0 and 1
|
||||||
|
tolerance: maximum distance from original points of polygon to approximated
|
||||||
|
area_threshold: if area of a polygon is less than this value, remove this small object
|
||||||
|
"""
|
||||||
|
polygons = []
|
||||||
|
# pad mask with 0 around borders
|
||||||
|
padded_mask = np.pad(mask, pad_width=1, mode='constant', constant_values=0)
|
||||||
|
contours = measure.find_contours(padded_mask, 0.5)
|
||||||
|
# Fix coordinates after padding
|
||||||
|
contours = np.subtract(contours, 1)
|
||||||
|
for contour in contours:
|
||||||
|
if not np.array_equal(contour[0], contour[-1]):
|
||||||
|
contour = np.vstack((contour, contour[0]))
|
||||||
|
contour = measure.approximate_polygon(contour, tolerance)
|
||||||
|
if len(contour) > 2:
|
||||||
|
contour = np.flip(contour, axis=1)
|
||||||
|
reshaped_contour = []
|
||||||
|
for xy in contour:
|
||||||
|
reshaped_contour.append(xy[0])
|
||||||
|
reshaped_contour.append(xy[1])
|
||||||
|
for rcontour in reshaped_contour:
|
||||||
|
if rcontour < 0:
|
||||||
|
rcontour = 0
|
||||||
|
# Check if area of a polygon is enough
|
||||||
|
rle = mask_util.frPyObjects([reshaped_contour], mask.shape[0], mask.shape[1])
|
||||||
|
area = mask_util.area(rle)
|
||||||
|
if sum(area) > area_threshold:
|
||||||
|
polygons.append(reshaped_contour)
|
||||||
|
return polygons
|
||||||
|
|
||||||
|
def fix_segments_intersections(polygons, height, width, img_name,
|
||||||
|
threshold=0.0, ratio_tolerance=0.001, area_threshold=1):
|
||||||
|
"""Find all intersected regions and crop contour for back object by objects which
|
||||||
|
are in front of the first one. It is related to a specialty of segmentation
|
||||||
|
in CVAT annotation. Intersection is calculated via function 'iou' from cocoapi
|
||||||
|
Args:
|
||||||
|
polygons: all objects on image represented as 2D array of objects' contours
|
||||||
|
height: height of image
|
||||||
|
width: width of image
|
||||||
|
img_name: name of image file
|
||||||
|
threshold: threshold of intersection over union of two objects.
|
||||||
|
By default is set to 0 and processes any two intersected objects
|
||||||
|
ratio_tolerance: used for situation when one object is fully or almost fully
|
||||||
|
inside another one and we don't want make "hole" in one of objects
|
||||||
|
"""
|
||||||
|
empty_polygon = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
|
||||||
|
|
||||||
|
for i, _ in enumerate(polygons):
|
||||||
|
rle_bottom = mask_util.frPyObjects([polygons[i]['points']], height, width)
|
||||||
|
segment_overlapped = False
|
||||||
|
for j in range(i + 1, len(polygons)):
|
||||||
|
rle_top = mask_util.frPyObjects([polygons[j]['points']], height, width)
|
||||||
|
iou = mask_util.iou(rle_bottom, rle_top, [0, 0])
|
||||||
|
area_top = sum(mask_util.area(rle_top))
|
||||||
|
area_bottom = sum(mask_util.area(rle_bottom))
|
||||||
|
if area_bottom == 0:
|
||||||
|
continue
|
||||||
|
area_ratio = area_top / area_bottom
|
||||||
|
sum_iou = sum(iou)
|
||||||
|
|
||||||
|
# If segment is fully inside another one, save this segment as is
|
||||||
|
if area_ratio - ratio_tolerance < sum_iou[0] < area_ratio + ratio_tolerance:
|
||||||
|
continue
|
||||||
|
# Check situation when bottom segment is fully inside top.
|
||||||
|
# It means that in annotation is mistake. Save this segment as is
|
||||||
|
if 1 / area_ratio - ratio_tolerance < sum_iou[0] < 1 / area_ratio + ratio_tolerance:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if sum_iou[0] > threshold:
|
||||||
|
segment_overlapped = True
|
||||||
|
bottom_mask = np.array(mask_util.decode(rle_bottom), dtype=np.uint8)
|
||||||
|
top_mask = np.array(mask_util.decode(rle_top), dtype=np.uint8)
|
||||||
|
|
||||||
|
bottom_mask = np.subtract(bottom_mask, top_mask)
|
||||||
|
bottom_mask[bottom_mask > 1] = 0
|
||||||
|
|
||||||
|
bottom_mask = np.sum(bottom_mask, axis=2)
|
||||||
|
bottom_mask = np.array(bottom_mask > 0, dtype=np.uint8)
|
||||||
|
polygons[i]['points'] = mask_to_polygon(bottom_mask, area_threshold=area_threshold)
|
||||||
|
# If some segment is empty, do small fix to avoid error in cocoapi function
|
||||||
|
if len(polygons[i]['points']) == 0:
|
||||||
|
polygons[i]['points'] = [empty_polygon]
|
||||||
|
rle_bottom = mask_util.frPyObjects(polygons[i]['points'], height, width)
|
||||||
|
if not segment_overlapped:
|
||||||
|
polygons[i]['points'] = [polygons[i]['points']]
|
||||||
|
|
||||||
|
output_polygons = []
|
||||||
|
for polygon in polygons:
|
||||||
|
poly_len = len(polygon['points'])
|
||||||
|
if poly_len != 0 and polygon['points'] != [empty_polygon]:
|
||||||
|
output_polygons.append(polygon)
|
||||||
|
|
||||||
|
return output_polygons
|
||||||
|
|
||||||
|
|
||||||
|
def polygon_area_and_bbox(polygon, height, width):
|
||||||
|
"""Calculate area of object's polygon and bounding box around it
|
||||||
|
Args:
|
||||||
|
polygon: objects contour represented as 2D array
|
||||||
|
height: height of object's region (use full image)
|
||||||
|
width: width of object's region (use full image)
|
||||||
|
"""
|
||||||
|
rle = mask_util.frPyObjects(polygon, height, width)
|
||||||
|
area = mask_util.area(rle)
|
||||||
|
bbox = mask_util.toBbox(rle)
|
||||||
|
bbox = [min(bbox[:, 0]),
|
||||||
|
min(bbox[:, 1]),
|
||||||
|
max(bbox[:, 0] + bbox[:, 2]) - min(bbox[:, 0]),
|
||||||
|
max(bbox[:, 1] + bbox[:, 3]) - min(bbox[:, 1])]
|
||||||
|
return area, bbox
|
||||||
|
|
||||||
|
|
||||||
|
def insert_license_data(result_annotation):
|
||||||
|
"""Fill license fields in annotation by blank data
|
||||||
|
Args:
|
||||||
|
result_annotation: output annotation in COCO representation
|
||||||
|
"""
|
||||||
|
result_annotation['licenses'].append({
|
||||||
|
'name': '',
|
||||||
|
'id': 0,
|
||||||
|
'url': ''
|
||||||
|
})
|
||||||
|
|
||||||
|
|
||||||
|
def insert_info_data(annotations, result_annotation):
|
||||||
|
"""Fill available information of annotation
|
||||||
|
Args:
|
||||||
|
xml_root: root for xml parser
|
||||||
|
result_annotation: output annotation in COCO representation
|
||||||
|
"""
|
||||||
|
version = annotations.data.version
|
||||||
|
description = annotations.meta['task']['name']
|
||||||
|
date = annotations.meta['dumped']
|
||||||
|
date = date.split(' ')[0]
|
||||||
|
year = date.split('-')[0]
|
||||||
|
result_annotation['info'] = {
|
||||||
|
'contributor': '',
|
||||||
|
'date_created': date,
|
||||||
|
'description': description,
|
||||||
|
'url': '',
|
||||||
|
'version': version,
|
||||||
|
'year': year
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def insert_categories_data(annotations, result_annotation):
|
||||||
|
"""Get labels from input annotation and fill categories field in output annotation
|
||||||
|
Args:
|
||||||
|
xml_root: root for xml parser
|
||||||
|
result_annotation: output annotation in COCO representation
|
||||||
|
labels_file: path to file with labels names.
|
||||||
|
If not defined, parse annotation to get labels names
|
||||||
|
"""
|
||||||
|
def get_categories(names, sort=False):
|
||||||
|
category_map = {}
|
||||||
|
categories = []
|
||||||
|
# Sort labels by its names to make the same order of ids for different annotations
|
||||||
|
if sort:
|
||||||
|
names.sort()
|
||||||
|
cat_id = 0
|
||||||
|
for name in names:
|
||||||
|
category_map[name] = cat_id
|
||||||
|
categories.append({'id': cat_id, 'name': name, 'supercategory': ''})
|
||||||
|
cat_id += 1
|
||||||
|
return category_map, categories
|
||||||
|
|
||||||
|
categories = []
|
||||||
|
category_map = {}
|
||||||
|
label_names = [label[1]["name"] for label in annotations.meta['task']['labels']]
|
||||||
|
|
||||||
|
category_map, categories = get_categories(label_names, sort=True)
|
||||||
|
|
||||||
|
result_annotation['categories'] = categories
|
||||||
|
return category_map
|
||||||
|
|
||||||
|
|
||||||
|
def insert_image_data(image, result_annotation):
|
||||||
|
"""Get data from input annotation for image and fill fields for this image in output annotation
|
||||||
|
Args:
|
||||||
|
image: dictionary with data for image from original annotation
|
||||||
|
result_annotation: output annotation in COCO representation
|
||||||
|
"""
|
||||||
|
new_img = {}
|
||||||
|
new_img['coco_url'] = ''
|
||||||
|
new_img['date_captured'] = ''
|
||||||
|
new_img['flickr_url'] = ''
|
||||||
|
new_img['license'] = 0
|
||||||
|
new_img['id'] = image.frame
|
||||||
|
new_img['file_name'] = image.name
|
||||||
|
new_img['height'] = image.height
|
||||||
|
new_img['width'] = image.width
|
||||||
|
result_annotation['images'].append(new_img)
|
||||||
|
|
||||||
|
|
||||||
|
def insert_annotation_data(image, category_map, segm_id, obj, result_annotation):
|
||||||
|
"""Get data from input annotation for object and fill fields for this object in output annotation
|
||||||
|
Args:
|
||||||
|
image: dictionary with data for image from input CVAT annotation
|
||||||
|
category_map: map for categories represented in the annotation {name: id}
|
||||||
|
segm_id: identificator of current object
|
||||||
|
obj: includes data for the object [label, polygon]
|
||||||
|
result_annotation: output annotation in COCO representation
|
||||||
|
"""
|
||||||
|
new_anno = {}
|
||||||
|
new_anno['category_id'] = category_map[obj['label']]
|
||||||
|
new_anno['id'] = segm_id
|
||||||
|
new_anno['image_id'] = image.frame
|
||||||
|
new_anno['iscrowd'] = 0
|
||||||
|
new_anno['segmentation'] = obj['points']
|
||||||
|
area, bbox = polygon_area_and_bbox(obj['points'], image.height, image.width)
|
||||||
|
new_anno['area'] = float(np.sum(area))
|
||||||
|
new_anno['bbox'] = bbox
|
||||||
|
result_annotation['annotations'].append(new_anno)
|
||||||
|
|
||||||
|
result_annotation = {
|
||||||
|
'licenses': [],
|
||||||
|
'info': {},
|
||||||
|
'categories': [],
|
||||||
|
'images': [],
|
||||||
|
'annotations': []
|
||||||
|
}
|
||||||
|
|
||||||
|
insert_license_data(result_annotation)
|
||||||
|
insert_info_data(annotations, result_annotation)
|
||||||
|
category_map = insert_categories_data(annotations, result_annotation)
|
||||||
|
|
||||||
|
segm_id = 0
|
||||||
|
for img in annotations.group_by_frame():
|
||||||
|
polygons = []
|
||||||
|
|
||||||
|
for shape in img.labeled_shapes:
|
||||||
|
if shape.type == 'polygon' or shape.type == 'rectangle':
|
||||||
|
polygon = {
|
||||||
|
'label': shape.label,
|
||||||
|
'points': shape.points,
|
||||||
|
'z_order': shape.z_order,
|
||||||
|
}
|
||||||
|
|
||||||
|
if shape.type == 'rectangle':
|
||||||
|
xtl = polygon['points'][0]
|
||||||
|
ytl = polygon['points'][1]
|
||||||
|
xbr = polygon['points'][2]
|
||||||
|
ybr = polygon['points'][3]
|
||||||
|
polygon['points'] = [xtl, ytl, xbr, ytl, xbr, ybr, xtl, ybr]
|
||||||
|
polygons.append(polygon)
|
||||||
|
|
||||||
|
polygons.sort(key=lambda x: int(x['z_order']))
|
||||||
|
|
||||||
|
# Create new image
|
||||||
|
insert_image_data(img, result_annotation)
|
||||||
|
polygons = fix_segments_intersections(polygons, img.height, img.width, img.name)
|
||||||
|
|
||||||
|
# Create new annotation for this image
|
||||||
|
for poly in polygons:
|
||||||
|
insert_annotation_data(img, category_map, segm_id, poly, result_annotation)
|
||||||
|
segm_id += 1
|
||||||
|
|
||||||
|
file_object.write(json.dumps(result_annotation, indent=2).encode())
|
||||||
|
file_object.flush()
|
||||||
|
|
||||||
|
# Try to load created annotation via cocoapi
|
||||||
|
try:
|
||||||
|
coco_loader.COCO(file_object.name)
|
||||||
|
except:
|
||||||
|
raise
|
||||||
Loading…
Reference in New Issue