#!/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