Spaces:
Build error
Build error
File size: 4,503 Bytes
463b952 |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
import numpy as np
import json
from pathlib import Path, PosixPath
from pycocotools.coco import COCO
def min_index(arr1, arr2):
"""
Find a pair of indexes with the shortest distance.
Args:
arr1: (N, 2).
arr2: (M, 2).
Return:
a pair of indexes (tuple)
"""
dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1)
return np.unravel_index(np.argmin(dis, axis=None), dis.shape)
def merge_multi_segment(segments):
"""
Merge multi segments to one list.
Find coordinates with min distance between each segment,
then connect these coordinates with one thin line to merge all
segments into one.
Args:
segments (List(List)): original segmentations in coco's json file
like [segmentation1, segmentation2, ...], where
each segmentation is a list of coordinates
"""
s = []
segments = [np.array(i).reshape(-1,2) for i in segments]
idx_list = [[] for _ in range(len(segments))]
# record the indexes with the min distance between each segment
for i in range(1, len(segments)):
idx1, idx2 = min_index(segments[i - 1, segments[i]])
idx_list[i - 1].append(idx1)
idx_list[i].append(idx2)
# use two round to connect all the segments
for k in range(2):
# forward connection
if k == 0:
for i, idx in enumerate(idx_list):
# middle segments have two indexes
# reverse the index of middle segments
if len(idx) == 2 and idx[0] > idx[1]:
idx = idx[::-1]
segments[i] = segments[i][::-1, :]
segments[i] = np.roll(segments[i], -idx[0], axis=0)
segments[i] = np.concatenate([segments[i], segments[i][:1]])
# deal with the first segment and the last one
if i in [0, len(idx_list) - 1]:
s.append(segments[i])
else:
idx = [0, idx[1] - idx[0]]
s.append(segments[i][idx[0]:idx[1] + 1])
else:
for i in range(len(idx_list) - 1, -1, -1):
if i not in [0, len(idx_list) - 1]:
idx = idx_list[i]
nidx = abs(idx[1] - idx[0])
s.append(segments[i][nidx:])
return s
def get_yolo_labels(path2json, use_segment=False):
if not isinstance(path2json, PosixPath):
path2json = Path(path2json)
path2labels = path2json.parents[0] / "labels"
path2labels.mkdir(parents=True, exist_ok=True)
coco = COCO(path2json)
img2anns = {}
for ann in coco.dataset['annotations']:
img_id = ann['image_id']
if img_id not in img2anns:
img2anns[img_id] = [ann]
else:
img2anns[img_id].append(ann)
id2img = {img["id"]: img for img in coco.dataset["images"]}
for img_id, anns in img2anns.items():
img = id2img[img_id]
h, w, f = img['height'], img['width'], img['file_name']
bboxes = []
segments = []
for ann in anns:
if ann['iscrowd']:
continue
# coco box format: [top left x, top left y, width, height]
box = np.array(ann['bbox'], dtype=np.float64)
box[:2] += box[2:] / 2 # center coordinates
box[[0, 2]] /= w # normalize x
box[[1, 3]] /= h # normalize y
if box[2] <= 0 or box[3] <= 0:
continue
cls = ann['category_id'] - 1
box = [cls] + box.tolist()
if box not in bboxes:
bboxes.append(box)
# segmentation?
if use_segment:
if len(ann['segmentation']) > 1:
s = merge_multi_segment(ann['segmentation'])
s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist()
else:
s = [j for i in ann['segmentation'] for j in i] # all segments concatenated
s = (np.array(s).reshape(-1, 2) / np.array([w, h])).reshape(-1).tolist()
s = [cls] + s
if s not in segments:
segments.append(s)
# write
with open((path2labels / f).with_suffix('.txt'), 'a') as file:
for i in range(len(bboxes)):
line = *(segments[i] if use_segment else bboxes[i]),
file.write(('%g ' * len(line)).rstrip() % line + '\n') |