Update app.py

app.py

@@ -1,185 +1,271 @@
-import spaces                # 必须最先 import，用于 ZeroGPU 装饰
+import os
 import cv2
 import numpy as np
-import torch
-from ultralytics import YOLO  # pip install ultralytics
 import gradio as gr
-import matplotlib.pyplot as plt
+import tempfile
+from mouse_tracker import MouseTrackerAnalyzer
 
-# 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())}")
+# 全局变量
+analyzer = None
+video_file_path = None
+model_file_path = None
+total_frames = 0
+output_path = None
 
-# 加载模型并指定分割任务
-model = YOLO("fst-v1.2-n.onnx", task="segment")
-if use_cuda:
-    try:
-        model.model.to("cuda")
-    except:
-        pass
-
-@spaces.GPU(duration=600)  # ZeroGPU 环境下执行该函数，超时 600s
-def analyze_video(video_path, num_mice, time_range, window_size_sec=1, fps=30):
-    """
-    分割 → 跟踪 → 计算挣扎强度，仅分析指定时间区间
-    返回：标注后视频 & 绘制的挣扎强度曲线 (matplotlib Figure)
-    """
-    # 打开视频并获取基本信息
+# 从视频中提取特定帧
+def extract_frame(video_path, frame_num):
+    """从视频中提取特定帧"""
+    if not video_path:
+        return None
+        
     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 = 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)
-
-    # 输出视频初始化
-    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
-    struggle_records = [[] for _ in range(num_mice)]
-    frame_idx = start_frame
-
-    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)
-        res     = next(results)
-
-        # 无检测帧处理
-        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_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_masks.append(m_res)
-        aligned_masks = np.array(aligned_masks)
-
-        # 计算质心 & ID 分配 (nearest-centroid)
-        curr_centroids = []
-        for m in aligned_masks:
-            ys, xs = np.where(m > 0)
-            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 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:
-                struggle_records[mid].append(None)
-            else:
-                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)
-
-            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)
-        frame_idx += 1
-
+    if not cap.isOpened():
+        return None
+        
+    # 设置帧位置
+    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
+    
+    # 读取帧
+    ret, frame = cap.read()
     cap.release()
-    out.release()
+    
+    if not ret:
+        return None
+    
+    # 转换为RGB格式
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    return frame_rgb
 
-    # 绘制挣扎曲线
-    win = int(window_size_sec * vid_fps)
-    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 = []
-        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_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")
-    ax.legend()
-
-    return out_path, fig
+# 选择视频文件
+def select_video(video_file):
+    global video_file_path, total_frames
+    
+    if not video_file:
+        return None, "Please select a video file", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+        
+    video_file_path = video_file
+    
+    # 获取视频总帧数
+    cap = cv2.VideoCapture(video_file_path)
+    if not cap.isOpened():
+        return None, "Cannot open video file", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+        
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    
+    # 提取第一帧
+    ret, first_frame = cap.read()
+    cap.release()
+    
+    if not ret:
+        return None, "Cannot read video frame", gr.Slider(minimum=0, maximum=0, value=0), gr.Slider(minimum=0, maximum=0, value=0)
+    
+    # 转为RGB
+    first_frame_rgb = cv2.cvtColor(first_frame, cv2.COLOR_BGR2RGB)
+        
+    # 更新帧滑块
+    start_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=0, step=1)
+    end_slider = gr.Slider(minimum=0, maximum=total_frames-1, value=total_frames-1, step=1)
+    
+    return first_frame_rgb, f"Video loaded successfully, total frames: {total_frames}", start_slider, end_slider
 
-# Gradio 前端
-with gr.Blocks(title="Mice Struggle Analysis") as demo:
-    gr.Markdown("上传视频，输入鼠标数量，选择分析时间范围，点击 Run")
-    with gr.Row():
-        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 select_model(model_file):
+    global model_file_path
+    
+    if model_file is None:
+        return "Please select a model file"
+        
+    model_file_path = model_file
+    return f"Model selected: {os.path.basename(model_file_path)}"
 
