Update app.py

app.py

@@ -20,308 +20,185 @@ except ImportError:
 # 全局变量
 analyzer = None
 video_file_path = None
-model_file_path = "./fst-v1.2-n.engine"  # 直接指定模型文件路径
+model_suffix = ".engine"  # 默认使用 TensorRT 格式
+model_base_name = "fst-v1.2-n"  # 模型基础名称，无后缀
 total_frames = 0
 output_path = None
 
+# 构造模型文件路径
+def get_model_file_path():
+    """根据用户选择的后缀返回完整模型文件路径"""
+    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)
     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
-    
-    # 转换为RGB格式
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    return frame_rgb
+    return cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
 
 # 选择视频文件
 def select_video(video_file):
     global video_file_path, total_frames
-    
     if not video_file:
-        return None, "请选择视频文件", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
-        
+        # 无视频时重置滑块和提示
+        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(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
-        
+        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()
     cap.release()
-    
     if not ret:
-        return None, "无法读取视频帧", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
-    
-    # 转为RGB
+        return None, "无法读取视频帧", gr.Slider(0,0,0), gr.Slider(0,0,0)
     first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
-        
-    # 更新帧滑块
-    start_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=0, step=1)
-    end_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=total_frames-1, step=1)
-    
-    model_status = f"使用模型: {os.path.basename(model_file_path)}"
-    return first_frame_rgb, f"视频加载成功，总帧数: {total_frames}. {model_status}", start_slider, end_slider
+    # 更新滑块范围
+    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
 
 # 预览帧
 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}"
 
-# 开始分析
-# 为HF Spaces环境添加GPU装饰器
-if is_spaces:
-    @spaces.GPU(duration=120)  # 申请GPU资源，持续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)
-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_impl(video, conf, iou, max_det, start_frame, end_frame, threshold):
-    global analyzer, output_path, model_file_path
-    
+    global analyzer, output_path
     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'
-        print(f"使用设备: {device}")
-        
-        # 确保模型文件存在
+        model_file_path = get_model_file_path()
         if not os.path.exists(model_file_path):
-            # 如果在Hugging Face Spaces环境中，尝试从Hub下载模型
             if is_spaces:
                 try:
-                    print(f"尝试从Hugging Face Hub下载模型: {os.path.basename(model_file_path)}")
                     model_file_path = hf_hub_download(
-                        repo_id="YOUR_HF_USERNAME/YOUR_REPO_NAME",  # 替换为您的仓库
-                        filename="weights/fst-v1.2-n.onnx"
+                        repo_id="YOUR_HF_USERNAME/YOUR_REPO_NAME",
+                        filename=f"weights/{model_base_name}{model_suffix}"
                     )
-                    print(f"模型已下载到: {model_file_path}")
                 except Exception as e:
-                    print(f"从Hub下载模型失败: {str(e)}")
+                    print(f"下载模型失败: {e}")
             else:
-                print(f"警告: 模型文件 {model_file_path} 不存在!")
-        
-        # 创建分析器
+                print(f"警告: 未找到模型文件 {model_file_path}")
         analyzer = MouseTrackerAnalyzer(
-            model_path=model_file_path, 
-            conf=conf, 
-            iou=iou, 
-            max_det=max_det, 
-            verbose=True  # 开启详细日志
+            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)}")
-            
-        print(f"处理视频: {video}")
-        print(f"输出路径: {output_path}")
-        print(f"参数: conf={conf}, iou={iou}, max_det={max_det}, threshold={threshold}")
-        print(f"使用模型: {model_file_path}")
-        
-        # 提取视频帧数范围并分析
+            print(f"进度: {progress}%，检测: {len(results)} 个对象")
         results = analyzer.process_video(
-            video_path=video, 
-            output_path=output_path, 
-            start_frame=start_frame, 
+            video_path=video,
+            output_path=output_path,
+            start_frame=start_frame,
             end_frame=end_frame,
             callback=progress_update
         )
-        
-        # 保存结果到CSV
-        print(f"保存结果到CSV: {csv_path}")
         analyzer.save_results(csv_path)
-        print(f"结果已保存到CSV，共 {len(analyzer.results)} 帧数据")
-        
-        # 生成分析图表
-        print("生成时间序列图...")
-        if len(analyzer.results) == 0:
-            print("警告: 没有可用于绘图的结果!")
-            plot_path = None
-        else:
+        plot_path = None
+        if analyzer.results:
             plot_path = analyzer.generate_time_series_plot()
