File size: 7,394 Bytes
c3e98fa 2ebf7ae c3e98fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 | """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
|