File size: 1,647 Bytes
3f3265f | 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 | # clean up the merged COCO dataset
from os.path import join as pjoin
import json
import pycocotools.mask as mask_util
ROOT = 'data/coco_merged/annotations'
JSONS = ['interaction.json', 'interactable.json', 'semantics.json']
# filter noise produced by labeling tool, denote mask in compressed RLE format
def fix_mask(data):
for anno in data['annotations']:
mask = mask_util.frPyObjects(anno['segmentation'], 540, 960)
mask = mask_util.decode(mask)
mask[:, -1] = mask[:, -2]
mask[-1:, :] = mask[-2, :]
RLE = mask_util.encode(mask)
bbox = mask_util.toBbox(RLE)
bbox = bbox.astype(int)
area = int(bbox[2] * bbox[3])
RLE['counts'] = RLE['counts'].decode('utf-8')
anno['segmentation'] = RLE
anno['bbox'] = bbox.tolist()
anno['area'] = area
# remove redundant annotations
def remove_redundant(data):
total = len(data['annotations'])
anno_list = []
for anno in data['annotations']:
anno['id']=0
if anno not in anno_list:
anno_list.append(anno)
for anno in anno_list:
anno['id'] = anno_list.index(anno) + 1
data['annotations'] = anno_list
print(f'{total - len(anno_list)} redundant annotations removed')
def main():
json_files = [pjoin(ROOT, j) for j in JSONS]
for j in json_files:
with open(j, 'r') as f:
data = json.load(f)
fix_mask(data)
remove_redundant(data)
new_j = j.replace('.json', '_cleaned.json')
with open(new_j, 'w') as f:
json.dump(data, f, indent=4)
if __name__ == '__main__':
main() |