parse_animalpose_dataset.py

import json
import os
import re
import time
import warnings

import cv2
import numpy as np
import xmltodict
from xtcocotools.coco import COCO

np.random.seed(0)


def list_all_files(root_dir, ext='.xml'):
    """List all files in the root directory and all its sub directories.

    :param root_dir: root directory
    :param ext: filename extension
    :return: list of files
    """
    files = []
    file_list = os.listdir(root_dir)
    for i in range(0, len(file_list)):
        path = os.path.join(root_dir, file_list[i])
        if os.path.isdir(path):
            files.extend(list_all_files(path))
        if os.path.isfile(path):
            if path.lower().endswith(ext):
                files.append(path)
    return files


def get_anno_info():
    keypoints_info = [
        'L_Eye',
        'R_Eye',
        'L_EarBase',
        'R_EarBase',
        'Nose',
        'Throat',
        'TailBase',
        'Withers',
        'L_F_Elbow',
        'R_F_Elbow',
        'L_B_Elbow',
        'R_B_Elbow',
        'L_F_Knee',
        'R_F_Knee',
        'L_B_Knee',
        'R_B_Knee',
        'L_F_Paw',
        'R_F_Paw',
        'L_B_Paw',
        'R_B_Paw',
    ]
    skeleton_info = [[1, 2], [1, 3], [2, 4], [1, 5], [2, 5], [5, 6], [6, 8],
                     [7, 8], [6, 9], [9, 13], [13, 17], [6, 10], [10, 14],
                     [14, 18], [7, 11], [11, 15], [15, 19], [7, 12], [12, 16],
                     [16, 20]]
    category_info = [{
        'supercategory': 'animal',
        'id': 1,
        'name': 'animal',
        'keypoints': keypoints_info,
        'skeleton': skeleton_info
    }]

    return keypoints_info, skeleton_info, category_info


def xml2coco_trainval(file_list, img_root, save_path, start_ann_id=0):
    """Save annotations in coco-format.

    :param file_list: list of data annotation files.
    :param img_root: the root dir to load images.
    :param save_path: the path to save transformed annotation file.
    :param start_ann_id: the starting point to count the annotation id.
    :param val_num: the number of annotated objects for validation.
    """
    images = []
    annotations = []
    img_ids = []
    ann_ids = []

    ann_id = start_ann_id

    name2id = {
        'L_Eye': 0,
        'R_Eye': 1,
        'L_EarBase': 2,
        'R_EarBase': 3,
        'Nose': 4,
        'Throat': 5,
        'TailBase': 6,
        'Withers': 7,
        'L_F_Elbow': 8,
        'R_F_Elbow': 9,
        'L_B_Elbow': 10,
        'R_B_Elbow': 11,
        'L_F_Knee': 12,
        'R_F_Knee': 13,
        'L_B_Knee': 14,
        'R_B_Knee': 15,
        'L_F_Paw': 16,
        'R_F_Paw': 17,
        'L_B_Paw': 18,
        'R_B_Paw': 19
    }
    for file in file_list:
        data_anno = xmltodict.parse(open(file).read())['annotation']

        img_id = int(data_anno['image'].split('_')[0] +
                     data_anno['image'].split('_')[1])

        if img_id not in img_ids:
            image_name = 'VOC2012/JPEGImages/' + data_anno['image'] + '.jpg'
            img = cv2.imread(os.path.join(img_root, image_name))

            image = {}
            image['id'] = img_id
            image['file_name'] = image_name
            image['height'] = img.shape[0]
            image['width'] = img.shape[1]

            images.append(image)
            img_ids.append(img_id)
        else:
            pass

        keypoint_anno = data_anno['keypoints']['keypoint']
        assert len(keypoint_anno) == 20

        keypoints = np.zeros([20, 3], dtype=np.float32)

        for kpt_anno in keypoint_anno:
            keypoint_name = kpt_anno['@name']
            keypoint_id = name2id[keypoint_name]

            visibility = int(kpt_anno['@visible'])

            if visibility == 0:
                continue
            else:
                keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
                keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
                keypoints[keypoint_id, 2] = 2

        anno = {}
        anno['keypoints'] = keypoints.reshape(-1).tolist()
        anno['image_id'] = img_id
        anno['id'] = ann_id
        anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))

        visible_bounds = data_anno['visible_bounds']
        anno['bbox'] = [
            float(visible_bounds['@xmin']),
            float(visible_bounds['@ymin']),
            float(visible_bounds['@width']),
            float(visible_bounds['@height'])
        ]
        anno['iscrowd'] = 0
        anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
        anno['category_id'] = 1

        annotations.append(anno)
        ann_ids.append(ann_id)
        ann_id += 1

    cocotype = {}

    cocotype['info'] = {}
    cocotype['info'][
        'description'] = 'AnimalPose dataset 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

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype['categories'] = category_info

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    json.dump(cocotype, open(save_path, 'w'), indent=4)
    print('number of images:', len(img_ids))
    print('number of annotations:', len(ann_ids))
    print(f'done {save_path}')


