app.py

import os
import gradio as gr
from ultralytics import YOLO
from fastapi import FastAPI
from PIL import Image
import torch
import spaces
import numpy as np
import cv2
from pathlib import Path
import tempfile
import imageio

# 从环境变量获取密码
APP_USERNAME = "admin"  # 用户名保持固定
APP_PASSWORD = os.getenv("APP_PASSWORD", "default_password")  # 从环境变量获取密码

app = FastAPI()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = YOLO('kunin-mice-pose.v0.1.4.pt')

# 定义认证状态
class AuthState:
    def __init__(self):
        self.is_logged_in = False

auth_state = AuthState()

def login(username, password):
    """登录验证"""
    if username == APP_USERNAME and password == APP_PASSWORD:
        auth_state.is_logged_in = True
        return gr.update(visible=False), gr.update(visible=True), "登录成功"
    return gr.update(visible=True), gr.update(visible=False), "用户名或密码错误"

@spaces.GPU(duration=300)
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: 每帧最大检测数量
    """
    if not auth_state.is_logged_in:
        return None, "请先登录"
    
    # 创建临时目录保存输出视频
    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) if process_seconds else int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    
    # 创建视频写入器
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(
        output_path, 
        fourcc, 
        fps, 
        (width, height)
    )
    
    # 计算基于分辨率的线宽
    base_size = min(width, height)
    line_thickness = max(1, int(base_size * 0.002))  # 0.2% 的最小边长

    # 设置推理参数并处理视频
    results = model.predict(
        source=video_path,
        device=device,
        conf=conf_threshold,
        save=False,
        show=False,
        stream=True,
        line_width=line_thickness,  # 线宽
        boxes=True,  # 显示边界框
        show_labels=True,
        show_conf=True,
        vid_stride=1,
        max_det=max_det,
        retina_masks=True,  # 更精细的显示
        verbose=False
    )
    
    # 处理结果
    frame_count = 0
    detection_info = []
    
    # 用于记录轨迹和热图数据
    all_positions = []
    heatmap = np.zeros((height, width), dtype=np.float32)
    
    for r in results:
        frame = r.plot()
        
        # 收集位置信息
        if hasattr(r, 'keypoints') and r.keypoints is not None:
            kpts = r.keypoints.data
            if isinstance(kpts, torch.Tensor):
                kpts = kpts.cpu().numpy()
            
            # 取第一个检测目标的第一个关键点（通常是头部）
            if kpts.shape == (1, 8, 3):  # [num_objects, num_keypoints, xyz]
                x, y = int(kpts[0, 0, 0]), int(kpts[0, 0, 1])  # 使用第一个关键点
                all_positions.append([x, y])
                
                # 更新热图，使用高斯核来平滑
                if 0 <= x < width and 0 <= y < height:
                    # 创建高斯核心点
                    sigma = 10  # 调整这个值来改变热点大小
                    kernel_size = 31  # 必须是奇数
                    temp_heatmap = np.zeros((height, width), dtype=np.float32)
                    temp_heatmap[y, x] = 1
                    temp_heatmap = cv2.GaussianBlur(temp_heatmap, (kernel_size, kernel_size), sigma)
                    heatmap += temp_heatmap
        
        # 收集检测信息
        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 process_seconds and frame_count >= total_frames:
            break
    
    # 释放视频写入器
    video_writer.release()
    
    # 生成分析报告
    confidences = [float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']]
    hist, bins = np.histogram(confidences, bins=5)
    
    confidence_report = "\n".join([
        f"置信度 {bins[i]:.2f}-{bins[i+1]:.2f}: {hist[i]:3d}个检测 ({hist[i]/len(confidences)*100:.1f}%)"
        for i in range(len(hist))
    ])
    
    report = f"""视频分析报告:
参数设置:
- 置信度阈值: {conf_threshold:.2f}
- 最大检测数量: {max_det}
- 处理时长: {process_seconds}秒

分析结果:
- 处理帧数: {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])}