-    def enable_slider(path):
-        cap = cv2.VideoCapture(path)
-        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=duration, value=(0,duration), disabled=False)
+# 预览帧
+def preview_frame(video_file, frame_num):
+    if not video_file:
+        return None, "Please select a video first"
+        
+    # 从视频提取帧
+    frame = extract_frame(video_file, frame_num)
+    if frame is None:
+        return None, "Cannot read specified frame"
+    
+    return frame, f"Frame {frame_num}"
 
-    video_input.change(fn=enable_slider, inputs=video_input, outputs=time_range)
+# 开始分析
+def start_analysis(video, model, conf, iou, max_det, start_frame, end_frame, threshold):
+    global analyzer, output_path
+    
+    if not video or not model:
+        return None, None, "Please select a video and model file"
+        
+    if start_frame >= end_frame:
+        return None, None, "Start frame must be less than end frame"
+    
+    # 创建输出路径
+    video_name = os.path.splitext(os.path.basename(video))[0]
+    output_path = os.path.join(os.path.dirname(video), f"{video_name}_out.mp4")
+    csv_path = os.path.join(os.path.dirname(video), f"{video_name}_results.csv")
+    
+    try:
+        # 创建分析器
+        analyzer = MouseTrackerAnalyzer(
+            model_path=model, 
+            conf=conf, 
+            iou=iou, 
+            max_det=max_det, 
+            verbose=True  # 开启详细日志
+        )
+        analyzer.struggle_threshold = threshold
+        
+        # 处理视频的进度回调
+        def progress_update(progress, frame, results):
+            print(f"Processing: {progress}%, Objects detected: {len(results)}")
+            
+        print(f"Processing video: {video}")
+        print(f"Output path: {output_path}")
+        print(f"Parameters: conf={conf}, iou={iou}, max_det={max_det}, threshold={threshold}")
+        
+        # 提取视频帧数范围并分析
+        results = analyzer.process_video(
+            video_path=video, 
+            output_path=output_path, 
+            start_frame=start_frame, 
+            end_frame=end_frame,
+            callback=progress_update
+        )
+        
+        # 保存结果到CSV
+        print(f"Saving results to CSV: {csv_path}")
+        analyzer.save_results(csv_path)
+        print(f"Results saved to CSV with {len(analyzer.results)} frames of data")
+        
+        # 生成分析图表
+        print("Generating time series plot...")
+        if len(analyzer.results) == 0:
+            print("WARNING: No results available for plotting!")
+            plot_path = None
+        else:
+            plot_path = analyzer.generate_time_series_plot()
+            if plot_path and os.path.exists(plot_path):
+                print(f"Plot generated and saved to: {plot_path}, size: {os.path.getsize(plot_path)/1024:.2f}KB")
+            else:
+                print(f"Failed to generate plot or plot file does not exist!")
+                plot_path = None
+        
+        # 检查输出文件是否存在
+        if os.path.exists(output_path):
+            file_size = os.path.getsize(output_path) / (1024 * 1024)  # MB
+            print(f"Output video size: {file_size:.2f}MB")
+            
+            # 处理debug帧
+            debug_frame_path = os.path.join(os.path.dirname(output_path), "debug_frame.jpg")
+            if os.path.exists(debug_frame_path):
+                print(f"Debug frame saved at: {debug_frame_path}")
+        
+        if plot_path and os.path.exists(plot_path):
+            print(f"Plot file exists at: {plot_path}, size: {os.path.getsize(plot_path)/1024:.2f}KB")
+            
+        # 确保返回正确的文件路径
+        status_message = "Analysis complete. "
+        
+        if os.path.exists(output_path):
+            status_message += f"Video saved."
+        else:
+            status_message += "WARNING: Output video not found. "
+            
+        if plot_path and os.path.exists(plot_path):
+            status_message += f" Time series plot generated."
+        else:
+            status_message += " WARNING: Failed to generate time series plot."
+            
+        status_message += f" Results saved to: {csv_path}"
+        
+        return output_path, plot_path, status_message
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        return None, None, f"Processing error: {str(e)}"
 
