Spaces:
Sleeping
Sleeping
| """Clinical image preprocessor for medical documents and prescriptions. | |
| Implements a robust preprocessing pipeline using OpenCV (with Pillow fallback): | |
| 1. RGB normalization | |
| 2. Grayscale conversion | |
| 3. Bilateral denoising (edge-preserving — preserves ink strokes) | |
| 4. CLAHE contrast enhancement | |
| 5. Morphological opening (gentle noise removal) | |
| 6. Min/max normalization | |
| Optional: | |
| 7. Deskew (not implemented in PASS 1 — flagged for PASS 2) | |
| Design principles from reference repository (sanketnaik99/doctors-prescription): | |
| - Grayscale + CLAHE + bilateral filter + morphological operations improve OCR | |
| on faded, low-contrast, or paper-grain-noisy prescription images. | |
| - Original bytes are NEVER modified. Returns new bytes. | |
| - Processed bytes are temporary and must not be persisted. | |
| IMPORTANT: This preprocessor does NOT call any OCR or AI model. | |
| It is independently testable with no external dependencies beyond OpenCV/Pillow. | |
| """ | |
| from __future__ import annotations | |
| import io | |
| import logging | |
| from dataclasses import dataclass | |
| import numpy as np | |
| from PIL import Image, ImageEnhance | |
| log = logging.getLogger(__name__) | |
| try: | |
| import cv2 | |
| _OPENCV_AVAILABLE = True | |
| except ImportError: | |
| cv2 = None # type: ignore[assignment] | |
| _OPENCV_AVAILABLE = False | |
| class PreprocessingResult: | |
| """Result of preprocessing a single image.""" | |
| processed_bytes: bytes | |
| mime_type: str # Always "image/png" after processing | |
| operations_applied: tuple[str, ...] # Ordered list of operations performed | |
| estimated_quality_improvement: float # Estimated improvement (0.0 = no change, 1.0 = max) | |
| original_size_bytes: int | |
| processed_size_bytes: int | |
| class ClinicalImagePreprocessor: | |
| """Preprocesses clinical document images for downstream OCR. | |
| Usage: | |
| preprocessor = ClinicalImagePreprocessor() | |
| result = preprocessor.preprocess_image(raw_bytes) | |
| # Use result.processed_bytes for OCR | |
| All operations are applied to an in-memory copy. Original bytes untouched. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| contrast_boost: float = 1.4, | |
| apply_clahe: bool = True, | |
| apply_morphology: bool = True, | |
| clahe_clip_limit: float = 2.0, | |
| clahe_grid_size: tuple[int, int] = (8, 8), | |
| bilateral_d: int = 5, | |
| bilateral_sigma_color: float = 50.0, | |
| bilateral_sigma_space: float = 50.0, | |
| ) -> None: | |
| self._contrast_boost = contrast_boost | |
| self._apply_clahe = apply_clahe | |
| self._apply_morphology = apply_morphology | |
| self._clahe_clip_limit = clahe_clip_limit | |
| self._clahe_grid_size = clahe_grid_size | |
| self._bilateral_d = bilateral_d | |
| self._bilateral_sigma_color = bilateral_sigma_color | |
| self._bilateral_sigma_space = bilateral_sigma_space | |
| def preprocess_image(self, image_bytes: bytes) -> PreprocessingResult: | |
| """Preprocess a clinical document image. | |
| Primary path: OpenCV (grayscale → bilateral → CLAHE → morph → normalize) | |
| Fallback: Pillow contrast enhancement only. | |
| Args: | |
| image_bytes: Raw image bytes (PNG, JPEG, or WebP). | |
| Returns: | |
| PreprocessingResult with processed_bytes in PNG format. | |
| """ | |
| operations: list[str] = [] | |
| original_size = len(image_bytes) | |
| # Open with Pillow first (handles multiple formats) | |
| try: | |
| pil_img = Image.open(io.BytesIO(image_bytes)) | |
| if pil_img.mode != "RGB": | |
| pil_img = pil_img.convert("RGB") | |
| operations.append("convert_to_rgb") | |
| except Exception as exc: | |
| log.warning("Pillow cannot open image: %s", exc) | |
| return PreprocessingResult( | |
| processed_bytes=image_bytes, | |
| mime_type="image/png", | |
| operations_applied=("error_passthrough",), | |
| estimated_quality_improvement=0.0, | |
| original_size_bytes=original_size, | |
| processed_size_bytes=original_size, | |
| ) | |
| np_rgb = np.array(pil_img) | |
| # Estimate pre-processing contrast as baseline | |
| pre_std = float(np.std(np.mean(np_rgb, axis=2))) | |
| if _OPENCV_AVAILABLE and cv2 is not None: | |
| result = self._opencv_pipeline(np_rgb, operations) | |
| if result is not None: | |
| post_std = float(np.std(result)) if result.ndim == 2 else float( | |
| np.std(np.mean(result, axis=2)) | |
| ) | |
| improvement = max(0.0, min(1.0, (post_std - pre_std) / max(1.0, pre_std))) | |
| # Encode back to PNG | |
| try: | |
| if result.ndim == 2: | |
| # Grayscale → back to RGB for downstream model compatibility | |
| rgb_out = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB) | |
| else: | |
| rgb_out = result | |
| success, encoded = cv2.imencode(".png", rgb_out) | |
| if success: | |
| processed_bytes = encoded.tobytes() | |
| return PreprocessingResult( | |
| processed_bytes=processed_bytes, | |
| mime_type="image/png", | |
| operations_applied=tuple(operations), | |
| estimated_quality_improvement=round(improvement, 3), | |
| original_size_bytes=original_size, | |
| processed_size_bytes=len(processed_bytes), | |
| ) | |
| except Exception as enc_err: | |
| log.warning("OpenCV encode failed: %s", enc_err) | |
| # Pillow fallback | |
| return self._pillow_fallback(pil_img, image_bytes, operations, pre_std, original_size) | |
| def _opencv_pipeline( | |
| self, np_rgb: np.ndarray, operations: list[str] | |
| ) -> np.ndarray | None: | |
| """Run the full OpenCV preprocessing pipeline. | |
| Returns the processed image array, or None on failure. | |
| """ | |
| try: | |
| # BGR for OpenCV | |
| cv_bgr = cv2.cvtColor(np_rgb, cv2.COLOR_RGB2BGR) # type: ignore[union-attr] | |
| # Grayscale | |
| gray = cv2.cvtColor(cv_bgr, cv2.COLOR_BGR2GRAY) # type: ignore[union-attr] | |
| operations.append("opencv_grayscale") | |
| # Bilateral denoising — preserves ink edges, smooths paper grain | |
| denoised = cv2.bilateralFilter( # type: ignore[union-attr] | |
| gray, | |
| d=self._bilateral_d, | |
| sigmaColor=self._bilateral_sigma_color, | |
| sigmaSpace=self._bilateral_sigma_space, | |
| ) | |
| operations.append("bilateral_denoise") | |
| # CLAHE — Contrast Limited Adaptive Histogram Equalization | |
| if self._apply_clahe: | |
| clahe = cv2.createCLAHE( # type: ignore[union-attr] | |
| clipLimit=self._clahe_clip_limit, | |
| tileGridSize=self._clahe_grid_size, | |
| ) | |
| enhanced = clahe.apply(denoised) | |
| operations.append("clahe_contrast_enhancement") | |
| else: | |
| enhanced = denoised | |
| # Morphological opening — gentle noise removal without eroding strokes | |
| if self._apply_morphology: | |
| kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (2, 2)) # type: ignore[union-attr] | |
| morphed = cv2.morphologyEx(enhanced, cv2.MORPH_OPEN, kernel) # type: ignore[union-attr] | |
| operations.append("morphological_open_2x2") | |
| else: | |
| morphed = enhanced | |
| # Normalize to full [0, 255] range | |
| normalized = cv2.normalize(morphed, None, 0, 255, cv2.NORM_MINMAX) # type: ignore[union-attr] | |
| operations.append("minmax_normalize") | |
| return normalized # type: ignore[return-value] | |
| except Exception as exc: | |
| log.warning("OpenCV pipeline failed: %s", exc) | |
| return None | |
| def _pillow_fallback( | |
| self, | |
| pil_img: Image.Image, | |
| original_bytes: bytes, | |
| operations: list[str], | |
| pre_std: float, | |
| original_size: int, | |
| ) -> PreprocessingResult: | |
| """Pillow-based fallback: simple contrast enhancement.""" | |
| try: | |
| enhancer = ImageEnhance.Contrast(pil_img) | |
| enhanced = enhancer.enhance(self._contrast_boost) | |
| operations.append(f"pillow_contrast_{self._contrast_boost}x") | |
| buf = io.BytesIO() | |
| enhanced.save(buf, format="PNG") | |
| processed_bytes = buf.getvalue() | |
| np_enhanced = np.array(enhanced.convert("L")) | |
| post_std = float(np.std(np_enhanced)) | |
| improvement = max(0.0, min(1.0, (post_std - pre_std) / max(1.0, pre_std))) | |
| return PreprocessingResult( | |
| processed_bytes=processed_bytes, | |
| mime_type="image/png", | |
| operations_applied=tuple(operations), | |
| estimated_quality_improvement=round(improvement, 3), | |
| original_size_bytes=original_size, | |
| processed_size_bytes=len(processed_bytes), | |
| ) | |
| except Exception as pil_err: | |
| log.warning("Pillow fallback failed: %s", pil_err) | |
| return PreprocessingResult( | |
| processed_bytes=original_bytes, | |
| mime_type="image/png", | |
| operations_applied=("error_passthrough",), | |
| estimated_quality_improvement=0.0, | |
| original_size_bytes=original_size, | |
| processed_size_bytes=original_size, | |
| ) | |
| __all__ = ["ClinicalImagePreprocessor", "PreprocessingResult"] | |