feat: kpt 4

HISTORY.md

@@ -0,0 +1,4 @@
+### 2024-12-28
+
+1. 使用 4 个关键点检测
+

app.py

@@ -5,7 +5,7 @@ from core.chessboard_detector import ChessboardDetector
 
 detector = ChessboardDetector(
     det_model_path="onnx/det/v1.onnx", 
-    pose_model_path="onnx/pose/v1.onnx",
+    pose_model_path="onnx/pose/4_v2.onnx",
     full_classifier_model_path="onnx/layout_recognition/v1.onnx"
 )
 

core/chessboard_detector.py

@@ -8,7 +8,7 @@ from .runonnx.rtmdet import RTMDET_ONNX
 from .runonnx.rtmpose import RTMPOSE_ONNX
 from .runonnx.full_classifier import FULL_CLASSIFIER_ONNX
 
-from core.helper_34 import extract_chessboard
+from core.helper_4_kpt import extract_chessboard
 
 class ChessboardDetector:
     def __init__(self, 

core/helper_34.py

@@ -1,316 +0,0 @@
-import numpy as np
-import cv2
-from typing import Tuple, List
-
-
-BONE_NAMES = [
-    "A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8",
-    "J0", "J1", "J2", "J3", "J4", "J5", "J6", "J7", "J8",
-    "B0", "C0", "D0", "E0", "F0", "G0", "H0", "I0",
-    "B8", "C8", "D8", "E8", "F8", "G8", "H8", "I8",
-]
-
-def check_keypoints(keypoints: np.ndarray):
-    """
-    检查关键点坐标是否正确
-    @param keypoints: 关键点坐标, shape 为 (34, 2)
-    """
-    if keypoints.shape != (34, 2):
-        raise Exception(f"keypoints shape error: {keypoints.shape}")
-
-
-def build_cells_xywh_by_cronners(corner_points: np.ndarray, padding: int = 3) -> np.ndarray:
-    """
-    根据 棋盘的 corner 点坐标 计算 每个位置的 xywh
-    @param corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2)
-    @param padding: 棋盘边框 padding
-
-    @return: 棋盘的 xywh, shape 为 (10, 9, 4), 4 为 center_x, center_y, w, h
-    """
-
-    if corner_points.shape != (4, 2):
-        raise Exception(f"corner_points shape error: {corner_points.shape}")
-    
-    top_left_xy = corner_points[0]
-    top_right_xy = corner_points[1]
-    bottom_left_xy = corner_points[2]
-    bottom_right_xy = corner_points[3]
-    
-    # 计算 每个框的 w 和 h
-    item_w = (top_right_xy[0] - top_left_xy[0]) / (9 - 1)
-    item_h = (bottom_left_xy[1] - top_left_xy[1]) / (10 - 1)
-
-    item_w = item_w
-    item_h = item_h
-
-    item_w_with_padding = item_w - padding * 2
-    item_h_with_padding = item_h - padding * 2
-    
-    # 计算 每个框的 center 坐标
-    cells_xywh = np.zeros((10, 9, 4))
-
-    for i in range(10):
-        for j in range(9):
-            center_x = top_left_xy[0] + item_w * j
-            center_y = top_left_xy[1] + item_h * i
-
-            cells_xywh[i, j] = [center_x, center_y, item_w_with_padding, item_h_with_padding]
-
-    return cells_xywh
-
-
-
-# todo: 需要优化
-def build_cells_xywh(keypoints: np.ndarray, width: int = 450, height: int = 500, padding: int = 3) -> np.ndarray:
-    """
-    @param keypoints: 关键点坐标, shape 为 (34, 2)
-    @param width: 棋盘宽度
-    @param height: 棋盘高度
-    @param padding: 棋盘边框 padding
-    @return: 棋盘的 xywh, shape 为 (10, 9, 4), 4 为 center_x, center_y, w, h
-    """
-    check_keypoints(keypoints)
-
-
-    # 生成 A0 到 J8 的坐标, 如 B1 坐标 为 A1-J1 与 B0-B8 的交集点
-    cells_xywh = np.zeros((10, 9, 4), dtype=np.int16)
-
-    # 遍历 full_points 的每个点，计算其坐标
-    for i in range(10):
-        for j in range(9):
-            # 计算 第 i 行 第 j 列 的坐标
-            row_name = chr(ord('A') + i)
-            col_name = str(j)
-            flag_name = f"{row_name}{col_name}"
-            if flag_name in BONE_NAMES:
-                # 计算 第 i 行 第 j 列 的坐标
-                cur_xy = keypoints[BONE_NAMES.index(flag_name)]
-                cells_xywh[i, j] = [cur_xy[0], cur_xy[1], 0, 0]
-            else:
-                # 计算 第 i 行 第 j 列 的坐标
-                row_start_name = f"{row_name}0"
-                row_end_name = f"{row_name}8"
-                
-                col_start_name = f"A{col_name}"
-                col_end_name = f"J{col_name}"
-
-                row_start_xy = keypoints[BONE_NAMES.index(row_start_name)]
-                row_end_xy = keypoints[BONE_NAMES.index(row_end_name)]
-
-                col_start_xy = keypoints[BONE_NAMES.index(col_start_name)]
-                col_end_xy = keypoints[BONE_NAMES.index(col_end_name)]
-
-                # 计算 row_start_xy 到 row_end_xy 的直线 与 col_start_xy 到 col_end_xy 的直线 的交点
-                # 使用参数方程法计算交点
-                x1, y1 = row_start_xy  # 横向直线起点
-                x2, y2 = row_end_xy    # 横向直线终点
-                x3, y3 = col_start_xy  # 纵向直线起点
-                x4, y4 = col_end_xy    # 纵向直线终点
-
-                # 计算交点坐标
-                # 使用克莱姆法则求解
-                denominator = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4)
-                
-                # 计算交点的 x 坐标
-                x = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denominator
-                # 计算交点的 y 坐标
-                y = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denominator
-                
-                cells_xywh[i, j] = [int(x), int(y), 0, 0]
-
-    # 计算每个点位的 wh
-    for i in range(10):
-        for j in range(9):
-            cur_xy = cells_xywh[i, j]
-            # 获取上下左右 4 个点, 根据 4 个点计算 wh， 宽高为 4 个点 计算出来的 x1y1x2y2 的距离 的 1/2
-            if i == 0:
-                # [i+1, j] 的 反向点
-                up_xy = 2 * cur_xy - cells_xywh[i+1, j]
-            else:
-                up_xy = cells_xywh[i - 1, j]
-            
-            if i == 9:
-                # [i-1, j] 的 反向点
-                down_xy = 2 * cur_xy - cells_xywh[i-1, j]
-            else:
-                down_xy = cells_xywh[i+1, j]
-
-            if j == 0:
-                left_xy = 2 * cur_xy - cells_xywh[i, j+1]
-            else:
-                left_xy = cells_xywh[i, j-1]
-
-            if j == 8:
-                right_xy = 2 * cur_xy - cells_xywh[i, j-1]
-            else:
-                right_xy = cells_xywh[i, j+1]
-
-            min_x = min(up_xy[0].tolist(), down_xy[0].tolist(), left_xy[0].tolist(), right_xy[0].tolist())
-            min_y = min(up_xy[1].tolist(), down_xy[1].tolist(), left_xy[1].tolist(), right_xy[1].tolist())
-
-            min_x += padding
-            min_y += padding
-
-            # 防止 min_x 和 min_y 为 0
-            min_x = max(min_x, 1)
-            min_y = max(min_y, 1)
-
-            max_x = max(up_xy[0].tolist(), down_xy[0].tolist(), left_xy[0].tolist(), right_xy[0].tolist())
-            max_y = max(up_xy[1].tolist(), down_xy[1].tolist(), left_xy[1].tolist(), right_xy[1].tolist())
-
-            max_x -= padding
-            max_y -= padding
-
-            # 防止 max_x 和 max_y 超出边界
-            max_x = min(max_x, width - 1)
-            max_y = min(max_y, height - 1)
-
-            w = (max_x - min_x) / 2
-            h = (max_y - min_y) / 2
-
-            cells_xywh[i, j] = [int(cur_xy[0]), int(cur_xy[1]), int(w), int(h)]
-
-    return cells_xywh
-
-
-def perspective_transform(
-        image: cv2.UMat, 
-        src_points: np.ndarray, 
-        keypoints: np.ndarray,
-        dst_size=(450, 500)) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
-    """
-    透视变换
-    @param image: 图片
-    @param src_points: 源点坐标
-    @param keypoints: 关键点坐标
-    @param dst_size: 目标尺寸 (width, height) 10 行 9 列
-
-    @return:
-        result: 透视变换后的图片
-        transformed_keypoints: 透视变换后的关键点坐标
-        corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
-    """
-
-    check_keypoints(keypoints)
-
-
-    # 源点和目标点
-    src = np.float32(src_points)
-    padding = 50
-    corner_points = np.float32([
-        # 左上角
-        [padding, padding], 
-        # 右上角
-        [dst_size[0]-padding, padding], 
-        # 左下角
-        [padding, dst_size[1]-padding], 
-        # 右下角
-        [dst_size[0]-padding, dst_size[1]-padding]])
-
-    # 计算透视变换矩阵
-    matrix = cv2.getPerspectiveTransform(src, corner_points)
-
-    # 执行透视变换
-    result = cv2.warpPerspective(image, matrix, dst_size)
-
-    # 重塑数组为要求的格式 (N,1,2)
-    keypoints_reshaped = keypoints.reshape(-1, 1, 2).astype(np.float32)
-    transformed_keypoints = cv2.perspectiveTransform(keypoints_reshaped, matrix)
-    # 转回原来的形状
-    transformed_keypoints = transformed_keypoints.reshape(-1, 2)
-
-    return result, transformed_keypoints, corner_points
-
-
-
-def get_board_corner_points(keypoints: np.ndarray) -> np.ndarray:
-    """
-    计算棋局四个边角的 points
-    @param keypoints: 关键点坐标, shape 为 (34, 2)
-    @return: 边角的坐标, shape 为 (4, 2)
-    """
-    check_keypoints(keypoints)
-
-    # 找到 A0 A8 J0 J8 的坐标 以及 A4 和 J4 的坐标
-    a0_index = BONE_NAMES.index("A0")
-    a8_index = BONE_NAMES.index("A8")
-    j0_index = BONE_NAMES.index("J0")
-    j8_index = BONE_NAMES.index("J8")
-
-    a0_xy = keypoints[a0_index]
-    a8_xy = keypoints[a8_index]
-    j0_xy = keypoints[j0_index]
-    j8_xy = keypoints[j8_index]
-
-    # 计算新的四个角点坐标
-    dst_points = np.array([
-        a0_xy,
-        a8_xy,
-        j0_xy,
-        j8_xy
-    ], dtype=np.float32)
-
-    return dst_points
-
-def extract_chessboard(img: cv2.UMat, keypoints: np.ndarray) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
-    """
-    提取棋盘信息
-    @param img: 图片
-    @param keypoints: 关键点坐标, shape 为 (34, 2)
-    @return:
-        transformed_image: 透视变换后的图片
-        transformed_keypoints: 透视变换后的关键点坐标
-        transformed_corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
-    """
-
-    check_keypoints(keypoints)
-
-    source_corner_points = get_board_corner_points(keypoints)
-
-    transformed_image, transformed_keypoints, transformed_corner_points = perspective_transform(img, source_corner_points, keypoints)
-
-    return transformed_image, transformed_keypoints, transformed_corner_points
-
-
-def collect_cells_images(image: cv2.UMat, cells_xywh: np.ndarray) -> List[List[np.ndarray]]:
-    """
-    收集 棋盘的 cells_xywh 对应的图片集合
-    """
-    width = image.shape[1]
-    height = image.shape[0]
-    crop_cells: List[List[np.ndarray]] = []
-
-    for i in range(10):
-        row_cells = []
-        for j in range(9):
-            x, y, w, h = cells_xywh[i, j]
-
-            x_0 = max(int(x-w/2), 0)
-            y_0 = max(int(y-h/2), 0)
-            x_1 = min(int(x+w/2), width-1)
-            y_1 = min(int(y+h/2), height-1)
-
-            crop_img = image[y_0:y_1, x_0:x_1]
-            row_cells.append(crop_img)
-        crop_cells.append(row_cells)
-
-    return crop_cells
-
-def draw_cells_box(image: cv2.UMat, cells_xywh: np.ndarray) -> cv2.UMat:
-    """
-    绘制 棋盘的 cells_xywh 对应的 矩形框
-    """
-    width = image.shape[1]
-    height = image.shape[0]
-    for i in range(10):
-        for j in range(9):
-            x, y, w, h = cells_xywh[i, j]
-
-            x_0 = max(int(x-w/2), 0)
-            y_0 = max(int(y-h/2), 0)
-            x_1 = min(int(x+w/2), width-1)
-            y_1 = min(int(y+h/2), height-1)
-
-            cv2.rectangle(image,(x_0, y_0), (x_1, y_1), (0, 0, 255), 1)
-
-    return image