-    run_button   = gr.Button("Run")
-    output_video = gr.Video(label="Annotated Video")
-    output_plot  = gr.Plot(label="Struggle Plot")
-    run_button.click(fn=analyze_video,
-                     inputs=[video_input, num_input, time_range],
-                     outputs=[output_video, output_plot])
+# 创建Gradio界面
+def create_interface():
+    with gr.Blocks(title="Mouse Struggle Analysis - Object Tracking") as app:
+        gr.Markdown("# Mouse Forced Swim Test Struggle Analysis (Object Tracking)")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # 视频和模型选择
+                video_input = gr.Video(label="Input Video")
+                model_input = gr.File(label="Model File (.pt format recommended)")
+                
+                # 参数设置
+                with gr.Row():
+                    conf = gr.Slider(minimum=0.1, maximum=0.9, value=0.25, step=0.05, label="Confidence Threshold")
+                    iou = gr.Slider(minimum=0.1, maximum=0.9, value=0.45, step=0.05, label="IoU Threshold")
+                
+                with gr.Row():
+                    max_det = gr.Slider(minimum=1, maximum=50, value=20, step=1, label="Max Detections")
+                    threshold = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.01, label="Struggle Threshold")
+                
+                # 帧选择
+                start_frame = gr.Slider(minimum=0, maximum=999999, value=0, step=1, label="Start Frame")
+                end_frame = gr.Slider(minimum=0, maximum=999999, value=999999, step=1, label="End Frame")
+                
+                # 预览按钮
+                preview_btn = gr.Button("Preview Frame")
+                
+                # 开始分析
+                start_btn = gr.Button("Start Analysis", variant="primary")
+                
+            with gr.Column(scale=2):
+                # 显示区域
+                with gr.Tab("Preview"):
+                    # 图像预览
+                    preview_image = gr.Image(label="Preview Image", type="numpy", height=400)
+                    status_text = gr.Textbox(label="Status", interactive=False)
+                    gr.Markdown("""
+                    ### Instructions:
+                    1. Select a video and model file (.pt format segmentation model like yolov8n-seg.pt recommended)
+                    2. Adjust parameters
+                       - Confidence Threshold: Minimum confidence for object detection, lower values detect more potential objects
+                       - IoU Threshold: For filtering overlapping detections
+                       - Max Detections: Maximum number of objects to detect per frame
+                       - Struggle Threshold: Minimum score to classify as struggle state
+                    3. Set frame range
+                    4. Click "Start Analysis" button
+                    
+                    The system will automatically track mice and analyze their struggle behavior, no need to manually define regions
+                    """)
+                
+                with gr.Tab("Results"):
+                    with gr.Row():
+                        output_video = gr.Video(label="Analysis Result Video")
+                        result_plot = gr.Image(label="Struggle Score Time Series")
+                    
+                    result_status = gr.Textbox(label="Analysis Status", interactive=False)
+        
+        # 绑定事件
+        video_input.change(select_video, inputs=[video_input], outputs=[preview_image, status_text, start_frame, end_frame])
+        model_input.change(select_model, inputs=[model_input], outputs=[status_text])
+        
+        preview_btn.click(preview_frame, inputs=[video_input, start_frame], outputs=[preview_image, status_text])
+        
+        start_btn.click(
+            start_analysis, 
+            inputs=[video_input, model_input, conf, iou, max_det, start_frame, end_frame, threshold], 
+            outputs=[output_video, result_plot, result_status]
+        )
+    
+    return app
 
+# 启动应用
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    # 清除可能干扰的代理设置
+    if 'http_proxy' in os.environ:
+        del os.environ['http_proxy']
+    if 'https_proxy' in os.environ:
+        del os.environ['https_proxy']
+    if 'all_proxy' in os.environ:
+        del os.environ['all_proxy']
+    
+    app = create_interface()
+    # 使用简化的启动配置
+    app.launch(server_name="127.0.0.1", server_port=7860, share=False)