diff --git a/json2yolo/__init__.py b/json2yolo/__init__.py new file mode 100644 index 0000000..e22c9d6 --- /dev/null +++ b/json2yolo/__init__.py @@ -0,0 +1,3 @@ +from json2yolo.general_json2yolo import convert_coco_json, convert_ath_json, convert_vott_json, convert_infolks_json + +__all__ = ["convert_coco_json", "convert_ath_json", "convert_vott_json", "convert_infolks_json"] diff --git a/general_json2yolo.py b/json2yolo/general_json2yolo.py similarity index 84% rename from general_json2yolo.py rename to json2yolo/general_json2yolo.py index f7a3989..745521e 100644 --- a/general_json2yolo.py +++ b/json2yolo/general_json2yolo.py @@ -4,7 +4,7 @@ import pandas as pd from PIL import Image -from utils import * +from json2yolo.utils import * # Convert INFOLKS JSON file into YOLO-format labels ---------------------------- @@ -295,72 +295,6 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91 with open((fn / f).with_suffix('.txt'), 'a') as file: file.write(('%g ' * len(line)).rstrip() % line + '\n') -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). - """ - # (N, M) - dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1) - index = np.unravel_index(np.argmin(dis, axis=None), dis.shape) - return index - - -def merge_multi_segment(segments): - """Merge multi segments to one list. - Find the 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,...], - 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 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][0:1]]) - # deal with the first segment and the last one - if i == 0 or i == (len(idx_list) - 1): - s.append(segments[i]) - # deal with the middle ones - else: - idx = [0, idx[1] - idx[0]] - s.append(segments[i][idx[0]:idx[1] + 1]) - - # backward connection - else: - for i in range(len(idx_list) - 1, -1, -1): - if i != 0 and i != (len(idx_list) - 1): - idx = idx_list[i] - nidx = abs(idx[1] - idx[0]) - s.append(segments[i][nidx:]) - return s - - if __name__ == '__main__': source = 'COCO' diff --git a/labelbox_json2yolo.py b/json2yolo/labelbox_json2yolo.py similarity index 98% rename from labelbox_json2yolo.py rename to json2yolo/labelbox_json2yolo.py index 843aaf8..4ef3d12 100644 --- a/labelbox_json2yolo.py +++ b/json2yolo/labelbox_json2yolo.py @@ -7,7 +7,7 @@ from PIL import Image from tqdm import tqdm -from utils import make_dirs +from json2yolo.utils import make_dirs def convert(file, zip=True): diff --git a/utils.py b/json2yolo/utils.py similarity index 71% rename from utils.py rename to json2yolo/utils.py index 2fcbf71..aad7f61 100644 --- a/utils.py +++ b/json2yolo/utils.py @@ -164,3 +164,69 @@ def coco91_to_coco80_class(): # converts 80-index (val2014) to 91-index (paper) 51, 52, 53, 54, 55, 56, 57, 58, 59, None, 60, None, None, 61, None, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, None, 73, 74, 75, 76, 77, 78, 79, None] return x + + +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). + """ + # (N, M) + dis = ((arr1[:, None, :] - arr2[None, :, :]) ** 2).sum(-1) + index = np.unravel_index(np.argmin(dis, axis=None), dis.shape) + return index + + +def merge_multi_segment(segments): + """Merge multi segments to one list. + Find the 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,...], + 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 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][0:1]]) + # deal with the first segment and the last one + if i == 0 or i == (len(idx_list) - 1): + s.append(segments[i]) + # deal with the middle ones + else: + idx = [0, idx[1] - idx[0]] + s.append(segments[i][idx[0]:idx[1] + 1]) + + # backward connection + else: + for i in range(len(idx_list) - 1, -1, -1): + if i != 0 and i != (len(idx_list) - 1): + idx = idx_list[i] + nidx = abs(idx[1] - idx[0]) + s.append(segments[i][nidx:]) + return s diff --git a/requirements.txt b/requirements.txt deleted file mode 100755 index a8a0038..0000000 --- a/requirements.txt +++ /dev/null @@ -1,9 +0,0 @@ -# pip install -r requirements.txt - -numpy -opencv-python>=4.1.2 -pandas -Pillow -pyYAML -requests -tqdm diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..c1b14e3 --- /dev/null +++ b/setup.py @@ -0,0 +1,18 @@ +from setuptools import setup + +setup(name='json2yolo', + version='0.0.1', + author='Ultralytics', + packages=['json2yolo'], + description='Dataset converter for yolo', + license='GPL', + install_requires=[ + "numpy", + "opencv-python>=4.1.2", + "pandas", + "Pillow", + "pyYAML", + "requests", + "tqdm", + ], +)