core/helper_4_kpt.py

@@ -0,0 +1,114 @@
+import numpy as np
+import cv2
+from typing import Tuple, List
+
+
+BONE_NAMES = [
+    "A0", "A8",
+    "J0", "J8",
+]
+
+def check_keypoints(keypoints: np.ndarray):
+    """
+    检查关键点坐标是否正确
+    @param keypoints: 关键点坐标, shape 为 (N, 2)
+    """
+    if keypoints.shape != (len(BONE_NAMES), 2):
+        raise Exception(f"keypoints shape error: {keypoints.shape}")
+def perspective_transform(
+        image: cv2.UMat, 
+        src_points: np.ndarray, 
+        keypoints: np.ndarray,
+        dst_size=(450, 500)) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
+    """
+    透视变换
+    @param image: 图片
+    @param src_points: 源点坐标
+    @param keypoints: 关键点坐标
+    @param dst_size: 目标尺寸 (width, height) 10 行 9 列
+
+    @return:
+        result: 透视变换后的图片
+        transformed_keypoints: 透视变换后的关键点坐标
+        corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
+    """
+
+    check_keypoints(keypoints)
+
+
+    # 源点和目标点
+    src = np.float32(src_points)
+    padding = 50
+    corner_points = np.float32([
+        # 左上角
+        [padding, padding], 
+        # 右上角
+        [dst_size[0]-padding, padding], 
+        # 左下角
+        [padding, dst_size[1]-padding], 
+        # 右下角
+        [dst_size[0]-padding, dst_size[1]-padding]])
+
+    # 计算透视变换矩阵
+    matrix = cv2.getPerspectiveTransform(src, corner_points)
+
+    # 执行透视变换
+    result = cv2.warpPerspective(image, matrix, dst_size)
+
+    # 重塑数组为要求的格式 (N,1,2)
+    keypoints_reshaped = keypoints.reshape(-1, 1, 2).astype(np.float32)
+    transformed_keypoints = cv2.perspectiveTransform(keypoints_reshaped, matrix)
+    # 转回原来的形状
+    transformed_keypoints = transformed_keypoints.reshape(-1, 2)
+
+    return result, transformed_keypoints, corner_points
+
+
+
+def get_board_corner_points(keypoints: np.ndarray) -> np.ndarray:
+    """
+    计算棋局四个边角的 points
+    @param keypoints: 关键点坐标, shape 为 (N, 2)
+    @return: 边角的坐标, shape 为 (4, 2)
+    """
+    check_keypoints(keypoints)
+
+    # 找到 A0 A8 J0 J8 的坐标 以及 A4 和 J4 的坐标
+    a0_index = BONE_NAMES.index("A0")
+    a8_index = BONE_NAMES.index("A8")
+    j0_index = BONE_NAMES.index("J0")
+    j8_index = BONE_NAMES.index("J8")
+
+    a0_xy = keypoints[a0_index]
+    a8_xy = keypoints[a8_index]
+    j0_xy = keypoints[j0_index]
+    j8_xy = keypoints[j8_index]
+
+    # 计算新的四个角点坐标
+    dst_points = np.array([
+        a0_xy,
+        a8_xy,
+        j0_xy,
+        j8_xy
+    ], dtype=np.float32)
+
+    return dst_points
+
+def extract_chessboard(img: cv2.UMat, keypoints: np.ndarray) -> Tuple[cv2.UMat, np.ndarray, np.ndarray]:
+    """
+    提取棋盘信息
+    @param img: 图片
+    @param keypoints: 关键点坐标, shape 为 (N, 2)
+    @return:
+        transformed_image: 透视变换后的图片
+        transformed_keypoints: 透视变换后的关键点坐标
+        transformed_corner_points: 棋盘的 corner 点坐标, shape 为 (4, 2) A0, A8, J0, J8
+    """
+
+    check_keypoints(keypoints)
+
+    source_corner_points = get_board_corner_points(keypoints)
+
+    transformed_image, transformed_keypoints, transformed_corner_points = perspective_transform(img, source_corner_points, keypoints)
+
+    return transformed_image, transformed_keypoints, transformed_corner_points

