add wflw2coco

tools/dataset_converters/wflw2coco.py

@@ -0,0 +1,117 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import json
+import math
+import os
+import os.path as osp
+
+import cv2
+import numpy as np
+
+
+def default_dump(obj):
+    """Convert numpy classes to JSON serializable objects."""
+    if isinstance(obj, (np.integer, np.floating, np.bool_)):
+        return obj.item()
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    else:
+        return obj
+
+
+def convert_wflw_to_coco(ann_file, out_file):
+    annotations = []
+    images = []
+    files = []
+    cnt = 0
+    image_cnt = 0
+
+    data_infos = open(ann_file).readlines()
+    data_infos = [x.strip().split() for x in data_infos]
+    for data in data_infos:
+        file_name = data[-1]
+        img_path = osp.join('data/wflw_raw/WFLW_images', file_name)
+        img = cv2.imread(img_path)
+
+        keypoints = []
+
+        coordinates = [data[i:i + 2] for i in range(0, 196, 2)]
+
+        for coordinate in coordinates:
+            x, y = coordinate[0], coordinate[1]
+            x, y = float(x), float(y)
+            keypoints.append([x, y, 1])
+        keypoints = np.array(keypoints)
+
+        x1, y1, _ = np.amin(keypoints, axis=0)
+        x2, y2, _ = np.amax(keypoints, axis=0)
+        w, h = x2 - x1, y2 - y1
+        scale = math.ceil(max(w, h)) / 200
+        w_new = w / scale
+        h_new = h / scale
+        center = [(x1 + x2) / 2, (y1 + y2) / 2]
+        x1_new = center[0] - w_new / 2
+        y1_new = center[1] - h_new / 2
+        bbox = [x1_new, y1_new, w_new, h_new]
+
+        image = {}
+        # check if the image already exists
+        if file_name in files:
+            image = images[files.index(file_name)]
+        else:
+            image['id'] = image_cnt
+            image['file_name'] = f'{file_name}'
+            image['height'] = img.shape[0]
+            image['width'] = img.shape[1]
+            image_cnt = image_cnt + 1
+            files.append(file_name)
+            images.append(image)
+
+        ann = {}
+        ann['keypoints'] = keypoints.reshape(-1).tolist()
+        ann['image_id'] = image['id']
+        ann['id'] = cnt
+        ann['num_keypoints'] = len(keypoints)
+        ann['bbox'] = bbox
+        ann['is_crowd'] = 0
+        ann['area'] = w * h
+        ann['category_id'] = 1
+        ann['center'] = center
+        ann['scale'] = scale
+        annotations.append(ann)
+        cnt = cnt + 1
+
+    cocotype = {}
+    cocotype['images'] = images
+    cocotype['annotations'] = annotations
+    cocotype['categories'] = [{
+        'supercategory': 'person',
+        'id': 1,
+        'name': 'face',
+        'keypoints': [],
+        'skeleton': []
+    }]
+
+    json.dump(
+        cocotype,
+        open(out_file, 'w'),
+        ensure_ascii=False,
+        default=default_dump)
+    print(f'done {out_file}')
+
+
+if __name__ == '__main__':
+    if not osp.exists('data/wflw_raw/annotations'):
+        os.makedirs('data/wflw_raw/annotations')
+    root_folder = 'data/wflw_raw'
+    train_file = 'list_98pt_rect_attr_train.txt'
+    test_file = 'list_98pt_rect_attr_test.txt'
+    ann_train_file = os.path.join(root_folder, 'WFLW_annotations',
+                                  'list_98pt_rect_attr_train_test', train_file)
+    ann_test_file = os.path.join(root_folder, 'WFLW_annotations',
+                                 'list_98pt_rect_attr_train_test', test_file)
+    out_train_file = os.path.join(root_folder, 'annotations',
+                                  'face_landmarks_wflw_train.json')
+    out_test_file = os.path.join(root_folder, 'annotations',
+                                 'face_landmarks_wflw_test.json')
+    convert_wflw_to_coco(ann_train_file, out_train_file)
+    convert_wflw_to_coco(ann_test_file, out_test_file)