Update app.py

app.py

@@ -1,6 +1,6 @@
 # filename: app.py
 
-import spaces  # 必须最先 import
+import spaces  # 必须最先 import，用于 ZeroGPU 装饰
 import os
 import cv2
 import numpy as np
@@ -15,24 +15,21 @@ print(f"CUDA available: {use_cuda}")
 if use_cuda:
     print(f"GPU Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-# === 2. 加载模型并显式指定 task ===
-# 避免无法猜测任务类型的警告，明确使用分割 (segment) 模式
-model = YOLO("fst-v1.2-n.onnx", task="segment")  # ONNX 模型需上传至空间
-
-# 若 CUDA 可用，迁移模型至 GPU
+# === 2. 加载模型 (显式指定 segmentation 任务 & 默认置信度阈值) ===
+#    conf 设置为 0.25，可根据实际降低到 0.1-0.2
+model = YOLO("fst-v1.2-n.onnx", task="segment")  # 避免自动猜测任务模式的警告 
 if use_cuda:
     try:
         model.model.to("cuda")
-    except Exception:
+    except:
         pass
 
-@spaces.GPU(duration=600)  # 调用时分配 GPU，超时 600s :contentReference[oaicite:3]{index=3}
+@spaces.GPU(duration=600)  # 在 ZeroGPU 上执行该函数，超时 600s :contentReference[oaicite:2]{index=2}
 def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     """
-    核心分析：分割 → 跟踪 → 计算“挣扎强度”
-    返回：标注后的视频路径 & 绘制好的挣扎曲线 (matplotlib Figure)
+    分割 → 跟踪 → “挣扎强度” 计算
+    返回：标注后视频文件路径 & matplotlib Figure
     """
-    # 视频读写设置
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     out_path = "output.mp4"
@@ -40,7 +37,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
 
-    # 跟踪初始化
+    # 跟踪状态
     prev_centroids   = [None] * num_mice
     prev_masks       = [None] * num_mice
     struggle_records = [[] for _ in range(num_mice)]
@@ -51,21 +48,33 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
         if not ret:
             break
 
-        # 分割推理（stream=True 加速），指定 device
+        # === 3. 分割推理 ===
         device = "cuda" if use_cuda else "cpu"
-        results = model(frame, stream=True, device=device)
+        results = model(frame, stream=True, device=device, conf=0.25)  # 指定置信度阈值 :contentReference[oaicite:3]{index=3}
         res     = next(results)
-        masks   = res.masks.data.cpu().numpy()  # [N, H, W]
 
-        # 计算质心 & 分配 ID（nearest-centroid）
+        # 空检测帧处理：res.masks 可能为 None
+        if res.masks is None or res.masks.data is None:
+            # 为每只鼠标补充 None 并直接写入原帧
+            for mid in range(num_mice):
+                struggle_records[mid].append(None)
+            out.write(frame)
+            frame_idx += 1
+            continue
+
+        masks = res.masks.data.cpu().numpy()  # [N, H, W]
+
+        # === 4. 质心计算 & ID 分配 (nearest-centroid) ===
         curr_centroids = []
         for m in masks:
             ys, xs = np.where(m > 0)
             curr_centroids.append(
-                (int(np.mean(xs)), int(np.mean(ys))) if len(xs)>0 else None
+                (int(np.mean(xs)), int(np.mean(ys))) if xs.size else None
             )
+
         assignments = [-1] * len(curr_centroids)
         unused_ids  = set(range(num_mice))
+        # 匹配已有 ID
         for i, c in enumerate(curr_centroids):
             if c is None: continue
             best_j, best_d = None, float("inf")
@@ -78,26 +87,31 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
             if best_j is not None and best_d < (50**2):
                 assignments[i] = best_j
                 unused_ids.remove(best_j)
+        # 分配新 ID
         for i in range(len(curr_centroids)):
             if assignments[i] < 0 and unused_ids:
                 assignments[i] = unused_ids.pop()
 
-        # 计算挣扎强度 & 可视化
+        # === 5. 计算挣扎强度 & 可视化叠加 ===
         for i, m in enumerate(masks):
             mid = assignments[i]
-            if mid < 0: continue
+            if mid < 0:
+                continue
             prev_m = prev_masks[mid]
             if prev_m is None:
                 struggle_records[mid].append(None)
             else:
+                # 异或统计像素差异
                 diff = int(np.logical_xor(prev_m, m).sum())
                 struggle_records[mid].append(diff)
-            # 叠加掩膜 & ID
+
+            # 掩膜叠加
             mask_rgb = np.stack([m*255 if c==1 else 0 for c in range(3)], axis=-1).astype(np.uint8)
             frame = cv2.addWeighted(frame, 1, mask_rgb, 0.5, 0)
             if curr_centroids[i]:
                 cv2.putText(frame, f"ID:{mid}", curr_centroids[i],
                             cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
+
             prev_centroids[mid] = curr_centroids[i]
             prev_masks[mid]     = m.copy()
 
@@ -107,7 +121,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     cap.release()
     out.release()
 
-    # 汇总 & 绘图
+    # === 6. 汇总 & 绘制挣扎曲线 ===
     win = int(window_size_sec * fps)
     fig, ax = plt.subplots(figsize=(8,4))
     times = np.arange(0, frame_idx, win) / fps
@@ -118,6 +132,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
         first = next((i for i,v in enumerate(rec) if v is not None), None)
         if first is not None:
             ax.axvspan(0, first/fps, alpha=0.3, color='gray')
+
     ax.set_xlabel("Time (s)")
     ax.set_ylabel("Struggle Intensity")
     ax.set_title("Mouse Struggle Over Time")
@@ -125,7 +140,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
     return out_path, fig
 
-# Gradio 前端
+# === 7. Gradio 前端 ===
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
     gr.Markdown("上传视频，输入鼠标数量，点击 Run。")
     with gr.Row():
@@ -139,5 +154,5 @@ with gr.Blocks(title="Mice Struggle Analysis") as demo:
                   outputs=[output_video, output_plot])
 
 if __name__ == "__main__":
-    # 去除不支持的 api_config 参数
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)  # ⚠️ 删除 api_config :contentReference[oaicite:4]{index=4}
+    # 不再使用 api_config，保持默认超时
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)