Update app.py

app.py

@@ -1,3 +1,6 @@
+下面是去除了多余注释标记、保证语法正确的关键代码片段及完整 `app.py`。我已删掉所有 `:contentReference[…]` 等非 Python 语法内容，并确保 `mask_rgb` 构造部分形状一致。
+
+```python
 # filename: app.py
 
 import spaces                # 必须最先 import，用于 ZeroGPU 装饰
@@ -8,26 +11,22 @@ from ultralytics import YOLO  # pip install ultralytics
 import gradio as gr
 import matplotlib.pyplot as plt
 
-# === 1. GPU 可用性检查 & 日志 ===
+# 1. GPU 可用性检查 & 日志
 use_cuda = torch.cuda.is_available()
-print(f"CUDA available: {use_cuda}")                # 输出 CUDA 状态
+print(f"CUDA available: {use_cuda}")
 if use_cuda:
     print(f"GPU Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-# === 2. 加载模型 & 指定分割任务 ===
-model = YOLO("fst-v1.2-n.onnx", task="segment")     # 明确 segment 任务，避免警告
+# 2. 加载模型并指定分割任务
+model = YOLO("fst-v1.2-n.onnx", task="segment")
 if use_cuda:
     try:
-        model.model.to("cuda")                      # 将模型迁移到 GPU
+        model.model.to("cuda")
     except:
         pass
 
 @spaces.GPU(duration=600)  # ZeroGPU 环境下执行该函数，超时 600s
 def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
-    """
-    分割 → 跟踪 → “挣扎强度” 计算
-    返回：标注后视频路径 & matplotlib Figure
-    """
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     out_path = "output.mp4"
@@ -45,12 +44,12 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
         if not ret:
             break
 
-        # 3. 分割推理
+        # 分割推理
         device = "cuda" if use_cuda else "cpu"
         results = model(frame, stream=True, device=device, conf=0.25)
         res     = next(results)
 
-        # 没有检测到掩膜时，全部记录 None 并写入原帧
+        # 无检测时直接写入并记录 None
         if res.masks is None or res.masks.data is None:
             for mid in range(num_mice):
                 struggle_records[mid].append(None)
@@ -61,21 +60,15 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
         # 原始掩膜 (N, H_model, W_model)
         masks = res.masks.data.cpu().numpy()
 
-        # 4. 对齐掩膜到原视频帧尺寸
+        # 对齐到视频帧尺寸
         aligned_masks = []
         for m in masks:
-            # 二值化掩膜 → uint8
-            m_uint8 = (m > 0).astype(np.uint8)        # 0/1
-            # 重采样到视频帧大小
-            m_resized = cv2.resize(
-                m_uint8,
-                (width, height),
-                interpolation=cv2.INTER_NEAREST        # 保持二值特性
-            )                                          
+            m_uint8 = (m > 0).astype(np.uint8)
+            m_resized = cv2.resize(m_uint8, (width, height), interpolation=cv2.INTER_NEAREST)
             aligned_masks.append(m_resized)
-        aligned_masks = np.array(aligned_masks)       # 形状: (N, height, width) :contentReference[oaicite:3]{index=3}
+        aligned_masks = np.array(aligned_masks)
 
-        # 5. 跟踪: 质心计算 & ID 分配 (nearest-centroid)
+        # 质心计算 & 分配 ID
         curr_centroids = []
         for m in aligned_masks:
             ys, xs = np.where(m > 0)
@@ -84,26 +77,31 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
             )
         assignments = [-1] * len(curr_centroids)
         unused_ids  = set(range(num_mice))
+
+        # 最近质心匹配
         for i, c in enumerate(curr_centroids):
-            if c is None: continue
+            if c is None:
+                continue
             best_j, best_d = None, float("inf")
             for j in unused_ids:
                 pc = prev_centroids[j]
-                if pc is None: continue
-                d = (c[0]-pc[0])**2 + (c[1]-pc[1])**2
+                if pc is None:
+                    continue
+                d = (c[0] - pc[0])**2 + (c[1] - pc[1])**2
                 if d < best_d:
                     best_j, best_d = j, d
-            if best_j is not None and best_d < (50**2):
+            if best_j is not None and best_d < 50**2:
                 assignments[i] = best_j
                 unused_ids.remove(best_j)
         for i in range(len(curr_centroids)):
             if assignments[i] < 0 and unused_ids:
                 assignments[i] = unused_ids.pop()
 
-        # 6. “挣扎强度”计算 & 掩膜叠加
+        # 计算挣扎强度 & 可视化叠加
         for i, m in enumerate(aligned_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)
@@ -111,14 +109,14 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
                 diff = int(np.logical_xor(prev_m, m).sum())
                 struggle_records[mid].append(diff)
 
-            # 关键：用 zeros_like 保证三通道形状一致
+            # 构建三通道掩膜，保证与 frame 形状一致
             mask_rgb = np.stack([
-                np.zeros_like(m),       # 通道 0
-                m * 255,                # 通道 1
-                np.zeros_like(m)        # 通道 2
-            ], axis=-1).astype(np.uint8) :contentReference[oaicite:4]{index=4}
-
+                np.zeros_like(m),      # 通道 0
+                m * 255,               # 通道 1
+                np.zeros_like(m)       # 通道 2
+            ], 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)
@@ -132,7 +130,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     cap.release()
     out.release()
 
-    # 7. 挣扎曲线汇总 & 绘制
+    # 汇总 & 绘制挣扎曲线
     win = int(window_size_sec * fps)
     fig, ax = plt.subplots(figsize=(8,4))
     times = np.arange(0, frame_idx, win) / fps
@@ -151,7 +149,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
     return out_path, fig
 
-# 8. Gradio 前端不变
+# Gradio 前端
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
     gr.Markdown("上传视频，输入鼠标数量，点击 Run。")
     with gr.Row():
@@ -166,3 +164,9 @@ with gr.Blocks(title="Mice Struggle Analysis") as demo:
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+```
+
+**说明：**  
+- 已彻底移除所有非 Python 语法的注释标记。  
+- `mask_rgb` 构造时三通道数组均为与 `m` 同形状，确保 `cv2.addWeighted` 能正常运行。  
+- 其他逻辑与之前保持一致，部署即可通过。