Create app.py

app.py

@@ -0,0 +1,149 @@
+# filename: app.py
+
+import cv2
+import numpy as np
+from ultralytics import YOLO  # pip install ultralytics :contentReference[oaicite:2]{index=2}
+import gradio as gr
+import matplotlib.pyplot as plt
+
+# 1. 加载已训练好的分割模型
+model = YOLO("yolo11n-seg.pt")  # 模型文件需手动上传至 Space :contentReference[oaicite:3]{index=3}
+
+def analyze_video(video_path, num_mice, window_size_sec=1, fps=30):
+    """
+    核心分析函数：对上传视频进行分割、跟踪与挣扎强度计算
+    返回：标注后的视频路径 & 挣扎强度曲线图（matplotlib Figure）
+    """
+    # 视频读取与输出配置
+    cap = cv2.VideoCapture(video_path)
+    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
+    out_path = "output.mp4"
+    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))
+
+    # 跟踪数据结构：每只鼠标保留上帧质心、掩膜
+    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
+
+        # 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 已经是二值化
+
+        # 计算当前帧每个实例的质心
+        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))
+        for i, c in enumerate(curr_centroids):
+            if c is None: 
+                continue
+            # 找到距离最近的上一帧质心
+            best_j, best_dist = None, float('inf')
+            for j in unused_prev:
+                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
+                assignments[i] = best_j
+                unused_prev.remove(best_j)
+        # 未匹配的实例新分配 ID
+        for i in range(len(curr_centroids)):
+            if assignments[i] == -1 and unused_prev:
+                assignments[i] = unused_prev.pop()
+
+        # 4. 计算“挣扎强度” & 叠加绘制
+        for i, m in enumerate(masks):
+            id_ = assignments[i]
+            if id_ is None or id_<0: 
+                continue
+            prev_m = prev_masks[id_]
+            if prev_m is None:
+                # 未检测到前，标记为 None
+                struggle_records[id_].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)
+            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()
+
+        out.write(frame)
+        frame_idx += 1
+
+    cap.release()
+    out.release()
+
+    # 5. 按时间窗口汇总并绘制
+    win_size = 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)
+
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel("Struggle Intensity")
+    ax.legend()
+    ax.set_title("Mouse Struggle Over Time")
+
+    return out_path, fig
+
+# 6. Gradio 接口
+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")
+    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])
+
+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})