File size: 1,716 Bytes
1da285f | 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 | 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/coco_merged/images/interactable'
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) |