Update app.py

app.py

@@ -1,29 +1,28 @@
 # filename: app.py
 
-import spaces  # 必须首 import，用于 ZeroGPU 装饰
-import os
+import spaces                # 必须最先 import，用于 ZeroGPU 装饰
 import cv2
 import numpy as np
 import torch
-from ultralytics import YOLO   # pip install ultralytics
+from ultralytics import YOLO  # pip install ultralytics
 import gradio as gr
 import matplotlib.pyplot as plt
 
 # === 1. GPU 可用性检查 & 日志 ===
 use_cuda = torch.cuda.is_available()
-print(f"CUDA available: {use_cuda}")
+print(f"CUDA available: {use_cuda}")                # 输出 CUDA 状态
 if use_cuda:
     print(f"GPU Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
-# === 2. 加载模型并显式指定 segmentation 任务 ===
-model = YOLO("fst-v1.2-n.onnx", task="segment")  # ONNX 模型需上传至 Space
+# === 2. 加载模型 & 指定分割任务 ===
+model = YOLO("fst-v1.2-n.onnx", task="segment")     # 明确 segment 任务，避免警告
 if use_cuda:
     try:
-        model.model.to("cuda")
+        model.model.to("cuda")                      # 将模型迁移到 GPU
     except:
         pass
 
-@spaces.GPU(duration=600)  # 在 ZeroGPU 上执行本函数，超时 600s
+@spaces.GPU(duration=600)  # ZeroGPU 环境下执行该函数，超时 600s
 def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     """
     分割 → 跟踪 → “挣扎强度” 计算
@@ -43,13 +42,15 @@ 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)
         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)
@@ -57,11 +58,26 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
             frame_idx += 1
             continue
 
-        masks = res.masks.data.cpu().numpy()  # [N, H, W]
+        # 原始掩膜 (N, H_model, W_model)
+        masks = res.masks.data.cpu().numpy()
 
-        # 质心计算与 ID 分配
-        curr_centroids = []
+        # 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        # 保持二值特性
+            )                                          
+            aligned_masks.append(m_resized)
+        aligned_masks = np.array(aligned_masks)       # 形状: (N, height, width) :contentReference[oaicite:3]{index=3}
+
+        # 5. 跟踪: 质心计算 & ID 分配 (nearest-centroid)
+        curr_centroids = []
+        for m in aligned_masks:
             ys, xs = np.where(m > 0)
             curr_centroids.append(
                 (int(np.mean(xs)), int(np.mean(ys))) if xs.size else None
@@ -84,8 +100,8 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
             if assignments[i] < 0 and unused_ids:
                 assignments[i] = unused_ids.pop()
 
-        # 计算挣扎强度 & 可视化叠加
-        for i, m in enumerate(masks):
+        # 6. “挣扎强度”计算 & 掩膜叠加
+        for i, m in enumerate(aligned_masks):
             mid = assignments[i]
             if mid < 0: continue
             prev_m = prev_masks[mid]
@@ -95,14 +111,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 保证三通道形状一致 ——  
+            # 关键：用 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)
+                np.zeros_like(m),       # 通道 0
+                m * 255,                # 通道 1
+                np.zeros_like(m)        # 通道 2
+            ], axis=-1).astype(np.uint8) :contentReference[oaicite:4]{index=4}
 
+            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)
@@ -116,7 +132,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
@@ -135,7 +151,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
     return out_path, fig
 
-# Gradio 前端部分保持不变
+# 8. Gradio 前端不变
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
     gr.Markdown("上传视频，输入鼠标数量，点击 Run。")
     with gr.Row():