# 使用paddlex训练好的模型对批量图片进行测试，统计过严和流出情况
# 需要修改模型路径，测试图片路径，图片格式
# 以人工判定结果命名图片(true_label，OK或者NG)

import os
import glob
import paddlex as pdx

model = pdx.load_model(r'D:\workspace\PaddleX\output\paddle_model\ppyolo\best_model')
# pdx.det.visualize(infer_img, result, threshold=0.3, save_dir='./')
infer_dict = {'过严':0, '流出':0, 'OK一致':0, 'NG一致':0, '其它':0}
total_image = glob.glob('test\*.bmp')
print('总共测试%d张图片' % len(total_image))
for image_path in total_image:
    result = model.predict(image_path)
    print(result[0]['category'], result[0]['score'])
    infer_label = result[0]['category'][:2]
    true_label = os.path.split(image_path)[1][:2]
    if true_label == 'OK' and infer_label == 'NG':
        infer_dict['过严'] += 1
        print('过严')
    elif true_label == 'NG' and infer_label == 'OK':
        infer_dict['流出'] += 1
        print('流出')
    elif true_label == 'NG' and infer_label == 'NG':
        infer_dict['NG一致'] += 1
        print('NG一致')
    elif true_label == 'OK' and infer_label == 'OK':
        infer_dict['OK一致'] += 1
        print('OK一致')
    else:
        infer_dict['其它'] += 1
        print('其它')
    print('--------------------------------')
print(infer_dict)
print('过严率：', infer_dict['过严'] / sum(infer_dict.values()))
print('流出率：', infer_dict['流出'] / sum(infer_dict.values()))
print('NG一致率：', infer_dict['NG一致'] / sum(infer_dict.values()))
print('OK一致率：', infer_dict['OK一致'] / sum(infer_dict.values()))