Update app.py

app.py

@@ -1,48 +1,54 @@
 import gradio as gr
 from ultralytics import YOLO
 from fastapi import FastAPI
-import cv2
+from PIL import Image
 import torch
 import spaces
 import numpy as np
+import cv2
 from pathlib import Path
+import tempfile
 
 app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = YOLO('kunin-mice-pose.v0.1.0.pt')  # 使用你的小鼠检测模型
+model = YOLO('kunin-mice-pose.v0.1.0.pt')
 
 @spaces.GPU
-def process_video(video_path, process_seconds=20):
-    # 创建临时输出路径
-    output_path = Path("temp_output.mp4")
+def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8):
+    """
+    处理视频并进行小鼠检测
+    Args:
+        video_path: 输入视频路径
+        process_seconds: 处理时长(秒)
+        conf_threshold: 置信度阈值(0-1)
+        max_det: 每帧最大检测数量
+    """
+    # 创建临时目录保存输出视频
+    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
+        output_path = tmp_file.name
     
     # 获取视频信息
     cap = cv2.VideoCapture(video_path)
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    total_frames = int(process_seconds * fps)
+    total_frames = int(process_seconds * fps) if process_seconds else int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
     
     # 创建视频写入器
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     video_writer = cv2.VideoWriter(
-        str(output_path), 
+        output_path, 
         fourcc, 
         fps, 
         (width, height)
     )
     
-    # 统计信息
-    frame_count = 0
-    total_detections = 0
-    max_mice = 0
-    detection_stats = []
-    
-    # 处理视频
+    # 设置推理参数并处理视频
     results = model.predict(
         source=video_path,
         device=device,
-        conf=0.5,
+        conf=conf_threshold,  # 使用用户设置的置信度阈值
         save=False,
         show=False,
         stream=True,
@@ -51,60 +57,109 @@ def process_video(video_path, process_seconds=20):
         show_labels=True,
         show_conf=True,
         vid_stride=1,
+        max_det=max_det,  # 使用用户设置的最大检测数量
     )
     
+    # 处理结果
+    frame_count = 0
+    detection_info = []
+    
     for r in results:
-        # 获取当前帧的检测结果
+        # 获取绘制了预测结果的帧
         frame = r.plot()
-        num_mice = len(r.boxes)
-        max_mice = max(max_mice, num_mice)
-        total_detections += num_mice
-        detection_stats.append(num_mice)
+        
+        # 收集检测信息
+        frame_info = {
+            "frame": frame_count + 1,
+            "count": len(r.boxes),
+            "detections": []
+        }
+        
+        for box in r.boxes:
+            conf = float(box.conf[0])
+            cls = int(box.cls[0])
+            cls_name = r.names[cls]
+            frame_info["detections"].append({
+                "class": cls_name,
+                "confidence": f"{conf:.2%}"
+            })
+        
+        detection_info.append(frame_info)
         
         # 写入视频
         video_writer.write(frame)
         
         frame_count += 1
-        if frame_count >= total_frames:
+        if process_seconds and frame_count >= total_frames:
             break
     
-    # 释放资源
+    # 释放视频写入器
     video_writer.release()
-    cap.release()
     
-    # 生成统计信息
-    avg_mice = total_detections / frame_count if frame_count > 0 else 0
-    output_text = f"""小鼠检测统计:
+    # 生成分析报告
+    report = f"""视频分析报告:
+参数设置:
+- 置信度阈值: {conf_threshold:.2f}
+- 最大检测数量: {max_det}
 - 处理时长: {process_seconds}秒
-- 总帧数: {frame_count}
-- 最大检测数量: {max_mice}只
-- 平均检测数量: {avg_mice:.1f}只
-- 检测帧率: {fps} FPS"""
-    
-    # 如果有检测到的帧
-    if detection_stats:
-        output_text += "\n\n帧检测分布:"
-        for count in range(max(detection_stats) + 1):
-            frames = detection_stats.count(count)
-            percentage = frames / frame_count * 100
-            output_text += f"\n{count}只小鼠: {frames}帧 ({percentage:.1f}%)"
+
+分析结果:
+- 处理帧数: {frame_count}
+- 平均每帧检测到的老鼠数: {np.mean([info['count'] for info in detection_info]):.1f}
+- 最大检测数: {max([info['count'] for info in detection_info])}
+- 最小检测数: {min([info['count'] for info in detection_info])}
+
+置信度分布:
+{np.histogram([float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']], bins=5)[0].tolist()}
+"""
     
-    return str(output_path), output_text
+    return output_path, report
 
 demo = gr.Interface(
     fn=process_video,
     inputs=[
         gr.Video(label="输入视频"),
-        gr.Slider(minimum=1, maximum=60, value=20, step=1, label="处理时长(秒)")
+        gr.Number(label="处理时长（秒，0表示处理整个视频）", value=20),
+        gr.Slider(
+            minimum=0.1,
+            maximum=1.0,
+            value=0.2,
+            step=0.05,
+            label="置信度阈值",
+            info="越高越严格，建议范围0.2-0.5"
+        ),
+        gr.Slider(
+            minimum=1,
+            maximum=10,
+            value=8,
+            step=1,
+            label="最大检测数量",
+            info="每帧最多检测的目标数量"
+        )
     ],
     outputs=[
         gr.Video(label="检测结果"),
-        gr.Textbox(label="检测统计")
+        gr.Textbox(label="分析报告")
     ],
-    title="🐁 小鼠行为检测系统",
-    description="上传视频来检测小鼠 | Upload a video to detect mice",
-    # examples=[["example.mp4", 20]],
-    # cache_examples=True
+    title="🐁 小鼠行为分析 (Mice Behavior Analysis)",
+    description="上传视频来检测和分析小鼠行为 | Upload a video to detect and analyze mice behavior",
+    article="""
+    ### 使用说明
+    1. 上传视频文件
+    2. 设置处理参数：
+       - 处理时长：需要分析的视频时长（秒）
+       - 置信度阈值：检测的置信度要求（越高越严格）
+       - 最大检测数量：每帧最多检测的目标数量
+    3. 等待处理完成
+    4. 查看检测结果视频和分析报告
+    
+    ### 注意事项
+    - 支持常见视频格式（mp4, avi 等）
+    - 建议视频分辨率不超过 1920x1080
+    - 处理时间与视频长度和分辨率相关
+    - 置信度建议范围：0.2-0.5
+    - 最大检测数量建议根据实际场景设置
+    """
 )
 
 if __name__ == "__main__":