置信度分布:
{confidence_report}
"""
    
    # 在生成轨迹图之前，添加异常点过滤和轨迹平滑
    def filter_trajectories(positions, width, height, max_jump_distance=100):
        """
        过滤轨迹中的异常点
        Args:
            positions: 位置列表 [[x1,y1], [x2,y2],...]
            width: 视频宽度
            height: 视频高度
            max_jump_distance: 允许的最大跳跃距离
        """
        if len(positions) < 3:
            return positions
        
        filtered_positions = []
        last_valid_pos = None
        
        for i, pos in enumerate(positions):
            x, y = pos
            
            # 检查点是否在有效范围内
            if not (0 <= x < width and 0 <= y < height):
                continue
            
            # 第一个有效点
            if last_valid_pos is None:
                filtered_positions.append(pos)
                last_valid_pos = pos
                continue
            
            # 计算与上一个点的距离
            distance = np.sqrt((x - last_valid_pos[0])**2 + (y - last_valid_pos[1])**2)
            
            if distance > max_jump_distance:
                # 如果距离太大，进行插值
                if len(filtered_positions) > 0:
                    # 寻找下一个有效点
                    next_valid_pos = None
                    for next_pos in positions[i:]:
                        nx, ny = next_pos
                        if (0 <= nx < width and 0 <= ny < height):
                            next_distance = np.sqrt((nx - last_valid_pos[0])**2 + (ny - last_valid_pos[1])**2)
                            if next_distance <= max_jump_distance:
                                next_valid_pos = next_pos
                                break
                    
                    if next_valid_pos is not None:
                        # 线性插值
                        steps = max(2, int(distance / max_jump_distance))
                        for j in range(1, steps):
                            alpha = j / steps
                            interp_x = int(last_valid_pos[0] * (1 - alpha) + next_valid_pos[0] * alpha)
                            interp_y = int(last_valid_pos[1] * (1 - alpha) + next_valid_pos[1] * alpha)
                            filtered_positions.append([interp_x, interp_y])
                        filtered_positions.append(next_valid_pos)
                        last_valid_pos = next_valid_pos
            else:
                filtered_positions.append(pos)
                last_valid_pos = pos
        
        # 使用移动平均平滑轨迹
        window_size = 5
        smoothed_positions = []
        
        if len(filtered_positions) >= window_size:
            # 添加开始的点
            smoothed_positions.extend(filtered_positions[:window_size//2])
            
            # 平滑中间的点
            for i in range(window_size//2, len(filtered_positions) - window_size//2):
                window = filtered_positions[i-window_size//2:i+window_size//2+1]
                smoothed_x = int(np.mean([p[0] for p in window]))
                smoothed_y = int(np.mean([p[1] for p in window]))
                smoothed_positions.append([smoothed_x, smoothed_y])
            
            # 添加结束的点
            smoothed_positions.extend(filtered_positions[-window_size//2:])
        else:
            smoothed_positions = filtered_positions
        
        return smoothed_positions

    # 修改轨迹图生成部分
    trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255  # 白色背景
    points = np.array(all_positions, dtype=np.int32)
    if len(points) > 1:
        # 过滤和平滑轨迹
        filtered_points = filter_trajectories(points.tolist(), width, height)
        points = np.array(filtered_points, dtype=np.int32)
        
        # 绘制轨迹线，使用渐变色
        for i in range(len(points) - 1):
            ratio = i / (len(points) - 1)
            # 使用从蓝到红的渐变色
            color = (
                int((1 - ratio) * 255),  # B
                50,                      # G
                int(ratio * 255)         # R
            )
            cv2.line(trajectory_img, tuple(points[i]), tuple(points[i + 1]), color, 2)
        
        # 绘制起点和终点
        cv2.circle(trajectory_img, tuple(points[0]), 8, (0, 255, 0), -1)  # 绿色起点
        cv2.circle(trajectory_img, tuple(points[-1]), 8, (0, 0, 255), -1)  # 红色终点
        
        # 每隔一定帧数添加方向箭头
        arrow_interval = max(len(points) // 20, 1)  # 控制箭头数量
        for i in range(0, len(points) - arrow_interval, arrow_interval):
            pt1 = tuple(points[i])
            pt2 = tuple(points[i + arrow_interval])
            # 计算箭头方向
            angle = np.arctan2(pt2[1] - pt1[1], pt2[0] - pt1[0])
            # 绘制箭头
            cv2.arrowedLine(trajectory_img, pt1, pt2, (100, 100, 100), 1, tipLength=0.2)
    
    # 处理热图
    if np.max(heatmap) > 0:  # 确保有数据
        heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
        heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), cv2.COLORMAP_JET)
        # 添加一些透明度
        alpha = 0.7
        heatmap_colored = cv2.addWeighted(heatmap_colored, alpha, np.full_like(heatmap_colored, 255), 1-alpha, 0)
    
    # 准备GIF动画帧
    trajectory_frames = []
    heatmap_frames = []
    
    # 创建基础图像
    base_trajectory = np.zeros((height, width, 3), dtype=np.uint8) + 255
    base_heatmap = np.zeros((height, width), dtype=np.float32)
    
    # 每N帧生成一个动画帧
    frame_interval = max(1, len(filtered_points) // 50)  # 控制GIF帧数
    
    for i in range(0, len(filtered_points), frame_interval):
        current_points = filtered_points[:i+1]
        
        # 轨迹动画帧 - 修改颜色方案，与静态图保持一致
        frame_trajectory = base_trajectory.copy()
        if len(current_points) > 1:
            points = np.array(current_points, dtype=np.int32)
            for j in range(len(points) - 1):
                ratio = j / (len(current_points) - 1)  # 修改这里，使用当前总长度
                color = (
                    int((1 - ratio) * 255),  # B
                    50,                       # G
                    int(ratio * 255)          # R
                )
                cv2.line(frame_trajectory, tuple(points[j]), tuple(points[j + 1]), color, 2)
            
            # 绘制当前位置点
            cv2.circle(frame_trajectory, tuple(points[-1]), 8, (0, 0, 255), -1)
        trajectory_frames.append(frame_trajectory)
        
        # 热图动画帧 - 修改颜色方案，使用与静态图相同的配色
        frame_heatmap = base_heatmap.copy()
        for x, y in current_points:
            if 0 <= x < width and 0 <= y < height:
                temp_heatmap = np.zeros((height, width), dtype=np.float32)
                temp_heatmap[y, x] = 1
                temp_heatmap = cv2.GaussianBlur(temp_heatmap, (31, 31), 10)
                frame_heatmap += temp_heatmap
        
        if np.max(frame_heatmap) > 0:
            frame_heatmap_norm = cv2.normalize(frame_heatmap, None, 0, 255, cv2.NORM_MINMAX)
            frame_heatmap_color = cv2.applyColorMap(frame_heatmap_norm.astype(np.uint8), cv2.COLORMAP_JET)
            # 添加白色背景，与静态图保持一致
            frame_heatmap_color = cv2.addWeighted(frame_heatmap_color, 0.7, np.full_like(frame_heatmap_color, 255), 0.3, 0)
            heatmap_frames.append(frame_heatmap_color)

    # 保存GIF动画 - 修改这部分
    trajectory_gif_path = output_path.replace('.mp4', '_trajectory.gif')  # 使用.gif后缀
    heatmap_gif_path = output_path.replace('.mp4', '_heatmap.gif')  # 使用.gif后缀
    
    imageio.mimsave(trajectory_gif_path, trajectory_frames, duration=50)  # 50ms per frame
    imageio.mimsave(heatmap_gif_path, heatmap_frames, duration=50)
    
    # 保存图像
    trajectory_path = output_path.replace('.mp4', '_trajectory.png')
    heatmap_path = output_path.replace('.mp4', '_heatmap.png')
    cv2.imwrite(trajectory_path, trajectory_img)
    cv2.imwrite(heatmap_path, heatmap_colored)
    
    return output_path, trajectory_path, heatmap_path, trajectory_gif_path, heatmap_gif_path, report

# 创建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("# 🐁 小鼠行为分析 (Mice Behavior Analysis)")
    
    # 登录界面
    with gr.Group() as login_interface:
        username = gr.Textbox(label="用户名")
        password = gr.Textbox(label="密码", type="password")
        login_button = gr.Button("登录")
        login_msg = gr.Textbox(label="消息", interactive=False)
    
    # 主界面
    with gr.Group(visible=False) as main_interface:
        gr.Markdown("上传视频来检测和分析小鼠行为 | Upload a video to detect and analyze mice behavior")
        
        with gr.Row():
            with gr.Column():
                video_input = gr.Video(label="输入视频")
                process_seconds = gr.Number(
                    label="处理时长（秒，0表示处理整个视频）", 
                    value=20
                )
                conf_threshold = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.2,
                    step=0.05,
                    label="置信度阈值",
                    info="越高越严格，建议范围0.2-0.5"
                )
                max_det = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=1,
                    step=1,
                    label="最大检测数量",
                    info="每帧最多检测的目标数量"
                )
                process_btn = gr.Button("开始处理")
            
            with gr.Column():
                video_output = gr.Video(label="检测结果")
                with gr.Row():
                    trajectory_output = gr.Image(label="运动轨迹")
                    trajectory_gif_output = gr.Image(label="轨迹动画")
                with gr.Row():
                    heatmap_output = gr.Image(label="热力图")
                    heatmap_gif_output = gr.Image(label="热力图动画")
                report_output = gr.Textbox(label="分析报告")
        
        gr.Markdown("""
        ### 使用说明
        1. 上传视频文件
        2. 设置处理参数：
           - 处理时长：需要分析的视频时长（秒
           - 置信度阈值：检测的置信度要求（越高越严格）
           - 最大检测数量：每帧最多检测的目标数量
        3. 等待处理完成
        4. 查看检测结果视频和分析报告
        
        ### 注意事项
        - 支持常见视频格式（mp4, avi 等）
        - 建议视频分辨率不超过 1920x1080
        - 处理时间与视频长度和分辨率相关
        - 置信度建议范围：0.2-0.5
        - 最大检测数量建议根据实际场景设置
        """)
    
    # 设置事件处理
    login_button.click(
        fn=login,
        inputs=[username, password],
        outputs=[login_interface, main_interface, login_msg]
    )
    
    process_btn.click(
        fn=process_video,
        inputs=[video_input, process_seconds, conf_threshold, max_det],
        outputs=[video_output, trajectory_output, heatmap_output, 
                trajectory_gif_output, heatmap_gif_output, report_output]
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)