preprocess_coco_person.py

import argparse
import json
import tqdm


def remove_link(cat, kp_a_name, kp_b_name):
    skeleton = cat['skeleton']
    keypoints = cat['keypoints']
    new_skeleton = []
    
    for link in skeleton:
        if (keypoints[link[0] - 1] == kp_a_name) and (keypoints[link[1] - 1] == kp_b_name):
            continue # skip
        if (keypoints[link[0] - 1] == kp_b_name) and (keypoints[link[1] - 1] == kp_a_name):
            continue # skip
        new_skeleton.append(link)
    cat['skeleton'] = new_skeleton

    
def add_link(cat, kp_a_name, kp_b_name):
    keypoints = cat['keypoints']
    cat['skeleton'].append([keypoints.index(kp_a_name) + 1, keypoints.index(kp_b_name) + 1])


def append_neck_keypoint(ann, cat):
    keypoints = cat['keypoints']
    kps = ann['keypoints']
    l_idx = 3 * keypoints.index('left_shoulder')
    r_idx = 3 * keypoints.index('right_shoulder')
    x_neck = round((kps[l_idx] + kps[r_idx]) / 2.0)
    y_neck = round((kps[l_idx + 1] + kps[r_idx + 1]) / 2.0)
    
    v_l = kps[l_idx + 2]
    v_r = kps[r_idx + 2]
    
    if v_l == 0 or v_r == 0:
        v_neck = 0
    elif v_l == 1 or v_r == 1:
        v_neck = 1
    else:
        v_neck = 2
    
    kps += [x_neck, y_neck, v_neck]
    
def get_cat(data, cat_name):
    return [c for c in data['categories'] if c['name'] == cat_name][0]

def get_anns(data, cat_id):
    return [a for a in data['annotations'] if a['category_id'] == cat_id]


if __name__ == '__main__':
    
    parser = argparse.ArgumentParser()
    parser.add_argument('input_annotation_file', type=str, help='Path to COCO annotation file')
    parser.add_argument('output_annotation_file', type=str, help='Path to COCO annotation file')
    args = parser.parse_args()
    
    print('Loading...')
    with open(args.input_annotation_file, 'r') as f:
        data = json.load(f)
    
    print('Preprocessing...')
    cat = get_cat(data, 'person')
    cat_id = cat['id']
    anns = get_anns(data, cat_id)

    for a in anns:
        append_neck_keypoint(a, cat)

    cat['keypoints'].append('neck')
    remove_link(cat, 'left_shoulder', 'right_shoulder')
    remove_link(cat, 'left_shoulder', 'left_hip')
    remove_link(cat, 'right_shoulder', 'right_hip')
    add_link(cat, 'neck', 'nose')
    add_link(cat, 'neck', 'left_shoulder')
    add_link(cat, 'neck', 'right_shoulder')
    add_link(cat, 'neck', 'left_hip')
    add_link(cat, 'neck', 'right_hip')
    
    print('Saving...')
    with open(args.output_annotation_file, 'w') as f:
        json.dump(data, f)