Update app.py

app.py

@@ -12,6 +12,7 @@ import tempfile
 import imageio
 from tqdm import tqdm
 import logging
+import torch.nn.functional as F
 
 # 新增: 配置logging
 logging.basicConfig(
@@ -49,7 +50,7 @@ def login(username, password):
 @spaces.GPU(duration=300)
 def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8):
     """
-    处理视频并进行小鼠检测
+    处理视频并进行��鼠检测
     Args:
         video_path: 输入视频路径
         process_seconds: 处理时长(秒)
@@ -167,7 +168,7 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     
     pbar.close()  # 关闭进度条
     video_writer.release()
-    logger.info(f"视频处理完成，共处理 {frame_count} 帧")
+    logger.info(f"视频处��完成，共处理 {frame_count} 帧")
 
     # 生成分析报告
     confidences = [float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']]
@@ -183,118 +184,114 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
 - 置信度阈值: {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: 允许��最大跳跃距离
-        """
+    def filter_trajectories_gpu(positions, width, height, max_jump_distance=100):
+        """GPU加速版本的轨迹过滤"""
         if len(positions) < 3:
             return positions
         
-        filtered_positions = []
-        last_valid_pos = None
+        # 转换为GPU张量
+        points = torch.tensor(positions, device=device, dtype=torch.float32)
         
-        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
+        # 计算相邻点之间的距离
+        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
+        
+        for i in range(len(valid_indices)-1):
+            curr_idx = valid_indices[i].item()
+            next_idx = valid_indices[i+1].item()
             
-            distance = np.sqrt((x - last_valid_pos[0])**2 + (y - last_valid_pos[1])**2)
+            filtered_points.append(points[curr_idx].tolist())
             
-            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
+            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())
         
-        window_size = 5
-        smoothed_positions = []
+        filtered_points.append(points[valid_indices[-1]].tolist())
         
-        if len(filtered_positions) >= window_size:
-            smoothed_positions.extend(filtered_positions[:window_size//2])
+        # 平滑处理
+        if len(filtered_points) >= 5:
+            points_tensor = torch.tensor(filtered_points, device=device)
+            kernel_size = 5
+            padding = kernel_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])
+            # 使用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()
             
-            smoothed_positions.extend(filtered_positions[-window_size//2:])
-        else:
-            smoothed_positions = filtered_positions
+            smoothed_points = torch.stack([smoothed_x, smoothed_y], dim=1)
+            return smoothed_points.cpu().numpy().tolist()
         
-        return smoothed_positions
+        return filtered_points
 
     # 修改轨迹图生成部分
-    trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255
+    trajectory_img = torch.ones((height, width, 3), device=device) * 255
     points = np.array(all_positions, dtype=np.int32)
     if len(points) > 1:
-        filtered_points = filter_trajectories(points.tolist(), width, height)
+        filtered_points = filter_trajectories_gpu(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 = (
+            color = torch.tensor([
                 int((1 - ratio) * 255),  # B
                 50,                      # G
                 int(ratio * 255)         # R
-            )
-            cv2.line(trajectory_img, tuple(points[i]), tuple(points[i + 1]), color, 2)
+            ], device=device)
+            
+            # 使用GPU绘制线段
+            pt1, pt2 = points[i], points[i + 1]
+            draw_line_gpu(trajectory_img, pt1, pt2, 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)
+        trajectory_img = trajectory_img.cpu().numpy().astype(np.uint8)
+    
+    if torch.cuda.is_available():
+        heatmap = torch.zeros((height, width), device=device)
+        for x, y in filtered_points:
+            if 0 <= x < width and 0 <= y < height:
+                temp_heatmap = torch.zeros((height, width), device=device)
+                temp_heatmap[y, x] = 1
+                # 使用GPU的高斯模糊
+                temp_heatmap = gaussian_blur_gpu(temp_heatmap, kernel_size=31, sigma=10)
+                heatmap += temp_heatmap
         
-        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)
+        heatmap = heatmap.cpu().numpy()
     
     trajectory_frames = []
     heatmap_frames = []
@@ -357,6 +354,62 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     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)
+    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 = pt1
+    x2, y2 = pt2
+    dx = abs(x2 - x1)
+    dy = abs(y2 - y1)
+    
+    if dx > dy:
+        steps = dx
+    else:
+        steps = dy
+    
+    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]
+    
+    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)
+                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
+
 # 创建 Gradio 界面
 with gr.Blocks() as demo:
     gr.Markdown("# 🐁 小鼠行为分析 (Mice Behavior Analysis)")