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 | |
| np.random.seed(0) | |
| def save_coco_anno(keypoints_all, | |
| annotated_all, | |
| imgs_all, | |
| keypoints_info, | |
| skeleton_info, | |
| dataset, | |
| img_root, | |
| save_path, | |
| start_img_id=0, | |
| start_ann_id=0): | |
| """Save annotations in coco-format. | |
| :param keypoints_all: keypoint annotations. | |
| :param annotated_all: images annotated or not. | |
| :param imgs_all: the array of images. | |
| :param keypoints_info: information about keypoint name. | |
| :param skeleton_info: information about skeleton connection. | |
| :param dataset: information about dataset name. | |
| :param img_root: the path to save images. | |
| :param save_path: the path to save transformed annotation file. | |
| :param start_img_id: the starting point to count the image id. | |
| :param start_ann_id: the starting point to count the annotation id. | |
| """ | |
| images = [] | |
| annotations = [] | |
| img_id = start_img_id | |
| ann_id = start_ann_id | |
| num_annotations, keypoints_num, _ = keypoints_all.shape | |
| for i in range(num_annotations): | |
| img = imgs_all[i] | |
| keypoints = np.concatenate( | |
| [keypoints_all[i], annotated_all[i][:, None] * 2], axis=1) | |
| min_x, min_y = np.min(keypoints[keypoints[:, 2] > 0, :2], axis=0) | |
| max_x, max_y = np.max(keypoints[keypoints[:, 2] > 0, :2], axis=0) | |
| anno = {} | |
| anno['keypoints'] = keypoints.reshape(-1).tolist() | |
| anno['image_id'] = img_id | |
| anno['id'] = ann_id | |
| anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0)) | |
| anno['bbox'] = [ | |
| float(min_x), | |
| float(min_y), | |
| float(max_x - min_x + 1), | |
| float(max_y - min_y + 1) | |
| ] | |
| anno['iscrowd'] = 0 | |
| anno['area'] = anno['bbox'][2] * anno['bbox'][3] | |
| anno['category_id'] = 1 | |
| annotations.append(anno) | |
| ann_id += 1 | |
| image = {} | |
| image['id'] = img_id | |
| image['file_name'] = f'{img_id}.jpg' | |
| image['height'] = img.shape[0] | |
| image['width'] = img.shape[1] | |
| images.append(image) | |
| img_id += 1 | |
| cv2.imwrite(os.path.join(img_root, image['file_name']), img) | |
| skeleton = np.concatenate( | |
| [np.arange(keypoints_num)[:, None], skeleton_info[:, 0][:, None]], | |
| axis=1) + 1 | |
| skeleton = skeleton[skeleton.min(axis=1) > 0] | |
| cocotype = {} | |
| cocotype['info'] = {} | |
| cocotype['info'][ | |
| 'description'] = 'DeepPoseKit-Data 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': 'animal', | |
| 'id': 1, | |
| 'name': dataset, | |
| 'keypoints': keypoints_info, | |
| 'skeleton': skeleton.tolist() | |
| }] | |
| os.makedirs(os.path.dirname(save_path), exist_ok=True) | |
| json.dump(cocotype, open(save_path, 'w'), indent=4) | |
| print('number of images:', img_id) | |
| print('number of annotations:', ann_id) | |
| print(f'done {save_path}') | |
| for dataset in ['fly', 'locust', 'zebra']: | |
| keypoints_info = [] | |
| if dataset == 'fly': | |
| keypoints_info = [ | |
| 'head', 'eyeL', 'eyeR', 'neck', 'thorax', 'abdomen', 'forelegR1', | |
| 'forelegR2', 'forelegR3', 'forelegR4', 'midlegR1', 'midlegR2', | |
| 'midlegR3', 'midlegR4', 'hindlegR1', 'hindlegR2', 'hindlegR3', | |
| 'hindlegR4', 'forelegL1', 'forelegL2', 'forelegL3', 'forelegL4', | |
| 'midlegL1', 'midlegL2', 'midlegL3', 'midlegL4', 'hindlegL1', | |
| 'hindlegL2', 'hindlegL3', 'hindlegL4', 'wingL', 'wingR' | |
| ] | |
| elif dataset == 'locust': | |
| keypoints_info = [ | |
| 'head', 'neck', 'thorax', 'abdomen1', 'abdomen2', 'anttipL', | |
| 'antbaseL', 'eyeL', 'forelegL1', 'forelegL2', 'forelegL3', | |
| 'forelegL4', 'midlegL1', 'midlegL2', 'midlegL3', 'midlegL4', | |
| 'hindlegL1', 'hindlegL2', 'hindlegL3', 'hindlegL4', 'anttipR', | |
| 'antbaseR', 'eyeR', 'forelegR1', 'forelegR2', 'forelegR3', | |
| 'forelegR4', 'midlegR1', 'midlegR2', 'midlegR3', 'midlegR4', | |
| 'hindlegR1', 'hindlegR2', 'hindlegR3', 'hindlegR4' | |
| ] | |
| elif dataset == 'zebra': | |
| keypoints_info = [ | |
| 'snout', 'head', 'neck', 'forelegL1', 'forelegR1', 'hindlegL1', | |
| 'hindlegR1', 'tailbase', 'tailtip' | |
| ] | |
| else: | |
| NotImplementedError() | |
| dataset_dir = f'data/DeepPoseKit-Data/datasets/{dataset}' | |
| with h5py.File( | |
| os.path.join(dataset_dir, 'annotation_data_release.h5'), 'r') as f: | |
| # List all groups | |
| annotations = np.array(f['annotations']) | |
| annotated = np.array(f['annotated']) | |
| images = np.array(f['images']) | |
| skeleton_info = np.array(f['skeleton']) | |
| annotation_num, kpt_num, _ = annotations.shape | |
| data_list = np.arange(0, annotation_num) | |
| np.random.shuffle(data_list) | |
| val_data_num = annotation_num // 10 | |
| train_data_num = annotation_num - val_data_num | |
| train_list = data_list[0:train_data_num] | |
| val_list = data_list[train_data_num:] | |
| img_root = os.path.join(dataset_dir, 'images') | |
| os.makedirs(img_root, exist_ok=True) | |
| save_coco_anno( | |
| annotations[train_list], annotated[train_list], images[train_list], | |
| keypoints_info, skeleton_info, dataset, img_root, | |
| os.path.join(dataset_dir, 'annotations', f'{dataset}_train.json')) | |
| save_coco_anno( | |
| annotations[val_list], | |
| annotated[val_list], | |
| images[val_list], | |
| keypoints_info, | |
| skeleton_info, | |
| dataset, | |
| img_root, | |
| os.path.join(dataset_dir, 'annotations', f'{dataset}_test.json'), | |
| start_img_id=train_data_num, | |
| start_ann_id=train_data_num) | |