Spaces:
Runtime error
Runtime error
File size: 1,894 Bytes
4a285f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 |
import argparse
import datetime
import json
import os
from PIL import Image
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/CIHP/annotations',
help="path to save coco-style annotation json file")
parser.add_argument("--test_img_dir", type=str, default='../data/CIHP/Testing/Images',
help="test image path")
return parser.parse_args()
args = get_arguments()
INFO = {
"description": args.dataset + "Dataset",
"url": "",
"version": "",
"year": 2020,
"contributor": "yunqiuxu",
"date_created": datetime.datetime.utcnow().isoformat(' ')
}
LICENSES = [
{
"id": 1,
"name": "",
"url": ""
}
]
CATEGORIES = [
{
'id': 1,
'name': 'person',
'supercategory': 'person',
},
]
def main(args):
coco_output = {
"info": INFO,
"licenses": LICENSES,
"categories": CATEGORIES,
"images": [],
"annotations": []
}
image_id = 1
for image_name in os.listdir(args.test_img_dir):
image = Image.open(os.path.join(args.test_img_dir, image_name))
image_info = pycococreatortools.create_image_info(
image_id, image_name, image.size
)
coco_output["images"].append(image_info)
image_id += 1
if not os.path.exists(os.path.join(args.json_save_dir)):
os.mkdir(os.path.join(args.json_save_dir))
with open('{}/{}.json'.format(args.json_save_dir, args.dataset), 'w') as output_json_file:
json.dump(coco_output, output_json_file)
if __name__ == "__main__":
main(args)
|