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)

class_AOI_name = {"bu_pi_pei":"1","fang_xiang_fan":"2","err.txt_c_not_f":"3", "shang_xi_bu_lia":"4"}

config_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/AD_dsxw_test61.py'
checkpoint_file_1 = '/home/shanwei-luo/userdata/mmdetection/work_dirs/AD_dsxw_test61/epoch_36.pth'
model_1 = init_detector(config_file_1, checkpoint_file_1, device='cuda:0')

dirs_path = '/home/shanwei-luo/userdata/datasets/dsxw_test_SMD12_2022_0301_0314_day/'
dirs_name = os.listdir(dirs_path)
for dir_name in dirs_name:
    imgs_ok_path = dirs_path+dir_name+"/ok/"
    imgs_ng_path = dirs_path+dir_name+"/ng/"

    imgs_ok = os.listdir(imgs_ok_path)
    imgs_ng = os.listdir(imgs_ng_path)

    count_label_ok = len(imgs_ok)
    count_label_ng = len(imgs_ng)
    #print(count_label_ok,count_label_ng)
    imgs_labels = []
    imgs_name = []
    for img in imgs_ok:
        img_name = img.split("@")
        if img_name[2] in class_AOI_name.keys():
            count_label_ok -= 1
            continue
        imgs_labels.append(0)
        imgs_name.append(imgs_ok_path+img)

    for img in imgs_ng:
        img_name = img.split("@")
        if img_name[2] in class_AOI_name.keys():
            count_label_ng -= 1
            continue
        imgs_labels.append(1)
        imgs_name.append(imgs_ng_path+img)

    #print(count_label_ok,count_label_ng, len(imgs_labels))
    #print("before infer")
    index = 0
    num = len(imgs_name)
    results_1 = []
    step = 256
    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("after infer")

    imgs_results_1 = []
    for result in results_1:
        res_predict = 0
        for i in result:
            for j in range(i.shape[0]):
                if i[j, 4]>0.015:
                    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(count_label_ok,count_label_ng)
    print(count_ok/count_label_ok, count_ng/count_label_ng)