Update app.py

app.py

@@ -1,20 +1,36 @@
 # filename: app.py
 
+import os
 import cv2
 import numpy as np
-from ultralytics import YOLO  # pip install ultralytics :contentReference[oaicite:2]{index=2}
+import torch
+from ultralytics import YOLO   # pip install ultralytics
 import gradio as gr
 import matplotlib.pyplot as plt
-
-# 1. 加载已训练好的分割模型
-model = YOLO("fst-v1.2-n.onnx")  # 模型文件需手动上传至 Space :contentReference[oaicite:3]{index=3}
-
+import spaces  # ZeroGPU 装饰器
+
+# === GPU 可用性检查 & 日志输出 ===
+use_cuda = torch.cuda.is_available()
+print(f"CUDA available: {use_cuda}")
+if use_cuda:
+    print(f"GPU Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
+# 1. 加载并迁移模型到 GPU/CPU
+model = YOLO("fst-v1.2-n.onnx")
+if use_cuda:
+    # ONNX 推理时可通过 .model（PyTorch）迁移，或在推理调用时指定 device="cuda"
+    try:
+        model.model.to("cuda")
+    except Exception:
+        pass
+
+@spaces.GPU(duration=600)  # 调用时分配 GPU，600s 超时 :contentReference[oaicite:4]{index=4}
 def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     """
-    核心分析函数：对上传视频进行分割、跟踪与挣扎强度计算
-    返回：标注后的视频路径 & 挣扎强度曲线图（matplotlib Figure）
+    对视频进行分割 → 跟踪 → '挣扎强度' 计算  
+    返回：标注后的视频路径 & 挣扎强度曲线 (matplotlib Figure)
     """
-    # 视频读取与输出配置
+    # 2. 视频读取 & 输出初始化
     cap = cv2.VideoCapture(video_path)
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")
     out_path = "output.mp4"
@@ -22,81 +38,65 @@ 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
-    # 时间序列数据：每只鼠标每帧的“挣扎程度”
+    # 3. 跟踪 & 记录结构
+    prev_centroids   = [None] * num_mice
+    prev_masks       = [None] * num_mice
     struggle_records = [[] for _ in range(num_mice)]
-
     frame_idx = 0
+
     while True:
         ret, frame = cap.read()
-        if not ret:
-            break
+        if not ret: break
 
-        # 2. 分割推理（stream=True 可加速）：
-        results = model(frame, stream=True, device='cpu')
-        # 取第一张结果
-        res = next(results)
-        masks = res.masks.data.cpu().numpy()  # shape: [N, H, W]
-        # 只保留 tag="mice" 的结果（假设模型只检测 mice 类）
-        # masks 已经是二值化
+        # 4. 分割推理 (stream=True 加速)，指定 device
+        device = "cuda" if use_cuda else "cpu"
+        results = model(frame, stream=True, device=device)
+        res     = next(results)
+        masks   = res.masks.data.cpu().numpy()  # [N, H, W]
 
-        # 计算当前帧每个实例的质心
+        # 5. 计算质心 & ID 分配 (nearest-centroid)
         curr_centroids = []
         for m in masks:
             ys, xs = np.where(m > 0)
-            if len(xs)==0:
-                curr_centroids.append(None)
-            else:
-                curr_centroids.append((int(np.mean(xs)), int(np.mean(ys))))
-
-        # 3. 质心匹配分配 ID
-        assignments = [-1]*len(curr_centroids)
-        unused_prev = set(range(num_mice))
+            curr_centroids.append(
+                (int(np.mean(xs)), int(np.mean(ys))) if len(xs)>0 else None
+            )
+        assignments = [-1] * len(curr_centroids)
+        unused_ids  = set(range(num_mice))
         for i, c in enumerate(curr_centroids):
-            if c is None: 
-                continue
-            # 找到距离最近的上一帧质心
-            best_j, best_dist = None, float('inf')
-            for j in unused_prev:
+            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 d < best_dist:
-                    best_j, best_dist = j, d
-            if best_j is not None and best_dist < (50**2):  # 距离阈值 50
+                if d < best_d:
+                    best_j, best_d = j, d
+            if best_j is not None and best_d < (50**2):
                 assignments[i] = best_j
-                unused_prev.remove(best_j)
-        # 未匹配的实例新分配 ID
+                unused_ids.remove(best_j)
         for i in range(len(curr_centroids)):
-            if assignments[i] == -1 and unused_prev:
-                assignments[i] = unused_prev.pop()
+            if assignments[i]<0 and unused_ids:
+                assignments[i] = unused_ids.pop()
 
-        # 4. 计算“挣扎强度” & 叠加绘制
+        # 6. 计算挣扎强度 & 渲染
         for i, m in enumerate(masks):
-            id_ = assignments[i]
-            if id_ is None or id_<0: 
-                continue
-            prev_m = prev_masks[id_]
+            mid = assignments[i]
+            if mid is None or mid<0: continue
+            prev_m = prev_masks[mid]
             if prev_m is None:
-                # 未检测到前，标记为 None
-                struggle_records[id_].append(None)
+                struggle_records[mid].append(None)
             else:
-                # XOR 统计像素差异
-                diff = np.logical_xor(prev_m, m).sum()
-                struggle_records[id_].append(int(diff))
-            # 叠加掩膜 & ID
-            color = (0,255,0)
-            mask_rgb = np.stack([m*color[c] for c in range(3)], axis=-1).astype(np.uint8)
-            frame = cv2.addWeighted(frame,1,mask_rgb,0.5,0)
+                diff = int(np.logical_xor(prev_m, m).sum())
+                struggle_records[mid].append(diff)
+            # 可视化 mask & 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:{id_}", curr_centroids[i],
-                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
-
-            # 更新上一帧数据
-            prev_centroids[id_] = curr_centroids[i]
-            prev_masks[id_]     = m.copy()
+                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()
 
         out.write(frame)
         frame_idx += 1
@@ -104,38 +104,31 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     cap.release()
     out.release()
 
-    # 5. 按时间窗口汇总并绘制
-    win_size = int(window_size_sec * fps)
+    # 7. 汇总 & 绘制挣扎曲线
+    win = int(window_size_sec * fps)
     fig, ax = plt.subplots(figsize=(8,4))
-    times = np.arange(0, frame_idx, win_size) / fps
-    for id_, records in enumerate(struggle_records):
-        # 将记录按窗口求和，None视为 0 或保持空白
-        sums = []
-        for w in range(len(times)):
-            segment = records[w*win_size:(w+1)*win_size]
-            # 把 None 当作 0，但在绘图时保留空白
-            vals = [v if v is not None else 0 for v in segment]
-            sums.append(sum(vals))
-        ax.plot(times, sums, label=f"Mouse {id_}")
-        # 标记 None 区间
-        first_detect = next((i for i,v in enumerate(records) if v is not None), None)
-        if first_detect:
-            ax.axvspan(0, first_detect/fps, color='grey', alpha=0.3)
-
+    times = np.arange(0, frame_idx, win) / fps
+    for mid, rec in enumerate(struggle_records):
+        sums = [sum(v if v is not None else 0 for v in rec[i*win:(i+1)*win])
+                for i in range(len(times))]
+        ax.plot(times, sums, label=f"Mouse {mid}")
+        first = next((i for i,v in enumerate(rec) if v is not None), None)
+        if first:
+            ax.axvspan(0, first/fps, alpha=0.3, color='gray')
     ax.set_xlabel("Time (s)")
     ax.set_ylabel("Struggle Intensity")
-    ax.legend()
     ax.set_title("Mouse Struggle Over Time")
+    ax.legend()
 
     return out_path, fig
 
-# 6. Gradio 接口
+# 8. Gradio 前端
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
-    gr.Markdown("上传实验视频，输入鼠标数量，点击 Run 开始分析。")
+    gr.Markdown("上传视频，输入鼠标数量，点击 Run。")
     with gr.Row():
         video_in = gr.Video(label="Input Video")
         num_in   = gr.Number(value=1, precision=0, label="Number of Mice")
-    run_btn = gr.Button("Run")
+    run_btn      = gr.Button("Run")
     output_video = gr.Video(label="Annotated Video")
     output_plot  = gr.Plot(label="Struggle Plot")
     run_btn.click(fn=analyze_video,
@@ -143,7 +136,6 @@ with gr.Blocks(title="Mice Struggle Analysis") as demo:
                   outputs=[output_video, output_plot])
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False, 
-                inbrowser=False,  
-                # Zero GPU 环境下设置 600s 超时
-                api_config={"timeout":600})
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False,
+                inbrowser=False,
+                api_config={"timeout":600})  # 保持 600s 超时 :contentReference[oaicite:5]{index=5}