File size: 1,850 Bytes
186701e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 |
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/')
|