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