core/{kpt_34_with_xanything.py → kpt_4_with_xanything.py}

@@ -3,7 +3,7 @@ import os
 import json
 import numpy as np
 
-from .helper_34 import BONE_NAMES
+from .helper_4_kpt import BONE_NAMES
 
 class Shape:
 
@@ -45,7 +45,7 @@ class KeyPoint(Shape):
         super().__init__(label, [point_xy], group_id, "point")
     
 class Rectangle(Shape):
-    def __init__(self, label="A1", xyxy=list[float, float, float, float], group_id=1):
+    def __init__(self, label="A0", xyxy=list[float, float, float, float], group_id=1):
         
         if len(xyxy) != 4:
             raise ValueError("xyxy 必须是一个包含 4 个元素的列表")
@@ -114,9 +114,9 @@ class Annotation:
         }
 
 
-def save_kpt_34_with_xanything(image_input: np.ndarray, image_ann_path, bbox: list[float, float, float, float], kpt_34: list[tuple[str, float, float]], save_dir: str):
+def save_kpt_4_with_xanything(image_input: np.ndarray, image_ann_path, bbox: list[float, float, float, float], kpt_4: list[tuple[str, float, float]], save_dir: str):
     """
-    保存 34 个关键点 和 一个 bbox 到 xanything 的 json 文件
+    保存 4 个关键点 和 一个 bbox 到 xanything 的 json 文件
     """
     x1, y1, x2, y2 = bbox
     x1, y1, x2, y2 = float(x1), float(y1), float(x2), float(y2)