def xml2coco_test(file_list, img_root, save_path, start_ann_id=0):
    """Save annotations in coco-format.

    :param file_list: list of data annotation files.
    :param img_root: the root dir to load images.
    :param save_path: the path to save transformed annotation file.
    :param start_ann_id: the starting point to count the annotation id.
    """
    images = []
    annotations = []
    img_ids = []
    ann_ids = []

    ann_id = start_ann_id

    name2id = {
        'L_eye': 0,
        'R_eye': 1,
        'L_ear': 2,
        'R_ear': 3,
        'Nose': 4,
        'Throat': 5,
        'Tail': 6,
        'withers': 7,
        'L_F_elbow': 8,
        'R_F_elbow': 9,
        'L_B_elbow': 10,
        'R_B_elbow': 11,
        'L_F_knee': 12,
        'R_F_knee': 13,
        'L_B_knee': 14,
        'R_B_knee': 15,
        'L_F_paw': 16,
        'R_F_paw': 17,
        'L_B_paw': 18,
        'R_B_paw': 19
    }

    cat2id = {'cat': 1, 'cow': 2, 'dog': 3, 'horse': 4, 'sheep': 5}

    for file in file_list:
        data_anno = xmltodict.parse(open(file).read())['annotation']

        category_id = cat2id[data_anno['category']]

        img_id = category_id * 1000 + int(
            re.findall(r'\d+', data_anno['image'])[0])

        assert img_id not in img_ids

        # prepare images
        image_name = os.path.join('animalpose_image_part2',
                                  data_anno['category'], data_anno['image'])
        img = cv2.imread(os.path.join(img_root, image_name))

        image = {}
        image['id'] = img_id
        image['file_name'] = image_name
        image['height'] = img.shape[0]
        image['width'] = img.shape[1]

        images.append(image)
        img_ids.append(img_id)

        # prepare annotations
        keypoint_anno = data_anno['keypoints']['keypoint']
        keypoints = np.zeros([20, 3], dtype=np.float32)

        for kpt_anno in keypoint_anno:
            keypoint_name = kpt_anno['@name']
            keypoint_id = name2id[keypoint_name]

            visibility = int(kpt_anno['@visible'])

            if visibility == 0:
                continue
            else:
                keypoints[keypoint_id, 0] = float(kpt_anno['@x'])
                keypoints[keypoint_id, 1] = float(kpt_anno['@y'])
                keypoints[keypoint_id, 2] = 2

        anno = {}
        anno['keypoints'] = keypoints.reshape(-1).tolist()
        anno['image_id'] = img_id
        anno['id'] = ann_id
        anno['num_keypoints'] = int(sum(keypoints[:, 2] > 0))

        visible_bounds = data_anno['visible_bounds']
        anno['bbox'] = [
            float(visible_bounds['@xmin']),
            float(visible_bounds['@xmax']
                  ),  # typo in original xml: should be 'ymin'
            float(visible_bounds['@width']),
            float(visible_bounds['@height'])
        ]
        anno['iscrowd'] = 0
        anno['area'] = float(anno['bbox'][2] * anno['bbox'][3])
        anno['category_id'] = 1

        annotations.append(anno)
        ann_ids.append(ann_id)
        ann_id += 1

    cocotype = {}

    cocotype['info'] = {}
    cocotype['info'][
        'description'] = 'AnimalPose dataset 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

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype['categories'] = category_info

    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    json.dump(cocotype, open(save_path, 'w'), indent=4)
    print('=========================================================')
    print('number of images:', len(img_ids))
    print('number of annotations:', len(ann_ids))
    print(f'done {save_path}')


