Update app.py

app.py

@@ -1,5 +1,6 @@
 # filename: app.py
 
+import spaces  # 必须最先 import
 import os
 import cv2
 import numpy as np
@@ -7,30 +8,31 @@ import torch
 from ultralytics import YOLO   # pip install ultralytics
 import gradio as gr
 import matplotlib.pyplot as plt
-import spaces  # ZeroGPU 装饰器
 
-# === GPU 可用性检查 & 日志输出 ===
+# === 1. 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")
+# === 2. 加载模型并显式指定 task ===
+# 避免无法猜测任务类型的警告，明确使用分割 (segment) 模式
+model = YOLO("fst-v1.2-n.onnx", task="segment")  # ONNX 模型需上传至空间
+
+# 若 CUDA 可用，迁移模型至 GPU
 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}
+@spaces.GPU(duration=600)  # 调用时分配 GPU，超时 600s :contentReference[oaicite:3]{index=3}
 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"
@@ -38,7 +40,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))
 
-    # 3. 跟踪 & 记录结构
+    # 跟踪初始化
     prev_centroids   = [None] * num_mice
     prev_masks       = [None] * num_mice
     struggle_records = [[] for _ in range(num_mice)]
@@ -46,15 +48,16 @@ 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
 
-        # 4. 分割推理 (stream=True 加速)，指定 device
+        # 分割推理（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)
+        # 计算质心 & 分配 ID（nearest-centroid）
         curr_centroids = []
         for m in masks:
             ys, xs = np.where(m > 0)
@@ -76,20 +79,20 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
                 assignments[i] = best_j
                 unused_ids.remove(best_j)
         for i in range(len(curr_centroids)):
-            if assignments[i]<0 and unused_ids:
+            if assignments[i] < 0 and unused_ids:
                 assignments[i] = unused_ids.pop()
 
-        # 6. 计算挣扎强度 & 渲染
+        # 计算挣扎强度 & 可视化
         for i, m in enumerate(masks):
             mid = assignments[i]
-            if mid is None or 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 & ID
+            # 叠加掩膜 & 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]:
@@ -104,7 +107,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
@@ -113,7 +116,7 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
                 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:
+        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")
@@ -122,7 +125,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():
@@ -136,6 +139,5 @@ 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,
-                api_config={"timeout":600})  # 保持 600s 超时 :contentReference[oaicite:5]{index=5}
+    # 去除不支持的 api_config 参数
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)  # ⚠️ 删除 api_config :contentReference[oaicite:4]{index=4}