publish v1

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+*.DS_Store

app.py

@@ -0,0 +1,107 @@
+import gradio as gr
+from core.chessboard_detector import ChessboardDetector
+
+detector = ChessboardDetector(
+    det_model_path="onnx/det/v1.onnx", 
+    pose_model_path="onnx/pose/v1.onnx",
+    full_classifier_model_path="onnx/layout_recognition/v1.onnx"
+)
+
+
+
+# 数据集路径
+dict_cate_names = {
+    '.': '.',
+    'x': 'x',
+    '红帅': 'K',
+    '红士': 'A',
+    '红相': 'B',
+    '红马': 'N',
+    '红车': 'R',
+    '红炮': 'C',
+    '红兵': 'P',
+
+    '黑将': 'k',
+    '黑仕': 'a',
+    '黑象': 'b',
+    '黑傌': 'n',
+    '黑車': 'r',
+    '黑砲': 'c',
+    '黑卒': 'p',
+}
+
+dict_cate_names_reverse = {v: k for k, v in dict_cate_names.items()}
+
+
+with gr.Blocks(
+    css="""
+        .image {
+            max-height: 512px;
+        }
+    """
+) as demo:
+    gr.Markdown("""
+                ## 棋盘检测, 棋子识别
+
+                步骤:  
+                    1. 流程分成两步，第一步检测边缘  
+                    2. 对整个棋盘画面进行棋子分类预测
+                """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(label="上传棋盘图片", type="numpy", elem_classes="image")
+
+        with gr.Column():
+            original_image_with_keypoints = gr.Image(
+                label="step1: 原图带关键点",
+                interactive=False,
+                visible=True,
+                elem_classes="image"
+            )
+
+
+    with gr.Row():
+        with gr.Column():   
+            transformed_image = gr.Image(
+                label="step2: 拉伸棋盘",
+                interactive=False,
+                visible=True,
+                elem_classes="image"
+            )
+
+        with gr.Column():
+            use_time = gr.Textbox(
+                label="用时",
+                interactive=False,
+                visible=True,
+            )
+            layout_pred_info = gr.Dataframe(
+                label="棋子识别",
+                interactive=False,
+                visible=True,
+            )
+
+
+    def detect_chessboard(image):
+        original_image_with_keypoints, transformed_image, cells_labels_str, scores, time_info = detector.pred_detect_board_and_classifier(image)
+
+        # 将 cells_labels 转换为 DataFrame
+        # cells_labels 通过  \n 分割
+        annotation_10_rows = [item for item in cells_labels_str.split("\n")]
+        # 将 annotation_10_rows 转换成为 10 行 9 列的二维数组
+        annotation_arr_10_9 = [list(item) for item in annotation_10_rows]
+
+        # 将 棋子类别 转换为 中文
+        annotation_arr_10_9 = [[dict_cate_names_reverse[item] for item in row] for row in annotation_arr_10_9]
+
+
+        return original_image_with_keypoints, transformed_image, annotation_arr_10_9, time_info
+
+    image_input.change(fn=detect_chessboard, 
+                       inputs=[image_input], 
+                       outputs=[original_image_with_keypoints, transformed_image, layout_pred_info, use_time])
+
+if __name__ == "__main__":
+    demo.launch()

core/chessboard_detector.py

@@ -0,0 +1,137 @@
+
+import time
+import numpy as np
+import cv2
+from typing import List, Tuple, Union
+from pandas import DataFrame
+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
+
+class ChessboardDetector:
+    def __init__(self, 
+                 det_model_path: str, 
+                 pose_model_path: str,
+                 full_classifier_model_path: str = None
+                 ):
+
+        self.det = RTMDET_ONNX(
+            model_path=det_model_path,
+        )
+
+
+        self.pose = RTMPOSE_ONNX(
+            model_path=pose_model_path,
+        )
+
+        if full_classifier_model_path is not None:
+            self.full_classifier = FULL_CLASSIFIER_ONNX(
+                model_path=full_classifier_model_path,
+            )
+
+        self.board_positions = []  # 存储棋盘位置坐标
+        self.current_image = None
+        self.current_filename = None
+    
+
+    # 检测中国象棋棋盘
+    def pred_detect_and_keypoints(self, image_bgr: Union[np.ndarray, None] = None) -> Tuple[List[int], float, List[List[int]], List[float]]:
+
+        xyxy, conf = self.det.pred(image_bgr)
+
+        # 预测关键点, 绘制关键点
+        keypoints, scores = self.pose.pred(image=image_bgr, bbox=xyxy)
+
+        return xyxy, conf, keypoints, scores
+    
+
+    def draw_pred_with_keypoints(self, image_rgb: Union[np.ndarray, None] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+        if image_rgb is None:
+            return None, None, None, None
+        
+        image_rgb = image_rgb.copy()
+
+        original_image = image_rgb.copy()
+
+        image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+
+        xyxy, conf, keypoints, scores = self.pred_detect_and_keypoints(image_bgr)
+
+        # 绘制棋盘框架
+        draw_image = self.det.draw_pred(image_rgb, xyxy, conf)
+
+        # 绘制关键点
+        draw_image = self.pose.draw_pred(img=draw_image, keypoints=keypoints, scores=scores)
+
+        # 融合 self.pose.bone_names 与 keypoints, 再转换成 DataFrame
+        keypoint_list = []
+        for bone_name, keypoint in zip(self.pose.bone_names, keypoints):
+            keypoint_list.append({"name": bone_name, "x": keypoint[0], "y": keypoint[1]})
+
+        keypoint_df = DataFrame(keypoint_list)
+
+        return draw_image, original_image, [xyxy], keypoint_df
+
+    # 拉伸棋盘 detect board, 然后预测
+    def extract_chessboard_and_classifier_layout(self, 
+                                                image_rgb: Union[np.ndarray, None] = None, 
+                                                keypoints: Union[np.ndarray, None] = None
+                                                ) -> Tuple[np.ndarray, List[List[str]], List[List[float]]]:
+        
+        # 提取棋盘, 绘制 每个位置的 范围信息
+        transformed_image, _transformed_keypoints, _corner_points = extract_chessboard(img=image_rgb, keypoints=keypoints)
+
+        transformed_image_copy = transformed_image.copy()
+   
+        # 预测每个位置的 棋子类别
+        _, _, scores, pred_result = self.full_classifier.pred(transformed_image_copy, is_rgb=True)
+
+
+        return transformed_image, pred_result, scores
+
+
+    # 检测棋盘 detect board
+    def pred_detect_board_and_classifier(self, 
+                                         image_rgb: Union[np.ndarray, None] = None, 
+                                        ) -> Tuple[np.ndarray, np.ndarray, str, List[List[float]], str]:
+
+        """
+        @param image_rgb: 输入的 RGB 图像
+        @return: 
+            - transformed_image_layout  # 拉伸棋盘
+            - original_image_with_keypoints  # 原图关键点
+            - layout_pred_info  # 每个位置的 棋子类别
+            - scores  # 每个位置的 置信度
+            - time_info  # 推理用时
+        """
+
+        if image_rgb is None:
+            return None, None, [], [], ""
+        
+        image_rgb_for_extract = image_rgb.copy()
+        
+        start_time = time.time()
+
+        image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
+
+        xyxy, conf, keypoints, scores = self.pred_detect_and_keypoints(image_bgr)
+
+        # 绘制棋盘框架
+        draw_image = self.det.draw_pred(image_rgb, xyxy, conf)
+
+        """
+        绘制 原图关键点
+        """
+        original_image_with_keypoints = self.pose.draw_pred(img=draw_image, keypoints=keypoints, scores=scores)
+
+        transformed_image, cells_labels, scores = self.extract_chessboard_and_classifier_layout(image_rgb=image_rgb_for_extract, keypoints=keypoints)
+
+
+        use_time = time.time() - start_time
+
+        time_info = f"推理用时: {use_time:.2f}s"
+
+        return original_image_with_keypoints, transformed_image, cells_labels, scores, time_info
+    

core/helper_34.py

@@ -0,0 +1,316 @@
+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_cls.py

@@ -0,0 +1,23 @@
+
+
+dict_cate_names = {
+    'point': '.',
+    'other': 'x',
+    'red_king': 'K',
+    'red_advisor': 'A',
+    'red_bishop': 'B',
+    'red_knight': 'N',
+    'red_rook': 'R',
+    'red_cannon': 'C',
+    'red_pawn': 'P',
+    'black_king': 'k',
+    'black_advisor': 'a',
+    'black_bishop': 'b',
+    'black_knight': 'n',
+    'black_rook': 'r',
+    'black_cannon': 'c',
+    'black_pawn': 'p',
+}
+
+
+full_cate_names = list(dict_cate_names.keys())

core/kpt_34_with_xanything.py

@@ -0,0 +1,197 @@
+import cv2
+import os
+import json
+import numpy as np
+
+from .helper_34 import BONE_NAMES
+
+class Shape:
+
+    @staticmethod
+    def init_from_dict(data: dict):
+        shape_ins = Shape(data["label"], data["points"], data["group_id"], data["shape_type"])
+
+        return shape_ins
+
+    def __init__(self, label="", points=None, group_id=1, shape_type=""):
+        self.label = label
+        self.score = None
+        self.points = points
+        self.group_id = group_id
+        self.description = ""
+        self.difficult = False
+        self.shape_type = shape_type
+        self.flags = {}
+        self.attributes = {}
+    
+    def to_dict(self):
+        return {
+            "label": self.label,
+            "score": self.score,
+            "points": self.points,
+            "group_id": self.group_id,
+            "description": self.description,
+            "difficult": self.difficult,
+            "shape_type": self.shape_type,
+            "flags": self.flags,
+            "attributes": self.attributes
+        }
+    
+class KeyPoint(Shape):
+    def __init__(self, label="", point_xy=list[float, float], group_id=1):
+        # 校验 point_xy 是否为 2 个元素的列表
+        if len(point_xy) != 2:
+            raise ValueError("point_xy 必须是一个包含 2 个元素的列表")
+        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):
+        
+        if len(xyxy) != 4:
+            raise ValueError("xyxy 必须是一个包含 4 个元素的列表")
+        
+        """
+        bbox [左上角坐标, 右上角坐标, 右下角坐标, 左下角坐标] [[x1,y1],[x2,y2],[x3,y3],[x4,y4]]
+        """
+        x1, y1, x2, y2 = xyxy
+
+        bbox = [
+            [x1, y1],
+            [x2, y1],
+            [x2, y2],
+            [x1, y2]
+        ]
+        
+        super().__init__(label, bbox, group_id, "rectangle")
+
+class Annotation:
+    @staticmethod
+    def init_from_dict(data: dict):
+        """
+        从 dict 初始化 Annotation 类
+        """
+
+        image_height = data["imageHeight"]
+        image_width = data["imageWidth"]
+
+        ann = Annotation(image_path=data["imagePath"], image_width=image_width, image_height=image_height)
+
+        for shape in data["shapes"]:
+            if shape["shape_type"] == "rectangle":
+                ann.add_shape(Rectangle.init_from_dict(shape))
+            elif shape["shape_type"] == "point":
+                ann.add_shape(KeyPoint.init_from_dict(shape))
+
+        return ann
+
+    def __init__(self, image_path="", image_width=-1, image_height=-1):
+        self.version = "2.4.4"
+        self.flags = {}
+        self.shapes = []
+        self.image_data = None
+
+        self.image_path = image_path
+        self.image_height = image_height
+        self.image_width = image_width
+    
+    def add_shape(self, shape: Rectangle | KeyPoint):
+        self.shapes.append(shape.to_dict())
+    
+    def to_dict(self):
+        if self.image_path == "":
+            raise ValueError("image_path 不能为空")
+        if self.image_height == -1 or self.image_width == -1:
+            raise ValueError("image_height 和 image_width 不能为 -1")
+
+        return {
+            "version": self.version,
+            "flags": self.flags,
+            "shapes": self.shapes,
+            "imagePath": self.image_path,
+            "imageData": self.image_data,
+            "imageHeight": self.image_height,
+            "imageWidth": self.image_width
+        }
+
+
+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):
+    """
+    保存 34 个关键点 和 一个 bbox 到 xanything 的 json 文件
+    """
+    x1, y1, x2, y2 = bbox
+    x1, y1, x2, y2 = float(x1), float(y1), float(x2), float(y2)
+
+
+    if image_input is None:
+        raise ValueError("image_input 不能为 None")
+    
+    image_height, image_width = image_input.shape[:2]
+
+
+    # image_ann_path 缺省 .json
+    if not image_ann_path.endswith(".json"):
+        image_ann_path = image_ann_path + ".json"
+
+    # 读取 image_ann_path 的 文件名
+    file_name = os.path.basename(image_ann_path)
+
+    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)]
+
+    for bone_name in BONE_NAMES:
+        x, y = kpt_34_dict[bone_name]
+        annotation.add_shape(KeyPoint(bone_name, [x, y]))
+
+    # 添加 bbox
+    annotation.add_shape(Rectangle("bbox", [x1, y1, x2, y2]))
+
+    ann_file_path = os.path.join(save_dir, file_name)
+    # 保存
+    with open(ann_file_path, "w") as f:
+        json.dump(annotation.to_dict(), f)
+
+    # 保存图片
+    image_input_rgb = image_input.copy()[:, :, ::-1]
+
+    # print('ann_file_path:', ann_file_path.replace(".json", ".jpg"))
+
+    cv2.imwrite(ann_file_path.replace(".json", ".jpg"), image_input_rgb)
+
+def read_xanything_to_json(json_path) -> tuple[list[tuple[str, float, float]], list[float, float, float, float]]:
+    """
+    读取 xanything 的 json 文件
+    """
+    data = {}
+    with open(json_path, "r") as f:
+        data = json.load(f)
+
+    # data 
+    annotation = Annotation.init_from_dict(data)
+
+    keypoints_34_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]]
+        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]] = []
+
+    for item in BONE_NAMES:
+        keypoints_34.append((item, keypoints_34_dict[item][0], keypoints_34_dict[item][1]))
+
+    return keypoints_34, bbox
+
+
+
+
+
+
+
+

