Spaces:
Sleeping
Sleeping
| """ | |
| preprocessor.py — Advanced Computer Vision preprocessing pipeline for documents and live camera photos. | |
| Features: | |
| - EXIF rotation normalization (vital for mobile/camera photos) | |
| - Auto-deskew via Hough line transform | |
| - Background shadow & illumination normalization | |
| - Unsharp masking for camera blur correction | |
| - CLAHE adaptive contrast enhancement | |
| - Bilateral edge-preserving denoising | |
| - Adaptive Otsu binarization | |
| - Perspective warp for angled captures | |
| """ | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image, ImageOps | |
| import math | |
| def fix_exif_orientation(pil_img: Image.Image) -> Image.Image: | |
| """Normalize image orientation using EXIF data (crucial for phone & webcam photos).""" | |
| try: | |
| return ImageOps.exif_transpose(pil_img) | |
| except Exception: | |
| return pil_img | |
| def pil_to_cv2(pil_img: Image.Image) -> np.ndarray: | |
| """Convert PIL Image → OpenCV BGR ndarray.""" | |
| arr = np.array(pil_img.convert("RGB")) | |
| return cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) | |
| def cv2_to_pil(cv2_img: np.ndarray) -> Image.Image: | |
| """Convert OpenCV BGR ndarray → PIL Image.""" | |
| rgb = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB) | |
| return Image.fromarray(rgb) | |
| def remove_shadows_and_normalize(cv2_img: np.ndarray) -> np.ndarray: | |
| """ | |
| Remove uneven background shadows and normalize illumination. | |
| Particularly effective for phone photos with cast shadows or glare. | |
| """ | |
| # Split into color planes | |
| rgb_planes = cv2.split(cv2_img) | |
| result_planes = [] | |
| for plane in rgb_planes: | |
| # Dilate to approximate background | |
| dilated_img = cv2.dilate(plane, np.ones((7, 7), np.uint8)) | |
| bg_img = cv2.medianBlur(dilated_img, 21) | |
| # Difference between plane and estimated background | |
| diff_img = 255 - cv2.absdiff(plane, bg_img) | |
| # Normalize to full dynamic range | |
| norm_img = cv2.normalize( | |
| diff_img, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1 | |
| ) | |
| result_planes.append(norm_img) | |
| return cv2.merge(result_planes) | |
| def sharpen_image(cv2_img: np.ndarray, strength: float = 1.3) -> np.ndarray: | |
| """Apply unsharp mask filter to sharpen slightly blurry camera captures.""" | |
| gaussian = cv2.GaussianBlur(cv2_img, (0, 0), 2.0) | |
| sharpened = cv2.addWeighted(cv2_img, 1.0 + strength, gaussian, -strength, 0) | |
| return sharpened | |
| def adjust_brightness_contrast(cv2_img: np.ndarray, brightness: int = 0, contrast: float = 1.0) -> np.ndarray: | |
| """ | |
| Adjust brightness (-100 to 100) and contrast (0.5 to 3.0). | |
| Allows user to brighten underexposed camera captures. | |
| """ | |
| adjusted = cv2.convertScaleAbs(cv2_img, alpha=contrast, beta=brightness) | |
| return adjusted | |
| def extreme_contrast_text(cv2_img: np.ndarray) -> np.ndarray: | |
| """ | |
| Extract faint, low-contrast ink on paper by computing adaptive local contrast. | |
| Guarantees text characters become distinct and visible. | |
| """ | |
| gray = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2GRAY) | |
| # Estimate background illumination | |
| kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 25)) | |
| bg = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel) | |
| # Background subtraction | |
| diff = 255 - cv2.absdiff(gray, bg) | |
| norm = cv2.normalize(diff, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1) | |
| # Adaptive enhancement | |
| clahe = cv2.createCLAHE(clipLimit=3.5, tileGridSize=(8, 8)) | |
| enhanced = clahe.apply(norm) | |
| return cv2.cvtColor(enhanced, cv2.COLOR_GRAY2BGR) | |
| def auto_deskew(img: np.ndarray) -> np.ndarray: | |
| """Detect text tilt and rotate image to align text horizontally.""" | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| gray = cv2.bitwise_not(gray) | |
| thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[1] | |
| # Find coordinates of white pixels | |
| coords = np.column_stack(np.where(thresh > 0)) | |
| if len(coords) < 10: | |
| return img | |
| # minAreaRect gives angle | |
| angle = cv2.minAreaRect(coords)[-1] | |
| if angle < -45: | |
| angle = -(90 + angle) | |
| else: | |
| angle = -angle | |
| # Only correct if tilt is moderate and significant | |
| if abs(angle) < 0.6 or abs(angle) > 40: | |
| return img | |
| h, w = img.shape[:2] | |
| center = (w // 2, h // 2) | |
| M = cv2.getRotationMatrix2D(center, angle, 1.0) | |
| rotated = cv2.warpAffine( | |
| img, M, (w, h), | |
| flags=cv2.INTER_CUBIC, | |
| borderMode=cv2.BORDER_REPLICATE | |
| ) | |
| return rotated | |
| def enhance_contrast(img: np.ndarray, clip_limit: float = 2.5) -> np.ndarray: | |
| """Apply CLAHE (Contrast Limited Adaptive Histogram Equalization).""" | |
| lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB) | |
| l, a, b = cv2.split(lab) | |
| clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8)) | |
| l_enhanced = clahe.apply(l) | |
| merged = cv2.merge([l_enhanced, a, b]) | |
| return cv2.cvtColor(merged, cv2.COLOR_LAB2BGR) | |
| def denoise(img: np.ndarray) -> np.ndarray: | |
| """Bilateral filter to remove camera noise while preserving text edges.""" | |
| return cv2.bilateralFilter(img, d=7, sigmaColor=50, sigmaSpace=50) | |
| def binarize(img: np.ndarray) -> np.ndarray: | |
| """Adaptive Gaussian thresholding for crisp black/white document text.""" | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| # Background shadow removal prior to binarization | |
| dilated = cv2.dilate(gray, np.ones((7, 7), np.uint8)) | |
| bg = cv2.medianBlur(dilated, 21) | |
| diff = 255 - cv2.absdiff(gray, bg) | |
| norm = cv2.normalize(diff, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1) | |
| binary = cv2.adaptiveThreshold( | |
| norm, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 21, 10 | |
| ) | |
| return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) | |
| def detect_document_corners(img: np.ndarray): | |
| """Detect the 4 corners of a document for perspective warp.""" | |
| gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
| blurred = cv2.GaussianBlur(gray, (5, 5), 0) | |
| edged = cv2.Canny(blurred, 50, 150) | |
| contours, _ = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) | |
| if not contours: | |
| return None | |
| contours = sorted(contours, key=cv2.contourArea, reverse=True)[:5] | |
| for c in contours: | |
| peri = cv2.arcLength(c, True) | |
| approx = cv2.approxPolyDP(c, 0.02 * peri, True) | |
| if len(approx) == 4 and cv2.contourArea(c) > (img.shape[0] * img.shape[1] * 0.15): | |
| return approx | |
| return None | |
| def order_points(pts): | |
| """Order points as: top-left, top-right, bottom-right, bottom-left.""" | |
| rect = np.zeros((4, 2), dtype="float32") | |
| s = pts.sum(axis=1) | |
| rect[0] = pts[np.argmin(s)] | |
| rect[2] = pts[np.argmax(s)] | |
| diff = np.diff(pts, axis=1) | |
| rect[1] = pts[np.argmin(diff)] | |
| rect[3] = pts[np.argmax(diff)] | |
| return rect | |
| def perspective_warp(img: np.ndarray, corners) -> np.ndarray: | |
| """Flatten a perspective-distorted document using 4-point transform.""" | |
| pts = corners.reshape(4, 2).astype("float32") | |
| rect = order_points(pts) | |
| tl, tr, br, bl = rect | |
| width_a = np.linalg.norm(br - bl) | |
| width_b = np.linalg.norm(tr - tl) | |
| max_width = max(int(width_a), int(width_b)) | |
| height_a = np.linalg.norm(tr - br) | |
| height_b = np.linalg.norm(tl - bl) | |
| max_height = max(int(height_a), int(height_b)) | |
| if max_width < 100 or max_height < 100: | |
| return img | |
| dst = np.array([ | |
| [0, 0], | |
| [max_width - 1, 0], | |
| [max_width - 1, max_height - 1], | |
| [0, max_height - 1] | |
| ], dtype="float32") | |
| M = cv2.getPerspectiveTransform(rect, dst) | |
| warped = cv2.warpPerspective(img, M, (max_width, max_height)) | |
| return warped | |
| def full_preprocess(pil_img: Image.Image, mode: str = "auto") -> tuple[Image.Image, dict]: | |
| """ | |
| Comprehensive document preprocessing pipeline. | |
| Modes: | |
| - 'auto': EXIF fix + deskew + shadow removal + CLAHE + sharpening + gentle denoise | |
| - 'camera_enhance': Specialized for live camera / phone photos (shadows, glare, blur compensation) | |
| - 'binarize': High contrast B&W text thresholding | |
| - 'warp': 4-corner perspective flattening | |
| - 'raw': Passthrough with EXIF rotation fixed | |
| """ | |
| # Always normalize EXIF orientation first | |
| fixed_pil = fix_exif_orientation(pil_img) | |
| steps_applied = ["exif_orientation_fixed"] | |
| if mode == "raw": | |
| return fixed_pil, {"steps": ["raw_passthrough"]} | |
| cv2_img = pil_to_cv2(fixed_pil) | |
| # Resolution scaling check: if too small (e.g. low-res webcam), upscale with bicubic | |
| h, w = cv2_img.shape[:2] | |
| if min(h, w) < 700: | |
| scale = 800.0 / min(h, w) | |
| cv2_img = cv2.resize(cv2_img, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_CUBIC) | |
| steps_applied.append("rescale_upscale") | |
| # Step 1: Perspective Warp if explicitly selected | |
| if mode == "warp": | |
| corners = detect_document_corners(cv2_img) | |
| if corners is not None: | |
| cv2_img = perspective_warp(cv2_img, corners) | |
| steps_applied.append("perspective_warp") | |
| else: | |
| steps_applied.append("warp_not_detected") | |
| # Step 2: Auto Deskew | |
| deskewed = auto_deskew(cv2_img) | |
| if not np.array_equal(deskewed, cv2_img): | |
| steps_applied.append("auto_deskew") | |
| cv2_img = deskewed | |
| # Step 3: Camera photo shadow removal & illumination flattening | |
| if mode in ("camera_enhance", "auto"): | |
| cv2_img = remove_shadows_and_normalize(cv2_img) | |
| steps_applied.append("shadow_illumination_norm") | |
| # Step 4: CLAHE Contrast Enhancement | |
| cv2_img = enhance_contrast(cv2_img, clip_limit=3.0 if mode == "camera_enhance" else 2.2) | |
| steps_applied.append("clahe_contrast") | |
| # Step 5: Unsharp masking to sharpen characters | |
| cv2_img = sharpen_image(cv2_img, strength=1.2 if mode == "camera_enhance" else 0.8) | |
| steps_applied.append("unsharp_mask_sharpen") | |
| # Step 6: Denoise | |
| cv2_img = denoise(cv2_img) | |
| steps_applied.append("bilateral_denoise") | |
| # Step 7: Binarization or Extreme Contrast if requested | |
| if mode == "binarize": | |
| cv2_img = binarize(cv2_img) | |
| steps_applied.append("adaptive_binarize") | |
| elif mode == "extreme_contrast": | |
| cv2_img = extreme_contrast_text(cv2_img) | |
| steps_applied.append("extreme_contrast_text") | |
| return cv2_to_pil(cv2_img), {"steps": steps_applied} | |