import sys
import os
import random
import numpy as np
import json
sys.path.append("/home/shanwei-luo/userdata/mmdetection")
from mmdet.apis import (async_inference_detector, inference_detector,
                        init_detector, show_result_pyplot)

data_coco = json.load(open('/home/shanwei-luo/userdata/pd_datasets/20220606_coco/annotations/val_cat_mode.json'))
data_name = data_coco["images"]
data_ann = data_coco['annotations']
boxes = {}
for res in data_ann:
    print(res)
    img_id_1 = res["image_id"]
    img_name = data_name[int(img_id_1)]["file_name"]
    #print(img_name)
    #print(res)
    bbox = res["bbox"]
    label = res["category_id"]
    bbox.append(int(label))
    if img_name in boxes.keys():
        boxes[img_name].append(bbox)
    else:
        boxes[img_name]=[]
        boxes[img_name].append(bbox)
#print(boxes)

config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_pd_test01/AD_pd_test01.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_pd_test01/epoch_18.pth'

img_path = '/home/shanwei-luo/userdata/pd_datasets/20220606_coco/images/'

model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:0')

imgs = os.listdir(img_path)
#img_id = random.randint(0, len(label_path))
imgs_labels = []
imgs_name = []
num_ng = 0
for i in range(len(data_name)):
    res_label = 0
    if data_name[i]["file_name"] in boxes.keys():
        res_label = 1
        num_ng += 1
    imgs_labels.append(res_label)
    imgs_name.append(img_path+data_name[i]["file_name"])
num_ok = len(data_name)-num_ng
print(len(imgs_labels), num_ok, num_ng)
imgs_name = imgs_name[:10]
print("before infer")
index = 0
num = len(imgs_name)
results_1 = []
step = 64
while index<num:
    index += step
    if index < num:
        results_1_tmp = inference_detector(model_1, imgs_name[index-step:index])
    else:
        results_1_tmp = inference_detector(model_1, imgs_name[index-step:num])
    results_1 += results_1_tmp
    print(len(results_1))
print("after infer")

#score_thrs = [0.01, 0.011, 0.012, 0.013, 0.014, 0.015, 0.016, 0.017, 0.018, 0.019, 0.02]
for score_thr in np.arange(0.01, 0.2, 0.01):
    imgs_results_1 = []
    for result in results_1:
        res_predict = 0
        print(len(result))
        for i in result:
            print(i.shape)
            for j in range(i.shape[0]):
                if i[j, 4]>score_thr:
                    res_predict = 1
        imgs_results_1.append(res_predict)

    count_ng = 0
    count_ok = 0
    for i in range(len(imgs_labels)):
        if imgs_labels[i]==0 and imgs_results_1[i]==0:
            count_ok += 1
        if imgs_labels[i]==1 and imgs_results_1[i]==1:
            count_ng += 1
        '''if imgs_labels[i]==1 and imgs_results_1[i]==0:
            print(imgs_name[i])'''
    print("score_thr:", score_thr, "  recall(ok):", count_ok/num_ok, "  recall(ng):", count_ng/num_ng)