-            if plot_path and os.path.exists(plot_path):
-                print(f"图表已生成并保存到: {plot_path}, 大小: {os.path.getsize(plot_path)/1024:.2f}KB")
-            else:
-                print(f"生成图表失败或图表文件不存在!")
-                plot_path = None
-        
-        # 检查输出文件是否存在
-        if os.path.exists(output_path):
-            file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
-            print(f"输出视频大小: {file_size:.2f}MB")
-            
-            # 处理debug帧
-            debug_frame_path = os.path.join(os.path.dirname(output_path), "debug_frame.jpg")
-            if os.path.exists(debug_frame_path):
-                print(f"调试帧保存在: {debug_frame_path}")
-        
-        if plot_path and os.path.exists(plot_path):
-            print(f"图表文件存在于: {plot_path}, 大小: {os.path.getsize(plot_path)/1024:.2f}KB")
-            
-        # 确保返回正确的文件路径
-        status_message = "分析完成。"
-        
-        if os.path.exists(output_path):
-            status_message += f"视频已保存。"
-        else:
-            status_message += "警告: 未找到输出视频。"
-            
-        if plot_path and os.path.exists(plot_path):
-            status_message += f" 时间序列图已生成。"
-        else:
-            status_message += " 警告: 生成时间序列图失败。"
-            
-        status_message += f" 结果已保存到: {csv_path}"
-        
-        return output_path, plot_path, status_message
+        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"处理错误: {str(e)}"
+        return None, None, f"处理错误: {e}"
 
-# 创建Gradio界面
+# HF Spaces 装饰器
+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)
+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)
+
+# 创建 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="输入视频")
-                
-                # 显示当前使用的模型和设备信息
-                device_info = "GPU" if torch.cuda.is_available() else "CPU"
-                model_info = gr.Textbox(
-                    label="系统信息", 
-                    value=f"使用模型: {os.path.basename(model_file_path)} | 计算设备: {device_info}", 
+                video_input = gr.Video(label="输入视频")  
+                # 新增模型格式选择下拉框
+                model_format = gr.Dropdown(
+                    label="模型格式",
+                    choices=[".onnx", ".engine", ".pt", ".mlpackage"],
+                    value=model_suffix,
+                    interactive=True
+                )
+                device_info = gr.Textbox(
+                    label="系统信息",
+                    value=f"设备: {'GPU' if torch.cuda.is_available() else 'CPU'}", 
                     interactive=False
                 )
-                
-                # 参数设置
                 with gr.Row():
-                    conf = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="置信度阈值")
-                    iou = gr.Slider(minimum=0.1, maximum=0.9, value=0.45, step=0.05, label="IoU阈值")
-                
+                    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(minimum=1, maximum=50, value=20, step=1, label="最大检测数")
-                    threshold = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.01, label="挣扎阈值")
-                
-                # 帧选择
-                start_frame = gr.Slider(minimum=0, maximum=999999, value=0, step=1, label="起始帧")
-                end_frame = gr.Slider(minimum=0, maximum=999999, value=999999, step=1, label="结束帧")
-                
-                # 预览按钮
+                    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)
-                    gr.Markdown("""
-                    ### 使用说明:
-                    1. 选择一个视频文件
-                    2. 调整参数
-                       - 置信度阈值: 对象检测的最低置信度，较低的值会检测更多潜在对象
-                       - IoU阈值: 用于过滤重叠检测
-                       - 最大检测数: 每帧检测的最大对象数
-                       - 挣扎阈值: 分类为挣扎状态的最低分数
-                    3. 设置帧范围
-                    4. 点击"开始分析"按钮
-                    
-                    系统将自动跟踪小鼠并分析其挣扎行为，无需手动定义区域
-                    """)
-                
                 with gr.Tab("结果"):
                     with gr.Row():
                         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])
-        
         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], 
+            start_analysis,
+            inputs=[video_input, conf, iou, max_det, start_frame, end_frame, threshold],
             outputs=[output_video, result_plot, result_status]
         )
-    
     return app
 
-# 启动应用
 if __name__ == "__main__":
-    # 清除可能干扰的代理设置
-    if 'http_proxy' in os.environ:
-        del os.environ['http_proxy']
-    if 'https_proxy' in os.environ:
-        del os.environ['https_proxy']
-    if 'all_proxy' in os.environ:
-        del os.environ['all_proxy']
-    
-    # 检查设备和模型
+    # 清除代理环境
+    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}")
-    
-    # 检查模型文件是否存在
-    if not os.path.exists(model_file_path):
-        print(f"警告: 模型文件 {model_file_path} 不存在!")
-    else:
-        print(f"使用模型: {model_file_path}")
-    
+    print(f"默认模型路径: {get_model_file_path()}")
     app = create_interface()
-    
-    # 根据环境决定启动方式
     if is_spaces:
-        # Hugging Face Spaces环境中的启动方式
         app.launch()
     else:
-        # 本地环境的启动方式
-        app.launch(server_name="0.0.0.0", server_port=7860, share=False) 
+        app.launch(server_name="0.0.0.0", server_port=7860, share=False)