add ImagePreprocessor and compute_image_hash to ocr module
This commit is contained in:
@@ -1,14 +1,13 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
from .engine import (
|
from .engine import (
|
||||||
ImagePreprocessor,
|
|
||||||
OCREngine,
|
OCREngine,
|
||||||
OCRResult,
|
OCRResult,
|
||||||
PaddleOCREngine,
|
PaddleOCREngine,
|
||||||
TesseractEngine,
|
TesseractEngine,
|
||||||
compute_image_hash,
|
|
||||||
)
|
)
|
||||||
from .parallel import ParallelProcessor, ProcessedImage
|
from .parallel import ParallelProcessor, ProcessedImage
|
||||||
|
from .preprocess import ImagePreprocessor, compute_image_hash
|
||||||
|
|
||||||
__all__ = [
|
__all__ = [
|
||||||
"OCREngine",
|
"OCREngine",
|
||||||
|
|||||||
119
src/ocr_pipeline/ocr/preprocess.py
Normal file
119
src/ocr_pipeline/ocr/preprocess.py
Normal file
@@ -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()
|
||||||
Reference in New Issue
Block a user