Spaces:
Sleeping
Sleeping
File size: 4,950 Bytes
c1b1f3a | 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 | import os
import cv2
import numpy as np
from PIL import Image
# Support PyMuPDF (fitz)
HAVE_FITZ = False
try:
import fitz
HAVE_FITZ = True
except ImportError:
pass
class PreprocessingPipeline:
def __init__(self, config: dict = None):
self.config = config or {}
self.do_grayscale = self.config.get("grayscale", True)
self.do_deskew = self.config.get("deskew", True)
self.do_denoise = self.config.get("denoise", True)
self.do_binarize = self.config.get("binarize", False)
self.target_width = self.config.get("target_width", 1000)
self.target_height = self.config.get("target_height", 1400)
self.pdf_dpi = self.config.get("pdf_dpi", 200)
def pdf_to_image(self, pdf_path: str, page_num: int = 0) -> np.ndarray:
"""Converts the specified page of a PDF file to a NumPy BGR image."""
if HAVE_FITZ:
doc = fitz.open(pdf_path)
if page_num < len(doc):
pix = doc[page_num].get_pixmap(dpi=self.pdf_dpi)
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
return cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
raise RuntimeError("PDF rendering requires PyMuPDF (fitz).")
def deskew(self, image: np.ndarray) -> np.ndarray:
"""Detects skew angle and rotates image to deskew."""
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
blur = cv2.GaussianBlur(gray, (9, 9), 0)
thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
# Use morphological dilation to find text line blocks
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (30, 5))
dilate = cv2.dilate(thresh, kernel, iterations=2)
contours, _ = cv2.findContours(dilate, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
angles = []
for c in contours:
if cv2.contourArea(c) > 1000:
rect = cv2.minAreaRect(c)
angle = rect[-1]
if angle < -45:
angle = -(90 + angle)
elif angle > 45:
angle = 90 - angle
if abs(angle) < 45:
angles.append(angle)
if not angles:
return image
median_angle = float(np.median(angles))
if abs(median_angle) < 0.5:
return image
(h, w) = image.shape[:2]
center = (w // 2, h // 2)
M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
rotated = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
return rotated
def denoise(self, image: np.ndarray) -> np.ndarray:
"""Applies median filter denoising."""
if len(image.shape) == 3:
return cv2.fastNlMeansDenoisingColored(image, None, 5, 5, 7, 21)
return cv2.medianBlur(image, 3)
def binarize(self, gray_image: np.ndarray) -> np.ndarray:
"""Applies adaptive Gaussian thresholding."""
return cv2.adaptiveThreshold(
gray_image, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
)
def resize(self, image: np.ndarray) -> np.ndarray:
"""Resizes image to target dimensions while maintaining aspect ratio."""
h, w = image.shape[:2]
scale = min(self.target_width / w, self.target_height / h)
new_w, new_h = int(w * scale), int(h * scale)
return cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
def process(self, image_or_path) -> np.ndarray:
"""Runs the complete configurable preprocessing pipeline."""
if isinstance(image_or_path, str):
if image_or_path.lower().endswith(".pdf"):
img = self.pdf_to_image(image_or_path)
else:
img = cv2.imread(image_or_path)
elif isinstance(image_or_path, np.ndarray):
img = image_or_path.copy()
elif isinstance(image_or_path, Image.Image):
img = cv2.cvtColor(np.array(image_or_path), cv2.COLOR_RGB2BGR)
else:
raise ValueError(f"Unsupported image input type: {type(image_or_path)}")
if img is None:
raise ValueError(f"Failed to load image from: {image_or_path}")
# 1. Deskew
if self.do_deskew:
img = self.deskew(img)
# 2. Denoise
if self.do_denoise:
img = self.denoise(img)
# 3. Grayscale
if self.do_grayscale and len(img.shape) == 3:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# 4. Adaptive Binarization
if self.do_binarize:
if len(img.shape) == 3:
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
img = self.binarize(img)
# 5. Resize
if self.target_width and self.target_height:
img = self.resize(img)
return img
|