Aman Nalakath
a0211155fa
- Slim core deps: move ML stack to optional extras (paddle/tables/figures/scientific/full) - Lazy settings proxy with config search paths (env var, cwd, user dir) - New commands: init, demo, setup [basic|full], first-run guard on run/watch - Engine-aware preprocessing: Paddle gets original image (fixes dark mode 0.83->0.95) - Results table shows Skipped count; lazy run-dir creation - kg_ocr marked experimental with extra, Docker defaults with OCR_PIPELINE_CONFIG - 25/25 tests, ruff clean
110 lines
4.0 KiB
Python
110 lines
4.0 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ocr_pipeline.config import settings
|
|
|
|
|
|
class ImagePreprocessor:
|
|
"""Creates an OCR-friendly derivative without mutating detector input images."""
|
|
|
|
def __init__(self) -> None:
|
|
self.config = settings.ocr.preprocess
|
|
self.max_dim = self.config.max_dimension
|
|
|
|
def load_image(self, path: str | Path) -> np.ndarray:
|
|
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:
|
|
height, width = image.shape[:2]
|
|
if max(height, width) <= self.max_dim:
|
|
return image
|
|
scale = self.max_dim / max(height, width)
|
|
return cv2.resize(
|
|
image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_AREA
|
|
)
|
|
|
|
def deskew(self, image: np.ndarray) -> np.ndarray:
|
|
if not self.config.deskew:
|
|
return image
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
_, foreground = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
|
coords = np.column_stack(np.where(foreground > 0))
|
|
if len(coords) < 20:
|
|
return image
|
|
angle = cv2.minAreaRect(coords)[-1]
|
|
angle = -(90 + angle) if angle < -45 else -angle
|
|
if abs(angle) <= 0.5:
|
|
return image
|
|
height, width = image.shape[:2]
|
|
matrix = cv2.getRotationMatrix2D((width // 2, height // 2), angle, 1.0)
|
|
return cv2.warpAffine(
|
|
image, matrix, (width, height), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE
|
|
)
|
|
|
|
def denoise(self, image: np.ndarray) -> np.ndarray:
|
|
return (
|
|
cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
|
|
if self.config.denoise
|
|
else image
|
|
)
|
|
|
|
def apply_clahe(self, image: np.ndarray) -> np.ndarray:
|
|
if not self.config.clahe:
|
|
return image
|
|
lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
|
lightness, a_channel, b_channel = cv2.split(lab)
|
|
lightness = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(lightness)
|
|
return cv2.cvtColor(cv2.merge((lightness, a_channel, b_channel)), cv2.COLOR_LAB2BGR)
|
|
|
|
def remove_lines(self, image: np.ndarray) -> np.ndarray:
|
|
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 = cv2.morphologyEx(
|
|
binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))
|
|
)
|
|
vertical = cv2.morphologyEx(
|
|
binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))
|
|
)
|
|
mask = cv2.bitwise_or(horizontal, vertical)
|
|
result = image.copy()
|
|
result[mask > 0] = (255, 255, 255)
|
|
return result
|
|
|
|
def adaptive_threshold(self, image: np.ndarray) -> np.ndarray:
|
|
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_image(self, image: np.ndarray) -> np.ndarray:
|
|
image = self.resize_if_needed(image)
|
|
image = self.deskew(image)
|
|
image = self.denoise(image)
|
|
image = self.apply_clahe(image)
|
|
image = self.remove_lines(image)
|
|
return self.adaptive_threshold(image)
|
|
|
|
def preprocess(self, path: str | Path) -> np.ndarray:
|
|
return self.preprocess_image(self.load_image(path))
|
|
|
|
|
|
def compute_image_hash(path: str | Path) -> str:
|
|
hasher = hashlib.sha256()
|
|
with Path(path).open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(8192), b""):
|
|
hasher.update(chunk)
|
|
return hasher.hexdigest()
|