core/runonnx/base_onnx.py

@@ -0,0 +1,88 @@
+import onnxruntime
+import numpy as np
+import cv2
+from abc import ABC, abstractmethod
+from typing import Any, Tuple, Union, List
+
+class BaseONNX(ABC):
+    def __init__(self, model_path: str, input_size: Tuple[int, int]):
+        """初始化ONNX模型基类
+
+        Args:
+            model_path (str): ONNX模型路径
+            input_size (tuple): 模型输入尺寸 (width, height)
+        """
+        self.session = onnxruntime.InferenceSession(model_path)
+        self.input_name = self.session.get_inputs()[0].name
+        self.input_size = input_size
+
+    def load_image(self, image: Union[cv2.UMat, str]) -> cv2.UMat:
+        """加载图像
+
+        Args:
+            image (Union[cv2.UMat, str]): 图像路径或cv2图像对象
+
+        Returns:
+            cv2.UMat: 加载的图像
+        """
+        if isinstance(image, str):
+            return cv2.imread(image)
+        return image.copy()
+
+    @abstractmethod
+    def preprocess_image(self, img_bgr: cv2.UMat, *args, **kwargs) -> np.ndarray:
+        """图像预处理抽象方法
+
+        Args:
+            img_bgr (cv2.UMat): BGR格式的输入图像
+            
+        Returns:
+            np.ndarray: 预处理后的图像
+        """
+        pass
+
+    @abstractmethod
+    def run_inference(self, image: np.ndarray) -> Any:
+        """运行推理的抽象方法
+
+        Args:
+            image (np.ndarray): 预处理后的输入图像
+
+        Returns:
+            Any: 模型输出结果
+        """
+        pass
+
+    @abstractmethod
+    def pred(self, image: Union[cv2.UMat, str], *args, **kwargs) -> Any:
+        """预测的抽象方法
+
+        Args:
+            image (Union[cv2.UMat, str]): 输入图像或图像路径
+
+        Returns:
+            Any: 预测结果
+        """
+        pass
+
+    @abstractmethod
+    def draw_pred(self, img: cv2.UMat, *args, **kwargs) -> cv2.UMat:
+        """绘制预测结果的抽象方法
+
+        Args:
+            img (cv2.UMat): 要绘制的图像
+
+        Returns:
+            cv2.UMat: 绘制结果后的图像
+        """
+        pass
+
+    
+    def check_images_list(self, images: List[Union[cv2.UMat, str, np.ndarray]]):
+        """
+        检查图像列表是否有效
+        """
+        for image in images:
+            if not isinstance(image, cv2.UMat) and not isinstance(image, str) and not isinstance(image, np.ndarray):
+                raise ValueError("The images must be a list of cv2.UMat or str or np.ndarray.")
+ 

