import json
import os
import pickle
import os.path as osp
import cv2
from IPython import embed
import copy

import torch.nn.functional as F
import torch

with open("arms.json","r") as f:
    arm_data = json.load(f)



def draw_box(img,box,img_name,arm=None,gt_box=None,box1=None):
    pt1 = (int(box[0]),int(box[1]))
    pt2 = (int(box[2]),int(box[3]))
    cv2.rectangle(img,pt1,pt2,(0,255,0),4) 
    if arm is not None:
        pt3 = (int(arm[0][0]),int(arm[0][1]))
        pt4 = (int(arm[1][0]),int(arm[1][1]))
        cv2.line(img,pt3,pt4,(0,0,255),5) 
    if gt_box is not None:
        pt5 = (int(gt_box[0]),int(gt_box[1]))
        pt6 = (int(gt_box[2]),int(gt_box[3]))
        cv2.rectangle(img,pt5,pt6,(255,0,0),4) 
    if box1 is not None: 
        pt7 = (int(box1[0]),int(box1[1]))
        pt8 = (int(box1[2]),int(box1[3]))
        cv2.rectangle(img,pt7,pt8,(255,255,0),4) 
    cv2.imwrite(img_name,img)



bbox_data = {}

i=0



for img_name in os.listdir('./yourefit/images'):
    pickle_file = osp.join('./yourefit/pickle',img_name[:-4]+'.p')
    pick = pickle.load(open(pickle_file, "rb" ))
    bbox_data[img_name[:-4]] = pick['bbox']



cos_sim_list=[]
for name in arm_data.keys():
    arm = arm_data[name]
    box = bbox_data[name]
    img_file = './yourefit/images/'+name+'.jpg'
    img = cv2.imread(img_file)
    print(name)
    draw_box(img,box,'./vis/gt_vis/'+name+'.jpg',arm)


    
    
    
    
    
    


