app.py

import os
import cv2
import numpy as np
import gradio as gr
import torch
from mouse_tracker import MouseTrackerAnalyzer
from huggingface_hub import hf_hub_download

# 检查是否在Hugging Face Spaces环境中
try:
    import spaces
    is_spaces = True
    print("检测到 Hugging Face Spaces 环境")
except ImportError:
    is_spaces = False
    print("在本地环境运行")

# 全局配置
model_base_name = "fst-v1.3-n"  # 模型基础名称，无后缀
total_frames = 0

# 根据后缀构造模型路径
def get_model_file_path(model_suffix):
    return f"./{model_base_name}{model_suffix}"

# 从视频中提取特定帧
def extract_frame(video_path, frame_num):
    if not video_path:
        return None
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        return None
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return None
    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

# 选择视频文件
def select_video(video_file, model_suffix):
    global total_frames
    if not video_file:
        return None, "请选择视频文件", gr.Slider(0,0,0), gr.Slider(0,0,0)
    total_frames = int(cv2.VideoCapture(video_file).get(cv2.CAP_PROP_FRAME_COUNT))
    # 读取首帧
    cap = cv2.VideoCapture(video_file)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return None, "无法读取视频帧", gr.Slider(0,0,0), gr.Slider(0,0,0)
    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # 更新滑块
    start = gr.Slider(minimum=0, maximum=total_frames-1, value=0, step=1)
    end = gr.Slider(minimum=0, maximum=total_frames-1, value=total_frames-1, step=1)
    status = f"视频加载成功，总帧数: {total_frames}. 使用模型: {os.path.basename(get_model_file_path(model_suffix))}"
    return frame_rgb, status, start, end

# 预览帧
def preview_frame(video_file, frame_num):
    if not video_file:
        return None, "请先选择视频文件"
    frame = extract_frame(video_file, frame_num)
    if frame is None:
        return None, "无法读取指定帧"
    return frame, f"帧 {frame_num}"

# 分析实现
def _start_analysis_impl(video, model_suffix, conf, iou, max_det, start_frame, end_frame, threshold):
    if not video:
        return None, None, "请选择视频文件"
    if start_frame >= end_frame:
        return None, None, "起始帧必须小于结束帧"
    # 构造路径
    video_name = os.path.splitext(os.path.basename(video))[0]
    output_path = os.path.join(os.path.dirname(video), f"{video_name}_out.mp4")
    csv_path = os.path.join(os.path.dirname(video), f"{video_name}_results.csv")
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    model_path = get_model_file_path(model_suffix)
    if not os.path.exists(model_path):
        if is_spaces:
            try:
                model_path = hf_hub_download(
                    repo_id="YOUR_HF_USERNAME/YOUR_REPO_NAME",
                    filename=f"weights/{model_base_name}{model_suffix}"
                )
            except Exception:
                print(f"下载模型失败: {model_path}")
        else:
            print(f"警告: 本地未找到模型文件 {model_path}")
    # 初始化分析器
    analyzer = MouseTrackerAnalyzer(
        model_path=model_path,
        conf=conf,
        iou=iou,
        max_det=max_det,
        verbose=True
    )
    analyzer.struggle_threshold = threshold
    # 运行分析
    analyzer.process_video(
        video_path=video,
        output_path=output_path,
        start_frame=start_frame,
        end_frame=end_frame,
        callback=lambda prog, frm, res: print(f"进度: {prog}% 检测: {len(res)} 项")
    )
    analyzer.save_results(csv_path)
    # 生成图表
    plot_path = None
    if analyzer.results:
        plot_path = analyzer.generate_time_series_plot()
    status = f"分析完成。视频: {output_path}, CSV: {csv_path}"
    if plot_path:
        status += f", 图表: {plot_path}"
    return output_path, plot_path, status

# HF Spaces GPU 装饰
if is_spaces:
    @spaces.GPU(duration=120)
    def start_analysis(video, model_suffix, conf, iou, max_det, start_frame, end_frame, threshold):
        return _start_analysis_impl(video, model_suffix, conf, iou, max_det, start_frame, end_frame, threshold)
else:
    def start_analysis(video, model_suffix, conf, iou, max_det, start_frame, end_frame, threshold):
        return _start_analysis_impl(video, model_suffix, conf, iou, max_det, start_frame, end_frame, threshold)

# 创建 Gradio 界面
def create_interface():
    with gr.Blocks(title="鼠强迫游泳挣扎度分析") as app:
        gr.Markdown("# 鼠强迫游泳测试挣扎度分析 (对象跟踪)")
        with gr.Row():
            with gr.Column(scale=1):
                video_input = gr.Video(label="输入视频")
                model_format = gr.Dropdown(
                    label="模型格式",
                    choices=[".onnx", ".engine", ".pt", ".mlpackage"],
                    value=".onnx",
                    interactive=True
                )
                device_info = gr.Textbox(
                    label="系统信息",
                    value=f"设备: {'GPU' if torch.cuda.is_available() else 'CPU'}",
                    interactive=False
                )
                conf = gr.Slider(0.1, 0.9, value=0.25, step=0.05, label="置信度阈值")
                iou = gr.Slider(0.1, 0.9, value=0.45, step=0.05, label="IoU阈值")
                max_det = gr.Slider(1, 50, value=20, step=1, label="最大检测数")
                threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="挣扎阈值")
                start_frame = gr.Slider(0, 999999, value=0, step=1, label="起始帧")
                end_frame = gr.Slider(0, 999999, value=999999, step=1, label="结束帧")
                preview_btn = gr.Button("预览帧")
                start_btn = gr.Button("开始分析", variant="primary")
            with gr.Column(scale=2):
                with gr.Tab("预览"):
                    preview_image = gr.Image(label="预览图像", type="numpy", height=400)
                    status_text = gr.Textbox(label="状态", interactive=False)
                with gr.Tab("结果"):
                    output_video = gr.Video(label="分析结果视频")
                    result_plot = gr.Image(label="挣扎分数时间序列")
                    result_status = gr.Textbox(label="分析状态", interactive=False)
        # 事件绑定，包含模型格式参数
        video_input.change(select_video, inputs=[video_input, model_format], outputs=[preview_image, status_text, start_frame, end_frame])
        preview_btn.click(preview_frame, inputs=[video_input, start_frame], outputs=[preview_image, status_text])
        start_btn.click(
            start_analysis,
            inputs=[video_input, model_format, conf, iou, max_det, start_frame, end_frame, threshold],
            outputs=[output_video, result_plot, result_status]
        )
    return app

if __name__ == "__main__":
    # 清理代理
    for key in ['http_proxy', 'https_proxy', 'all_proxy']:
        os.environ.pop(key, None)
    print(f"设备: {'GPU' if torch.cuda.is_available() else 'CPU'}")
    print(f"默认模型路径: {get_model_file_path('.onnx')}")
    app = create_interface()
    if is_spaces:
        app.launch()
    else:
        app.launch(server_name="0.0.0.0", server_port=7860, share=False)