import sys
import os
import random
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_test50/AD_dsxw_test50.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test50/epoch_50.pth'

img_path = '/home/shanwei-luo/userdata/datasets/dsxw_test_SMD11_2201_2202/images/'
label_path = '/home/shanwei-luo/userdata/datasets/dsxw_test_SMD11_2201_2202/labels/'

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

imgs = os.listdir(img_path)
labels = os.listdir(label_path)
#img_id = random.randint(0, len(label_path))
label_ng = len(labels)
label_ok = len(imgs)-label_ng
print(label_ok, label_ng)
imgs_labels = []
imgs_name = []
#imgs = imgs[:40]
for img in imgs:
    label = img[:-3]+'txt'
    res_label = 0
    if label in labels:
        res_label = 1
    imgs_labels.append(res_label)
    imgs_name.append(img_path+img)
print(len(imgs_labels))
print("before infer")
index = 0
num = len(imgs_name)
results_1 = []
results_2 = []
step = 32
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")

imgs_results_1 = []
for result in results_1:
    res_predict = 0
    for i in result:
        if i.shape[0]>0:
            res_predict = max(max(i[:,4]), res_predict)
    print(res_predict)
    imgs_results_1.append(res_predict)

print(len(imgs_results_1))

f = open("CI_SMD11.txt", "w")
for i in range(len(imgs_labels)):
    print(imgs[i], imgs_results_1[i])
    f.write(imgs[i]+" "+str(imgs_results_1[i])+"\n")

f.close()