Spaces:
Runtime error
Runtime error
fifa-tryon-demo
/
Self-Correction-Human-Parsing-for-ACGPN
/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) | |