Convert the COCO RLE format to YOLOv5/v8 segmentation format. (Sourcery refactored)

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,133 @@ 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(f'kpt_shape: [{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]
+        p2 = contour[i + 1] if i < num - 1 else 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 or distance < distance_min:
+                distance_min = distance
+                idx1 = i
+                idx2 = j
+    return idx1, idx2
+
+def merge_contours(contour1, contour2, idx1, idx2):
+    contour = [contour1[i] for i in list(range(0, idx1 + 1))]
+    contour.extend(contour2[i] for i in list(range(idx2, len(contour2))))
+    contour.extend(contour2[i] for i in list(range(0, idx2 + 1)))
+    contour.extend(contour1[i] for i in list(range(idx1, len(contour1))))
+    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 = [
+        contour for contour in contours_parent if len(contour) != 0
+    ]
+    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
+    return mask2polygon(m)
 
 def min_index(arr1, arr2):
     """Find a pair of indexes with the shortest distance. 
@@ -386,7 +517,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',