Update app.py

app.py

@@ -16,7 +16,7 @@ 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.2.pt')
+model = YOLO('kunin-mice-pose.v0.1.3.pt')
 
 # 定义认证状态
 class AuthState:
@@ -173,10 +173,95 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
 {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)
@@ -269,7 +354,7 @@ with gr.Blocks() as demo:
         ### 使用说明
         1. 上传视频文件
         2. 设置处理参数：
-           - 处理时长：需要分析的视频时长（秒）
+           - 处理时长：需要分析的视频时长（秒��
            - 置信度阈值：检测的置信度要求（越高越严格）
            - 最大检测数量：每帧最多检测的目标数量
         3. 等待处理完成
@@ -297,4 +382,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)