def split_train_val(work_dir, trainval_file, train_file, val_file,
                    val_ann_num):
    """Split train-val json file into training and validation files.

    :param work_dir: path to load train-val json file, and save split files.
    :param trainval_file: The input json file combining both train and val.
    :param trainval_file: The output json file for training.
    :param trainval_file: The output json file for validation.
    :param val_ann_num: the number of validation annotations.
    """

    coco = COCO(os.path.join(work_dir, trainval_file))

    img_list = list(coco.imgs.keys())
    np.random.shuffle(img_list)

    count = 0

    images_train = []
    images_val = []
    annotations_train = []
    annotations_val = []

    for img_id in img_list:
        ann_ids = coco.getAnnIds(img_id)

        if count + len(ann_ids) <= val_ann_num:
            # for validation
            count += len(ann_ids)
            images_val.append(coco.imgs[img_id])
            for ann_id in ann_ids:
                annotations_val.append(coco.anns[ann_id])

        else:
            images_train.append(coco.imgs[img_id])
            for ann_id in ann_ids:
                annotations_train.append(coco.anns[ann_id])

    if count == val_ann_num:
        print(f'We have found {count} annotations for validation.')
    else:
        warnings.warn(
            f'We only found {count} annotations, instead of {val_ann_num}.')

    cocotype_train = {}
    cocotype_val = {}

    keypoints_info, skeleton_info, category_info = get_anno_info()

    cocotype_train['info'] = {}
    cocotype_train['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype_train['info']['version'] = '1.0'
    cocotype_train['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype_train['info']['date_created'] = time.strftime(
        '%Y/%m/%d', time.localtime())
    cocotype_train['images'] = images_train
    cocotype_train['annotations'] = annotations_train
    cocotype_train['categories'] = category_info

    json.dump(
        cocotype_train,
        open(os.path.join(work_dir, train_file), 'w'),
        indent=4)
    print('=========================================================')
    print('number of images:', len(images_train))
    print('number of annotations:', len(annotations_train))
    print(f'done {train_file}')

    cocotype_val['info'] = {}
    cocotype_val['info'][
        'description'] = 'AnimalPose dataset Generated by MMPose Team'
    cocotype_val['info']['version'] = '1.0'
    cocotype_val['info']['year'] = time.strftime('%Y', time.localtime())
    cocotype_val['info']['date_created'] = time.strftime(
        '%Y/%m/%d', time.localtime())
    cocotype_val['images'] = images_val
    cocotype_val['annotations'] = annotations_val
    cocotype_val['categories'] = category_info

    json.dump(
        cocotype_val, open(os.path.join(work_dir, val_file), 'w'), indent=4)
    print('=========================================================')
    print('number of images:', len(images_val))
    print('number of annotations:', len(annotations_val))
    print(f'done {val_file}')


dataset_dir = 'data/animalpose/'

# We choose the images from PascalVOC for train + val
# In total, train+val: 3608 images, 5117 annotations
xml2coco_trainval(
    list_all_files(os.path.join(dataset_dir, 'PASCAL2011_animal_annotation')),
    dataset_dir,
    os.path.join(dataset_dir, 'annotations', 'animalpose_trainval.json'),
    start_ann_id=1000000)

# train: 2798 images, 4000 annotations
# val: 810 images, 1117 annotations
split_train_val(
    os.path.join(dataset_dir, 'annotations'),
    'animalpose_trainval.json',
    'animalpose_train.json',
    'animalpose_val.json',
    val_ann_num=1117)

# We choose the remaining 1000 images for test
# 1000 images, 1000 annotations
xml2coco_test(
    list_all_files(os.path.join(dataset_dir, 'animalpose_anno2')),
    dataset_dir,
    os.path.join(dataset_dir, 'annotations', 'animalpose_test.json'),
    start_ann_id=0)