import tkinter as tk
from tkinter import filedialog
from PIL import Image, ImageTk, ImageDraw
from ultralytics import YOLO

class YoloPredictorApp:
    def __init__(self, master):
        self.master = master
        self.master.title("YOLO Predictor")

        # 加载预训练模型
        self.model = YOLO(
            '/Users/sunmingzheng/Downloads/20241026/yolov8/ultralytics-main/runs/detect/train12/weights/best.pt')

        self.show_rain_effect = False
        self.default_image_path = "/Users/sunmingzheng/Desktop/rain.jpg"

        self.original_image = None  # 用于保存原始图像

        # 类别名称列表
        self.names = [
            "xiansuwu","xiansushiwu","xiansusanshi","xiansusishi","xiansuwushi","xiansuliushi","xiansuqishi","xiansubashi",
            "jinzhizhixingzuozhuan","huicherangxing","jinzhizhixing","jinzhizuozhuan","jinzhizuozhuanyouzhuan","jinzhiyouzhuan",
            "jinzhichaoche","jinzhidiaotou","jinzhijidongche","jinzhimingdi","xiankuansanmi","xianzhongshidun","zhixingyouzhuan",
            "zhixing","zuozhuan","zuozhuanyouzhuan","youzhuan","kaozuoxingshi","kaoyouxingshi","huandaoxingshi","jidongchexingshi",
            "mingdi","feijidongchexingshi","diaotou","zuoyouraoxing","youceraoxing","zhuyiweixian","zhuyixingren","zhuyifeijidongche",
            "zhuyiertong","xiangyoujiwan","xiangzuojiwan","xiadoupo","shangdoupo","manxing","youTxingjiaocha","zuoTxingjiaocha",
            "cunzhuang","fanxiangwanlu","wurenkanshoutielukou","shigong","lianxuwanlu","yourenkanshoutiedaokou","shiguyifaluduan",
            "tingcherangxing","jinzhitongxing","jinzhitingfang","jinzhishiru","jiansurangxing","tingchejiancha"
        ]

        # 创建主容器
        main_frame = tk.Frame(self.master)
        main_frame.pack(fill=tk.BOTH, expand=True)

        # 左侧图片区域
        self.image_label = tk.Label(main_frame)
        self.image_label.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 信息标签区域
        info_frame = tk.Frame(main_frame, width=200)
        info_frame.pack(side=tk.RIGHT, fill=tk.BOTH)

        # 各种信息标签
        label_options = {"font": ("Helvetica", 16), "anchor": "center"}
        self.label_cls = tk.Label(info_frame, text="标志类别: None", **label_options)
        self.label_speed = tk.Label(info_frame, text="速度: 40 km/h", **label_options)
        self.label_rain = tk.Label(info_frame, text="是否有雨: No", **label_options)
        self.label_roadblock = tk.Label(info_frame, text="是否有临时路障: No", **label_options)
        self.label_special_vehicle = tk.Label(info_frame, text="是否有特种车辆: No", **label_options)

        # 显示标签
        for label in [self.label_cls, self.label_speed, self.label_rain, self.label_roadblock, self.label_special_vehicle]:
            label.pack(fill=tk.BOTH, pady=20, padx=20, expand=True)

        # 创建菜单栏
        menu = tk.Menu(self.master)
        self.master.config(menu=menu)

        # 文件菜单
        file_menu = tk.Menu(menu)
        menu.add_cascade(label="File", menu=file_menu)
        file_menu.add_command(label="Choose Image", command=self.choose_image)
        file_menu.add_separator()
        file_menu.add_command(label="Predict", command=self.predict)
        file_menu.add_command(label="Add Rain Effect", command=self.add_rain_effect)
        file_menu.add_command(label="Remove Rain Effect", command=self.remove_rain_effect)

    # 选择图片
    def choose_image(self):
        file_path = filedialog.askopenfilename(filetypes=[("Image files", "*.jpg *.jpeg *.png")])
        if file_path:
            self.image = Image.open(file_path)
            self.image = self.image.resize((600, 600))
            self.original_image = self.image.copy()  # 保存原始图像的副本

            self.photo = ImageTk.PhotoImage(self.image)

            # 将图片显示在窗口左侧
            self.image_label.config(image=self.photo)
            self.image_label.image = self.photo

    # 进行预测
    def predict(self):
        if hasattr(self, 'image'):
            results = self.model(self.image)
            draw = ImageDraw.Draw(self.image)
            for r in results:
                bbox = r.boxes.xyxy.tolist()
                cls = r.boxes.cls.data.tolist()
                bbox = [int(coord) for coord in bbox[0]]
                draw.rectangle(bbox, outline="red", width=2)
                cls_index = int(cls[0])  # 类别索引
                class_name = self.names[cls_index] if cls_index < len(self.names) else "未知类别"
                draw.text((bbox[0], bbox[1]), class_name, fill="red", font_size=20)

                # 更新标志类别
                self.label_cls.config(text=f"标志类别: {class_name}")

            self.photo = ImageTk.PhotoImage(self.image)
            self.image_label.config(image=self.photo)
            self.image_label.image = self.photo

    # 添加雨滴效果
    def add_rain_effect(self):
        if self.image is not None:
            default_image = Image.open(self.default_image_path)
            default_image = default_image.resize(self.image.size)

            self.show_rain_effect = True
            if self.show_rain_effect:
                self.image = Image.blend(self.image, default_image, 0.5)

                # 更新速度和是否有雨
                self.label_speed.config(text="速度: 30 km/h")
                self.label_rain.config(text="是否有雨: Yes")

            self.photo = ImageTk.PhotoImage(self.image)
            self.image_label.config(image=self.photo)

    # 移除雨滴效果
    def remove_rain_effect(self):
        if self.original_image is not None:
            self.image = self.original_image.copy()  # 恢复到原始图像

            # 更新速度和是否有雨
            self.label_speed.config(text="速度: 40 km/h")
            self.label_rain.config(text="是否有雨: No")

            self.photo = ImageTk.PhotoImage(self.image)
            self.image_label.config(image=self.photo)

if __name__ == "__main__":
    root = tk.Tk()
    app = YoloPredictorApp(root)
    root.mainloop()