Spaces:
Sleeping
Sleeping
| 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 | |