mmpose/tools/dataset_converters/lapa2coco.py

# 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')