Update app.py

app.py

@@ -1,6 +1,6 @@
 # filename: app.py
 
-import spaces  # 必须最先 import，用于 ZeroGPU 装饰
+import spaces  # 必须首 import，用于 ZeroGPU 装饰
 import os
 import cv2
 import numpy as np
@@ -15,20 +15,19 @@ print(f"CUDA available: {use_cuda}")
 if use_cuda:
     print(f"GPU Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-# === 2. 加载模型 (显式指定 segmentation 任务 & 默认置信度阈值) ===
-#    conf 设置为 0.25，可根据实际降低到 0.1-0.2
-model = YOLO("fst-v1.2-n.onnx", task="segment")  # 避免自动猜测任务模式的警告 
+# === 2. 加载模型并显式指定 segmentation 任务 ===
+model = YOLO("fst-v1.2-n.onnx", task="segment")  # ONNX 模型需上传至 Space
 if use_cuda:
     try:
         model.model.to("cuda")
     except:
         pass
 
-@spaces.GPU(duration=600)  # 在 ZeroGPU 上执行该函数，超时 600s :contentReference[oaicite:2]{index=2}
+@spaces.GPU(duration=600)  # 在 ZeroGPU 上执行本函数，超时 600s
 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")
@@ -37,7 +36,6 @@ 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)]
@@ -45,17 +43,14 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
     while True:
         ret, frame = cap.read()
-        if not ret:
-            break
+        if not ret: break
 
-        # === 3. 分割推理 ===
         device = "cuda" if use_cuda else "cpu"
-        results = model(frame, stream=True, device=device, conf=0.25)  # 指定置信度阈值 :contentReference[oaicite:3]{index=3}
+        results = model(frame, stream=True, device=device, conf=0.25)
         res     = next(results)
 
-        # 空检测帧处理：res.masks 可能为 None
+        # 若无分割结果，则记录 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)
@@ -64,17 +59,15 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
         masks = res.masks.data.cpu().numpy()  # [N, H, W]
 
-        # === 4. 质心计算 & ID 分配 (nearest-centroid) ===
+        # 质心计算与 ID 分配
         curr_centroids = []
         for m in masks:
             ys, xs = np.where(m > 0)
             curr_centroids.append(
                 (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")
@@ -87,27 +80,29 @@ 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)
 
-            # 掩膜叠加
-            mask_rgb = np.stack([m*255 if c==1 else 0 for c in range(3)], axis=-1).astype(np.uint8)
+            # —— 关键修改：用 zeros_like 保证三通道形状一致 ——  
+            mask_rgb = np.stack([
+                np.zeros_like(m),         # 通道 0
+                m.astype(np.uint8)*255,   # 通道 1
+                np.zeros_like(m)          # 通道 2
+            ], axis=-1).astype(np.uint8)  # 避免标量与数组形状不匹配 :contentReference[oaicite:2]{index=2}
             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)
@@ -121,7 +116,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
@@ -140,7 +135,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
     return out_path, fig
 
-# === 7. Gradio 前端 ===
+# Gradio 前端部分保持不变
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
     gr.Markdown("上传视频，输入鼠标数量，点击 Run。")
     with gr.Row():
@@ -154,5 +149,4 @@ 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)