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

config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test70/AD_dsxw_test70.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test70/epoch_38.pth'

img_path = '/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/ok/'
dis_path1 = '/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/ok_score/'
dis_path2 = '/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/v1/dsxw_train/images/'

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

imgs = os.listdir(img_path)
imgs_name = []
for img in imgs:
    imgs_name.append(img_path+img)
print(len(imgs_name))
print("before infer")
index = 0
num = len(imgs_name)
results_1 = []
step = 128
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")
result_score = []
for result in results_1:
    max_v = 0
    for i in result:
        for j in range(i.shape[0]):
            if max_v < i[j, 4]:
                max_v = i[j, 4]
    result_score.append(max_v)
np_result = np.array(result_score)
print(np_result.shape)
np.save('/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/result.npy', np_result)

score_thr = 0.07
for k_v in np.arange(0.5, 1, 0.05):
    count = 0
    index = 0
    for result in results_1:
        hard = 0
        for i in result:
            for j in range(i.shape[0]):
                '''if i[j, 4]>=score_thr-score_thr*k_v and i[j, 4]<=score_thr+(1-score_thr)*k_v:
                    hard = 1'''
                if i[j, 4]>=k_v:
                    hard = 1
        if hard==1:
            count+=1
            #print(img_path+imgs[index], dis_path1+imgs[index])
            #shutil.copy(img_path+imgs[index], dis_path1+imgs[index])
            #shutil.copy(img_path+imgs[index], dis_path2+imgs[index])
        index += 1
    print(k_v, count)

np_result = np.load('/home/shanwei-luo/userdata/datasets/PCBA_dataset_v15_MLOPS/result.npy')
result_score = np_result.tolist()
for k_v in np.arange(0.5, 1, 0.05):
    count = 0
    index = 0
    for result in result_score:
        if result>=k_v:
            count +=1
        index += 1
    print(k_v, count)