core/runonnx/full_classifier.py

@@ -0,0 +1,203 @@
+import onnxruntime
+import numpy as np
+import cv2
+from typing import Tuple, List, Union
+from .base_onnx import BaseONNX
+
+
+def center_crop(image: np.ndarray, target_size: Tuple[int, int]) -> np.ndarray:
+    """
+    Center crop the image to the target size.
+
+    Args:
+        image (np.ndarray): The input image.
+        target_size (Tuple[int, int]): The desired output size (height, width).
+
+    Returns:
+        np.ndarray: The cropped image.
+    """
+    h, w, _ = image.shape
+    target_w, target_h = target_size
+
+    center_x = w // 2
+    center_y = h // 2
+
+    start_x = int(center_x - target_w // 2)
+    start_y = int(center_y - target_h // 2)
+
+    cropped_image = image[start_y:start_y + target_h, start_x:start_x + target_w]
+
+    return cropped_image
+
+
+dict_cate_names = {
+    'point': '.',
+    'other': 'x',
+    'red_king': 'K',
+    'red_advisor': 'A',
+    'red_bishop': 'B',
+    'red_knight': 'N',
+    'red_rook': 'R',
+    'red_cannon': 'C',
+    'red_pawn': 'P',
+    'black_king': 'k',
+    'black_advisor': 'a',
+    'black_bishop': 'b',
+    'black_knight': 'n',
+    'black_rook': 'r',
+    'black_cannon': 'c',
+    'black_pawn': 'p',
+}
+
+class FULL_CLASSIFIER_ONNX(BaseONNX):
+
+    label_2_short = dict_cate_names
+
+    classes_labels = list(dict_cate_names.keys())
+
+    def __init__(self, 
+                 model_path,
+                 # 输入图片大小
+                 input_size=(280, 315), # (w, h)
+                 # 图片裁剪大小
+                 crop_size=(400, 450), # (w, h)
+                 ):
+        super().__init__(model_path, input_size)
+
+        self.crop_size = crop_size
+
+
+    def preprocess_image(self, img_bgr: cv2.UMat, is_rgb: bool = True):
+
+        if not is_rgb:
+            img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+        else:
+            img_rgb = img_bgr
+
+        if img_rgb.shape[:2] != self.crop_size:
+            # 调整图片大小 执行 center crop
+            img_rgb = center_crop(img_rgb, self.crop_size) # dst_size = (w, h)
+
+        # resize 到 input_size
+        img_rgb = cv2.resize(img_rgb, self.input_size)
+
+        # normalize mean and std
+        img = (img_rgb - np.array([ 123.675, 116.28, 103.53])) / np.array([58.395, 57.12, 57.375])
+
+        img = img.astype(np.float32)
+        # 转换为浮点型并归一化
+        # img = img.astype(np.float32) / 255.0
+        
+        # 调整维度顺序 (H,W,C) -> (C,H,W)
+        img = np.transpose(img, (2, 0, 1))
+        
+        # 添加 batch 维度
+        img = np.expand_dims(img, axis=0)
+
+        return img
+    
+
+    def run_inference(self, image: np.ndarray) -> np.ndarray:
+        """
+        Run inference on the image.
+
+        Args:
+            image (np.ndarray): The image to run inference on.
+
+        Returns:
+            tuple: A tuple containing the detection results and labels.
+        """
+        # 运行推理
+        outputs, = self.session.run(None, {self.input_name: image})
+        
+        return outputs
+
+    def pred(self, image: List[Union[cv2.UMat, str]], is_rgb: bool = True) -> Tuple[List[List[str]], List[List[str]], List[List[float]], str]:
+        """
+        Predict the detection results of the image.
+
+        Args:
+            image (cv2.UMat, str): The image to predict.
+
+        Returns:
+          
+        """
+        if isinstance(image, str):
+            img_bgr = cv2.imread(image)
+            is_rgb = False
+        else:
+            img_bgr = image.copy()
+
+        image = self.preprocess_image(img_bgr, is_rgb)
+
+        labels = self.run_inference(image)
+
+        # 校验 labels 的 shape
+        assert labels.shape[1:] == (90, 16)
+
+        # shape (90, 16)
+        first_batch_labels = labels[0]
+
+        # 获取置信度最高的标签
+        # list[int]
+        label_indexes = np.argmax(first_batch_labels, axis=-1).tolist()
+
+        # 将标签索引转换为标签
+        # list[str]
+        label_names = [self.classes_labels[index] for index in label_indexes]
+
+        # list[str]
+        label_short = [self.label_2_short[name] for name in label_names]
+
+        # 获取置信度, 根据  first_batch_labels 和 label_indexes
+        confidence = first_batch_labels[np.arange(first_batch_labels.shape[0]), label_indexes]
+
+        label_names_10x9 = [label_names[i*9:(i+1)*9] for i in range(10)]
+        label_short_10x9 = [label_short[i*9:(i+1)*9] for i in range(10)]
+        confidence_10x9 = [confidence[i*9:(i+1)*9] for i in range(10)]
+
+
+        layout_str = "\n".join(["".join(row) for row in label_short_10x9])
+
+        return label_names_10x9, label_short_10x9, confidence_10x9, layout_str
+    
+    def draw_pred(self, image: cv2.UMat, label_index: int, label_name: str, label_short: str, confidence: float) -> cv2.UMat:
+
+        # 在图像上绘制预测结果
+        cv2.putText(image, f"{label_short} {confidence:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
+
+        return image
+    
+
+    def draw_pred_with_result(self, image: cv2.UMat, results: List[Tuple[int, str, str, float]], cells_xyxy: np.ndarray, is_rgb: bool = True) -> cv2.UMat:
+
+        assert len(results) == cells_xyxy.shape[0]
+
+        if not is_rgb:
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+        for i, (label_index, label_name, label_short, confidence) in enumerate(results):
+            # 确保坐标是整数类型
+            x1, y1, x2, y2 = map(int, cells_xyxy[i])
+
+            if label_name.startswith('red'):
+                color = (180, 105, 255) # 粉红色
+            elif label_name.startswith('black'):
+                color = (0, 100, 50) # 黑色
+            else:
+                color = (0, 0, 255) # 蓝色
+
+            
+            if confidence < 0.5:
+                # yellow 
+                color = (255, 255, 0)
+
+            # confidence:.2f 仅保留两位小数 移除
+
+            label_str = f"{label_short} {confidence:.2f}" if confidence < 0.9 else f"{label_short}"
+
+            cv2.putText(image, label_str, (x1 + 8, y2 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.3, color, 1)
+
+        return image
+
+

core/runonnx/rtmdet.py

@@ -0,0 +1,117 @@
+
+
+import numpy as np
+import cv2
+from typing import Tuple, List, Union
+from .base_onnx import BaseONNX
+
+
+class RTMDET_ONNX(BaseONNX):
+
+    def __init__(self, model_path, input_size=(640, 640)):
+        super().__init__(model_path, input_size)
+
+    def preprocess_image(self, img_bgr: cv2.UMat):
+        # 调整图片大小
+        img_bgr = cv2.resize(img_bgr, self.input_size)
+
+        # normalize mean and std
+        img = (img_bgr - np.array([103.53, 116.28, 123.675])) / np.array([57.375, 57.12, 58.395])
+
+        img = img.astype(np.float32)
+        # 转换为浮点型并归一化
+        # img = img.astype(np.float32) / 255.0
+        
+        # 调整维度顺序 (H,W,C) -> (C,H,W)
+        img = np.transpose(img, (2, 0, 1))
+        
+        # 添加 batch 维度
+        img = np.expand_dims(img, axis=0)
+
+        return img
+    
+
+    def run_inference(self, image: np.ndarray):
+        """
+        Run inference on the image.
+
+        Args:
+            image (np.ndarray): The image to run inference on.
+
+        Returns:
+            tuple: A tuple containing the detection results and labels.
+        """
+        # 运行推理
+        outputs = self.session.run(None, {self.input_name: image})
+
+        """
+        dets: 检测框 [batch, num_dets, [x1, y1, x2, y2, conf]] （[batch, num_dets, Reshape(dets_dim_2)]）
+        labels: 标签 [batch,num_dets]
+        """
+        dets, labels = outputs
+
+        return dets, labels
+
+    def pred(self, image: List[Union[cv2.UMat, str]]) -> Tuple[List[int], float]:
+        """
+        Predict the detection results of the image.
+
+        Args:
+            image (cv2.UMat, str): The image to predict.
+
+        Returns:
+           xyxy (list[int, int, int, int]): The detection results.
+           conf (float): The confidence of the detection results.
+        """
+        if isinstance(image, str):
+            img_bgr = cv2.imread(image)
+        else:
+            img_bgr = image.copy()
+
+        original_w, original_h = img_bgr.shape[1], img_bgr.shape[0]
+
+        image = self.preprocess_image(img_bgr)
+        dets, labels = self.run_inference(image)
+
+        # 获取置信度最高的检测框
+        # dets = dets[0][0]
+        # labels = labels[0][0]
+
+        x1, y1, x2, y2, conf = dets[0][0]
+
+        xyxy = [x1, y1, x2, y2]
+
+        xyxy = self.transform_xyxy_to_original(xyxy, original_w, original_h)
+
+        return xyxy, conf
+    
+    def transform_xyxy_to_original(self, xyxy, original_w, original_h) -> List[int]:
+        """
+        将检测框从输入图像的尺寸转换为原始图像的尺寸
+        """
+        x1, y1, x2, y2 = xyxy
+
+        input_w, input_h = self.input_size
+        ratio_w, ratio_h = original_w / input_w, original_h / input_h
+
+        x1, y1, x2, y2 = x1 * ratio_w, y1 * ratio_h, x2 * ratio_w, y2 * ratio_h
+        # 转换为整数
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
+
+        return [x1, y1, x2, y2]
+    
+    def draw_pred(self, img: cv2.UMat, xyxy: List[int], conf: float, is_rgb: bool = True) -> cv2.UMat:
+        """
+        Draw the detection results on the image.
+        """
+
+        if not is_rgb:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        x1, y1, x2, y2 = xyxy
+
+        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
+        cv2.putText(img, f"{conf:.2f}", (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+
+        return img
+    

core/runonnx/rtmpose.py

@@ -0,0 +1,356 @@
+import numpy as np
+import cv2
+from typing import Tuple, List, Union
+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",
+    ]
+
+    def __init__(self, model_path, input_size=(256, 256), padding=1.25):
+        super().__init__(model_path, input_size)
+        self.padding = padding
+
+    
+    def get_bbox_center_scale(self, bbox: List[int]):
+        """Convert bounding box to center and scale.
+        
+        The center is the coordinates of the bbox center, and the scale is the 
+        bbox width and height normalized by the padding factor.
+        
+        Args:
+            bbox: Bounding box in format [x1, y1, x2, y2]
+            
+        Returns:
+            tuple: A tuple containing:
+                - center (numpy.ndarray): Center coordinates [x, y]
+                - scale (numpy.ndarray): Scale [width, height] 
+        """
+        
+        # Get bbox center
+        x1, y1, x2, y2 = bbox
+        center = np.array([(x1 + x2) / 2.0, (y1 + y2) / 2.0])
+        
+        # Get bbox scale (width and height)
+        w = x2 - x1
+        h = y2 - y1
+        
+        # Convert to scaled width/height
+        scale = np.array([w, h]) * self.padding
+        
+        return center, scale
+
+
+    @staticmethod
+    def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+        """Rotate a point by an angle.
+
+        Args:
+            pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+            angle_rad (float): rotation angle in radian
+
+        Returns:
+            np.ndarray: Rotated point in shape (2, )
+        """
+
+        sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+        rot_mat = np.array([[cs, -sn], [sn, cs]])
+        return rot_mat @ pt
+
+
+    @staticmethod
+    def _get_3rd_point(a: np.ndarray, b: np.ndarray):
+        """To calculate the affine matrix, three pairs of points are required. This
+        function is used to get the 3rd point, given 2D points a & b.
+
+        The 3rd point is defined by rotating vector `a - b` by 90 degrees
+        anticlockwise, using b as the rotation center.
+
+        Args:
+            a (np.ndarray): The 1st point (x,y) in shape (2, )
+            b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+        Returns:
+            np.ndarray: The 3rd point.
+        """
+        direction = a - b
+        c = b + np.r_[-direction[1], direction[0]]
+        return c
+
+
+    @staticmethod
+    def get_warp_matrix(
+        center: np.ndarray,
+        scale: np.ndarray,
+        rot: float,
+        output_size: Tuple[int, int],
+        shift: Tuple[float, float] = (0., 0.),
+        inv: bool = False,
+        fix_aspect_ratio: bool = True,
+    ) -> np.ndarray:
+        """Calculate the affine transformation matrix that can warp the bbox area
+        in the input image to the output size.
+
+        Args:
+            center (np.ndarray[2, ]): Center of the bounding box (x, y).
+            scale (np.ndarray[2, ]): Scale of the bounding box
+                wrt [width, height].
+            rot (float): Rotation angle (degree).
+            output_size (np.ndarray[2, ] | list(2,)): Size of the
+                destination heatmaps.
+            shift (0-100%): Shift translation ratio wrt the width/height.
+                Default (0., 0.).
+            inv (bool): Option to inverse the affine transform direction.
+                (inv=False: src->dst or inv=True: dst->src)
+            fix_aspect_ratio (bool): Whether to fix aspect ratio during transform.
+                Defaults to True.
+
+        Returns:
+            np.ndarray: A 2x3 transformation matrix
+        """
+        assert len(center) == 2
+        assert len(scale) == 2
+        assert len(output_size) == 2
+        assert len(shift) == 2
+
+        shift = np.array(shift)
+        src_w, src_h = scale[:2]
+        dst_w, dst_h = output_size[:2]
+
+        rot_rad = np.deg2rad(rot)
+        src_dir = RTMPOSE_ONNX._rotate_point(np.array([src_w * -0.5, 0.]), rot_rad)
+        dst_dir = np.array([dst_w * -0.5, 0.])
+
+        src = np.zeros((3, 2), dtype=np.float32)
+        src[0, :] = center + scale * shift
+        src[1, :] = center + src_dir + scale * shift
+
+        dst = np.zeros((3, 2), dtype=np.float32)
+        dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+        dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+
+        if fix_aspect_ratio:
+            src[2, :] = RTMPOSE_ONNX._get_3rd_point(src[0, :], src[1, :])
+            dst[2, :] = RTMPOSE_ONNX._get_3rd_point(dst[0, :], dst[1, :])
+        else:
+            src_dir_2 = RTMPOSE_ONNX._rotate_point(np.array([0., src_h * -0.5]), rot_rad)
+            dst_dir_2 = np.array([0., dst_h * -0.5])
+            src[2, :] = center + src_dir_2 + scale * shift
+            dst[2, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir_2
+
+        if inv:
+            warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+        else:
+            warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+        return warp_mat
+
+
+    def get_warp_size_with_input_size(self, 
+                                      bbox_center: List[int], 
+                                      bbox_scale: List[int],
+                                      inv: bool = False,
+                                      ):
+        """
+        获取仿射变换矩阵的输出尺寸
+        """
+
+        w, h = self.input_size
+        warp_size = self.input_size
+        
+        # 修正长宽比
+        scale_w, scale_h = bbox_scale
+        aspect_ratio = w / h
+        if scale_w > scale_h * aspect_ratio:
+            bbox_scale = [scale_w, scale_w / aspect_ratio]
+        else:
+            bbox_scale = [scale_h * aspect_ratio, scale_h]
+        
+        # 计算仿射变换矩阵 确保数据类型正确
+        center = np.array(bbox_center, dtype=np.float32)
+        scale = np.array(bbox_scale, dtype=np.float32)
+        
+        rot = 0.0  # 不考虑旋转
+        
+        warp_mat = self.get_warp_matrix(center, scale, rot, output_size=warp_size, inv=inv)
+
+        return warp_mat
+
+    def topdown_affine(self, img: cv2.UMat, bbox_center: List[int], bbox_scale: List[int]):
+        """简化版的 top-down 仿射变换函数
+        
+        Args:
+        img: 输入图像
+        
+        Returns:
+            变换后的图像
+        """
+        
+        warp_mat = self.get_warp_size_with_input_size(bbox_center, bbox_scale)
+
+        # 应用仿射变换
+        dst_img = cv2.warpAffine(img, warp_mat, self.input_size, flags=cv2.INTER_LINEAR)
+        
+        return dst_img
+    
+
+    # 获取每个关键点的最优预测位置
+    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,)
+        
+
+        input_w, input_h = self.input_size
+
+        # 将索引转换为实际坐标 (0-1之间)
+        x_coords = x_indices / (input_w * 2)  # 归一化到0-1
+        y_coords = y_indices / (input_h * 2)
+        
+        # 组合成坐标对
+        keypoints = np.stack([x_coords, y_coords], axis=1)  # (34, 2)
+        
+        # 获取每个点的置信度分数
+        scores = np.max(simcc_x[0], axis=1) * np.max(simcc_y[0], axis=1)
+        
+        return keypoints, scores
+
+
+
+    def preprocess_image(self, img_bgr: cv2.UMat, bbox_center: List[int], bbox_scale: List[int]):
+
+        """
+        预处理图像
+
+        Args:
+            img_bgr (cv2.UMat): 输入图像
+            bbox_center (list[int, int]): 边界框中心坐标 [x, y]
+            bbox_scale (list[int, int]): 边界框尺度 [w, h]
+        """
+
+        affine_img_bgr = self.topdown_affine(img_bgr, bbox_center, bbox_scale)
+
+        # 转RGB并进行归一化
+        affine_img_rgb = cv2.cvtColor(affine_img_bgr, cv2.COLOR_BGR2RGB)
+        # normalize mean and std
+        affine_img_rgb_norm = (affine_img_rgb - np.array([123.675, 116.28, 103.53])) / np.array([58.395, 57.12, 57.375])
+        # 转换为浮点型并归一化
+        img = affine_img_rgb_norm.astype(np.float32)
+        # 调整维度顺序 (H,W,C) -> (C,H,W)
+        img = np.transpose(img, (2, 0, 1))
+        # 添加 batch 维度
+        img = np.expand_dims(img, axis=0)
+
+        return img
+    
+
+    def run_inference(self, image: np.ndarray):
+        """
+        Run inference on the image.
+
+        Args:
+            image (np.ndarray): The image to run inference on.
+
+        Returns:
+            tuple: A tuple containing the detection results and labels.
+        """
+        # 运行推理
+        outputs = self.session.run(None, {self.input_name: image})
+        """
+        simcc_x: float32[batch,MatMulsimcc_x_dim_1,512]
+        simcc_y: float32[batch,MatMulsimcc_x_dim_1,512]
+        """
+        simcc_x, simcc_y = outputs
+
+        return simcc_x, simcc_y
+
+    def pred(self, image: List[Union[cv2.UMat, str]], bbox: List[int]) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Predict the keypoints results of the image.
+
+        Args:
+            image (str | cv2.UMat): The image to predict.
+            bbox (list[int, int, int, int]): The bounding box to predict.
+
+        Returns:
+            keypoints (np.ndarray): The predicted keypoints.
+            scores (np.ndarray): The predicted scores.
+        """
+        if isinstance(image, str):
+            img_bgr = cv2.imread(image)
+        else:
+            img_bgr = image.copy()
+
+        bbox_center, bbox_scale = self.get_bbox_center_scale(bbox)
+
+        image = self.preprocess_image(img_bgr, bbox_center, bbox_scale)
+        simcc_x, simcc_y  = self.run_inference(image)
+
+        # 获取SimCC预测的最大值位置，返回关键点坐标��置信度分数
+        # 对应 width 和 height 为  input_size 的归一化，即 （256,256）
+        keypoints, scores = self.get_simcc_maximum(simcc_x, simcc_y)
+
+        # 将预测的关键点坐标从模型输出尺寸映射回原图尺寸
+        keypoints = self.transform_keypoints_to_original(keypoints, bbox_center, bbox_scale, self.input_size)
+
+        return keypoints, scores
+    
+    def transform_keypoints_to_original(self, keypoints, center, scale, output_size):
+        """
+        将预测的关键点坐标从模型输出尺寸映射回原图尺寸
+        
+        Args:
+            keypoints: 预测的关键点坐标 [N, 2]
+            center: bbox中心点 [x, y]
+            scale: bbox尺度 [w, h]
+            output_size: 模型输入尺寸 (w, h)
+
+        Returns:
+            np.ndarray: 转换后的关键点坐标 [N, 2]
+        """
+        target_coords = keypoints.copy()
+    
+        # 将0-1的预测坐标转换为像素坐标， 256*256
+        target_coords[:, 0] = target_coords[:, 0] * output_size[0]
+        target_coords[:, 1] = target_coords[:, 1] * output_size[1]
+        
+        # 计算仿射变换矩阵
+        warp_mat = self.get_warp_size_with_input_size(center, scale, inv=True)
+        
+        # 转换为齐次坐标
+        ones = np.ones((len(target_coords), 1))
+        target_coords_homogeneous = np.hstack([target_coords, ones])
+        
+        # 应用逆变换
+        original_keypoints = target_coords_homogeneous @ warp_mat.T
+        
+        return original_keypoints
+    
+    def draw_pred(self, img: cv2.UMat, keypoints: np.ndarray, scores: np.ndarray, is_rgb: bool = True) -> cv2.UMat:
+        """
+        Draw the keypoints results on the image.
+        """
+
+        if not is_rgb:
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+        # 获取 随机的  34 中颜色
+        colors = np.random.randint(0, 256, (34, 3))
+
+        for i, (point, score) in enumerate(zip(keypoints, scores)):
+            if score > 0.3:  # 设置置信度阈值
+                x, y = map(int, point)
+                # 使用不同颜色标注不同的关键点
+                color = colors[i]
+
+                cv2.circle(img, (x, y), 5, (int(color[0]), int(color[1]), int(color[2])), -1)
+                # 添加关键点索引标注
+                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)
+        return img
+    

requirements.txt

@@ -0,0 +1 @@
+opencv-python