import gradio as gr
import cv2
import time
from ultralytics import YOLO
from paddleocr import PaddleOCR
import numpy as np
import detect_tools as tools
from imgTest import get_license_result
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# 加载YOLOv8检测模型
model_path = 'models/best.pt'
yolo_model = YOLO(model_path, task='detect')

# 加载车牌识别模型
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)


def Car_detection():
    def predict_image(query_image):
        start_time = time.time()

        # 图像预处理
        img = cv2.cvtColor(query_image, cv2.COLOR_BGR2RGB)
        print(f"Image preprocessing time: {time.time() - start_time:.2f}s")

        # 使用YOLOv8检测车辆和车牌位置
        yolo_start = time.time()
        results = yolo_model(img)[0]
        yolo_output = img.copy()  # 复制原图像用于显示YOLO结果
        location_list = results.boxes.xyxy.tolist()
        print(f"YOLO detection time: {time.time() - yolo_start:.2f}s")

        # 处理每个检测到的车牌区域
        license_numbers = []
        for location in location_list:
            x1, y1, x2, y2 = list(map(int, location))
            crop_img = img[y1:y2, x1:x2]

            # 使用PaddleOCR识别车牌号
            license_num, confidence = get_license_result(ocr, crop_img)
            if license_num:
                license_numbers.append(license_num)
            else:
                license_numbers.append("无法识别")

            # 在YOLO结果图上绘制检测框
            cv2.rectangle(yolo_output, (x1, y1), (x2, y2), (0, 255, 0), 2)

        return yolo_output, "\n".join(license_numbers)

    title = "<h1 align='center'>基于Opencv图像处理的车牌检测与识别</h1>"
    description = "上传一张包含车辆的图像，系统将检测车辆并识别车牌号码"
    # examples = [['images/car.jpg'], ['images/car.png'], ['images/car_test.jpg']]

    with gr.Blocks() as demo:
        gr.Markdown(title)
        gr.Markdown(description)
        with gr.Row():
            with gr.Column(scale=1):
                img = gr.components.Image(label="上传图片")
                btn = gr.Button("点击检测与识别")
            with gr.Column(scale=1):
                out_1 = gr.components.Image(label="YOLO定位结果:", height="auto")
                out_2 = gr.components.Textbox(label="车牌识别结果:", type="text", lines=6)

        inputs = [img]
        outputs = [out_1, out_2]
        btn.click(fn=predict_image, inputs=inputs, outputs=outputs)
        # gr.Examples(examples, inputs=inputs)

    return demo


if __name__ == "__main__":
    with gr.TabbedInterface(
            [Car_detection()],
            ["Opencv车牌检测与识别"],
    ) as demo:
        demo.launch(show_api=False, inbrowser=True)