diff --git a/src/ocr_pipeline/ocr/__init__.py b/src/ocr_pipeline/ocr/__init__.py index 7a1acc5..baa4bf0 100644 --- a/src/ocr_pipeline/ocr/__init__.py +++ b/src/ocr_pipeline/ocr/__init__.py @@ -1,14 +1,13 @@ from __future__ import annotations from .engine import ( - ImagePreprocessor, OCREngine, OCRResult, PaddleOCREngine, TesseractEngine, - compute_image_hash, ) from .parallel import ParallelProcessor, ProcessedImage +from .preprocess import ImagePreprocessor, compute_image_hash __all__ = [ "OCREngine", diff --git a/src/ocr_pipeline/ocr/preprocess.py b/src/ocr_pipeline/ocr/preprocess.py new file mode 100644 index 0000000..74075a9 --- /dev/null +++ b/src/ocr_pipeline/ocr/preprocess.py @@ -0,0 +1,119 @@ +from __future__ import annotations +import hashlib +import cv2 +import numpy as np + +from ocr_pipeline.config import settings +from ocr_pipeline.utils.logging import get_logger + +logger = get_logger(__name__) + + +class ImagePreprocessor: + """Image preprocessing for OCR optimization.""" + + def __init__(self): + self.config = settings.ocr.preprocess + self.max_dim = self.config.max_dimension + + def load_image(self, path: str) -> np.ndarray: + """Load image from file path.""" + image = cv2.imread(str(path)) + if image is None: + raise ValueError(f"Failed to load image: {path}") + return image + + def resize_if_needed(self, image: np.ndarray) -> np.ndarray: + """Resize image if it exceeds max dimension.""" + h, w = image.shape[:2] + if max(h, w) > self.max_dim: + scale = self.max_dim / max(h, w) + new_w, new_h = int(w * scale), int(h * scale) + image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA) + return image + + def deskew(self, image: np.ndarray) -> np.ndarray: + """Correct skew in the image.""" + if not self.config.deskew: + return image + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + coords = np.column_stack(np.where(gray > 0)) + if len(coords) == 0: + return image + angle = cv2.minAreaRect(coords)[-1] + if angle < -45: + angle = -(90 + angle) + else: + angle = -angle + if abs(angle) > 0.5: + h, w = image.shape[:2] + center = (w // 2, h // 2) + M = cv2.getRotationMatrix2D(center, angle, 1.0) + image = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE) + return image + + def denoise(self, image: np.ndarray) -> np.ndarray: + """Remove noise from image.""" + if not self.config.denoise: + return image + return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21) + + def apply_clahe(self, image: np.ndarray) -> np.ndarray: + """Apply CLAHE for contrast enhancement.""" + if not self.config.clahe: + return image + lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) + l, a, b = cv2.split(lab) + clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) + l = clahe.apply(l) + lab = cv2.merge((l, a, b)) + return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR) + + def remove_lines(self, image: np.ndarray) -> np.ndarray: + """Remove horizontal and vertical lines (tables, grids).""" + if not self.config.remove_lines: + return image + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) + horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1)) + vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40)) + horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel) + vertical_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel) + lines = cv2.bitwise_or(horizontal_lines, vertical_lines) + contours, _ = cv2.findContours(lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + mask = np.ones(image.shape[:2], dtype=np.uint8) * 255 + cv2.drawContours(mask, contours, -1, 0, thickness=cv2.FILLED) + result = image.copy() + result[mask == 0] = [255, 255, 255] + return result + + def adaptive_threshold(self, image: np.ndarray) -> np.ndarray: + """Apply adaptive thresholding for binarization.""" + if not self.config.adaptive_threshold: + return image + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + binary = cv2.adaptiveThreshold( + gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, + cv2.THRESH_BINARY, 11, 2 + ) + return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) + + def preprocess(self, path: str) -> np.ndarray: + """Full preprocessing pipeline.""" + image = self.load_image(path) + image = self.resize_if_needed(image) + image = self.deskew(image) + image = self.denoise(image) + image = self.apply_clahe(image) + image = self.remove_lines(image) + image = self.adaptive_threshold(image) + return image + + +def compute_image_hash(path: str) -> str: + """Compute SHA256 hash of image file for deduplication.""" + hasher = hashlib.sha256() + with open(path, "rb") as f: + for chunk in iter(lambda: f.read(8192), b""): + hasher.update(chunk) + return hasher.hexdigest()