Spaces:
Sleeping
Sleeping
| # -------------------------------------------------------- | |
| # SiamMask | |
| # Licensed under The MIT License | |
| # Written by Qiang Wang (wangqiang2015 at ia.ac.cn) | |
| # -------------------------------------------------------- | |
| from pycocotools.coco import COCO | |
| from os.path import join | |
| import json | |
| dataDir = '.' | |
| for data_subset in ['val2017', 'train2017']: | |
| dataset = dict() | |
| annFile = '{}/annotations/instances_{}.json'.format(dataDir, data_subset) | |
| coco = COCO(annFile) | |
| n_imgs = len(coco.imgs) | |
| for n, img_id in enumerate(coco.imgs): | |
| print('subset: {} image id: {:04d} / {:04d}'.format(data_subset, n, n_imgs)) | |
| img = coco.loadImgs(img_id)[0] | |
| annIds = coco.getAnnIds(imgIds=img['id'], iscrowd=None) | |
| anns = coco.loadAnns(annIds) | |
| crop_base_path = join(data_subset, img['file_name'].split('/')[-1].split('.')[0]) | |
| if len(anns) > 0: | |
| dataset[crop_base_path] = dict() | |
| for track_id, ann in enumerate(anns): | |
| rect = ann['bbox'] | |
| if rect[2] <= 0 or rect[3] <= 0: # lead nan error in cls. | |
| continue | |
| bbox = [rect[0], rect[1], rect[0]+rect[2]-1, rect[1]+rect[3]-1] # x1,y1,x2,y2 | |
| dataset[crop_base_path]['{:02d}'.format(track_id)] = {'000000': bbox} | |
| print('save json (dataset), please wait 20 seconds~') | |
| json.dump(dataset, open('{}.json'.format(data_subset), 'w'), indent=4, sort_keys=True) | |
| print('done!') | |