Spaces:
Runtime error
Runtime error
| """ | |
| preprocessing.py | |
| ---------------- | |
| Full radiology image preprocessing pipeline. | |
| Covers ALL modalities: X-ray, CT, MRI, Ultrasound, PET, Mammography, | |
| Fluoroscopy, Nuclear Medicine, DEXA, Angiography. | |
| Supports DICOM (.dcm), PNG, JPG. | |
| """ | |
| import io | |
| import logging | |
| from pathlib import Path | |
| from typing import Optional | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| logger = logging.getLogger(__name__) | |
| class RadiologyPreprocessor: | |
| TARGET_SIZE = (224, 224) | |
| # ── Window presets for every modality + region combination ────────── | |
| WINDOW_PRESETS = { | |
| # X-RAY | |
| "xray": {"wc": 40, "ww": 400}, | |
| "xray_chest": {"wc": -600, "ww": 1500}, | |
| "xray_abdomen": {"wc": 60, "ww": 400}, | |
| "xray_pelvis": {"wc": 40, "ww": 400}, | |
| "xray_spine": {"wc": 40, "ww": 400}, | |
| "xray_knee": {"wc": 40, "ww": 400}, | |
| "xray_shoulder": {"wc": 40, "ww": 400}, | |
| "xray_hand": {"wc": 40, "ww": 350}, | |
| "xray_foot": {"wc": 40, "ww": 350}, | |
| "xray_wrist": {"wc": 40, "ww": 350}, | |
| "xray_ankle": {"wc": 40, "ww": 350}, | |
| "xray_hip": {"wc": 40, "ww": 400}, | |
| "xray_skull": {"wc": 35, "ww": 80}, | |
| "xray_neck": {"wc": 40, "ww": 300}, | |
| # CT | |
| "ct": {"wc": 40, "ww": 400}, | |
| "ct_chest": {"wc": -600, "ww": 1500}, | |
| "ct_abdomen": {"wc": 50, "ww": 350}, | |
| "ct_pelvis": {"wc": 50, "ww": 350}, | |
| "ct_brain": {"wc": 35, "ww": 80}, | |
| "ct_brain_bone": {"wc": 400, "ww": 2000}, | |
| "ct_brain_subdural": {"wc": 75, "ww": 215}, | |
| "ct_spine": {"wc": 400, "ww": 1800}, | |
| "ct_spine_soft": {"wc": 50, "ww": 350}, | |
| "ct_liver": {"wc": 60, "ww": 160}, | |
| "ct_lung": {"wc": -600, "ww": 1500}, | |
| "ct_mediastinum": {"wc": 50, "ww": 350}, | |
| "ct_bone": {"wc": 400, "ww": 1800}, | |
| "ct_angiography": {"wc": 300, "ww": 600}, | |
| "ct_cardiac": {"wc": 50, "ww": 350}, | |
| "ct_neck": {"wc": 50, "ww": 350}, | |
| "ct_kidney": {"wc": 50, "ww": 350}, | |
| "ct_pancreas": {"wc": 50, "ww": 250}, | |
| "ct_adrenal": {"wc": 35, "ww": 150}, | |
| # MRI | |
| "mri": {"wc": 128, "ww": 256}, | |
| "mri_brain": {"wc": 128, "ww": 256}, | |
| "mri_spine": {"wc": 128, "ww": 256}, | |
| "mri_knee": {"wc": 128, "ww": 256}, | |
| "mri_shoulder": {"wc": 128, "ww": 256}, | |
| "mri_abdomen": {"wc": 128, "ww": 256}, | |
| "mri_pelvis": {"wc": 128, "ww": 256}, | |
| "mri_cardiac": {"wc": 128, "ww": 256}, | |
| "mri_breast": {"wc": 128, "ww": 256}, | |
| "mri_liver": {"wc": 128, "ww": 256}, | |
| "mri_prostate": {"wc": 128, "ww": 256}, | |
| "mri_neck": {"wc": 128, "ww": 256}, | |
| "mri_hip": {"wc": 128, "ww": 256}, | |
| "mri_ankle": {"wc": 128, "ww": 256}, | |
| "mri_wrist": {"wc": 128, "ww": 256}, | |
| "mri_fetal": {"wc": 128, "ww": 256}, | |
| # ULTRASOUND | |
| "ultrasound": {"wc": 128, "ww": 256}, | |
| "ultrasound_abdomen": {"wc": 128, "ww": 256}, | |
| "ultrasound_pelvis": {"wc": 128, "ww": 256}, | |
| "ultrasound_thyroid": {"wc": 128, "ww": 256}, | |
| "ultrasound_breast": {"wc": 128, "ww": 256}, | |
| "ultrasound_cardiac": {"wc": 128, "ww": 256}, | |
| "ultrasound_obstetric": {"wc": 128, "ww": 256}, | |
| "ultrasound_vascular": {"wc": 128, "ww": 256}, | |
| "ultrasound_neck": {"wc": 128, "ww": 256}, | |
| "ultrasound_liver": {"wc": 128, "ww": 256}, | |
| "ultrasound_kidney": {"wc": 128, "ww": 256}, | |
| "ultrasound_testicular": {"wc": 128, "ww": 256}, | |
| "ultrasound_neonatal": {"wc": 128, "ww": 256}, | |
| # PET / NUCLEAR MEDICINE | |
| "pet": {"wc": 128, "ww": 256}, | |
| "pet_brain": {"wc": 128, "ww": 256}, | |
| "pet_body": {"wc": 128, "ww": 256}, | |
| "nuclear medicine": {"wc": 128, "ww": 256}, | |
| "spect": {"wc": 128, "ww": 256}, | |
| # MAMMOGRAPHY | |
| "mammography": {"wc": 1500, "ww": 6000}, | |
| "mammography_breast": {"wc": 1500, "ww": 6000}, | |
| # FLUOROSCOPY | |
| "fluoroscopy": {"wc": 40, "ww": 400}, | |
| "fluoroscopy_gi": {"wc": 40, "ww": 400}, | |
| # ANGIOGRAPHY | |
| "angiography": {"wc": 300, "ww": 600}, | |
| "angiography_cardiac": {"wc": 300, "ww": 600}, | |
| "angiography_cerebral": {"wc": 300, "ww": 600}, | |
| # DEXA (Bone Density) | |
| "dexa": {"wc": 128, "ww": 256}, | |
| # Default fallback | |
| "default": {"wc": 127, "ww": 255}, | |
| } | |
| # ── Modality-specific enhancement settings ─────────────────────────── | |
| ENHANCEMENT_SETTINGS = { | |
| "xray": {"clahe": True, "clahe_clip": 3.0, "denoise": False}, | |
| "ct": {"clahe": False, "clahe_clip": 2.0, "denoise": False}, | |
| "mri": {"clahe": True, "clahe_clip": 2.0, "denoise": True}, | |
| "ultrasound": {"clahe": True, "clahe_clip": 4.0, "denoise": True}, | |
| "pet": {"clahe": False, "clahe_clip": 2.0, "denoise": False}, | |
| "mammography": {"clahe": True, "clahe_clip": 2.5, "denoise": True}, | |
| "fluoroscopy": {"clahe": True, "clahe_clip": 3.0, "denoise": False}, | |
| "angiography": {"clahe": True, "clahe_clip": 2.0, "denoise": False}, | |
| "dexa": {"clahe": False, "clahe_clip": 2.0, "denoise": False}, | |
| "default": {"clahe": True, "clahe_clip": 2.0, "denoise": False}, | |
| } | |
| def __init__(self, modality: str = "xray", body_region: str = "chest"): | |
| self.modality = modality.lower().replace("-", " ").replace("_", " ") | |
| self.body_region = body_region.lower().replace("-", " ") | |
| self._window_preset = self._select_window_preset() | |
| self._enhancement = self._select_enhancement() | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def preprocess(self, file_bytes: bytes, filename: str) -> np.ndarray: | |
| """Returns float32 numpy array (3, 224, 224) normalised to [0, 1].""" | |
| ext = Path(filename).suffix.lower() | |
| if ext == ".dcm": | |
| image = self._load_dicom(file_bytes) | |
| elif ext in {".png", ".jpg", ".jpeg"}: | |
| image = self._load_standard_image(file_bytes) | |
| else: | |
| raise ValueError(f"Unsupported format: {ext}. Use .dcm, .png, or .jpg") | |
| image = self._apply_modality_normalisation(image) | |
| image = self._apply_enhancement(image) | |
| image = self._resize(image) | |
| return self._to_tensor(image) | |
| def preprocess_pil(self, file_bytes: bytes, filename: str) -> Image.Image: | |
| """Returns a PIL Image for model input.""" | |
| tensor = self.preprocess(file_bytes, filename) | |
| arr = (tensor.transpose(1, 2, 0) * 255).astype(np.uint8) | |
| return Image.fromarray(arr) | |
| # ------------------------------------------------------------------ | |
| # Private | |
| # ------------------------------------------------------------------ | |
| def _select_window_preset(self) -> dict: | |
| # Try modality_region key first | |
| mod_key = self.modality.replace(" ", "_") | |
| reg_key = self.body_region.replace(" ", "_") | |
| combined = f"{mod_key}_{reg_key}" | |
| if combined in self.WINDOW_PRESETS: | |
| return self.WINDOW_PRESETS[combined] | |
| if mod_key in self.WINDOW_PRESETS: | |
| return self.WINDOW_PRESETS[mod_key] | |
| return self.WINDOW_PRESETS["default"] | |
| def _select_enhancement(self) -> dict: | |
| mod_key = self.modality.replace(" ", "_").split("_")[0] | |
| return self.ENHANCEMENT_SETTINGS.get( | |
| mod_key, self.ENHANCEMENT_SETTINGS["default"] | |
| ) | |
| def _load_dicom(self, file_bytes: bytes) -> np.ndarray: | |
| try: | |
| import pydicom | |
| except ImportError: | |
| raise RuntimeError("Install pydicom: pip install pydicom") | |
| ds = pydicom.dcmread(io.BytesIO(file_bytes)) | |
| pixel_array = ds.pixel_array.astype(np.float32) | |
| slope = float(getattr(ds, "RescaleSlope", 1)) | |
| intercept = float(getattr(ds, "RescaleIntercept", 0)) | |
| pixel_array = pixel_array * slope + intercept | |
| # Try DICOM tags first, fall back to preset | |
| wc = float(getattr(ds, "WindowCenter", | |
| self._window_preset["wc"])) | |
| ww = float(getattr(ds, "WindowWidth", | |
| self._window_preset["ww"])) | |
| try: | |
| import pydicom.multival | |
| if isinstance(wc, pydicom.multival.MultiValue): | |
| wc = float(wc[0]) | |
| if isinstance(ww, pydicom.multival.MultiValue): | |
| ww = float(ww[0]) | |
| except Exception: | |
| pass | |
| return self._apply_windowing(pixel_array, wc, ww) | |
| def _load_standard_image(self, file_bytes: bytes) -> np.ndarray: | |
| arr = np.frombuffer(file_bytes, dtype=np.uint8) | |
| img = cv2.imdecode(arr, cv2.IMREAD_GRAYSCALE) | |
| if img is None: | |
| raise ValueError("Could not decode image.") | |
| return img.astype(np.float32) | |
| def _apply_windowing(pixel_array, wc, ww) -> np.ndarray: | |
| lower = wc - ww / 2 | |
| upper = wc + ww / 2 | |
| windowed = np.clip(pixel_array, lower, upper) | |
| windowed = (windowed - lower) / (ww + 1e-8) | |
| return (windowed * 255).astype(np.float32) | |
| def _apply_modality_normalisation(self, image: np.ndarray) -> np.ndarray: | |
| p2, p98 = np.percentile(image, [2, 98]) | |
| image = np.clip(image, p2, p98) | |
| image = (image - p2) / (p98 - p2 + 1e-8) * 255.0 | |
| return image.astype(np.float32) | |
| def _apply_enhancement(self, image: np.ndarray) -> np.ndarray: | |
| """Apply CLAHE and/or denoising based on modality.""" | |
| img_uint8 = np.clip(image, 0, 255).astype(np.uint8) | |
| if self._enhancement.get("denoise"): | |
| img_uint8 = cv2.fastNlMeansDenoising(img_uint8, h=10) | |
| if self._enhancement.get("clahe"): | |
| clip = self._enhancement.get("clahe_clip", 2.0) | |
| clahe = cv2.createCLAHE(clipLimit=clip, tileGridSize=(8, 8)) | |
| img_uint8 = clahe.apply(img_uint8) | |
| return img_uint8.astype(np.float32) | |
| def _resize(self, image: np.ndarray) -> np.ndarray: | |
| h, w = self.TARGET_SIZE | |
| return cv2.resize(image, (w, h), interpolation=cv2.INTER_LANCZOS4) | |
| def _to_tensor(self, image: np.ndarray) -> np.ndarray: | |
| normalised = image / 255.0 | |
| return np.stack([normalised, normalised, normalised], axis=0).astype(np.float32) | |