Update app.py

app.py

@@ -4,84 +4,118 @@ import cv2
 from ultralytics import YOLO
 from pathlib import Path
 import tempfile
+import os
 
-# 初始化模型占位符
-model = YOLO('yolov8n.pt')  # 默认加载官方模型
+# 初始化默认模型
+model = YOLO('yolov8n.pt')  # 自动下载基础模型
 
 def load_custom_model(model_path):
     global model
-    model = YOLO(model_path)
-    return "模型加载成功！"
+    try:
+        model = YOLO(model_path.name)  # 适配HuggingFace的文件对象
+        return "✅ 模型加载成功！"
+    except Exception as e:
+        return f"❌ 加载失败：{str(e)}"
 
-def detect_media(input_type, input_data):
-    if input_type == "camera":
-        # 摄像头实时检测
-        frame = input_data
-        results = model.predict(source=frame, stream=True)
-        annotated_frame = results[0].plot()
-        return annotated_frame[:, :, ::-1]  # BGR转RGB
+def process_frame(frame, input_type):
+    # 统一处理函数
+    results = model.predict(
+        source=frame,
+        verbose=False,     # 关闭控制台输出
+        device="cpu",      # 适配免费Space环境
+        conf=0.5          # 置信度阈值
+    )
+    annotated = results[0].plot()
+    return cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
+
+def video_pipeline(input_video):
+    # 视频处理优化
+    temp_dir = tempfile.mkdtemp()
+    output_path = os.path.join(temp_dir, "output.mp4")
+    
+    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))
+    
+    writer = cv2.VideoWriter(
+        output_path,
+        cv2.VideoWriter_fourcc(*'mp4v'),
+        fps, 
+        (width, height)
+    )
     
-    elif input_type == "image":
-        # 图片检测
-        results = model.predict(source=input_data)
-        return results[0].plot()[:, :, ::-1]
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret: break
+        writer.write(process_frame(frame, "video"))
     
-    elif input_type == "video":
-        # 视频检测处理
-        cap = cv2.VideoCapture(input_data)
-        output_frames = []
-        
-        while cap.isOpened():
-            ret, frame = cap.read()
-            if not ret: break
-            
-            results = model.predict(source=frame)
-            annotated_frame = results[0].plot()
-            output_frames.append(annotated_frame)
-        
-        # 生成临时输出视频
-        output_path = str(Path(tempfile.gettempdir()) / "output.mp4")
-        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), 
-                            30, (annotated_frame.shape[1], annotated_frame.shape[0]))
-        for f in output_frames:
-            out.write(f)
-        out.release()
-        
-        return output_path
+    cap.release()
+    writer.release()
+    return output_path
 
 # Gradio界面布局
-with gr.Blocks(title="YOLOv8检测系统") as demo:
-    gr.Markdown("# 🚀 YOLOv8多功能检测系统")
+with gr.Blocks(title="YOLOv8检测系统", css=".gradio-container {background: #f0f0f0}") as demo:
+    gr.Markdown("""# 🚀 YOLOv8多功能检测系统
+                *欢迎上传自定义模型或使用默认模型进行检测*""")
     
     with gr.Tab("⚙️ 模型管理"):
-        model_upload = gr.File(label="上传模型文件(.pt)", file_types=[".pt"])
-        load_btn = gr.Button("加载模型", variant="primary")
-        model_status = gr.Textbox(label="模型状态")
+        with gr.Row():
+            model_upload = gr.UploadButton(
+                "上传模型文件 (.pt)",
+                file_types=[".pt"],
+                variant="primary"
+            )
+            model_status = gr.Textbox(label="状态", interactive=False)
+        gr.Examples(
+            examples=["yolov8n.pt", "yolov8s.pt"],
+            inputs=model_upload,
+            label="示例模型"
+        )
     
-    with gr.Tab("📷 实时摄像头"):
-        webcam = gr.Image(source="webcam", streaming=True, label="摄像头画面")
-        cam_output = gr.Image(label="检测结果", interactive=False)
-        webcam.stream(fn=lambda x: detect_media("camera", x), 
-                     inputs=webcam, outputs=cam_output)
+    with gr.Tab("📷 实时检测"):
+        with gr.Row():
+            webcam = gr.Webcam(label="摄像头输入", mirror=True)
+            cam_output = gr.Image(label="检测结果")
+        webcam.stream(
+            fn=lambda x: process_frame(x, "camera"),
+            inputs=webcam,
+            outputs=cam_output,
+            show_progress="hidden"
+        )
     
     with gr.Tab("🖼️ 图片检测"):
-        img_input = gr.Image(type="filepath", label="上传图片")
+        with gr.Row():
+            img_input = gr.Image(type="filepath", sources=["upload"], label="输入图片")
+            img_output = gr.Image(label="检测结果")
         img_btn = gr.Button("开始检测", variant="primary")
-        img_output = gr.Image(label="检测结果")
+        img_btn.click(
+            fn=lambda x: process_frame(cv2.imread(x), "image"),
+            inputs=img_input,
+            outputs=img_output
+        )
     
     with gr.Tab("🎥 视频检测"):
-        vid_input = gr.Video(label="上传视频")
+        with gr.Row():
+            vid_input = gr.Video(label="输入视频", sources=["upload"])
+            vid_output = gr.Video(label="处理结果")
         vid_btn = gr.Button("处理视频", variant="primary")
-        vid_output = gr.Video(label="处理结果")
+        vid_btn.click(
+            fn=video_pipeline,
+            inputs=vid_input,
+            outputs=vid_output
+        )
     
-    # 事件绑定
-    load_btn.click(fn=load_custom_model, inputs=model_upload, outputs=model_status)
-    img_btn.click(fn=lambda x: detect_media("image", x), inputs=img_input, outputs=img_output)
-    vid_btn.click(fn=lambda x: detect_media("video", x), inputs=vid_input, outputs=vid_output)
+    # 模型加载事件
+    model_upload.upload(
+        fn=load_custom_model,
+        inputs=model_upload,
+        outputs=model_status
+    )
 
 if __name__ == "__main__":
-    demo.queue(concurrency_count=3).launch(
+    demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=True
+        show_error=True
     )