mirror of https://github.com/open-mmlab/mmyolo.git
59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
import os.path as osp
|
|
|
|
import mmcv
|
|
import mmengine
|
|
|
|
|
|
def convert_balloon_to_coco(ann_file, out_file, image_prefix):
|
|
|
|
data_infos = mmengine.load(ann_file)
|
|
|
|
annotations = []
|
|
images = []
|
|
obj_count = 0
|
|
for idx, v in enumerate(mmengine.track_iter_progress(data_infos.values())):
|
|
filename = v['filename']
|
|
img_path = osp.join(image_prefix, filename)
|
|
height, width = mmcv.imread(img_path).shape[:2]
|
|
|
|
images.append(
|
|
dict(id=idx, file_name=filename, height=height, width=width))
|
|
|
|
for _, obj in v['regions'].items():
|
|
assert not obj['region_attributes']
|
|
obj = obj['shape_attributes']
|
|
px = obj['all_points_x']
|
|
py = obj['all_points_y']
|
|
poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
|
|
poly = [p for x in poly for p in x]
|
|
|
|
x_min, y_min, x_max, y_max = (min(px), min(py), max(px), max(py))
|
|
|
|
data_anno = dict(
|
|
image_id=idx,
|
|
id=obj_count,
|
|
category_id=0,
|
|
bbox=[x_min, y_min, x_max - x_min, y_max - y_min],
|
|
area=(x_max - x_min) * (y_max - y_min),
|
|
segmentation=[poly],
|
|
iscrowd=0)
|
|
annotations.append(data_anno)
|
|
obj_count += 1
|
|
|
|
coco_format_json = dict(
|
|
images=images,
|
|
annotations=annotations,
|
|
categories=[{
|
|
'id': 0,
|
|
'name': 'balloon'
|
|
}])
|
|
mmengine.dump(coco_format_json, out_file)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
convert_balloon_to_coco('data/balloon/train/via_region_data.json',
|
|
'data/balloon/train.json', 'data/balloon/train/')
|
|
convert_balloon_to_coco('data/balloon/val/via_region_data.json',
|
|
'data/balloon/val.json', 'data/balloon/val/')
|