videoTest.py

import tkinter as tk
from tkinter import filedialog, ttk
import cv2
import numpy as np
from PIL import Image, ImageTk
from ultralytics import YOLO
import threading
import time

class CowDetectorApp:
    def __init__(self, root):
        self.root = root
        self.root.title("牛识别计数系统")
        self.root.geometry("1000x700")
        self.root.configure(bg="#f0f0f0")
        
        # 模型加载
        self.model = None
        self.model_loaded = False
        
        # 视频相关变量
        self.video_path = None
        self.cap = None
        self.video_thread = None
        self.running = False
        self.paused = False
        self.current_frame = None
        
        # 创建界面组件
        self._create_widgets()
        
        # 尝试加载模型
        self._load_model()
    
    def _create_widgets(self):
        # 顶部控制栏
        control_frame = tk.Frame(self.root, bg="#e0e0e0", padx=10, pady=10)
        control_frame.pack(fill=tk.X)
        
        # 加载视频按钮
        self.load_btn = tk.Button(control_frame, text="加载视频", command=self.load_video,
                                 bg="#4CAF50", fg="white", padx=10, pady=5, font=("SimHei", 10))
        self.load_btn.pack(side=tk.LEFT, padx=5)
        
        # 播放/暂停按钮
        self.play_btn = tk.Button(control_frame, text="播放", command=self.toggle_play,
                                 bg="#2196F3", fg="white", padx=10, pady=5, font=("SimHei", 10), state=tk.DISABLED)
        self.play_btn.pack(side=tk.LEFT, padx=5)
        
        # 视频路径显示
        self.video_path_var = tk.StringVar()
        self.video_path_label = tk.Label(control_frame, textvariable=self.video_path_var, 
                                        bg="#e0e0e0", font=("SimHei", 10), wraplength=600)
        self.video_path_label.pack(side=tk.LEFT, padx=10, fill=tk.X, expand=True)
        
        # 中间视频显示区域
        video_frame = tk.Frame(self.root, bg="black")
        video_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)
        
        self.video_label = tk.Label(video_frame, bg="black")
        self.video_label.pack(fill=tk.BOTH, expand=True)
        
        # 底部状态栏
        status_frame = tk.Frame(self.root, bg="#e0e0e0", height=50)
        status_frame.pack(fill=tk.X, side=tk.BOTTOM, pady=10)
        
        # 计数显示
        self.count_var = tk.StringVar(value="牛数量: 0")
        count_label = tk.Label(status_frame, textvariable=self.count_var, 
                             font=("SimHei", 16, "bold"), bg="#e0e0e0", fg="#ff5722")
        count_label.pack(side=tk.LEFT, padx=20)
        
        # 模型状态显示
        self.model_status_var = tk.StringVar(value="模型状态: 未加载")
        model_status_label = tk.Label(status_frame, textvariable=self.model_status_var, 
                                    font=("SimHei", 12), bg="#e0e0e0")
        model_status_label.pack(side=tk.RIGHT, padx=20)
    
    def _load_model(self):
        """加载YOLO模型"""
        try:
            # 尝试加载模型
            self.model = YOLO('runs/detect/train14/weights/best.pt')  # 使用用户提供的模型文件
            self.model_loaded = True
            self.model_status_var.set("模型状态: 已加载")
        except Exception as e:
            self.model_status_var.set(f"模型加载失败: {str(e)}")
            print(f"模型加载错误: {e}")
    
    def load_video(self):
        """加载视频文件"""
        if not self.model_loaded:
            tk.messagebox.showerror("错误", "模型未加载成功，无法处理视频")
            return
            
        self.video_path = filedialog.askopenfilename(
            filetypes=[("视频文件", "*.mp4;*.avi;*.mov;*.mkv")]
        )
        
        if self.video_path:
            # 停止当前播放的视频
            self.stop_video()
            
            # 更新界面
            self.video_path_var.set(self.video_path)
            self.play_btn.config(state=tk.NORMAL)
            
            # 加载视频并显示第一帧
            self.cap = cv2.VideoCapture(self.video_path)
            ret, frame = self.cap.read()
            if ret:
                self._display_frame(frame, 0)
    
    def toggle_play(self):
        """切换播放/暂停状态"""
        if not self.running:
            # 开始播放
            self.running = True
            self.paused = False
            self.play_btn.config(text="暂停")
            self.video_thread = threading.Thread(target=self._process_video)
            self.video_thread.daemon = True
            self.video_thread.start()
        else:
            # 切换暂停状态
            self.paused = not self.paused
            self.play_btn.config(text="继续" if self.paused else "暂停")
    
    def stop_video(self):
        """停止视频播放"""
        self.running = False
        self.paused = False
        if self.cap:
            self.cap.release()
            self.cap = None
        self.play_btn.config(text="播放", state=tk.DISABLED)
    
    def _process_video(self):
        """处理视频帧并进行检测"""
        if not self.cap:
            return
            
        fps = self.cap.get(cv2.CAP_PROP_FPS)
        delay = 1.0 / fps if fps > 0 else 0.033  # 约30fps
        
        while self.running:
            if not self.paused:
                ret, frame = self.cap.read()
                if not ret:
                    # 视频播放完毕
                    self.running = False
                    self.root.after(0, lambda: self.play_btn.config(text="播放"))
                    break
                
                # 检测牛并计数
                results = self.model(frame)
                count = len(results[0].boxes)  # 牛的数量
                
                # 绘制检测结果
                annotated_frame = results[0].plot()
                
                # 更新界面
                self.root.after(0, lambda f=annotated_frame, c=count: self._display_frame(f, c))
            
            # 控制播放速度
            time.sleep(delay)
    
    def _display_frame(self, frame, count):
        """在界面上显示帧和计数"""
        # 更新计数
        self.count_var.set(f"牛数量: {count}")
        
        # 转换为Tkinter可用的格式
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        h, w = frame.shape[:2]
        
        # 调整大小以适应窗口
        window_width = self.video_label.winfo_width()
        window_height = self.video_label.winfo_height()
        
        if window_width > 10 and window_height > 10:  # 确保窗口已初始化
            # 计算缩放比例
            scale = min(window_width / w, window_height / h)
            new_w = int(w * scale)
            new_h = int(h * scale)
            frame = cv2.resize(frame, (new_w, new_h))
        
        # 显示图像
        image = Image.fromarray(frame)
        photo = ImageTk.PhotoImage(image=image)
        self.video_label.config(image=photo)
        self.video_label.image = photo  # 保持引用

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