def split_coco(coco_path, val_ratio, shuffle):
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
train_coco = create_coco(coco)
val_coco = create_coco(coco)
img_info_list, anns_list = find_all_img_anns(coco)
indexs = list(range(len(img_info_list)))
if shuffle:
random.shuffle(indexs)
for ii, i in enumerate(indexs):
img_info = img_info_list[i]
anns = anns_list[i]
if ii < len(img_info_list) * (1 - val_ratio):
train_coco = insert_img_anns(train_coco, img_info, anns)
else:
val_coco = insert_img_anns(val_coco, img_info, anns)
save_path = osp.join(osp.dirname(coco_path), 'train.json')
with open(save_path, 'w') as f:
f.write(json.dumps(train_coco, indent=4, sort_keys=True))
save_path = osp.join(osp.dirname(coco_path), 'val.json')
with open(save_path, 'w') as f:
f.write(json.dumps(val_coco, indent=4, sort_keys=True))
def coco_offline_aug(coco_path, img_root, output, repeats=1):
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
new_coco = create_coco(coco)
save_dir = osp.join(output, 'images')
os.makedirs(save_dir, exist_ok=True)
for img_info in coco['images']:
anns = find_anns(coco, img_info)
img = cv2.imread(osp.join(img_root, img_info['file_name']))
img_name = osp.splitext(osp.basename(img_info['file_name']))[0]
anns = find_anns(coco, img_info)
for ri in range(repeats):
aug_img, aug_anns = aug_img_anns(img.copy(), anns)
iname = img_name + '_%05d.png' % ri
cv2.imwrite(osp.join(save_dir, iname), aug_img)
print(osp.join(save_dir, iname))
img_info = {
'file_name': osp.join('images', iname),
'width': aug_img.shape[1],
'height': aug_img.shape[0],
}
insert_img_anns(new_coco, img_info, aug_anns)
save_path = osp.join(output, 'coco_aug.json')
# print(new_coco)
with open(save_path, 'w') as f:
f.write(json.dumps(new_coco, indent=4))
def crop_coco_image(coco_path, img_root, output, img_size, steps):
print(img_size)
save_path = osp.join(output, 'images')
os.makedirs(save_path, exist_ok=True)
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
new_coco = create_coco(coco)
for img_info in coco['images']:
if img_info['width'] <= img_size[0] and img_info['height'] <= img_size[
1]:
anns = find_anns(coco, img_info)
insert_img_anns(new_coco, img_info, anns)
continue
print(osp.join(img_root, img_info['file_name']))
img = cv2.imread(osp.join(img_root, img_info['file_name']))
img_name = osp.splitext(osp.basename(img_info['file_name']))[0]
anns = find_anns(coco, img_info)
imgs_split, anns_split = crop_img_ann(img, anns, img_size, steps)
for si, (img, anns) in enumerate(zip(imgs_split, anns_split)):
iname = img_name + '_%05d.png' % si
cv2.imwrite(osp.join(save_path, iname), img)
print(osp.join('images', iname))
img_info = {
'file_name': osp.join('images', iname),
'width': img.shape[1],
'height': img.shape[0],
}
insert_img_anns(new_coco, img_info, anns)
save_path = osp.join(output, 'coco_split.json')
with open(save_path, 'w') as f:
f.write(json.dumps(new_coco, indent=4))
def crop_coco_instance(coco_path, img_root, output, square=False):
save_path = osp.join(output, 'images')
os.makedirs(save_path, exist_ok=True)
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
new_coco = create_coco(coco)
img_info_list, anns_list = find_all_img_anns(coco)
for img_info, anns in zip(img_info_list, anns_list):
if len(anns) == 0:
continue
print(osp.join(img_root, img_info['file_name']))
img = cv2.imread(osp.join(img_root, img_info['file_name']))
img_name = osp.splitext(osp.basename(img_info['file_name']))[0]
for ai, ann in enumerate(anns):
iname = img_name + '_%05d.png' % ai
seg = np.float32(ann['segmentation'])
p = seg.reshape(-1, 2).transpose(1, 0).astype(np.int32)
x1 = p[0].min()
x2 = p[0].max()
y1 = p[1].min()
y2 = p[1].max()
# square crop - it may be better for segmentation
if square:
center = ((x1 + x2) // 2, (y1 + y2) // 2)
r = max(x2 - x1, y2 - y1) // 2
x1 = center[0] - r
x2 = center[0] + r
y1 = center[1] - r
y2 = center[1] + r
x1_ = max(random.randint(x1 - 50, x1 - 10), 0)
y1_ = max(random.randint(y1 - 50, y1 - 10), 0)
x2_ = min(random.randint(x2 + 10, x2 + 50), img.shape[1])
y2_ = min(random.randint(y2 + 10, y2 + 50), img.shape[0])
cut = img[y1_:y2_, x1_:x2_]
seg[:, ::2] -= x1_
seg[:, 1::2] -= y1_
ann['bbox'] = np.array([x1 - x1_, y1 - y1_, x2 - x1,
y2 - y1]).tolist()
ann['segmentation'] = seg.tolist()
cv2.imwrite(osp.join(save_path, iname), cut)
print(osp.join('images', iname))
img_info = {
'file_name': osp.join('images', iname),
'width': cut.shape[1],
'height': cut.shape[0],
}
insert_img_anns(new_coco, img_info, [ann])
save_path = osp.join(output, 'coco_instance.json')
with open(save_path, 'w') as f:
f.write(json.dumps(new_coco, indent=4))
def yolo2coco(img_root):
coco = create_coco()
cls_txt = osp.join(img_root, 'class_names.txt')
if osp.exists(cls_txt):
with open(cls_txt, 'r') as f:
classes = [c for c in f.read().split('\n') if c]
for ci in range(len(coco['categories'])):
if ci < len(classes):
coco['categories'][ci]['name'] = classes[ci]
coco['categories'] = [
c for c in coco['categories'] if c['name'] is not None
]
img_list = os.listdir(osp.join(img_root, 'images'))
for img in img_list:
ih, iw, _ = cv2.imread(osp.join(img_root, 'images', img)).shape
label_name = osp.join(img_root, 'labels',
''.join(*(osp.splitext(img)[:-1])) + '.txt')
anns = []
if osp.exists(label_name):
labels = np.loadtxt(label_name, delimiter=' ', ndmin=2)
for (c, x, y, w, h) in labels:
x *= iw
y *= ih
w *= iw
h *= ih
x1 = x - w / 2
y1 = y - h / 2
x2 = x + w / 2
y2 = y + h / 2
anns.append({
'area': w * h,
'bbox': [x1, y1, w, h],
'category_id': int(c),
'iscrowd': 0,
# mask, 矩形是从左上角点按顺时针的四个顶点
'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
})
img_info = {
'file_name': osp.join('images', img),
'width': iw,
'height': ih
}
insert_img_anns(coco, img_info, anns)
with open(osp.join(img_root, 'coco.json'), 'w') as f:
f.write(json.dumps(coco, indent=4, sort_keys=True))
def merge_coco(path):
merge_coco = None
coco_list = os.listdir(path)
coco_list = [c for c in coco_list if osp.splitext(c)[-1] == '.json']
for coco_path in coco_list:
coco_path = osp.join(path, coco_path)
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
if merge_coco is None:
merge_coco = create_coco(coco)
img_info_list, anns_list = find_all_img_anns(coco)
for i in range(len(img_info_list)):
img_info = img_info_list[i]
anns = anns_list[i]
merge_coco = insert_img_anns(merge_coco, img_info, anns)
save_path = osp.join(osp.dirname(coco_path), 'merge.json')
with open(save_path, 'w') as f:
f.write(json.dumps(merge_coco, indent=4, sort_keys=True))
def coco_seg2points(coco_path, img_root, output):
save_path = osp.join(output, 'images')
os.makedirs(save_path, exist_ok=True)
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(coco)
new_coco = create_coco(coco)
for img_info in coco['images']:
anns = find_anns(coco, img_info)
if len(anns) == 0:
continue
print(osp.join(img_root, img_info['file_name']))
new_anns = []
for ai, ann in enumerate(anns):
points = np.int32(ann['segmentation']).reshape(-1, 2)
for p in points:
x, y = p.tolist()
x1 = max(0, x - 5)
y1 = max(0, y - 5)
x2 = min(img_info['width'], x + 5)
y2 = min(img_info['height'], y + 5)
w = x2 - x1
h = y2 - y1
new_anns.append({
'area':
w * h,
'bbox': [x1, y1, x2 - x1, y2 - y1],
'category_id':
ann['category_id'],
'iscrowd':
0,
'segmentation': [[x1, y1, x2, y1, x2, y2, x1, y2]]
})
insert_img_anns(new_coco, img_info, new_anns)
save_path = osp.join(output, 'coco_seg2points.json')
with open(save_path, 'w') as f:
f.write(json.dumps(new_coco, indent=4))
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))