Update app.py

app.py

@@ -9,18 +9,7 @@ import numpy as np
 import cv2
 from pathlib import Path
 import tempfile
-import imageio
 from tqdm import tqdm
-import logging
-import torch.nn.functional as F
-
-# 新增: 配置logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s',
-    datefmt='%Y-%m-%d %H:%M:%S'
-)
-logger = logging.getLogger(__name__)
 
 # 从环境变量获取密码
 APP_USERNAME = "admin"  # 用户名保持固定
@@ -28,7 +17,9 @@ APP_PASSWORD = os.getenv("APP_PASSWORD", "default_password")  # 从环境变量
 
 app = FastAPI()
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"使用设备: {device}")
 model = YOLO('kunin-mice-pose.v0.1.5n.pt')
+print("模型加载完成")
 
 # 定义认证状态
 class AuthState:
@@ -39,38 +30,26 @@ auth_state = AuthState()
 
 def login(username, password):
     """登录验证"""
-    logger.info(f"用户尝试登录: {username}")
     if username == APP_USERNAME and password == APP_PASSWORD:
         auth_state.is_logged_in = True
-        logger.info("登录成功")
         return gr.update(visible=False), gr.update(visible=True), "登录成功"
-    logger.warning("登录失败：用户名或密码错误")
     return gr.update(visible=True), gr.update(visible=False), "用户名或密码错误"
 
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=120)
 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: 每帧最大检测数量
     """
-    logger.info(f"开始处理视频: {video_path}")
-    logger.info(f"参数设置 - 处理时长: {process_seconds}秒, 置信度阈值: {conf_threshold}, 最大检测数: {max_det}")
+    print("开始处理视频...")
     
     if not auth_state.is_logged_in:
-        logger.warning("用户未登录，拒绝访问")
         return None, "请先登录"
     
-    # 创建临时目录保存输出视频
-    logger.info("创建临时输出目录")
+    print("创建临时输出文件...")
     with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as tmp_file:
         output_path = tmp_file.name
     
-    # 获取视频信息
-    logger.info("读取视频信息")
+    print("读取视频信息...")
     cap = cv2.VideoCapture(video_path)
     fps = int(cap.get(cv2.CAP_PROP_FPS))
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -78,9 +57,9 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     total_frames = int(process_seconds * fps) if process_seconds else int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
     cap.release()
     
-    logger.info(f"视频信息 - FPS: {fps}, 分辨率: {width}x{height}, 总帧数: {total_frames}")
+    print(f"视频信息: {width}x{height} @ {fps}fps, 总帧数: {total_frames}")
     
-    # 创建视频写入器
+    print("初始化视频写入器...")
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     video_writer = cv2.VideoWriter(
         output_path, 
@@ -89,11 +68,10 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         (width, height)
     )
     
-    # 计算基于分辨率的线宽
     base_size = min(width, height)
-    line_thickness = max(1, int(base_size * 0.002))  # 0.2% 的最小边长
+    line_thickness = max(1, int(base_size * 0.002))
 
-    logger.info("开始YOLO模型推理")
+    print("开始YOLO推理...")
     results = model.predict(
         source=video_path,
         device=device,
@@ -108,28 +86,26 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         vid_stride=1,
         max_det=max_det,
         retina_masks=True,
-        verbose=False  # 关闭YOLO默认日志输出
+        verbose=False
     )
     
-    logger.info("开始处理检测结果")
     frame_count = 0
     detection_info = []
     all_positions = []
     heatmap = np.zeros((height, width), dtype=np.float32)
     
-    # 新增: 创建进度条
-    pbar = tqdm(total=total_frames, desc="处理视频", unit="帧")
+    print("处理检测结果...")
+    progress_bar = tqdm(total=total_frames, desc="处理帧")
     
     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]
+            if kpts.shape == (1, 8, 3):
                 x, y = int(kpts[0, 0, 0]), int(kpts[0, 0, 1])
                 all_positions.append([x, y])
                 
@@ -141,7 +117,6 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
                     temp_heatmap = cv2.GaussianBlur(temp_heatmap, (kernel_size, kernel_size), sigma)
                     heatmap += temp_heatmap
         
-        # 收集检测信息
         frame_info = {
             "frame": frame_count + 1,
             "count": len(r.boxes),
@@ -161,16 +136,17 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         video_writer.write(frame)
         
         frame_count += 1
-        pbar.update(1)  # 更新进度条
+        progress_bar.update(1)
         
         if process_seconds and frame_count >= total_frames:
             break
     
-    pbar.close()  # 关闭进度条
+    progress_bar.close()
+    print("视频处理完成")
+    
     video_writer.release()
-    logger.info(f"视频处理完成，共处理 {frame_count} 帧")
-
-    # 生成分析报告
+    
+    print("生成分析报告...")
     confidences = [float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']]
     hist, bins = np.histogram(confidences, bins=5)
     
@@ -193,267 +169,113 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
 {confidence_report}
 """
     
