Update app.py

app.py

@@ -1,86 +1,109 @@
 import gradio as gr
 from ultralytics import YOLO
 from fastapi import FastAPI
-from PIL import Image
+import cv2
 import torch
 import spaces
 import numpy as np
+from pathlib import Path
 
 app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-# 移除 .to(device)，在预测时指定设备
-model = YOLO('nailong_yolo11.onnx')
+model = YOLO('kunin-mice-pose.v0.1.0.pt')  # 使用你的小鼠检测模型
 
 @spaces.GPU
-def predict(img):
-    # 将输入图像转换为PIL Image对象
-    input_image = Image.fromarray(img)
+def process_video(video_path, process_seconds=20):
+    # 创建临时输出路径
+    output_path = Path("temp_output.mp4")
     
-    # 保持长宽比的情况下调整尺寸
-    w, h = input_image.size
-    scale = min(640/w, 640/h)
-    new_w, new_h = int(w * scale), int(h * scale)
-    if scale != 1:
-        input_image = input_image.resize((new_w, new_h), Image.LANCZOS)
+    # 获取视频信息
+    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)
     
-    # 转换为numpy数组并进行预测
-    img_array = np.array(input_image)
+    # 创建视频写入器
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    video_writer = cv2.VideoWriter(
+        str(output_path), 
+        fourcc, 
+        fps, 
+        (width, height)
+    )
     
-    # 在predict时指定device
-    results = model.predict(img_array, device=device)
-    result = results[0]
+    # 统计信息
+    frame_count = 0
+    total_detections = 0
+    max_mice = 0
+    detection_stats = []
     
-    # 获取预测结果
-    result_img = result.plot()
+    # 处理视频
+    results = model.predict(
+        source=video_path,
+        device=device,
+        conf=0.5,
+        save=False,
+        show=False,
+        stream=True,
+        line_width=2,
+        show_boxes=True,
+        show_labels=True,
+        show_conf=True,
+        vid_stride=1,
+    )
     
-    # 处理检测信息
-    info = {
-        "detected": len(result.boxes) > 0,
-        "count": len(result.boxes),
-        "detections": []
-    }
+    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)
+        
+        # 写入视频
+        video_writer.write(frame)
+        
+        frame_count += 1
+        if frame_count >= total_frames:
+            break
     
-    if info["detected"]:
-        # 获取每个检测框的信息
-        for box in result.boxes:
-            conf = float(box.conf[0])
-            cls = int(box.cls[0])
-            cls_name = result.names[cls]
-            
-            detection_info = {
-                "class": cls_name,
-                "confidence": f"{conf:.2%}"
-            }
-            info["detections"].append(detection_info)
+    # 释放资源
+    video_writer.release()
+    cap.release()
     
-    # 生成输出文本
-    output_text = f"""检测结果:
-- 是否检测到目标: {'是' if info['detected'] else '否'}
-- 检测到的目标数量: {info['count']}"""
+    # 生成统计信息
+    avg_mice = total_detections / frame_count if frame_count > 0 else 0
+    output_text = f"""小鼠检测统计:
+- 处理时长: {process_seconds}秒
+- 总帧数: {frame_count}
+- 最大检测数量: {max_mice}只
+- 平均检测数量: {avg_mice:.1f}只
+- 检测帧率: {fps} FPS"""
     
-    if info["detections"]:
-        output_text += "\n- 详细信息:"
-        for idx, det in enumerate(info["detections"], 1):
-            output_text += f"\n  目标 {idx}: {det['class']} (置信度: {det['confidence']})"
+    # 如果有检测到的帧
+    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}%)"
     
-    # 如果需要将结果图像缩放回原始尺寸
-    if scale != 1:
-        result_img = Image.fromarray(result_img)
-        result_img = result_img.resize((w, h), Image.LANCZOS)
-        result_img = np.array(result_img)
-    
-    return result_img, output_text
+    return str(output_path), output_text
 
 demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(label="输入图片"),
+    fn=process_video,
+    inputs=[
+        gr.Video(label="输入视频"),
+        gr.Slider(minimum=1, maximum=60, value=20, step=1, label="处理时长(秒)")
+    ],
     outputs=[
-        gr.Image(label="检测结果", type="numpy"),
-        gr.Textbox(label="检测信息")
+        gr.Video(label="检测结果"),
+        gr.Textbox(label="检测统计")
     ],
-    title="🐉 奶龙杀手 (NailongKiller)",
-    description="上传图片来检测奶龙 | Upload an image to detect Nailong",
-    # examples=[["example1.jpg"]],
+    title="🐁 小鼠行为检测系统",
+    description="上传视频来检测小鼠 | Upload a video to detect mice",
+    # examples=[["example.mp4", 20]],
     # cache_examples=True
 )