| import os |
| import cv2 |
| import json |
| import argparse |
| from tqdm import tqdm |
|
|
| ANNOTATION_FILE = '../../dataset/data/coco_merged/annotations/interactable.json' |
| |
| IMAGE_DIR = '../../dataset/data/visualize/result_ours_i' |
| INTERACT_FILE = './guided_interact.json' |
| OUTPUT_DIR = './output' |
|
|
| parser = argparse.ArgumentParser(description='Visualize interact points') |
| parser.add_argument('-a', '--annotation', type=str, default=ANNOTATION_FILE, help='Annotation file') |
| parser.add_argument('-i', '--image', type=str, default=IMAGE_DIR, help='Original image directory') |
| parser.add_argument('-int', '--interact', type=str, default=INTERACT_FILE, help='Interact points to visualize') |
| parser.add_argument('-o', '--output', type=str, default=OUTPUT_DIR, help='Output directory') |
|
|
|
|
| def main(args): |
| os.makedirs(args.output, exist_ok=True) |
| with open(args.annotation, 'r') as f: |
| anno_data = json.load(f) |
| with open(args.interact, 'r') as f: |
| interact_data = json.load(f) |
| for img in tqdm(anno_data['images'], desc='Visualizing interact points', total=len(anno_data['images'])): |
| img_id = img['id'] |
| img_path = os.path.join('./result_ours_i', img['file_name']) |
| if not os.path.exists(img_path): |
| continue |
| cimg = cv2.imread(img_path) |
| for k in interact_data: |
| for interact in interact_data[k]: |
| if interact['img_id'] == img_id: |
| cv2.circle(cimg, (interact['X'], interact['Y']), 5, (0, 0, 100 + 125 // int(k)), -1) |
| cv2.imwrite(os.path.join(args.output, img['file_name']), cimg) |
|
|
| if __name__ == '__main__': |
| args = parser.parse_args() |
| main(args) |