Update app.py

app.py

@@ -100,25 +100,49 @@ 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}")
+            
             for kpts in r.keypoints:
                 if isinstance(kpts, torch.Tensor):
                     kpts = kpts.cpu().numpy()
-                # 打印关键点格式以便调试
-                # print(f"Keypoints shape: {kpts.shape}")
+                print(f"Single keypoints shape: {kpts.shape}")  # 打印形状
+                print(f"Single keypoints data: {kpts}")  # 打印数据
                 
                 # 确保关键点数据是正确的格式
                 if isinstance(kpts, np.ndarray):
-                    if len(kpts.shape) == 2:  # [n_points, 2/3]
-                        if kpts.shape[0] > 0:  # 确保有关键点
-                            # 使用第一个关键点
-                            if kpts.shape[1] >= 2:  # 确保有x,y坐标
-                                x, y = kpts[0, 0], kpts[0, 1]
+                    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:
-                                        heatmap[y, x] += 1
+                                        # 创建高斯核心点
+                                        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
         
         # 收集检测信息
         frame_info = {
@@ -177,11 +201,19 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255  # 白色背景
     points = np.array(all_positions, dtype=np.int32)
     if len(points) > 1:
-        # 绘制轨迹线
-        cv2.polylines(trajectory_img, [points], False, (0, 0, 255), 2)
+        # 绘制轨迹线，使用渐变色
+        for i in range(len(points) - 1):
+            ratio = i / (len(points) - 1)
+            color = (
+                int((1 - ratio) * 255),  # B
+                0,                       # 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]), 5, (0, 255, 0), -1)  # 绿色起点
-        cv2.circle(trajectory_img, tuple(points[-1]), 5, (255, 0, 0), -1)  # 红色终点
+        cv2.circle(trajectory_img, tuple(points[0]), 8, (0, 255, 0), -1)  # 绿色起点
+        cv2.circle(trajectory_img, tuple(points[-1]), 8, (255, 0, 0), -1)  # 红色终点
     
     # 生成热图
     heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)