Update app.py

app.py

@@ -1,36 +1,95 @@
 import gradio as gr
 import cv2
-import numpy as np
+import tempfile
 from ultralytics import YOLO
-from importlib.metadata import version
+from pathlib import Path
 
-# 动态组件选择
-GRADIO_VERSION = version("gradio")
-if int(GRADIO_VERSION.split('.')[0]) >= 4 and int(GRADIO_VERSION.split('.')[1]) >= 12:
-    Webcam = gr.Webcam
-else:
-    class Webcam(gr.Image):
-        def __init__(self, ​**kwargs):
-            kwargs.update(source="webcam", streaming=True)
-            super().__init__(**kwargs)
+# 全局变量存储当前模型
+current_model = None
 
-# YOLO模型初始化
-model = None
+def load_model(model_path):
+    global current_model
+    try:
+        current_model = YOLO(model_path)
+        return "模型加载成功！"
+    except Exception as e:
+        return f"模型加载失败：{str(e)}"
 
-def load_model(model_file):
-    global model
-    model = YOLO(model_file.name)
-    return "✅ 模型加载成功"
+def detect_image(input_image, conf_threshold):
+    if current_model is None:
+        raise gr.Error("请先上传模型文件")
+    
+    results = current_model(input_image, conf=conf_threshold)
+    plotted = results[0].plot()
+    return plotted[:, :, ::-1]  # BGR转RGB
+
+def detect_video(input_video, conf_threshold):
+    if current_model is None:
+        raise gr.Error("请先上传模型文件")
+    
+    cap = cv2.VideoCapture(input_video)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    # 创建临时输出文件
+    temp_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+    out = cv2.VideoWriter(temp_file.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        results = current_model(frame, conf=conf_threshold)
+        plotted = results[0].plot()
+        out.write(plotted)
+    
+    cap.release()
+    out.release()
+    return temp_file.name
+
+def detect_webcam(camera_input, conf_threshold):
+    if current_model is None:
+        raise gr.Error("请先上传模型文件")
+    
+    results = current_model(camera_input, conf=conf_threshold)
+    plotted = results[0].plot()
+    return plotted[:, :, ::-1]  # BGR转RGB
 
 with gr.Blocks() as demo:
-    with gr.Tab("📷 实时检测"):
-        webcam = Webcam(label="摄像头")
-        output = gr.Image()
-        webcam.stream(
-            fn=lambda x: model.predict(x)[0].plot(),
-            inputs=webcam,
-            outputs=output
-        )
+    gr.Markdown("# YOLOv8 自定义模型检测系统")
+    
+    with gr.Row():
+        model_input = gr.File(label="上传模型文件 (.pt)", type="filepath")
+        model_status = gr.Textbox(label="模型状态", interactive=False)
+    
+    model_input.upload(fn=load_model, inputs=model_input, outputs=model_status)
+    
+    with gr.Tabs():
+        with gr.TabItem("图片检测"):
+            with gr.Row():
+                img_input = gr.Image(label="输入图片", type="filepath")
+                img_output = gr.Image(label="检测结果")
+            img_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
+            img_button = gr.Button("执行检测")
+        
+        with gr.TabItem("视频检测"):
+            with gr.Row():
+                video_input = gr.Video(label="输入视频")
+                video_output = gr.Video(label="检测结果")
+            video_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
+            video_button = gr.Button("执行检测")
+        
+        with gr.TabItem("实时摄像头"):
+            webcam_input = gr.Image(label="摄像头画面", source="webcam", streaming=True)
+            webcam_output = gr.Image(label="检测结果", streaming=True)
+            webcam_conf = gr.Slider(0, 1, value=0.5, label="置信度阈值")
+    
+    # 绑定事件处理
+    img_button.click(fn=detect_image, inputs=[img_input, img_conf], outputs=img_output)
+    video_button.click(fn=detect_video, inputs=[video_input, video_conf], outputs=video_output)
+    webcam_input.stream(fn=detect_webcam, inputs=[webcam_input, webcam_conf], outputs=webcam_output)
 
 if __name__ == "__main__":
     demo.launch()