@@ -137,12 +137,12 @@ def save_kpt_34_with_xanything(image_input: np.ndarray, image_ann_path, bbox: li
 
     annotation = Annotation(file_name, image_width, image_height)
 
-    kpt_34_dict = {}
-    for bone_name, x, y in kpt_34:
-        kpt_34_dict[bone_name] = [float(x), float(y)]
+    kpt_4_dict = {}
+    for bone_name, x, y in kpt_4:
+        kpt_4_dict[bone_name] = [float(x), float(y)]
 
     for bone_name in BONE_NAMES:
-        x, y = kpt_34_dict[bone_name]
+        x, y = kpt_4_dict[bone_name]
         annotation.add_shape(KeyPoint(bone_name, [x, y]))
 
     # 添加 bbox
@@ -171,22 +171,22 @@ def read_xanything_to_json(json_path) -> tuple[list[tuple[str, float, float]], l
     # data 
     annotation = Annotation.init_from_dict(data)
 
-    keypoints_34_dict: dict[str, list[float, float]] = {}
+    keypoints_4_dict: dict[str, list[float, float]] = {}
     # x1, y1, x2, y2
     bbox: list[float, float, float, float] = [] 
 
     for shape in annotation.shapes:
         if shape["shape_type"] == "point":
-            keypoints_34_dict[shape["label"]] = [shape["points"][0][0], shape["points"][0][1]]
+            keypoints_4_dict[shape["label"]] = [shape["points"][0][0], shape["points"][0][1]]
         elif shape["shape_type"] == "rectangle":
             bbox = [shape["points"][0][0], shape["points"][0][1], shape["points"][2][0], shape["points"][2][1]]
 
-    keypoints_34: list[tuple[str, float, float]] = []
+    keypoints_4: list[tuple[str, float, float]] = []
 
     for item in BONE_NAMES:
-        keypoints_34.append((item, keypoints_34_dict[item][0], keypoints_34_dict[item][1]))
+        keypoints_4.append((item, keypoints_4_dict[item][0], keypoints_4_dict[item][1]))
 
-    return keypoints_34, bbox
+    return keypoints_4, bbox
 
 
 

core/runonnx/rtmpose.py

@@ -6,16 +6,35 @@ from .base_onnx import BaseONNX
 class RTMPOSE_ONNX(BaseONNX):
 
     bone_names = [
-        "A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8",
-        "J0", "J1", "J2", "J3", "J4", "J5", "J6", "J7", "J8",
-        "B0", "C0", "D0", "E0", "F0", "G0", "H0", "I0",
-        "B8", "C8", "D8", "E8", "F8", "G8", "H8", "I8",
+        "A0", "A8",
+        "J0", "J8",
     ]
 
-    def __init__(self, model_path, input_size=(256, 256), padding=1.25):
+
+    skeleton_links = [
+        "A0-A8",
+        "A8-J8",
+        "J8-J0",
+        "J0-A0",
+    ]
+
+    def __init__(self, 
+                model_path, input_size=(256, 256),
+                padding=1.25,
+                bone_names=None,
+                skeleton_links=None,
+                ):
         super().__init__(model_path, input_size)
         self.padding = padding
 
+        if bone_names is not None:
+            self.bone_names = bone_names
+
+        if skeleton_links is not None:
+            self.skeleton_links = skeleton_links
+
+        self.bone_colors = np.random.randint(0, 256, (len(self.bone_names), 3))
+
     
     def get_bbox_center_scale(self, bbox: List[int]):
         """Convert bounding box to center and scale.
@@ -202,8 +221,8 @@ class RTMPOSE_ONNX(BaseONNX):
     def get_simcc_maximum(self, simcc_x, simcc_y):
         
         # 在最后一维上找到最大值的索引
-        x_indices = np.argmax(simcc_x[0], axis=1)  # (34,)
-        y_indices = np.argmax(simcc_y[0], axis=1)  # (34,)
+        x_indices = np.argmax(simcc_x[0], axis=1)  # (N,)
+        y_indices = np.argmax(simcc_y[0], axis=1)  # (N,)
         
 
         input_w, input_h = self.input_size
@@ -213,7 +232,7 @@ class RTMPOSE_ONNX(BaseONNX):
         y_coords = y_indices / (input_h * 2)
         
         # 组合成坐标对
-        keypoints = np.stack([x_coords, y_coords], axis=1)  # (34, 2)
+        keypoints = np.stack([x_coords, y_coords], axis=1)  # (N, 2)
         
         # 获取每个点的置信度分数
         scores = np.max(simcc_x[0], axis=1) * np.max(simcc_y[0], axis=1)
@@ -339,8 +358,7 @@ class RTMPOSE_ONNX(BaseONNX):
         if not is_rgb:
             img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
 
-        # 获取 随机的  34 中颜色
-        colors = np.random.randint(0, 256, (34, 3))
+        colors = self.bone_colors
 
         for i, (point, score) in enumerate(zip(keypoints, scores)):
             if score > 0.3:  # 设置置信度阈值
@@ -352,5 +370,25 @@ class RTMPOSE_ONNX(BaseONNX):
                 # 添加关键点索引标注
                 cv2.putText(img, self.bone_names[i], (x+5, y+5), 
                             cv2.FONT_HERSHEY_SIMPLEX, 1.0, (int(color[0]), int(color[1]), int(color[2])), 1)
+                
+                # 绘制 关节连接线
+        for link in self.skeleton_links:
+            start_bone, end_bone = link.split("-")
+
+            start_index = self.bone_names.index(start_bone)
+            end_index = self.bone_names.index(end_bone)
+
+            start_keypoint = keypoints[start_index]
+            end_keypoint = keypoints[end_index]
+            link_color = colors[start_index]
+
+            # 绘制连线
+            if scores[start_index] > 0.3 and scores[end_index] > 0.3:
+                start_point = tuple(map(int, start_keypoint))
+                end_point = tuple(map(int, end_keypoint))
+                cv2.line(img, start_point, end_point, 
+                        (int(link_color[0]), int(link_color[1]), int(link_color[2])), 
+                        thickness=2)
+                
         return img