-    def filter_trajectories_gpu(positions, width, height, max_jump_distance=100):
-        """GPU加速版本的轨迹过滤"""
+    def filter_trajectories(positions, width, height, max_jump_distance=100):
         if len(positions) < 3:
             return positions
         
-        # 转换为GPU张量
-        points = torch.tensor(positions, device=device, dtype=torch.float32)
-        
-        # 计算相邻点之间的距离
-        diffs = points[1:] - points[:-1]
-        distances = torch.norm(diffs, dim=1)
-        
-        # 找出需要插值的位置
-        mask = distances > max_jump_distance
-        valid_indices = (~mask).nonzero().squeeze()
-        
-        if len(valid_indices) < 2:
-            return positions
-        
-        # 使用GPU进行插值
-        filtered_points = []
-        last_valid_idx = 0
+        filtered_positions = []
+        last_valid_pos = None
         
-        for i in range(len(valid_indices)-1):
-            curr_idx = valid_indices[i].item()
-            next_idx = valid_indices[i+1].item()
+        for i, pos in enumerate(positions):
+            x, y = pos
             
-            filtered_points.append(points[curr_idx].tolist())
+            if not (0 <= x < width and 0 <= y < height):
+                continue
             
-            if next_idx - curr_idx > 1:
-                # 线性插值
-                steps = max(2, int((next_idx - curr_idx)))
-                interp_points = torch.linspace(0, 1, steps)
-                start_point = points[curr_idx]
-                end_point = points[next_idx]
-                
-                interpolated = start_point[None] * (1 - interp_points[:, None]) + \
-                             end_point[None] * interp_points[:, None]
-                
-                filtered_points.extend(interpolated[1:-1].tolist())
+            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
         
-        filtered_points.append(points[valid_indices[-1]].tolist())
+        window_size = 5
+        smoothed_positions = []
         
