"""Note OCR pipeline — extract text from Note regions. Supports Vietnamese and English using PaddleOCR with dual-language detection. Lines are sorted top-to-bottom for natural reading order. Includes image preprocessing (upscale, denoise, binarize) for low-res inputs. Usage: from scripts.inference.ocr_note import NoteOCR ocr = NoteOCR() text = ocr.extract(crop_image) """ from __future__ import annotations import cv2 import numpy as np class NoteOCR: """Extract text from Note crop images using PaddleOCR or EasyOCR.""" def __init__(self, langs: list[str] | None = None, backend: str = "auto"): """Initialize OCR engines. Args: langs: Language codes to use. Defaults to ["vi", "en"]. backend: "paddle", "easyocr", or "auto" (try paddle first). """ self.langs = langs or ["vi", "en"] self.backend = backend self.engines: dict = {} if backend == "auto": try: self._init_paddle() self.backend = "paddle" except Exception: self._init_easyocr() self.backend = "easyocr" elif backend == "paddle": self._init_paddle() else: self._init_easyocr() print(f" NoteOCR backend: {self.backend}") def _init_paddle(self): from paddleocr import PaddleOCR for lang in self.langs: self.engines[lang] = PaddleOCR(lang=lang) def _init_easyocr(self): import easyocr # EasyOCR uses different lang codes: vi, en self.engines["easyocr"] = easyocr.Reader(self.langs, gpu=True) def _run_ocr(self, image: np.ndarray, lang: str) -> list[dict]: """Run OCR with one language, return structured line results.""" if self.backend == "easyocr": return self._run_easyocr(image) engine = self.engines[lang] try: result = engine.ocr(image, cls=True) except TypeError: result = engine.ocr(image) if not result or not result[0]: return [] lines = [] for detection in result[0]: bbox = detection[0] # [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] text = detection[1][0] confidence = detection[1][1] y_center = sum(pt[1] for pt in bbox) / 4 x_center = sum(pt[0] for pt in bbox) / 4 lines.append({ "text": text, "confidence": float(confidence), "y_center": y_center, "x_center": x_center, "bbox": bbox, }) return lines def _run_easyocr(self, image: np.ndarray) -> list[dict]: """Run EasyOCR, return structured line results.""" reader = self.engines["easyocr"] result = reader.readtext(image) lines = [] for bbox, text, confidence in result: # EasyOCR bbox: [[x1,y1],[x2,y1],[x2,y2],[x1,y2]] y_center = sum(pt[1] for pt in bbox) / 4 x_center = sum(pt[0] for pt in bbox) / 4 lines.append({ "text": text, "confidence": float(confidence), "y_center": y_center, "x_center": x_center, "bbox": bbox, }) return lines def _sort_lines(self, lines: list[dict], line_height_threshold: float = 15.0) -> list[dict]: """Sort lines in reading order: top-to-bottom, left-to-right. Groups lines that are on the same vertical level (within threshold) and sorts them left-to-right within each group. """ if not lines: return lines # Sort by y first lines.sort(key=lambda l: l["y_center"]) # Group into rows (lines within threshold are same row) rows: list[list[dict]] = [] current_row = [lines[0]] for line in lines[1:]: if abs(line["y_center"] - current_row[0]["y_center"]) < line_height_threshold: current_row.append(line) else: rows.append(current_row) current_row = [line] rows.append(current_row) # Sort each row left-to-right sorted_lines = [] for row in rows: row.sort(key=lambda l: l["x_center"]) sorted_lines.extend(row) return sorted_lines def _preprocess(self, image: np.ndarray, target_height: int = 1000) -> np.ndarray: """Preprocess image for better OCR: upscale, denoise, sharpen. Engineering drawings are often low-res scans. Upscaling + sharpening significantly improves character recognition, especially for Vietnamese diacritics. """ h, w = image.shape[:2] # Upscale small images (most impactful improvement) if h < target_height: scale = target_height / h image = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_CUBIC) # Convert to grayscale for processing if len(image.shape) == 3: gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) else: gray = image # Denoise gray = cv2.fastNlMeansDenoising(gray, h=10) # Adaptive threshold binarization (handles uneven lighting in scans) binary = cv2.adaptiveThreshold( gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 15, 8 ) # Convert back to BGR (PaddleOCR expects 3-channel) result = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) return result def extract(self, image: np.ndarray, min_confidence: float = 0.1) -> dict: """Extract text from a Note crop image. Args: image: BGR numpy array of the cropped Note region. min_confidence: Minimum confidence threshold per line. Returns: dict with keys: - text: joined text string - lines: list of {text, confidence, bbox} - language: detected primary language """ best_result = {"text": "", "lines": [], "language": "unknown", "avg_confidence": 0.0} # Try both raw and preprocessed image, keep best result images_to_try = [image, self._preprocess(image)] for img in images_to_try: for lang in self.langs: lines = self._run_ocr(img, lang) lines = [l for l in lines if l["confidence"] >= min_confidence] if not lines: continue avg_conf = sum(l["confidence"] for l in lines) / len(lines) if avg_conf > best_result["avg_confidence"]: sorted_lines = self._sort_lines(lines) best_result = { "text": "\n".join(l["text"] for l in sorted_lines), "lines": [ { "text": l["text"], "confidence": round(l["confidence"], 4), "bbox": l["bbox"], } for l in sorted_lines ], "language": lang, "avg_confidence": round(avg_conf, 4), } return best_result