Spaces:
Running
on
Zero
Running
on
Zero
Virtual-Try-On
/
preprocess
/humanparsing
/mhp_extension
/coco_style_annotation_creator
/human_to_coco.py
import argparse | |
import datetime | |
import json | |
import os | |
from PIL import Image | |
import numpy as np | |
import pycococreatortools | |
def get_arguments(): | |
parser = argparse.ArgumentParser(description="transform mask annotation to coco annotation") | |
parser.add_argument("--dataset", type=str, default='CIHP', help="name of dataset (CIHP, MHPv2 or VIP)") | |
parser.add_argument("--json_save_dir", type=str, default='../data/msrcnn_finetune_annotations', | |
help="path to save coco-style annotation json file") | |
parser.add_argument("--use_val", type=bool, default=False, | |
help="use train+val set for finetuning or not") | |
parser.add_argument("--train_img_dir", type=str, default='../data/instance-level_human_parsing/Training/Images', | |
help="train image path") | |
parser.add_argument("--train_anno_dir", type=str, | |
default='../data/instance-level_human_parsing/Training/Human_ids', | |
help="train human mask path") | |
parser.add_argument("--val_img_dir", type=str, default='../data/instance-level_human_parsing/Validation/Images', | |
help="val image path") | |
parser.add_argument("--val_anno_dir", type=str, | |
default='../data/instance-level_human_parsing/Validation/Human_ids', | |
help="val human mask path") | |
return parser.parse_args() | |
def main(args): | |
INFO = { | |
"description": args.split_name + " Dataset", | |
"url": "", | |
"version": "", | |
"year": 2019, | |
"contributor": "xyq", | |
"date_created": datetime.datetime.utcnow().isoformat(' ') | |
} | |
LICENSES = [ | |
{ | |
"id": 1, | |
"name": "", | |
"url": "" | |
} | |
] | |
CATEGORIES = [ | |
{ | |
'id': 1, | |
'name': 'person', | |
'supercategory': 'person', | |
}, | |
] | |
coco_output = { | |
"info": INFO, | |
"licenses": LICENSES, | |
"categories": CATEGORIES, | |
"images": [], | |
"annotations": [] | |
} | |
image_id = 1 | |
segmentation_id = 1 | |
for image_name in os.listdir(args.train_img_dir): | |
image = Image.open(os.path.join(args.train_img_dir, image_name)) | |
image_info = pycococreatortools.create_image_info( | |
image_id, image_name, image.size | |
) | |
coco_output["images"].append(image_info) | |
human_mask_name = os.path.splitext(image_name)[0] + '.png' | |
human_mask = np.asarray(Image.open(os.path.join(args.train_anno_dir, human_mask_name))) | |
human_gt_labels = np.unique(human_mask) | |
for i in range(1, len(human_gt_labels)): | |
category_info = {'id': 1, 'is_crowd': 0} | |
binary_mask = np.uint8(human_mask == i) | |
annotation_info = pycococreatortools.create_annotation_info( | |
segmentation_id, image_id, category_info, binary_mask, | |
image.size, tolerance=10 | |
) | |
if annotation_info is not None: | |
coco_output["annotations"].append(annotation_info) | |
segmentation_id += 1 | |
image_id += 1 | |
if not os.path.exists(args.json_save_dir): | |
os.makedirs(args.json_save_dir) | |
if not args.use_val: | |
with open('{}/{}_train.json'.format(args.json_save_dir, args.split_name), 'w') as output_json_file: | |
json.dump(coco_output, output_json_file) | |
else: | |
for image_name in os.listdir(args.val_img_dir): | |
image = Image.open(os.path.join(args.val_img_dir, image_name)) | |
image_info = pycococreatortools.create_image_info( | |
image_id, image_name, image.size | |
) | |
coco_output["images"].append(image_info) | |
human_mask_name = os.path.splitext(image_name)[0] + '.png' | |
human_mask = np.asarray(Image.open(os.path.join(args.val_anno_dir, human_mask_name))) | |
human_gt_labels = np.unique(human_mask) | |
for i in range(1, len(human_gt_labels)): | |
category_info = {'id': 1, 'is_crowd': 0} | |
binary_mask = np.uint8(human_mask == i) | |
annotation_info = pycococreatortools.create_annotation_info( | |
segmentation_id, image_id, category_info, binary_mask, | |
image.size, tolerance=10 | |
) | |
if annotation_info is not None: | |
coco_output["annotations"].append(annotation_info) | |
segmentation_id += 1 | |
image_id += 1 | |
with open('{}/{}_trainval.json'.format(args.json_save_dir, args.split_name), 'w') as output_json_file: | |
json.dump(coco_output, output_json_file) | |
coco_output_val = { | |
"info": INFO, | |
"licenses": LICENSES, | |
"categories": CATEGORIES, | |
"images": [], | |
"annotations": [] | |
} | |
image_id_val = 1 | |
segmentation_id_val = 1 | |
for image_name in os.listdir(args.val_img_dir): | |
image = Image.open(os.path.join(args.val_img_dir, image_name)) | |
image_info = pycococreatortools.create_image_info( | |
image_id_val, image_name, image.size | |
) | |
coco_output_val["images"].append(image_info) | |
human_mask_name = os.path.splitext(image_name)[0] + '.png' | |
human_mask = np.asarray(Image.open(os.path.join(args.val_anno_dir, human_mask_name))) | |
human_gt_labels = np.unique(human_mask) | |
for i in range(1, len(human_gt_labels)): | |
category_info = {'id': 1, 'is_crowd': 0} | |
binary_mask = np.uint8(human_mask == i) | |
annotation_info = pycococreatortools.create_annotation_info( | |
segmentation_id_val, image_id_val, category_info, binary_mask, | |
image.size, tolerance=10 | |
) | |
if annotation_info is not None: | |
coco_output_val["annotations"].append(annotation_info) | |
segmentation_id_val += 1 | |
image_id_val += 1 | |
with open('{}/{}_val.json'.format(args.json_save_dir, args.split_name), 'w') as output_json_file_val: | |
json.dump(coco_output_val, output_json_file_val) | |
if __name__ == "__main__": | |
args = get_arguments() | |
main(args) | |