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