Update app.py

app.py

@@ -3,12 +3,24 @@ 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环境")
+except ImportError:
+    is_spaces = False
+    print("在本地环境运行")
 
 # 全局变量
 analyzer = None
 video_file_path = None
-model_file_path = "./fst-v1.2-n.onnx"  # 直接指定模型文件路径
+model_file_path = "weights/fst-v1.2-n.onnx"  # 直接指定模型文件路径
 total_frames = 0
 output_path = None
 
@@ -82,7 +94,17 @@ def preview_frame(video_file, frame_num):
     return frame, f"帧 {frame_num}"
 
 # 开始分析
-def start_analysis(video, conf, iou, max_det, start_frame, end_frame, threshold):
+# 为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
     
     if not video:
@@ -97,6 +119,26 @@ def start_analysis(video, conf, iou, max_det, start_frame, end_frame, threshold)
     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}")
+        
+        # 确保模型文件存在
+        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"
+                    )
+                    print(f"模型已下载到: {model_file_path}")
+                except Exception as e:
+                    print(f"从Hub下载模型失败: {str(e)}")
+            else:
+                print(f"警告: 模型文件 {model_file_path} 不存在!")
+        
         # 创建分析器
         analyzer = MouseTrackerAnalyzer(
             model_path=model_file_path, 
@@ -187,8 +229,13 @@ def create_interface():
                 # 只保留视频选择，移除模型选择
                 video_input = gr.Video(label="输入视频")
                 
-                # 显示当前使用的模型
-                model_info = gr.Textbox(label="模型信息", value=f"使用模型: {os.path.basename(model_file_path)}", interactive=False)
+                # 显示当前使用的模型和设备信息
+                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}", 
+                    interactive=False
+                )
                 
                 # 参数设置
                 with gr.Row():
@@ -251,6 +298,17 @@ def create_interface():
 
 # 启动应用
 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']
+    
+    # 检查设备和模型
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"使用设备: {device}")
     
     # 检查模型文件是否存在
     if not os.path.exists(model_file_path):
@@ -259,5 +317,11 @@ if __name__ == "__main__":
         print(f"使用模型: {model_file_path}")
     
     app = create_interface()
-    # 使用简化的启动配置
-    app.launch(server_name="0.0.0.0", server_port=7860, share=True) 
+    
+    # 根据环境决定启动方式
+    if is_spaces:
+        # Hugging Face Spaces环境中的启动方式
+        app.launch()
+    else:
+        # 本地环境的启动方式
+        app.launch(server_name="127.0.0.1", server_port=7860, share=False)