ifly/sign_recognition/scripts/sign_recognition_node.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Service-gated full-frame OCR confirmation for a factory sign."""

import os
import subprocess

import cv2
import rospy
from cv_bridge import CvBridge, CvBridgeError
from sensor_msgs.msg import Image
from std_msgs.msg import String
from std_srvs.srv import SetBool, SetBoolResponse


class SignRecognitionNode:
    def __init__(self):
        rospy.init_node("sign_recognition_node")
        self._bridge = CvBridge()
        self._enabled = bool(rospy.get_param("~enabled", False))
        self._image_topic = rospy.get_param("~image_topic", "/usb_cam/image_raw")
        self._flip_horizontal = bool(rospy.get_param("~flip_horizontal", True))
        self._keywords = tuple(rospy.get_param("~keywords", ["食品", "日用品", "电子"]))
        self._interval = float(rospy.get_param("~interval_seconds", 0.50))
        self._max_frames = int(rospy.get_param("~max_frames", 10))
        self._scale = float(rospy.get_param("~scale", 2.0))
        self._language = rospy.get_param("~language", "chi_sim")
        self._psm = int(rospy.get_param("~psm", 6))
        self._tesseract_timeout = float(rospy.get_param("~tesseract_timeout_seconds", 4.0))
        self._debug_save = bool(rospy.get_param("~debug_save", False))
        self._debug_dir = rospy.get_param("~debug_dir", "/tmp/sign_debug")
        if self._max_frames <= 0:
            raise ValueError("max_frames must be positive")
        if self._tesseract_timeout <= 0.0:
            raise ValueError("tesseract_timeout_seconds must be positive")
        if self._debug_save:
            os.makedirs(self._debug_dir, exist_ok=True)

        self._last_process_time = 0.0
        self._frames_processed = 0
        self._result_pub = rospy.Publisher("/sign_recognition", String, queue_size=1)
        self._status_pub = rospy.Publisher(
            "/sign_recognition/status", String, queue_size=1, latch=True
        )
        self._debug_pub = rospy.Publisher(
            "/sign_recognition/debug_image", Image, queue_size=1
        )
        self._image_sub = rospy.Subscriber(
            self._image_topic, Image, self._image_callback, queue_size=1, buff_size=2 ** 24
        )
        self._enable_service = rospy.Service(
            "/sign_recognition/set_enabled", SetBool, self._set_enabled_callback
        )
        self._publish_status("OCR_WAITING_FOR_IMAGE" if self._enabled else "OCR_IDLE")
        rospy.loginfo(
            "sign OCR ready: enabled=%s image_topic=%s max_frames=%d",
            self._enabled, self._image_topic, self._max_frames,
        )

    def _publish_status(self, status):
        self._status_pub.publish(String(data=status))

    def _reset_attempt(self):
        self._last_process_time = 0.0
        self._frames_processed = 0

    def _set_enabled_callback(self, request):
        self._enabled = bool(request.data)
        self._reset_attempt()
        self._publish_status("OCR_WAITING_FOR_IMAGE" if self._enabled else "OCR_IDLE")
        state = "enabled" if self._enabled else "disabled"
        rospy.loginfo("sign OCR %s", state)
        return SetBoolResponse(success=True, message="sign OCR %s" % state)

    @staticmethod
    def _clean_text(raw_text):
        return raw_text.replace(" ", "").replace("\n", "").replace("\u3000", "")

    def _preprocess_variants(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        enlarged = cv2.resize(
            gray, None, fx=self._scale, fy=self._scale, interpolation=cv2.INTER_CUBIC
        )
        equalized = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(enlarged)
        _threshold, binary = cv2.threshold(
            equalized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU
        )
        return (("gray", enlarged), ("binary", binary))

    def _run_tesseract(self, image):
        success, encoded = cv2.imencode(".png", image)
        if not success:
            raise RuntimeError("cannot encode OCR image")
        try:
            completed = subprocess.run(
                ["tesseract", "stdin", "stdout", "-l", self._language,
                 "--psm", str(self._psm)],
                input=encoded.tobytes(), stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                timeout=self._tesseract_timeout, check=False,
            )
        except (OSError, subprocess.TimeoutExpired) as error:
            raise RuntimeError("tesseract execution failed: %s" % error)
        if completed.returncode != 0:
            raise RuntimeError(completed.stderr.decode("utf-8", errors="replace").strip())
        return completed.stdout.decode("utf-8", errors="replace")

    def _match_category(self, cleaned):
        # OCR often splits Chinese characters or inserts unrelated characters.
        # A single full-frame OCR result is sufficient if it contains every character of one
        # factory category, regardless of order or intervening characters.
        for category in self._keywords:
            if all(character in cleaned for character in category):
                return category
        # The electronics factory sign may show "生产" rather than "电子".
        if "电子" in self._keywords and all(character in cleaned for character in "生产"):
            return "电子"
        return None

    def _extract_keyword(self, frame):
        last_variant = None
        for variant_name, variant in self._preprocess_variants(frame):
            last_variant = variant
            raw_text = self._run_tesseract(variant)
            cleaned = self._clean_text(raw_text)
            rospy.loginfo("[OCR %s] raw=%r cleaned=%r", variant_name, raw_text, cleaned)
            category = self._match_category(cleaned)
            if category is not None:
                return category, last_variant
        return None, last_variant

    def _publish_debug(self, image):
        if image is None:
            return
        try:
            self._debug_pub.publish(self._bridge.cv2_to_imgmsg(image, encoding="mono8"))
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "OCR debug image conversion failed: %s", error)
        if self._debug_save:
            cv2.imwrite(
                os.path.join(self._debug_dir, "ocr_%03d.jpg" % self._frames_processed), image
            )

    def _finish(self, result, status):
        self._enabled = False
        self._result_pub.publish(String(data=result))
        self._publish_status(status)

    def _image_callback(self, message):
        if not self._enabled:
            return
        now = rospy.Time.now().to_sec()
        if now - self._last_process_time < self._interval:
            return
        self._last_process_time = now
        try:
            frame = self._bridge.imgmsg_to_cv2(message, desired_encoding="bgr8")
        except CvBridgeError as error:
            rospy.logerr_throttle(5.0, "OCR image conversion failed: %s", error)
            return
        if self._flip_horizontal:
            frame = cv2.flip(frame, 1)
        if frame.size == 0 or frame.shape[0] < 8 or frame.shape[1] < 8:
            rospy.logwarn_throttle(2.0, "OCR ignored an empty or too-small image")
            return

        self._frames_processed += 1
        try:
            keyword, debug_image = self._extract_keyword(frame)
        except (RuntimeError, OSError) as error:
            rospy.logerr_throttle(5.0, "OCR call failed: %s", error)
            keyword, debug_image = None, None
        self._publish_debug(debug_image)

        if keyword is not None:
            rospy.loginfo("factory OCR confirmed from one full image frame: %s", keyword)
            self._finish(keyword, "OCR_CONFIRMED type=%s" % keyword)
            return
        else:
            self._publish_status(
                "OCR_READING frame=%d/%d no_keyword" % (
                    self._frames_processed, self._max_frames
                )
            )

        if self._frames_processed >= self._max_frames:
            rospy.logwarn("factory OCR failed after %d ROI frames", self._frames_processed)
            self._finish("UNKNOWN", "OCR_FAILED no_consensus")


if __name__ == "__main__":
    try:
        SignRecognitionNode()
        rospy.spin()
    except (rospy.ROSInterruptException, ValueError):
        pass