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 kmeans_anchor(coco_path, n_clusters, img_size):
wh_list = []
with open(coco_path, 'r') as f:
coco = f.read()
coco = json.loads(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]
for ann in anns:
wh_list.append([
ann['bbox'][2] / img_info['width'],
ann['bbox'][3] / img_info['height']
])
wh_list = np.float32(wh_list)
estimator = KMeans(n_clusters=n_clusters)
estimator.fit(wh_list)
anchors = np.float32(estimator.cluster_centers_)
anchors[:, 0] *= img_size[0]
anchors[:, 1] *= img_size[1]
area = anchors[:, 0] * anchors[:, 1]
output = 'anchors: '
for i in range(n_clusters):
index = np.argmax(area)
area[index] = 0
output += '[%d,%d], ' % (int(anchors[index][0]), int(
anchors[index][1]))
print(output[:-2])
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 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 coco2labelme(coco_path, outputs):
os.makedirs(outputs, exist_ok=True)
with open(coco_path, 'r') as f:
coco = json.loads(f.read())
img_info_list, anns_list = find_all_img_anns(coco)
for img_info, anns in zip(img_info_list, anns_list):
labelme_json = copy.deepcopy(DEFAULT_LABELME)
labelme_json['imageHeight'] = img_info['height']
labelme_json['imageWidth'] = img_info['width']
shapes = []
for ann in anns:
shape = copy.deepcopy(DEFAULT_SHAPE)
points = ann['segmentation'][0]
points = list(zip(points[::2], points[1::2]))
#label = coco['categories'][ann['category_id']]['name']
label = coco['categories'][ann['category_id'] - 1][
'name'] # 注意有的ann['category_id']是从1开始的，而coco['categories']索引是0，这样需要减1
if len(points) == 1:
shape_type = 'point'
elif len(points) == 2:
shape_type = 'line'
elif len(points) <= 0:
continue
else:
shape_type = 'polygon'
shape['points'] = points
shape['shape_type'] = shape_type
shape['label'] = label
shapes.append(shape)
labelme_json['shapes'] = shapes
labelme_json['imagePath'] = img_info[
'file_name'] # osp.relpath(osp.join(osp.dirname(coco_path), img_info['file_name']), outputs)
with open(
osp.join(
outputs,
osp.splitext(osp.basename(img_info['file_name']))[0] +
'.json'), 'w') as f:
f.write(json.dumps(labelme_json, indent=4, sort_keys=True))