Task1/qr_code_recognition/scripts/qr_npu_detector.py

#!/usr/bin/env python3
"""RK3588 NPU QR-code box detector used by qr_scan_node.py.

The model is a single-class YOLOv5 export with a 640x640 RGB input and a
25200x6 output.  Detection is explicitly disabled by default so navigation
never spends NPU time on QR detection.
"""

from __future__ import print_function

import logging
import time

import cv2
import numpy as np


class RKNNQrDetector(object):
    """Run a single-class YOLOv5 QR detector through RKNNLite."""

    def __init__(
        self,
        model_file,
        input_size=640,
        confidence_threshold=0.35,
        nms_threshold=0.45,
        min_interval_seconds=0.20,
    ):
        self.model_file = str(model_file)
        self.input_size = int(input_size)
        self.confidence_threshold = float(confidence_threshold)
        self.nms_threshold = float(nms_threshold)
        self.min_interval_seconds = float(min_interval_seconds)
        self.enabled = False
        self._rknn = None
        self._last_inference_time = float("-inf")

    @property
    def loaded(self):
        return self._rknn is not None

    def prepare(self):
        """Initialize RKNN from the node's main thread before rospy.spin()."""
        self._ensure_loaded()

    def set_enabled(self, enabled):
        # Runtime creation is deliberately not done from a rospy service
        # callback thread: RKNNLite.init_runtime can block there on RK3588.
        if enabled and not self.loaded:
            raise RuntimeError("二维码NPU运行时尚未在主线程初始化")
        if enabled:
            self._last_inference_time = float("-inf")
        self.enabled = bool(enabled)

    def close(self):
        self.enabled = False
        if self._rknn is not None:
            self._rknn.release()
            self._rknn = None

    def _ensure_loaded(self):
        if self._rknn is not None:
            return
        try:
            # RKNNLite 1.5 changes logging level names during import.
            # Delay it until rospy has finished configuring its log handlers.
            from rknnlite.api import RKNNLite
            # RKNNLite 1.5 renames standard logging levels (for example
            # INFO -> I). rospy's roslogging only accepts standard names,
            # therefore restore both name-to-level and level-to-name maps.
            for level, name in (
                (logging.CRITICAL, "CRITICAL"),
                (logging.ERROR, "ERROR"),
                (logging.WARNING, "WARNING"),
                (logging.INFO, "INFO"),
                (logging.DEBUG, "DEBUG"),
                (logging.NOTSET, "NOTSET"),
            ):
                logging.addLevelName(level, name)
            logging._nameToLevel.update({
                "FATAL": logging.FATAL,
                "WARN": logging.WARNING,
            })
        except ImportError:
            raise RuntimeError(
                "未找到rknnlite；请使用venv3.9启动二维码NPU节点"
            )
        rknn = RKNNLite()
        result = rknn.load_rknn(self.model_file)
        if result != 0:
            raise RuntimeError("加载RKNN二维码模型失败，错误码%d" % result)
        result = rknn.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
        if result != 0:
            rknn.release()
            raise RuntimeError("初始化RK3588 NPU失败，错误码%d" % result)
        self._rknn = rknn

    def _letterbox(self, image):
        height, width = image.shape[:2]
        size = self.input_size
        scale = min(float(size) / float(width), float(size) / float(height))
        resized_width = int(round(width * scale))
        resized_height = int(round(height * scale))
        resized = cv2.resize(
            image, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR
        )
        pad_x = (size - resized_width) // 2
        pad_y = (size - resized_height) // 2
        padded = cv2.copyMakeBorder(
            resized,
            pad_y,
            size - resized_height - pad_y,
            pad_x,
            size - resized_width - pad_x,
            cv2.BORDER_CONSTANT,
            value=(114, 114, 114),
        )
        return cv2.cvtColor(padded, cv2.COLOR_BGR2RGB), scale, pad_x, pad_y

    def _nms_indices(self, boxes, scores):
        if not boxes:
            return []
        xywh_boxes = []
        for x1, y1, x2, y2 in boxes:
            xywh_boxes.append([
                int(round(x1)), int(round(y1)),
                int(round(max(0.0, x2 - x1))),
                int(round(max(0.0, y2 - y1))),
            ])
        indices = cv2.dnn.NMSBoxes(
            xywh_boxes,
            scores,
            self.confidence_threshold,
            self.nms_threshold,
        )
        if len(indices) == 0:
            return []
        return np.asarray(indices).reshape(-1).tolist()

    def detect(self, image):
        """Return (x1, y1, x2, y2, score) boxes in original-image pixels."""
        if not self.enabled or self._rknn is None:
            return []
        now = time.monotonic()
        if now - self._last_inference_time < self.min_interval_seconds:
            return []
        self._last_inference_time = now

        model_image, scale, pad_x, pad_y = self._letterbox(image)
        outputs = self._rknn.inference(
            inputs=[model_image],
            data_format="nhwc",
        )
        if len(outputs) != 1:
            raise RuntimeError("二维码NPU输出数量异常: %d" % len(outputs))
        prediction = np.asarray(outputs[0]).squeeze()
        if prediction.ndim != 2 or prediction.shape[1] != 6:
            raise RuntimeError(
                "二维码NPU输出形状异常: %s" % (np.asarray(outputs[0]).shape,)
            )

        scores = prediction[:, 4] * prediction[:, 5]
        selected = np.where(scores >= self.confidence_threshold)[0]
        boxes = []
        candidate_scores = []
        for index in selected:
            center_x, center_y, box_width, box_height = prediction[index, :4]
            boxes.append((
                center_x - box_width / 2.0,
                center_y - box_height / 2.0,
                center_x + box_width / 2.0,
                center_y + box_height / 2.0,
            ))
            candidate_scores.append(float(scores[index]))

        image_height, image_width = image.shape[:2]
        detections = []
        for index in self._nms_indices(boxes, candidate_scores):
            x1, y1, x2, y2 = boxes[index]
            x1 = max(0.0, min(float(image_width), (x1 - pad_x) / scale))
            y1 = max(0.0, min(float(image_height), (y1 - pad_y) / scale))
            x2 = max(0.0, min(float(image_width), (x2 - pad_x) / scale))
            y2 = max(0.0, min(float(image_height), (y2 - pad_y) / scale))
            if x2 - x1 >= 8.0 and y2 - y1 >= 8.0:
                detections.append((x1, y1, x2, y2, candidate_scores[index]))
        return detections

    @staticmethod
    def expanded_crop(image, detection, expand_ratio, scale):
        """Return an expanded, optionally enlarged crop for pyzbar decoding."""
        x1, y1, x2, y2 = detection[:4]
        image_height, image_width = image.shape[:2]
        expand_x = (x2 - x1) * float(expand_ratio)
        expand_y = (y2 - y1) * float(expand_ratio)
        left = max(0, int(round(x1 - expand_x)))
        top = max(0, int(round(y1 - expand_y)))
        right = min(image_width, int(round(x2 + expand_x)))
        bottom = min(image_height, int(round(y2 + expand_y)))
        if right <= left or bottom <= top:
            return None
        crop = image[top:bottom, left:right]
        if float(scale) > 1.0:
            crop = cv2.resize(
                crop,
                None,
                fx=float(scale),
                fy=float(scale),
                interpolation=cv2.INTER_CUBIC,
            )
        return crop