diff --git a/general_json2yolo.py b/general_json2yolo.py index f1ada2c..688b570 100644 --- a/general_json2yolo.py +++ b/general_json2yolo.py @@ -5,6 +5,7 @@ import pandas as pd from PIL import Image from collections import defaultdict +from pycocotools import mask from utils import * @@ -250,7 +251,7 @@ def convert_ath_json(json_dir): # dir contains json annotations and images print(f'Done. Output saved to {Path(dir).absolute()}') -def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91to80=False): +def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, use_keypoints=False, cls91to80=False): save_dir = make_dirs() # output directory coco80 = coco91_to_coco80_class() @@ -268,18 +269,24 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91 for ann in data['annotations']: imgToAnns[ann['image_id']].append(ann) + if use_keypoints: + show_kpt_shape_flip_idx(data) + # Write labels file for img_id, anns in tqdm(imgToAnns.items(), desc=f'Annotations {json_file}'): img = images['%g' % img_id] h, w, f = img['height'], img['width'], img['file_name'] - + f = f.split('/')[-1] + bboxes = [] segments = [] + keypoints = [] for ann in anns: - if ann['iscrowd']: - continue # The COCO box format is [top left x, top left y, width, height] - box = np.array(ann['bbox'], dtype=np.float64) + if len(ann['bbox']) == 0: + box = bbox_from_keypoints(ann) + else: + box = np.array(ann['bbox'], dtype=np.float64) box[:2] += box[2:] / 2 # xy top-left corner to center box[[0, 2]] /= w # normalize x box[[1, 3]] /= h # normalize y @@ -290,8 +297,12 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91 box = [cls] + box.tolist() if box not in bboxes: bboxes.append(box) - # Segments if use_segments: + if len(ann['segmentation']) == 0: + segments.append([]) + continue + if isinstance(ann['segmentation'], dict): + ann['segmentation'] = rle2polygon(ann['segmentation']) if len(ann['segmentation']) > 1: s = merge_multi_segment(ann['segmentation']) s = (np.concatenate(s, axis=0) / np.array([w, h])).reshape(-1).tolist() @@ -301,13 +312,149 @@ def convert_coco_json(json_dir='../coco/annotations/', use_segments=False, cls91 s = [cls] + s if s not in segments: segments.append(s) + if use_keypoints: + if 'keypoints' not in ann: + keypoints.append([]) + continue + else: + k = (np.array(ann['keypoints']).reshape(-1, 3) / np.array([w, h, 1])).reshape(-1).tolist() + k = box + k + keypoints.append(k) # Write with open((fn / f).with_suffix('.txt'), 'a') as file: for i in range(len(bboxes)): - line = *(segments[i] if use_segments else bboxes[i]), # cls, box or segments + if use_keypoints: + line = *(keypoints[i]), # cls, box, keypoints + else: + line = *(segments[i] if use_segments and len(segments[i]) > 0 else bboxes[i]), # cls, box or segments file.write(('%g ' * len(line)).rstrip() % line + '\n') +def bbox_from_keypoints(ann): + if 'keypoints' not in ann: + return + k = np.array(ann['keypoints']).reshape(-1, 3) + x_list, y_list, v_list = zip(*k) + box = [min(x_list), min(y_list), max(x_list) - min(x_list), max(y_list) - min(y_list)] + return np.array(box, dtype=np.float64) + +def show_kpt_shape_flip_idx(data): + for category in data['categories']: + if 'keypoints' not in category: + continue + keypoints = category['keypoints'] + num = len(keypoints) + print('kpt_shape: [' + str(num) + ', 3]') + flip_idx = list(range(num)) + for i, name in enumerate(keypoints): + name = name.lower() + left_pos = name.find('left') + if left_pos < 0: + continue + name_right = name.replace('left', 'right') + for j, namej in enumerate(keypoints): + namej = namej.lower() + if namej == name_right: + flip_idx[i] = j + flip_idx[j] = i + break + print('flip_idx: [' + ', '.join(str(x) for x in flip_idx) + ']') + + +def is_clockwise(contour): + value = 0 + num = len(contour) + for i, point in enumerate(contour): + p1 = contour[i] + if i < num - 1: + p2 = contour[i + 1] + else: + p2 = contour[0] + value += (p2[0][0] - p1[0][0]) * (p2[0][1] + p1[0][1]); + return value < 0 + +def get_merge_point_idx(contour1, contour2): + idx1 = 0 + idx2 = 0 + distance_min = -1 + for i, p1 in enumerate(contour1): + for j, p2 in enumerate(contour2): + distance = pow(p2[0][0] - p1[0][0], 2) + pow(p2[0][1] - p1[0][1], 2); + if distance_min < 0: + distance_min = distance + idx1 = i + idx2 = j + elif distance < distance_min: + distance_min = distance + idx1 = i + idx2 = j + return idx1, idx2 + +def merge_contours(contour1, contour2, idx1, idx2): + contour = [] + for i in list(range(0, idx1 + 1)): + contour.append(contour1[i]) + for i in list(range(idx2, len(contour2))): + contour.append(contour2[i]) + for i in list(range(0, idx2 + 1)): + contour.append(contour2[i]) + for i in list(range(idx1, len(contour1))): + contour.append(contour1[i]) + contour = np.array(contour) + return contour + +def merge_with_parent(contour_parent, contour): + if not is_clockwise(contour_parent): + contour_parent = contour_parent[::-1] + if is_clockwise(contour): + contour = contour[::-1] + idx1, idx2 = get_merge_point_idx(contour_parent, contour) + return merge_contours(contour_parent, contour, idx1, idx2) + +def mask2polygon(image): + contours, hierarchies = cv2.findContours(image, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS) + contours_approx = [] + polygons = [] + for contour in contours: + epsilon = 0.001 * cv2.arcLength(contour, True) + contour_approx = cv2.approxPolyDP(contour, epsilon, True) + contours_approx.append(contour_approx) + + contours_parent = [] + for i, contour in enumerate(contours_approx): + parent_idx = hierarchies[0][i][3] + if parent_idx < 0 and len(contour) >= 3: + contours_parent.append(contour) + else: + contours_parent.append([]) + + for i, contour in enumerate(contours_approx): + parent_idx = hierarchies[0][i][3] + if parent_idx >= 0 and len(contour) >= 3: + contour_parent = contours_parent[parent_idx] + if len(contour_parent) == 0: + continue + contours_parent[parent_idx] = merge_with_parent(contour_parent, contour) + + contours_parent_tmp = [] + for contour in contours_parent: + if len(contour) == 0: + continue + contours_parent_tmp.append(contour) + + polygons = [] + for contour in contours_parent_tmp: + polygon = contour.flatten().tolist() + polygons.append(polygon) + return polygons + +def rle2polygon(segmentation): + if isinstance(segmentation["counts"], list): + segmentation = mask.frPyObjects(segmentation, *segmentation["size"]) + m = mask.decode(segmentation) + m[m > 0] = 255 + polygons = mask2polygon(m) + return polygons def min_index(arr1, arr2): """Find a pair of indexes with the shortest distance. @@ -386,7 +533,8 @@ def delete_dsstore(path='../datasets'): if source == 'COCO': convert_coco_json('../datasets/coco/annotations', # directory with *.json use_segments=True, - cls91to80=True) + use_keypoints=False, + cls91to80=False) elif source == 'infolks': # Infolks https://infolks.info/ convert_infolks_json(name='out',