app.py

import gradio as gr
import numpy as np
import cv2
from ultralytics import YOLO
from pathlib import Path
import tempfile

# 初始化模型占位符
model = YOLO('yolov8n.pt')  # 默认加载官方模型

def load_custom_model(model_path):
    global model
    model = YOLO(model_path)
    return "模型加载成功！"

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
    
    elif input_type == "image":
        # 图片检测
        results = model.predict(source=input_data)
        return results[0].plot()[:, :, ::-1]
    
    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

# Gradio界面布局
with gr.Blocks(title="YOLOv8检测系统") 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.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("🖼️ 图片检测"):
        img_input = gr.Image(type="filepath", label="上传图片")
        img_btn = gr.Button("开始检测", variant="primary")
        img_output = gr.Image(label="检测结果")
    
    with gr.Tab("🎥 视频检测"):
        vid_input = gr.Video(label="上传视频")
        vid_btn = gr.Button("处理视频", variant="primary")
        vid_output = gr.Video(label="处理结果")
    
    # 事件绑定
    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)

if __name__ == "__main__":
    demo.queue(concurrency_count=3).launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )