# coding:utf-8
from ultralytics import YOLO
import cv2
import detect_tools as tools
from PIL import ImageFont
from paddleocr import PaddleOCR
import matplotlib.pyplot as plt
import os

def get_license_result(ocr, image):
    """
    image:输入的车牌截取照片
    输出，车牌号与置信度
    """
    result = ocr.ocr(image, cls=True)[0]
    if result:
        license_name, conf = result[0][1]
        if '·' in license_name:
            license_name = license_name.replace('·', '')
        return license_name, conf
    else:
        return None, None

# 获取图片地址
img_path = "images"
files = os.listdir(img_path)
s = []
for file in files:
    temp = img_path+"/"+file
    s.append(temp)
fontC = ImageFont.truetype("Fonts/msyhbd.ttc", 50, 0, encoding="utf-8")

# 加载ocr模型
cls_model_dir = 'paddleModels/whl/cls/ch_ppocr_mobile_v2.0_cls_infer'
rec_model_dir = 'paddleModels/whl/rec/ch/ch_PP-OCRv4_rec_infer'
ocr = PaddleOCR(use_angle_cls=False, lang="ch", det=False, cls_model_dir=cls_model_dir, rec_model_dir=rec_model_dir)

# 所需加载的模型目录
path = 'runs/detect/train/weights/best.pt'
# 加载预训练模型
# conf	0.25	object confidence threshold for detection
# iou	0.7	int.ersection over union (IoU) threshold for NMS
model = YOLO(path, task='detect')
# model = YOLO(path, task='detect',conf=0.5)

#检测图片
i = 0
for now_img in s:
    now_img = tools.img_cvread(now_img)
    results = model(now_img)[0]
    location_list = results.boxes.xyxy.tolist()
    if len(location_list) >= 1:
        location_list = [list(map(int, e)) for e in location_list]
        # 截取每个车牌区域的照片
        license_imgs = []
        for each in location_list:
            x1, y1, x2, y2 = each
            cropImg = now_img[y1:y2, x1:x2]
            license_imgs.append(cropImg)
        # 车牌识别结果
        lisence_res = []
        conf_list = []
        for each in license_imgs:
            license_num, conf = get_license_result(ocr, each)
            if license_num:
                lisence_res.append(license_num)
                conf_list.append(conf)
            else:
                lisence_res.append('无法识别')
                conf_list.append(0)
        for text, box in zip(lisence_res, location_list):
            now_img = tools.drawRectBox(now_img, box, text, fontC)
            now_img = cv2.resize(now_img, dsize=None, fx=0.5, fy=0.5, interpolation=cv2.INTER_LINEAR)
            # cv2.imshow("YOLOv8 Detection", now_img)
            cv2.imwrite(f"./result/{i}.jpg", now_img)
            i += 1