aman/kg-scr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

improve the OCR pipeline processing and outputs formatting

Browse Source
This commit is contained in:
Aman Nalakath
2026-07-19 19:54:53 +02:00
parent d375db47c3
commit 1cf1c6eeab
30 changed files with 3838 additions and 3900 deletions
Download Patch File Download Diff File Expand all files Collapse all files

98
src/ocr_pipeline/ocr/preprocess.py
View File

@@ -1,119 +1,95 @@
from __future__ import annotations
import hashlib
from pathlib import Path
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."""
"""Creates an OCR-friendly derivative without mutating detector input images."""
def __init__(self):
def __init__(self) -> None:
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."""
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:
"""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
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:
"""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:
_, 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]
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
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:
"""Remove noise from image."""
if not self.config.denoise:
return image
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21) if self.config.denoise else image
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)
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:
"""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)
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]
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
)
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)
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)
image = self.adaptive_threshold(image)
return 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) -> str:
"""Compute SHA256 hash of image file for deduplication."""
def compute_image_hash(path: str | Path) -> str:
hasher = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(8192), b""):
with Path(path).open("rb") as handle:
for chunk in iter(lambda: handle.read(8192), b""):
hasher.update(chunk)
return hasher.hexdigest()