Task1/qr_code_recognition/scripts/qr_scan_node.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Undistort camera frames, detect QR boxes on RK3588 NPU, and decode on demand."""

from pathlib import Path

import cv2
import numpy as np
import requests
import rospy
import yaml
from cv_bridge import CvBridge
from pyzbar.pyzbar import decode
from sensor_msgs.msg import Image
from std_msgs.msg import Float32MultiArray, String
from std_srvs.srv import SetBool, SetBoolResponse, Trigger, TriggerResponse

from qr_npu_detector import RKNNQrDetector


class QRCodeScanner(object):
    """Three explicit task modes: idle, discovery, and decode."""

    def __init__(self):
        rospy.init_node("qr_code_scanner", anonymous=True)
        self.bridge = CvBridge()
        default_calibration = Path(__file__).resolve().parent.parent / "config" / "head_camera.yaml"
        self.calibration_file = Path(rospy.get_param("~calibration_file", str(default_calibration)))
        self.camera_matrix = None
        self.distortion_coefficients = None
        self.calibration_width = None
        self.calibration_height = None
        self._load_calibration()

        self.scanned_results = []
        self.scanned_urls = set()
        self.decode_enabled = False
        default_model = Path(__file__).resolve().parent.parent / "models" / "qr_detector.rknn"
        self.npu_crop_expand_ratio = float(rospy.get_param("~npu_crop_expand_ratio", 0.30))
        self.npu_crop_scale = float(rospy.get_param("~npu_crop_scale", 2.0))
        self.npu_detector = RKNNQrDetector(
            model_file=rospy.get_param("~npu_model_file", str(default_model)),
            input_size=rospy.get_param("~npu_input_size", 640),
            confidence_threshold=rospy.get_param("~npu_confidence_threshold", 0.35),
            nms_threshold=rospy.get_param("~npu_nms_threshold", 0.45),
            min_interval_seconds=rospy.get_param("~npu_min_interval_seconds", 0.20),
        )
        # RKNNLite must initialize on this main thread, not in the SetBool
        # service callback that task1 invokes at center C.
        self.npu_detector.prepare()
        rospy.loginfo("二维码RKNN运行时已在主线程初始化，当前仍为idle模式")
        rospy.on_shutdown(self.npu_detector.close)

        self.image_sub = rospy.Subscriber("/usb_cam/image_raw", Image, self.image_callback, queue_size=1)
        self.result_pub = rospy.Publisher("/qr_scan_result", String, queue_size=10)
        self.detection_pub = rospy.Publisher("/qr_scan/npu_detection", Float32MultiArray, queue_size=1)
        self.debug_image_pub = rospy.Publisher("/qr_scan/debug_image", Image, queue_size=1)
        self.reset_service = rospy.Service("/qr_scan/reset", Trigger, self.reset_callback)
        self.npu_enable_service = rospy.Service(
            "/qr_scan/set_npu_enabled", SetBool, self.npu_enable_callback
        )
        self.decode_enable_service = rospy.Service(
            "/qr_scan/set_decode_enabled", SetBool, self.decode_enable_callback
        )
        rospy.loginfo("QR Code Scanner节点已启动（idle模式）")

    def _load_calibration(self):
        try:
            with self.calibration_file.open("r", encoding="utf-8") as stream:
                calibration = yaml.safe_load(stream)
            self.calibration_width = int(calibration["image_width"])
            self.calibration_height = int(calibration["image_height"])
            self.camera_matrix = np.asarray(
                calibration["camera_matrix"]["data"], dtype=np.float64
            ).reshape(3, 3)
            self.distortion_coefficients = np.asarray(
                calibration["distortion_coefficients"]["data"], dtype=np.float64
            ).reshape(-1, 1)
        except (OSError, KeyError, TypeError, ValueError, yaml.YAMLError) as error:
            rospy.logfatal("无法加载相机标定文件 %s: %s", self.calibration_file, error)
            raise
        rospy.loginfo("已加载相机标定: %s (%dx%d, %s)", self.calibration_file,
                      self.calibration_width, self.calibration_height,
                      calibration.get("distortion_model", "unknown"))

    def reset_callback(self, _request):
        self.scanned_results = []
        self.scanned_urls.clear()
        self.decode_enabled = False
        self.npu_detector.set_enabled(False)
        rospy.loginfo("二维码扫描记录已重置，已切换到idle模式")
        return TriggerResponse(success=True, message="qr scan state reset")

    def npu_enable_callback(self, request):
        try:
            self.npu_detector.set_enabled(request.data)
        except Exception as error:
            rospy.logerr("无法切换二维码NPU检测: %s", error)
            return SetBoolResponse(success=False, message=str(error))
        state = "开启" if request.data else "关闭"
        rospy.loginfo("二维码NPU检测已%s", state)
        return SetBoolResponse(success=True, message="NPU QR detection %s" % state)

    def decode_enable_callback(self, request):
        self.decode_enabled = bool(request.data)
        state = "decode" if self.decode_enabled else "discovery/idle"
        rospy.loginfo("二维码URL解码已%s（当前模式%s）", "开启" if self.decode_enabled else "关闭", state)
        return SetBoolResponse(success=True, message="QR decode %s" % state)

    def _handle_decoded_url(self, url, source):
        if url in self.scanned_urls:
            return
        self.scanned_urls.add(url)
        rospy.loginfo("%s识别到二维码URL: %s", source, url)
        try:
            response = requests.get(url, timeout=5, proxies={"http": None, "https": None})
            json_data = response.json()
            if json_data.get("code") == 200:
                product_name = json_data.get("result", "未知")
                rospy.loginfo("识别到货品: %s", product_name)
                self.scanned_results.append(product_name)
                self.result_pub.publish(String(data=product_name))
            else:
                rospy.logwarn("JSON返回错误: %s", json_data)
        except Exception as error:
            rospy.logerr("请求URL失败: %s", error)

    @staticmethod
    def _draw_pyzbar_polygon(frame, obj):
        points = obj.polygon
        if len(points) > 4:
            hull = cv2.convexHull(np.array(points, dtype=np.float32))
            hull = list(map(tuple, np.squeeze(hull)))
        else:
            hull = points
        for index in range(len(hull)):
            cv2.line(frame, hull[index], hull[(index + 1) % len(hull)], (0, 255, 0), 3)
        if hull:
            cv2.putText(frame, "QR", tuple(hull[0]), cv2.FONT_HERSHEY_SIMPLEX,
                        0.7, (0, 255, 0), 2)

    def _publish_best_detection(self, detections):
        if not detections:
            return
        x1, y1, x2, y2, score = max(detections, key=lambda item: item[4])
        message = Float32MultiArray()
        message.data = [
            (x1 + x2) / 2.0, (y1 + y2) / 2.0,
            x1, y1, x2, y2, score,
        ]
        self.detection_pub.publish(message)

    def image_callback(self, data):
        try:
            frame = self.bridge.imgmsg_to_cv2(data, "bgr8")
        except Exception as error:
            rospy.logerr("图像转换失败: %s", error)
            return
        # Idle must be genuinely lightweight: do not undistort, decode, run
        # inference, or publish debug frames until task1 explicitly enables
        # discovery or decode mode.  This also keeps service callbacks
        # responsive when the manager turns NPU detection on at center C.
        if not self.decode_enabled and not self.npu_detector.enabled:
            return

        if (frame.shape[1], frame.shape[0]) != (self.calibration_width, self.calibration_height):
            rospy.logwarn_throttle(5.0, "图像分辨率%dx%d与标定%dx%d不一致，跳过该帧",
                                   frame.shape[1], frame.shape[0],
                                   self.calibration_width, self.calibration_height)
            return

        frame = cv2.undistort(frame, self.camera_matrix, self.distortion_coefficients)
        decoded_objects = decode(frame) if self.decode_enabled else []
        for obj in decoded_objects:
            self._draw_pyzbar_polygon(frame, obj)
            try:
                self._handle_decoded_url(obj.data.decode("utf-8"), "pyzbar全图")
            except UnicodeDecodeError:
                continue

        detections = []
        if self.npu_detector.enabled:
            try:
                detections = self.npu_detector.detect(frame)
            except Exception as error:
                rospy.logerr_throttle(5.0, "二维码NPU推理失败: %s", error)
            self._publish_best_detection(detections)

        for detection in detections:
            x1, y1, x2, y2, score = detection
            crop_objects = []
            if self.decode_enabled:
                crop = self.npu_detector.expanded_crop(
                    frame, detection, self.npu_crop_expand_ratio, self.npu_crop_scale
                )
                crop_objects = decode(crop) if crop is not None else []
            color = (0, 255, 0) if crop_objects else (0, 165, 255)
            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), color, 2)
            cv2.putText(frame, "QR NPU %.2f" % score,
                        (int(x1), max(20, int(y1) - 6)),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
            for obj in crop_objects:
                try:
                    self._handle_decoded_url(obj.data.decode("utf-8"), "NPU裁剪pyzbar")
                except UnicodeDecodeError:
                    continue

        self.debug_image_pub.publish(self.bridge.cv2_to_imgmsg(frame, "bgr8"))
        if len(self.scanned_results) >= 3:
            rospy.loginfo_throttle(2.0, "===== 已识别商品: %s =====", self.scanned_results)

    def run(self):
        rospy.spin()


if __name__ == "__main__":
    try:
        QRCodeScanner().run()
    except rospy.ROSInterruptException:
        pass