-        # 平滑处理
-        if len(filtered_points) >= 5:
-            points_tensor = torch.tensor(filtered_points, device=device)
-            kernel_size = 5
-            padding = kernel_size // 2
+        if len(filtered_positions) >= window_size:
+            smoothed_positions.extend(filtered_positions[:window_size//2])
             
-            # 使用1D卷积进行平滑
-            weights = torch.ones(1, 1, kernel_size, device=device) / kernel_size
-            smoothed_x = F.conv1d(
-                points_tensor[:, 0].view(1, 1, -1), 
-                weights, 
-                padding=padding
-            ).squeeze()
-            smoothed_y = F.conv1d(
-                points_tensor[:, 1].view(1, 1, -1), 
-                weights, 
-                padding=padding
-            ).squeeze()
+            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_points = torch.stack([smoothed_x, smoothed_y], dim=1)
-            return smoothed_points.cpu().numpy().tolist()
+            smoothed_positions.extend(filtered_positions[-window_size//2:])
+        else:
+            smoothed_positions = filtered_positions
         
-        return filtered_points
+        return smoothed_positions
 
-    # 修改轨迹图生成部分
-    trajectory_img = torch.ones((height, width, 3), device=device, dtype=torch.float32)
+    print("生成轨迹图...")
+    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_gpu(points.tolist(), width, height)
+        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 = torch.tensor([
-                int((1 - ratio) * 255),  # B
-                50,                      # G
-                int(ratio * 255)         # R
-            ], device=device, dtype=torch.float32)
-            
-            # 使用GPU绘制线段
-            pt1, pt2 = points[i], points[i + 1]
-            draw_line_gpu(trajectory_img, pt1, pt2, color, 2)
+            color = (
+                int((1 - ratio) * 255),
+                50,
+                int(ratio * 255)
+            )
+            cv2.line(trajectory_img, tuple(points[i]), tuple(points[i + 1]), color, 2)
         
-        trajectory_img = trajectory_img.cpu().numpy().astype(np.uint8)
-    
-    # 修改热力图生成部分
-    def generate_heatmap(points, width, height):
-        """简化版本的热力图生成函数"""
-        logger.info(f"开始生成热力图，共 {len(points)} 个点")
-        
-        # 使用numpy生成热力图
-        heatmap = np.zeros((height, width), dtype=np.float32)
-        
-        for point in points:
-            try:
-                # 确保是整数坐标
-                x = min(max(0, int(point[0])), width-1)
-                y = min(max(0, int(point[1])), height-1)
-                
-                # 在点位置标记1
-                heatmap[y, x] = 1.0
-                
-                # 直接使用OpenCV的高斯模糊
-                temp_heatmap = cv2.GaussianBlur(heatmap, (31, 31), 10)
-                heatmap = temp_heatmap
-                
-            except Exception as e:
-                logger.warning(f"处理点 {point} 时出错: {e}")
-                continue
+        cv2.circle(trajectory_img, tuple(points[0]), 8, (0, 255, 0), -1)
+        cv2.circle(trajectory_img, tuple(points[-1]), 8, (0, 0, 255), -1)
         
-        logger.info("热力图生成完成")
-        return heatmap
-
-    # 生成热力图
-    logger.info("开始生成热力图")
-    heatmap = generate_heatmap(filtered_points, width, height)
-
-    # 归一化和着色
+        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)
+    
+    print("生成热力图...")
     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)
-        # 添加白色背景
-        heatmap_colored = cv2.addWeighted(heatmap_colored, 0.7, np.full_like(heatmap_colored, 255), 0.3, 0)
-    else:
-        # 如果没有检测到点，返回白色图像
-        heatmap_colored = np.full((height, width, 3), 255, dtype=np.uint8)
-
-    logger.info("热力图生成完成")
-    
-    trajectory_frames = []
-    heatmap_frames = []
-    
-    base_trajectory = np.zeros((height, width, 3), dtype=np.uint8) + 255
-    base_heatmap = np.zeros((height, width), dtype=np.float32)
-    
-    frame_interval = max(1, len(filtered_points) // 30)
-    
-    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),
-                    50,
-                    int(ratio * 255)
-                )
-                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)
-
-    logger.info("开始生成轨迹图和热力图")
-    trajectory_gif_path = output_path.replace('.mp4', '_trajectory.gif')
-    heatmap_gif_path = output_path.replace('.mp4', '_heatmap.gif')
-    
-    imageio.mimsave(trajectory_gif_path, trajectory_frames, duration=50)
-    imageio.mimsave(heatmap_gif_path, heatmap_frames, duration=50)
+        alpha = 0.7
+        heatmap_colored = cv2.addWeighted(heatmap_colored, alpha, np.full_like(heatmap_colored, 255), 1-alpha, 0)
     
+    print("保存结果图像...")
     trajectory_path = output_path.replace('.mp4', '_trajectory.png')
     heatmap_path = output_path.replace('.mp4', '_heatmap.png')
     cv2.imwrite(trajectory_path, trajectory_img)
-    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)
-        heatmap_colored = cv2.addWeighted(heatmap_colored, 0.7, np.full_like(heatmap_colored, 255), 0.3, 0)
     cv2.imwrite(heatmap_path, heatmap_colored)
     
