Update app.py

app.py

@@ -2,37 +2,29 @@ import os
 import cv2
 import numpy as np
 import gradio as gr
-import tempfile
 import torch
 from mouse_tracker import MouseTrackerAnalyzer
-import huggingface_hub
 from huggingface_hub import hf_hub_download
 
 # 检查是否在Hugging Face Spaces环境中
 try:
     import spaces
     is_spaces = True
-    print("检测到Hugging Face Spaces环境")
+    print("检测到 Hugging Face Spaces 环境")
 except ImportError:
     is_spaces = False
     print("在本地环境运行")
 
-# 全局变量
-analyzer = None
-video_file_path = None
-model_suffix = ".onnx"  # 默认使用 TensorRT 格式
+# 全局配置
 model_base_name = "fst-v1.2-n"  # 模型基础名称，无后缀
 total_frames = 0
-output_path = None
 
-# 构造模型文件路径
-def get_model_file_path():
-    """根据用户选择的后缀返回完整模型文件路径"""
+# 根据后缀构造模型路径
+def get_model_file_path(model_suffix):
     return f"./{model_base_name}{model_suffix}"
 
 # 从视频中提取特定帧
 def extract_frame(video_path, frame_num):
-    """从视频中提取指定帧号并返回 RGB 图像"""
     if not video_path:
         return None
     cap = cv2.VideoCapture(video_path)
@@ -46,26 +38,23 @@ def extract_frame(video_path, frame_num):
     return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
 # 选择视频文件
-def select_video(video_file):
-    global video_file_path, total_frames
+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)
-    video_file_path = video_file
-    cap = cv2.VideoCapture(video_file_path)
-    if not cap.isOpened():
-        return None, "无法打开视频文件", gr.Slider(0,0,0), gr.Slider(0,0,0)
-    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    ret, first_frame = cap.read()
+        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)
-    first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
-    # 更新滑块范围
+    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())}"
-    return first_frame_rgb, status, start, end
+    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):
@@ -76,94 +65,87 @@ def preview_frame(video_file, frame_num):
         return None, "无法读取指定帧"
     return frame, f"帧 {frame_num}"
 
-# 分析逻辑实现
-def _start_analysis_impl(video, conf, iou, max_det, start_frame, end_frame, threshold):
-    global analyzer, output_path
+# 分析实现
+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")
-    try:
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        model_file_path = get_model_file_path()
-        if not os.path.exists(model_file_path):
-            if is_spaces:
-                try:
-                    model_file_path = hf_hub_download(
-                        repo_id="YOUR_HF_USERNAME/YOUR_REPO_NAME",
-                        filename=f"weights/{model_base_name}{model_suffix}"
-                    )
-                except Exception as e:
-                    print(f"下载模型失败: {e}")
-            else:
-                print(f"警告: 未找到模型文件 {model_file_path}")
-        analyzer = MouseTrackerAnalyzer(
-            model_path=model_file_path,
-            conf=conf,
-            iou=iou,
-            max_det=max_det,
-            verbose=True
-        )
-        analyzer.struggle_threshold = threshold
-        def progress_update(progress, frame, results):
-            print(f"进度: {progress}%，检测: {len(results)} 个对象")
-        results = analyzer.process_video(
-            video_path=video,
-            output_path=output_path,
-            start_frame=start_frame,
-            end_frame=end_frame,
-            callback=progress_update
-        )
-        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
-    except Exception as e:
-        import traceback
-        traceback.print_exc()
-        return None, None, f"处理错误: {e}"
+    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 装饰器
+# HF Spaces GPU 装饰
 if is_spaces:
     @spaces.GPU(duration=120)
-    def start_analysis(video, conf, iou, max_det, start_frame, end_frame, threshold):
-        return _start_analysis_impl(video, conf, iou, max_det, start_frame, end_frame, threshold)
+    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, conf, iou, max_det, start_frame, end_frame, threshold):
-        return _start_analysis_impl(video, conf, iou, max_det, start_frame, end_frame, threshold)
+    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:
+    with gr.Blocks(title="鼠强迫游泳挣扎度分析") as app:
         gr.Markdown("# 鼠强迫游泳测试挣扎度分析 (对象跟踪)")
         with gr.Row():
             with gr.Column(scale=1):
-                video_input = gr.Video(label="输入视频")  
-                # 新增模型格式选择下拉框
+                video_input = gr.Video(label="输入视频")
                 model_format = gr.Dropdown(
                     label="模型格式",
                     choices=[".onnx", ".engine", ".pt", ".mlpackage"],
-                    value=model_suffix,
+                    value=".onnx",
                     interactive=True
                 )
                 device_info = gr.Textbox(
                     label="系统信息",
-                    value=f"设备: {'GPU' if torch.cuda.is_available() else 'CPU'}", 
+                    value=f"设备: {'GPU' if torch.cuda.is_available() else 'CPU'}",
                     interactive=False
                 )
-                with gr.Row():
-                    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阈值")
-                with gr.Row():
-                    max_det = gr.Slider(1, 50, value=20, step=1, label="最大检测数")
-                    threshold = gr.Slider(0, 1, value=0.3, step=0.01, label="挣扎阈值")
+                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("预览帧")
@@ -173,30 +155,25 @@ def create_interface():
                     preview_image = gr.Image(label="预览图像", type="numpy", height=400)
                     status_text = gr.Textbox(label="状态", interactive=False)
                 with gr.Tab("结果"):
-                    with gr.Row():
-                        output_video = gr.Video(label="分析结果视频")
-                        result_plot = gr.Image(label="挣扎分数时间序列")
+                    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], outputs=[preview_image, status_text, start_frame, end_frame])
+        # 事件绑定，包含模型格式参数
+        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])
-        # 传递模型格式到全局
-        model_format.change(lambda fmt: setattr(globals(), 'model_suffix', fmt) or fmt, inputs=[model_format], outputs=[])        
         start_btn.click(
             start_analysis,
-            inputs=[video_input, conf, iou, max_det, start_frame, end_frame, threshold],
+            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 p in ['http_proxy', 'https_proxy', 'all_proxy']:
-        os.environ.pop(p, None)
-    # 日志设备和模型信息
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    print(f"使用设备: {device}")
-    print(f"默认模型路径: {get_model_file_path()}")
+    # 清理代理
+    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()