Update app.py

app.py

@@ -100,49 +100,24 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         
         # 收集位置信息
         if hasattr(r, 'keypoints') and r.keypoints is not None:
-            # 打印关键点对象信息
-            print(f"Keypoints type: {type(r.keypoints)}")
-            print(f"Keypoints data: {r.keypoints}")
+            kpts = r.keypoints.data
+            if isinstance(kpts, torch.Tensor):
+                kpts = kpts.cpu().numpy()
             
-            for kpts in r.keypoints:
-                if isinstance(kpts, torch.Tensor):
-                    kpts = kpts.cpu().numpy()
-                print(f"Single keypoints shape: {kpts.shape}")  # 打印形状
-                print(f"Single keypoints data: {kpts}")  # 打印数据
+            # 取第一个检测目标的第一个关键点（通常是头部）
+            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 isinstance(kpts, np.ndarray):
-                    if len(kpts.shape) == 3:  # [num_objects, num_keypoints, 3]
-                        for obj_kpts in kpts:
-                            if len(obj_kpts) > 0:
-                                x, y = obj_kpts[0][:2]  # 使用第一个关键点的x,y坐标
-                                if isinstance(x, (int, float)) and isinstance(y, (int, float)):
-                                    x, y = int(x), int(y)
-                                    all_positions.append([x, y])
-                                    # 更新热图，使用高斯核来平滑
-                                    if 0 <= x < width and 0 <= y < height:
-                                        # 创建高斯核心点
-                                        sigma = 5  # 调整这个值来改变热点大小
-                                        kernel_size = 15  # 必须是奇数
-                                        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
-                    elif len(kpts.shape) == 2:  # [num_keypoints, 3]
-                        if len(kpts) > 0:
-                            x, y = kpts[0][:2]  # 使用第一个关键点的x,y坐标
-                            if isinstance(x, (int, float)) and isinstance(y, (int, float)):
-                                x, y = int(x), int(y)
-                                all_positions.append([x, y])
-                                # 更新热图，使用高斯核来平滑
-                                if 0 <= x < width and 0 <= y < height:
-                                    # 创建高斯核心点
-                                    sigma = 5  # 调整这个值来改变热点大小
-                                    kernel_size = 15  # 必须是奇数
-                                    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
+                # 更新热图，使用高斯核来平滑
+                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 = {
@@ -204,20 +179,35 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         # 绘制轨迹线，使用渐变色
         for i in range(len(points) - 1):
             ratio = i / (len(points) - 1)
+            # 使用从蓝到红的渐变色
             color = (
                 int((1 - ratio) * 255),  # B
-                0,                       # G
+                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, (255, 0, 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)
     
-    # 生成热图
-    heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
-    heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), cv2.COLORMAP_JET)
+    # 处理热图
+    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)
     
     # 保存图像
     trajectory_path = output_path.replace('.mp4', '_trajectory.png')