-    logger.info("轨迹图和热力图生成完成")
-    logger.info("开始生成GIF动画")
-    imageio.mimsave(trajectory_gif_path, trajectory_frames, duration=50)
-    imageio.mimsave(heatmap_gif_path, heatmap_frames, duration=50)
-    logger.info("GIF动画生成完成")
-
-    logger.info("所有处理完成，准备返回结果")
-    return output_path, trajectory_path, heatmap_path, trajectory_gif_path, heatmap_gif_path, report
-
-def gaussian_blur_gpu(tensor, kernel_size=31, sigma=10):
-    """GPU版本的高斯模糊"""
-    channels = 1
-    kernel = get_gaussian_kernel2d(kernel_size, sigma).to(device)
-    kernel = kernel.view(1, 1, kernel_size, kernel_size)
-    tensor = tensor.view(1, 1, tensor.shape[0], tensor.shape[1])
-    
-    return F.conv2d(tensor, kernel, padding=kernel_size//2).squeeze()
-
-def get_gaussian_kernel2d(kernel_size, sigma):
-    """生成2D高斯核"""
-    kernel_x = torch.linspace(-kernel_size//2, kernel_size//2, kernel_size)
-    x, y = torch.meshgrid(kernel_x, kernel_x, indexing='ij')
-    kernel = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * sigma ** 2))
-    return kernel / kernel.sum()
-
-def draw_line_gpu(image, pt1, pt2, color, thickness=1):
-    """GPU版本的线段绘制"""
-    x1, y1 = map(int, pt1)  # 确保是整数
-    x2, y2 = map(int, pt2)  # 确保是整数
-    dx = abs(x2 - x1)
-    dy = abs(y2 - y1)
-    
-    # 防止除零错误
-    steps = max(dx, dy)
-    if steps == 0:
-        # 如果是同一个点，直接画点
-        if 0 <= x1 < image.shape[1] and 0 <= y1 < image.shape[0]:
-            image[y1, x1] = color
-        return
-    
-    x_inc = (x2 - x1) / steps
-    y_inc = (y2 - y1) / steps
-    
-    x = x1
-    y = y1
-    
-    points = torch.zeros((int(steps) + 1, 2), device=device)
-    for i in range(int(steps) + 1):
-        points[i] = torch.tensor([x, y])
-        x += x_inc
-        y += y_inc
-    
-    points = points.long()  # 转换为整数类型
-    valid_points = (points[:, 0] >= 0) & (points[:, 0] < image.shape[1]) & \
-                  (points[:, 1] >= 0) & (points[:, 1] < image.shape[0])
-    points = points[valid_points]
-    
-    color = color.to(image.dtype)
-    
-    if thickness > 1:
-        for dx in range(-thickness//2, thickness//2 + 1):
-            for dy in range(-thickness//2, thickness//2 + 1):
-                offset_points = points + torch.tensor([dx, dy], device=device, dtype=torch.long)
-                valid_offset = (offset_points[:, 0] >= 0) & (offset_points[:, 0] < image.shape[1]) & \
-                             (offset_points[:, 1] >= 0) & (offset_points[:, 1] < image.shape[0])
-                offset_points = offset_points[valid_offset]
-                image[offset_points[:, 1], offset_points[:, 0]] = color
-    else:
-        image[points[:, 1], points[:, 0]] = color
+    print("处理完成!")
+    return output_path, trajectory_path, heatmap_path, report
 
 # 创建 Gradio 界面
 with gr.Blocks() as demo:
-    gr.Markdown("# 🐁 小鼠行为分析 (Mice Behavior Analysis)")
+    gr.Markdown("# 🐭 小鼠行为分析 (Mice Behavior Analysis)")
     
     with gr.Group() as login_interface:
         username = gr.Textbox(label="用户名")
@@ -472,7 +294,7 @@ with gr.Blocks() as demo:
                     value=20
                 )
                 conf_threshold = gr.Slider(
-                    minimum=0.0,
+                    minimum=0.1,
                     maximum=1.0,
                     value=0.2,
                     step=0.05,
@@ -485,7 +307,7 @@ with gr.Blocks() as demo:
                     value=1,
                     step=1,
                     label="最大检测数量",
-                    info="每帧���多检测的目标数量"
+                    info="每帧最多检测的目标数量"
                 )
                 process_btn = gr.Button("开始处理")
             
@@ -493,17 +315,14 @@ with gr.Blocks() as demo:
                 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. 等待处理完成
@@ -526,20 +345,8 @@ with gr.Blocks() as demo:
     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]
+        outputs=[video_output, trajectory_output, heatmap_output, report_output]
     )
 
 if __name__ == "__main__":
-    try:
-        # GPU相关操作
-        if torch.cuda.is_available():
-            logger.info("使用GPU进行轨迹和热力图计算")
-            # ... GPU操作 ...
-        else:
-            logger.info("使用CPU进行轨迹和热力图计算")
-            # ... CPU操作 ...
-    except Exception as e:
-        logger.error(f"处理轨迹和热力图时出错: {str(e)}")
-        raise
     demo.launch(server_name="0.0.0.0", server_port=7860)