mirror of https://github.com/open-mmlab/mmpose
112 lines
3.2 KiB
Python
112 lines
3.2 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import json
|
|
import os
|
|
import os.path as osp
|
|
import time
|
|
|
|
import cv2
|
|
import mmengine
|
|
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_labpa_to_coco(ann_dir, out_file):
|
|
annotations = []
|
|
images = []
|
|
cnt = 0
|
|
|
|
if 'trainval' in ann_dir:
|
|
ann_dir_list = ['train', 'val']
|
|
else:
|
|
ann_dir_list = [ann_dir]
|
|
|
|
for tv in ann_dir_list:
|
|
ann_dir = 'data/LaPa/' + tv
|
|
landmark_dir = osp.join(ann_dir, 'landmarks')
|
|
ann_list = os.listdir(landmark_dir)
|
|
|
|
img_dir = osp.join(ann_dir, 'images')
|
|
|
|
for idx, ann_file in enumerate(mmengine.track_iter_progress(ann_list)):
|
|
cnt += 1
|
|
ann_path = osp.join(landmark_dir, ann_file)
|
|
file_name = ann_file[:-4] + '.jpg'
|
|
img_path = osp.join(img_dir, file_name)
|
|
data_info = open(ann_path).readlines()
|
|
|
|
img = cv2.imread(img_path)
|
|
|
|
keypoints = []
|
|
for line in data_info[1:]:
|
|
x, y = line.strip().split(' ')
|
|
x, y = float(x), float(y)
|
|
keypoints.append([x, y, 2])
|
|
keypoints = np.array(keypoints)
|
|
|
|
x1, y1, _ = np.amin(keypoints, axis=0)
|
|
x2, y2, _ = np.amax(keypoints, axis=0)
|
|
w, h = x2 - x1, y2 - y1
|
|
bbox = [x1, y1, w, h]
|
|
|
|
image = {}
|
|
image['id'] = cnt
|
|
image['file_name'] = f'{tv}/images/{file_name}'
|
|
image['height'] = img.shape[0]
|
|
image['width'] = img.shape[1]
|
|
images.append(image)
|
|
|
|
ann = {}
|
|
ann['keypoints'] = keypoints.reshape(-1).tolist()
|
|
ann['image_id'] = cnt
|
|
ann['id'] = cnt
|
|
ann['num_keypoints'] = len(keypoints)
|
|
ann['bbox'] = bbox
|
|
ann['iscrowd'] = 0
|
|
ann['area'] = int(ann['bbox'][2] * ann['bbox'][3])
|
|
ann['category_id'] = 1
|
|
|
|
annotations.append(ann)
|
|
|
|
cocotype = {}
|
|
|
|
cocotype['info'] = {}
|
|
cocotype['info']['description'] = 'LaPa Generated by MMPose Team'
|
|
cocotype['info']['version'] = 1.0
|
|
cocotype['info']['year'] = time.strftime('%Y', time.localtime())
|
|
cocotype['info']['date_created'] = time.strftime('%Y/%m/%d',
|
|
time.localtime())
|
|
|
|
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/LaPa/annotations'):
|
|
os.makedirs('data/LaPa/annotations')
|
|
for tv in ['val', 'test', 'train', 'trainval']:
|
|
print(f'processing {tv}')
|
|
convert_labpa_to_coco(tv, f'data/LaPa/annotations/lapa_{tv}.json')
|