feat: 模型优化

HISTORY.md

@@ -11,3 +11,8 @@
 
 1. 使用 4 个关键点检测
 
+
+### 2025-01-25
+
+1. 修改 pose 模型, 不再需要 bbox 输入
+

app.py

@@ -4,8 +4,7 @@ import os
 from core.chessboard_detector import ChessboardDetector
 
 detector = ChessboardDetector(
-    det_model_path="onnx/det/v2.onnx", 
-    pose_model_path="onnx/pose/4_v2.onnx",
+    pose_model_path="onnx/pose/4_v3.onnx",
     full_classifier_model_path="onnx/layout_recognition/nano_v1.onnx"
 )
 
@@ -58,15 +57,14 @@ with gr.Blocks(css="""
     gr.Markdown("""
                 ## 棋盘检测, 棋子识别
 
+                features: 轻量化模型
+
                 x 表示 有遮挡位置   
                 . 表示 棋盘上的普通交叉点  
 
                 步骤:  
-                    1. 流程分成两步，第一步检测边缘  
-                    2. 对整个棋盘画面进行棋子分类预测
-                
-                ### log
-                2025-01-24 模型优化 200M -> 30M
+                    1. 流程分成两步，第一步 keypoints 检测
+                    2. 拉伸棋盘，并预测棋子
                 """
     )
     with gr.Row():
@@ -105,7 +103,8 @@ with gr.Blocks(css="""
 
     with gr.Row():
         with gr.Column():
-            gr.Examples(full_examples, inputs=[image_input], label="示例视频、图片")
+            gr.Examples(
+                full_examples, inputs=[image_input], label="示例图片",  examples_per_page=15,)
 
 
     def detect_chessboard(image):

core/chessboard_detector.py

@@ -4,32 +4,24 @@ 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_4_kpt import extract_chessboard
 
 class ChessboardDetector:
-    def __init__(self, 
-                 det_model_path: str, 
+    def __init__(self,
                  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.full_classifier = FULL_CLASSIFIER_ONNX(
+            model_path=full_classifier_model_path,
+        )
 
         self.board_positions = []  # 存储棋盘位置坐标
         self.current_image = None
@@ -37,19 +29,21 @@ class ChessboardDetector:
     
 
     # 检测中国象棋棋盘
-    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)
+    def pred_keypoints(self, image_bgr: Union[np.ndarray, None] = None) -> Tuple[List[List[int]], List[float]]:
 
         # 预测关键点, 绘制关键点
-        keypoints, scores = self.pose.pred(image=image_bgr, bbox=xyxy)
 
-        return xyxy, conf, keypoints, scores
+        width, height = image_bgr.shape[:2]
+        bbox = [0, 0, width, height]
+
+        keypoints, scores = self.pose.pred(image=image_bgr, bbox=bbox)
+
+        return 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
+            return None, None, None
         
         image_rgb = image_rgb.copy()
 
@@ -57,13 +51,10 @@ class ChessboardDetector:
 
         image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
 
-        xyxy, conf, keypoints, scores = self.pred_detect_and_keypoints(image_bgr)
+        keypoints, scores = self.pred_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)
+        draw_image = self.pose.draw_pred(img=image_rgb, keypoints=keypoints, scores=scores)
 
         # 融合 self.pose.bone_names 与 keypoints, 再转换成 DataFrame
         keypoint_list = []
@@ -72,7 +63,7 @@ class ChessboardDetector:
 
         keypoint_df = DataFrame(keypoint_list)
 
-        return draw_image, original_image, [xyxy], keypoint_df
+        return draw_image, original_image, keypoint_df
 
     # 拉伸棋盘 detect board, 然后预测
     def extract_chessboard_and_classifier_layout(self, 
@@ -111,22 +102,24 @@ class ChessboardDetector:
             return None, None, [], [], ""
         
         image_rgb_for_extract = image_rgb.copy()
+        image_rgb_for_draw = 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)
+        try:
+            image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
 
-        # 绘制棋盘框架
-        draw_image = self.det.draw_pred(image_rgb, xyxy, conf)
+            keypoints, scores = self.pred_keypoints(image_bgr)
 
-        """
-        绘制 原图关键点
-        """
-        original_image_with_keypoints = self.pose.draw_pred(img=draw_image, keypoints=keypoints, scores=scores)
+            """
+            绘制 原图关键点
+            """
+            original_image_with_keypoints = self.pose.draw_pred(img=image_rgb_for_draw, keypoints=keypoints, scores=scores)
 
-        transformed_image, cells_labels, scores = self.extract_chessboard_and_classifier_layout(image_rgb=image_rgb_for_extract, keypoints=keypoints)
+            transformed_image, cells_labels, scores = self.extract_chessboard_and_classifier_layout(image_rgb=image_rgb_for_extract, keypoints=keypoints)
+        except Exception as e:
+            print("检测棋盘失败", e)
+            return None, None, None, None, ""
 
 
         use_time = time.time() - start_time

core/runonnx/rtmdet.py

@@ -1,117 +0,0 @@
-
-
-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

@@ -378,7 +378,7 @@ class RTMPOSE_ONNX(BaseONNX):
                 else:
                     text = f"{self.bone_names[i]}"
                 cv2.putText(img, text, (x+5, y+5), 
-                            cv2.FONT_HERSHEY_SIMPLEX, 1.0, (int(color[0]), int(color[1]), int(color[2])), 1)
+                            cv2.FONT_HERSHEY_SIMPLEX, 1.0, (int(color[0]), int(color[1]), int(color[2])), 2)
                 
                 # 绘制 关节连接线
         for link in self.skeleton_links: