core/poker_detector.py

import time
import numpy as np
import cv2
from typing import List, Tuple, Union
from .runonnx.common_detection import COMMON_DETECTION_ONNX

class PokerDetector:
    def __init__(self,
                 model_path: str,
                 ):

        self.poker_detection = COMMON_DETECTION_ONNX(
            model_path=model_path,
            labels=['A', '2', '3', '4', '5', '6', '7', '8', '9', '10', 'J', 'Q', 'K', 'R', 'B'],
        )
    
    # 检测棋盘 detect board
    def pred_and_draw(self, image_rgb: Union[np.ndarray, None] = None) -> Tuple[Union[np.ndarray, None], str]:

        if image_rgb is None:
            return None, ""
        
        start_time = time.time()

        try:
            image_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)
            origin_boxes, filtered_scores, label_names = self.poker_detection.pred(image=image_bgr)
            
            # draw
            image_rgb_with_pred = self.poker_detection.draw_pred(image_rgb, boxes=origin_boxes, scores=filtered_scores, labels=label_names)

        except Exception as e:
            print("检测失败2", e)
            return None, "检测失败2"

        use_time = time.time() - start_time

        time_info = f"推理用时: {use_time:.2f}s"

        return image_rgb_with_pred, time_info