Update app.py

app.py

@@ -1,5 +1,3 @@
-# filename: app.py
-
 import spaces                # 必须最先 import，用于 ZeroGPU 装饰
 import cv2
 import numpy as np
@@ -8,117 +6,119 @@ from ultralytics import YOLO  # pip install ultralytics
 import gradio as gr
 import matplotlib.pyplot as plt
 
-# 1. GPU 可用性检查 & 日志
-use_cuda = torch.cuda.is_available()
+# GPU 可用性检查 & 日志
+evice_is_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())}")
 
-# 2. 加载模型并指定分割任务
-model = YOLO("fst-v1.2-n.onnx", task="segment")
+# 加载模型并指定分割任务
+yolo_model = YOLO("fst-v1.2-n.onnx", task="segment")
 if use_cuda:
     try:
-        model.model.to("cuda")
+        yolo_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):
+@spaces.GPU(duration=600)  # 在 ZeroGPU 上运行，超时 600s
+def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
+    """
+    分割 → 跟踪 → 计算“挣扎强度”，只分析指定时间范围
+    返回：标注后视频 & 挣扎曲线 Figure
+    """
+    # 打开视频，获取基本信息
     cap = cv2.VideoCapture(video_path)
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    out_path = "output.mp4"
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    vid_fps = cap.get(cv2.CAP_PROP_FPS) or fps
     width  = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    out = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
+    start_s, end_s = time_range
+    start_frame = int(start_s * vid_fps)
+    end_frame   = int(end_s * vid_fps)
+
+    # 跳转到起始帧
+    cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
+
+    # 输出视频初始化
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out_path = "output.mp4"
+    out = cv2.VideoWriter(out_path, fourcc, vid_fps, (width, height))
 
-    prev_centroids   = [None] * num_mice
-    prev_masks       = [None] * num_mice
+    prev_centroids = [None] * num_mice
+    prev_masks     = [None] * num_mice
     struggle_records = [[] for _ in range(num_mice)]
-    frame_idx = 0
+    frame_idx = start_frame
 
-    while True:
+    while frame_idx <= end_frame:
         ret, frame = cap.read()
         if not ret:
             break
 
-        # 分割推理
-        device = "cuda" if use_cuda else "cpu"
-        results = model(frame, stream=True, device=device, conf=0.25)
+        # 分割推理\        device = "cuda" if use_cuda else "cpu"
+        results = yolo_model(frame, stream=True, device=device, conf=0.25)
         res     = next(results)
 
-        # 无检测时直接写入并记录 None
-        if res.masks is None or res.masks.data is None:
+        # 处理无检测帧\        if res.masks is None or res.masks.data is None:
             for mid in range(num_mice):
                 struggle_records[mid].append(None)
             out.write(frame)
             frame_idx += 1
             continue
 
-        # 原始掩膜 (N, H_model, W_model)
-        masks = res.masks.data.cpu().numpy()
+        masks = res.masks.data.cpu().numpy()  # (N, H_model, W_model)
 
-        # 对齐到视频帧尺寸
-        aligned_masks = []
+        # 对齐掩膜至帧尺寸
+        aligned = []
         for m in masks:
-            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)
-
-        # 质心计算 & 分配 ID
-        curr_centroids = []
-        for m in aligned_masks:
+            m_bin = (m > 0).astype(np.uint8)
+            m_res = cv2.resize(m_bin, (width, height), interpolation=cv2.INTER_NEAREST)
+            aligned.append(m_res)
+        aligned = np.array(aligned)
+
+        # 计算质心 & 分配 ID
+        curr_cent = []
+        for m in aligned:
             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))
-
-        # 最近质心匹配
-        for i, c in enumerate(curr_centroids):
-            if c is None:
-                continue
-            best_j, best_d = None, float("inf")
-            for j in unused_ids:
+            curr_cent.append((int(xs.mean()), int(ys.mean())) if xs.size else None)
+        assign = [-1] * len(curr_cent)
+        unused = set(range(num_mice))
+        for i, c in enumerate(curr_cent):
+            if c is None: continue
+            best_j, best_d = None, float('inf')
+            for j in unused:
                 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:
-                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()
-
-        # 计算挣扎强度 & 可视化叠加
-        for i, m in enumerate(aligned_masks):
-            mid = assignments[i]
-            if mid < 0:
-                continue
-            prev_m = prev_masks[mid]
-            if prev_m is None:
+                assign[i] = best_j
+                unused.remove(best_j)
+        for i in range(len(curr_cent)):
+            if assign[i] < 0 and unused:
+                assign[i] = unused.pop()
+
+        # 计算挣扎强度 & 叠加
+        for i, m in enumerate(aligned):
+            mid = assign[i]
+            if mid < 0: continue
+            pm = prev_masks[mid]
+            if pm is None:
                 struggle_records[mid].append(None)
             else:
-                diff = int(np.logical_xor(prev_m, m).sum())
+                diff = int(np.logical_xor(pm, m).sum())
                 struggle_records[mid].append(diff)
 
-            # 构建三通道掩膜，保证与 frame 形状一致
             mask_rgb = np.stack([
-                np.zeros_like(m),      # 通道 0
-                m * 255,               # 通道 1
-                np.zeros_like(m)       # 通道 2
+                np.zeros_like(m),
+                m * 255,
+                np.zeros_like(m)
             ], axis=-1).astype(np.uint8)
             frame = cv2.addWeighted(frame, 1, mask_rgb, 0.5, 0)
+            if curr_cent[i]:
+                cv2.putText(frame, f"ID:{mid}", curr_cent[i], cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
 
-            if curr_centroids[i]:
-                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_centroids[mid] = curr_cent[i]
             prev_masks[mid]     = m.copy()
 
         out.write(frame)
@@ -127,18 +127,16 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
     cap.release()
     out.release()
 
-    # 汇总 & 绘制挣扎曲线
-    win = int(window_size_sec * fps)
+    # 绘制挣扎曲线
+    win = int(window_size_sec * vid_fps)
     fig, ax = plt.subplots(figsize=(8,4))
-    times = np.arange(0, frame_idx, win) / fps
+    times = np.arange(start_s, end_s, win/vid_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))]
+        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 is not None:
-            ax.axvspan(0, first/fps, alpha=0.3, color='gray')
-
+            ax.axvspan(start_s, start_s+first/vid_fps, alpha=0.3, color='gray')
     ax.set_xlabel("Time (s)")
     ax.set_ylabel("Struggle Intensity")
     ax.set_title("Mouse Struggle Over Time")
@@ -148,16 +146,27 @@ def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
 
 # 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")
+        video_in   = gr.Video(label="Input Video")
+        num_in     = gr.Number(value=1, precision=0, label="Number of Mice")
+        time_range = gr.RangeSlider(label="Analysis Time Range (s)", minimum=0, maximum=1, value=(0,1), step=1, disabled=True)
+
+    # 上传视频后激活滑块并设置最大值
+    def get_video_duration(path):
+        cap = cv2.VideoCapture(path)
+        fps = cap.get(cv2.CAP_PROP_FPS) or fps
+        frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        dur = int(frames / fps)
+        cap.release()
+        return gr.update(maximum=dur, value=(0, dur), disabled=False)
+
+    video_in.change(fn=get_video_duration, inputs=video_in, outputs=time_range)
+
     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,
-                  inputs=[video_in, num_in],
-                  outputs=[output_video, output_plot])
+    run_btn.click(fn=analyze_video, inputs=[video_in, num_in, time_range], outputs=[output_video, output_plot])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)