Update app.py

app.py

@@ -7,36 +7,36 @@ import gradio as gr
 import matplotlib.pyplot as plt
 
 # GPU 可用性检查 & 日志
-evice_is_cuda = torch.cuda.is_available()
+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())}")
 
 # 加载模型并指定分割任务
-yolo_model = YOLO("fst-v1.2-n.onnx", task="segment")
+model = YOLO("fst-v1.2-n.onnx", task="segment")
 if use_cuda:
     try:
-        yolo_model.model.to("cuda")
+        model.model.to("cuda")
     except:
         pass
 
-@spaces.GPU(duration=600)  # 在 ZeroGPU 上运行，超时 600s
+@spaces.GPU(duration=600)  # ZeroGPU 环境下执行该函数，超时 600s
 def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
     """
-    分割 → 跟踪 → 计算“挣扎强度”，只分析指定时间范围
-    返回：标注后视频 & 挣扎曲线 Figure
+    分割 → 跟踪 → 计算挣扎强度，仅分析指定时间区间
+    返回：标注后视频 & 绘制的挣扎强度曲线 (matplotlib Figure)
     """
-    # 打开视频，获取基本信息
+    # 打开视频并获取基本信息
     cap = cv2.VideoCapture(video_path)
     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))
     start_s, end_s = time_range
-    start_frame = int(start_s * vid_fps)
-    end_frame   = int(end_s * vid_fps)
+    start_frame = min(int(start_s * vid_fps), total_frames)
+    end_frame   = min(int(end_s * vid_fps), total_frames)
 
-    # 跳转到起始帧
+    # 跳转到指定起始帧
     cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame)
 
     # 输出视频初始化
@@ -44,8 +44,8 @@ def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
     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 = start_frame
 
@@ -54,71 +54,79 @@ def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
         if not ret:
             break
 
-        # 分割推理\        device = "cuda" if use_cuda else "cpu"
-        results = yolo_model(frame, stream=True, device=device, conf=0.25)
+        # 分割推理
+        device = "cuda" if use_cuda else "cpu"
+        results = model(frame, stream=True, device=device, conf=0.25)
         res     = next(results)
 
-        # 处理无检测帧\        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
 
+        # 获取并对齐掩膜至帧尺寸
         masks = res.masks.data.cpu().numpy()  # (N, H_model, W_model)
-
-        # 对齐掩膜至帧尺寸
-        aligned = []
+        aligned_masks = []
         for m in 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)
+            aligned_masks.append(m_res)
+        aligned_masks = np.array(aligned_masks)
 
-        # 计算质心 & 分配 ID
-        curr_cent = []
-        for m in aligned:
+        # 计算质心 & ID 分配 (nearest-centroid)
+        curr_centroids = []
+        for m in aligned_masks:
             ys, xs = np.where(m > 0)
-            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:
+            curr_centroids.append((int(xs.mean()), int(ys.mean())) 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:
                 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:
-                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:
+                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:
                 struggle_records[mid].append(None)
             else:
-                diff = int(np.logical_xor(pm, m).sum())
-                struggle_records[mid].append(diff)
+                struggle = int(np.logical_xor(prev_m, m).sum())
+                struggle_records[mid].append(struggle)
 
+            # 构建三通道掩膜
             mask_rgb = np.stack([
                 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)
 
-            prev_centroids[mid] = curr_cent[i]
+            centroid = curr_centroids[i]
+            if centroid:
+                cv2.putText(frame, f"ID:{mid}", centroid,
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
+
+            prev_centroids[mid] = curr_centroids[i]
             prev_masks[mid]     = m.copy()
 
         out.write(frame)
@@ -132,11 +140,15 @@ def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
     fig, ax = plt.subplots(figsize=(8,4))
     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 = []
+        for i in range(len(times)):
+            segment = rec[i*win:(i+1)*win]
+            sums.append(sum(v if v is not None else 0 for v in segment))
         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(start_s, start_s+first/vid_fps, alpha=0.3, color='gray')
+        first_valid = next((i for i,v in enumerate(rec) if v is not None), None)
+        if first_valid is not None:
+            ax.axvspan(start_s, start_s+first_valid/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,25 +160,26 @@ def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
 with gr.Blocks(title="Mice Struggle Analysis") as demo:
     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")
-        time_range = gr.RangeSlider(label="Analysis Time Range (s)", minimum=0, maximum=1, value=(0,1), step=1, disabled=True)
+        video_input = gr.Video(label="Input Video")
+        num_input   = 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):
+    def enable_slider(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)
+        vid_fps = cap.get(cv2.CAP_PROP_FPS) or fps
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        duration = total_frames / vid_fps
         cap.release()
-        return gr.update(maximum=dur, value=(0, dur), disabled=False)
+        return gr.update(maximum=duration, value=(0,duration), disabled=False)
 
-    video_in.change(fn=get_video_duration, inputs=video_in, outputs=time_range)
+    video_input.change(fn=enable_slider, inputs=video_input, outputs=time_range)
 
-    run_btn      = gr.Button("Run")
+    run_button   = 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, time_range], outputs=[output_video, output_plot])
+    run_button.click(fn=analyze_video,
+                     inputs=[video_input, num_input, time_range],
+                     outputs=[output_video, output_plot])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False)