forked from WoodsGao/cv_utils
-
Notifications
You must be signed in to change notification settings - Fork 0
/
labelme2coco.py
executable file
·100 lines (92 loc) · 3.06 KB
/
labelme2coco.py
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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
#!/usr/bin/python3
import argparse
import json
import os
import os.path as osp
import cv2
import numpy as np
from coco_utils import create_coco, insert_img_anns
POINTS_WH = 30
def labelme2coco(path, img_root='images'):
if not img_root:
img_root = path
coco = create_coco()
categories = []
if osp.isfile(path):
files = [path]
path = osp.dirname(path)
else:
files = [
osp.join(path, f) for f in os.listdir(path)
if osp.splitext(f)[-1] == '.json'
]
for data in files:
with open(data, 'r') as f:
data = json.loads(f.read())
img_info = {
'file_name':
osp.relpath(osp.abspath(osp.join(img_root, data['imagePath'])),
'./'),
'width':
data['imageWidth'],
'height':
data['imageHeight'],
}
anns = []
for shapes in data['shapes']:
label = shapes['label']
if label not in categories:
coco['categories'][len(categories)]['name'] = label
categories.append(label)
cid = categories.index(label)
points = np.float32(shapes['points']).reshape(-1).tolist()
# polygon
if shapes['shape_type'] == 'polygon':
xs = points[::2]
ys = points[1::2]
x1 = min(xs)
y1 = min(ys)
w = max(xs) - x1
h = max(ys) - y1
anns.append({
'area': w * h,
'bbox': [x1, y1, w, h],
'category_id': cid,
'iscrowd': 0,
'segmentation': [points]
})
elif shapes['shape_type'] == 'line':
x, y = points[::2], points[1::2]
x1 = min(x) - POINTS_WH
y1 = min(y) - POINTS_WH
x2 = max(x) + POINTS_WH
y2 = max(y) + POINTS_WH
w = x2 - x1
h = y2 - y1
kps = []
for x_, y_ in zip(x, y):
kps.append(x_)
kps.append(y_)
kps.append(2)
anns.append({
'area': w * h,
'bbox': [x1, y1, w, h],
'category_id': cid,
'iscrowd': 0,
'segmentation': [[x1, y1, x1, y2, x2, y2, x2, y1]],
"keypoints": kps,
"num_keypoints": len(x)
})
coco = insert_img_anns(coco, img_info, anns)
coco['categories'] = [
c for c in coco['categories'] if c['name'] is not None
]
save_path = osp.join(path, '../coco.json')
with open(save_path, 'w') as f:
f.write(json.dumps(coco, indent=4))
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('path', type=str)
parser.add_argument('--img-root', default='images', type=str)
opt = parser.parse_args()
labelme2coco(opt.path, opt.img_root)