Spaces:
Build error
Build error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import json | |
| import os | |
| import time | |
| import cv2 | |
| import h5py | |
| import numpy as np | |
| mat_files = ['COFW_train_color.mat', 'COFW_test_color.mat'] | |
| dataset_dir = 'data/cofw/' | |
| image_root = os.path.join(dataset_dir, 'images/') | |
| annotation_root = os.path.join(dataset_dir, 'annotations/') | |
| os.makedirs(image_root, exist_ok=True) | |
| os.makedirs(annotation_root, exist_ok=True) | |
| cnt = 0 | |
| for mat_file in mat_files: | |
| mat = h5py.File(os.path.join(dataset_dir, mat_file), 'r') | |
| if 'train' in mat_file: | |
| imgs = mat['IsTr'] | |
| pts = mat['phisTr'] | |
| bboxes = mat['bboxesTr'] | |
| is_train = True | |
| json_file = 'cofw_train.json' | |
| else: | |
| imgs = mat['IsT'] | |
| pts = mat['phisT'] | |
| bboxes = mat['bboxesT'] | |
| is_train = False | |
| json_file = 'cofw_test.json' | |
| images = [] | |
| annotations = [] | |
| num = pts.shape[1] | |
| for idx in range(0, num): | |
| cnt += 1 | |
| img = np.array(mat[imgs[0, idx]]).transpose() | |
| keypoints = pts[:, idx].reshape(3, -1).transpose() | |
| # 2 for valid and 1 for occlusion | |
| keypoints[:, 2] = 2 - keypoints[:, 2] | |
| # matlab 1-index to python 0-index | |
| keypoints[:, :2] -= 1 | |
| bbox = bboxes[:, idx] | |
| # check nonnegativity | |
| bbox[bbox < 0] = 0 | |
| keypoints[keypoints < 0] = 0 | |
| image = {} | |
| image['id'] = cnt | |
| image['file_name'] = f'{str(cnt).zfill(6)}.jpg' | |
| image['height'] = img.shape[0] | |
| image['width'] = img.shape[1] | |
| cv2.imwrite( | |
| os.path.join(image_root, image['file_name']), | |
| cv2.cvtColor(img, cv2.COLOR_RGB2BGR)) | |
| images.append(image) | |
| anno = {} | |
| anno['keypoints'] = keypoints.reshape(-1).tolist() | |
| anno['image_id'] = cnt | |
| anno['id'] = cnt | |
| anno['num_keypoints'] = len(keypoints) # all keypoints are labelled | |
| anno['bbox'] = bbox.tolist() | |
| anno['iscrowd'] = 0 | |
| anno['area'] = anno['bbox'][2] * anno['bbox'][3] | |
| anno['category_id'] = 1 | |
| annotations.append(anno) | |
| cocotype = {} | |
| cocotype['info'] = {} | |
| cocotype['info']['description'] = 'COFW 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': [] | |
| }] | |
| ann_path = os.path.join(annotation_root, json_file) | |
| json.dump(cocotype, open(ann_path, 'w')) | |
| print(f'done {ann_path}') | |