Spaces:
Runtime error
Runtime error
| """ | |
| predict.py | |
| ---------- | |
| Full inference pipeline: | |
| 1. Image loading -- PNG/JPG/TIF (original) | |
| -- DICOM (.dcm) with Hounsfield Unit windowing [NEW] | |
| -- NIfTI (.nii / .nii.gz) with slice selection [NEW] | |
| 2. Preprocessing -- resize to 256x256, normalize to [0,1] | |
| 3. Segmentation -- Encoder-Decoder / U-Net lung mask | |
| 4. Visualization -- overlay, heatmap | |
| 5. Metrics -- Dice, IoU, Precision, Recall | |
| 6. Nodule detect -- detect_nodules() on masked lung region [NEW] | |
| """ | |
| import os | |
| import numpy as np | |
| from PIL import Image | |
| from skimage.transform import resize | |
| IMG_SIZE = 256 | |
| # =========================================================================== | |
| # CT WINDOWING (clinical standard for Lung Window) | |
| # =========================================================================== | |
| LUNG_WINDOW = dict(center=-600, width=1500) # best for lung parenchyma | |
| MEDIASTINUM_W = dict(center=50, width=400) # soft tissue / vessels | |
| BONE_WINDOW = dict(center=400, width=1800) # ribs / spine | |
| def apply_window(hu_array, center=-600, width=1500): | |
| """ | |
| Convert raw Hounsfield Unit values to [0, 255] uint8 using a clinical window. | |
| Radiologists choose a Window Center (WC) and Window Width (WW) to make | |
| specific tissue types visible: | |
| - Lung Window (WC=-600, WW=1500): lung tissue, nodules, air spaces | |
| - Mediastinum (WC= 50, WW= 400): heart, vessels, lymph nodes | |
| - Bone Window (WC= 400, WW=1800): ribs, vertebrae | |
| Any HU below (WC - WW/2) maps to 0 (black). | |
| Any HU above (WC + WW/2) maps to 255 (white). | |
| Values inside are linearly scaled. | |
| Args: | |
| hu_array : np.ndarray, raw Hounsfield Units (any shape) | |
| center : int, window center in HU | |
| width : int, window width in HU | |
| Returns: | |
| windowed : np.ndarray, same shape, uint8, range [0, 255] | |
| """ | |
| lo = center - width / 2.0 | |
| hi = center + width / 2.0 | |
| clipped = np.clip(hu_array.astype(np.float32), lo, hi) | |
| windowed = ((clipped - lo) / (hi - lo) * 255.0).astype(np.uint8) | |
| return windowed | |
| # =========================================================================== | |
| # FILE LOADERS | |
| # =========================================================================== | |
| def load_dicom(path): | |
| """ | |
| Load a single DICOM slice (.dcm) and return: | |
| - hu_slice : np.ndarray [H, W], float32, Hounsfield Units | |
| - display : np.ndarray [H, W], uint8, lung-windowed grayscale | |
| - metadata : dict with patient info and scan parameters | |
| Hounsfield Unit conversion: | |
| HU = pixel_value * RescaleSlope + RescaleIntercept | |
| This is the standard DICOM linear rescaling defined in the file header. | |
| Raises ImportError if pydicom is not installed. | |
| """ | |
| try: | |
| import pydicom | |
| except ImportError: | |
| raise ImportError( | |
| "pydicom is required for DICOM support.\n" | |
| "Install it with: pip install pydicom" | |
| ) | |
| ds = pydicom.dcmread(path) | |
| pixels = ds.pixel_array.astype(np.float32) | |
| # Apply DICOM rescale to get Hounsfield Units | |
| slope = float(getattr(ds, "RescaleSlope", 1.0)) | |
| intercept = float(getattr(ds, "RescaleIntercept", 0.0)) | |
| hu_slice = pixels * slope + intercept | |
| # Extract clinical metadata | |
| metadata = { | |
| "PatientID" : str(getattr(ds, "PatientID", "Unknown")), | |
| "StudyDate" : str(getattr(ds, "StudyDate", "Unknown")), | |
| "Modality" : str(getattr(ds, "Modality", "CT")), | |
| "SliceThickness" : float(getattr(ds, "SliceThickness", 0.0)), | |
| "PixelSpacing" : list(getattr(ds, "PixelSpacing", [1.0, 1.0])), | |
| "KVP" : float(getattr(ds, "KVP", 120.0)), | |
| "InstanceNumber" : int(getattr(ds, "InstanceNumber", 0)), | |
| "ImageShape" : list(pixels.shape), | |
| } | |
| display = apply_window(hu_slice, **LUNG_WINDOW) | |
| return hu_slice, display, metadata | |
| def load_nifti(path, slice_index=None): | |
| """ | |
| Load a NIfTI volume (.nii or .nii.gz) and return a single axial slice. | |
| NIfTI is the standard format for research datasets (LUNA16, MSD, etc.). | |
| The volume is stored as a 3D array [H, W, D] where D is the depth (number | |
| of axial slices). We extract one slice, apply lung windowing, and return it. | |
| Args: | |
| path : str, path to .nii or .nii.gz file | |
| slice_index : int or None — which axial slice to extract. | |
| If None, uses the middle slice. | |
| Returns: | |
| hu_slice : np.ndarray [H, W], float32 | |
| display : np.ndarray [H, W], uint8, lung-windowed | |
| volume_info : dict with shape and slice count | |
| n_slices : int, total number of slices in the volume | |
| Raises ImportError if nibabel is not installed. | |
| """ | |
| try: | |
| import nibabel as nib | |
| except ImportError: | |
| raise ImportError( | |
| "nibabel is required for NIfTI support.\n" | |
| "Install it with: pip install nibabel" | |
| ) | |
| nii = nib.load(path) | |
| volume = nii.get_fdata(dtype=np.float32) # [H, W, D] or [H, W, D, T] | |
| # Squeeze any time dimension | |
| if volume.ndim == 4: | |
| volume = volume[:, :, :, 0] | |
| n_slices = volume.shape[2] | |
| if slice_index is None: | |
| slice_index = n_slices // 2 | |
| slice_index = max(0, min(slice_index, n_slices - 1)) | |
| hu_slice = volume[:, :, slice_index] | |
| display = apply_window(hu_slice, **LUNG_WINDOW) | |
| volume_info = { | |
| "VolumeShape" : list(volume.shape), | |
| "NSlices" : n_slices, | |
| "SliceIndex" : slice_index, | |
| "VoxelSize" : list(nii.header.get_zooms()[:3]), | |
| "Affine" : nii.affine.tolist(), | |
| } | |
| return hu_slice, display, volume_info, n_slices | |
| def load_simpleitk_volume(path): | |
| """ | |
| Load a .mhd / .raw volumetric CT file (LUNA16 format) using SimpleITK. | |
| SimpleITK reads the full 3D volume and handles all internal coordinate | |
| transforms. The output is a numpy array in [D, H, W] order (depth-first). | |
| Args: | |
| path : str, path to .mhd file (the .raw file must be in the same folder) | |
| Returns: | |
| volume : np.ndarray [D, H, W], float32, Hounsfield Units | |
| spacing : tuple (z_mm, y_mm, x_mm), voxel spacing in millimetres | |
| origin : tuple (x, y, z), world coordinates of the first voxel | |
| Raises ImportError if SimpleITK is not installed. | |
| """ | |
| try: | |
| import SimpleITK as sitk | |
| except ImportError: | |
| raise ImportError( | |
| "SimpleITK is required for .mhd/.raw support.\n" | |
| "Install it with: pip install SimpleITK" | |
| ) | |
| itk_img = sitk.ReadImage(path) | |
| volume = sitk.GetArrayFromImage(itk_img).astype(np.float32) # [D, H, W] | |
| spacing = itk_img.GetSpacing() # (x, y, z) mm | |
| origin = itk_img.GetOrigin() | |
| return volume, spacing, origin | |
| # =========================================================================== | |
| # PREPROCESSING (unchanged for PNG/JPG/TIF — matches training pipeline) | |
| # =========================================================================== | |
| def preprocess_image(image): | |
| """ | |
| Preprocess any image input (PIL, numpy, or already-windowed uint8 array) | |
| into the model input format. | |
| Training pipeline match: | |
| resize(img, (256, 256, 1), preserve_range=True) then img /= 255.0 | |
| Args: | |
| image : PIL Image or np.ndarray (2D grayscale or 3D RGB/RGBA) | |
| Returns: | |
| model_input : np.ndarray [1, 256, 256, 1], float32, range [0, 1] | |
| display : np.ndarray [256, 256], uint8, contrast-stretched | |
| """ | |
| if isinstance(image, Image.Image): | |
| img = np.array(image) | |
| else: | |
| img = np.asarray(image, dtype=np.float32) | |
| # Convert to 2D grayscale | |
| if img.ndim == 3: | |
| if img.shape[2] in (3, 4): | |
| img = (0.299 * img[:, :, 0] | |
| + 0.587 * img[:, :, 1] | |
| + 0.114 * img[:, :, 2]).astype(np.float32) | |
| elif img.shape[2] == 1: | |
| img = img[:, :, 0].astype(np.float32) | |
| elif img.ndim == 2: | |
| img = img.astype(np.float32) | |
| else: | |
| raise ValueError(f"Unsupported image shape: {img.shape}") | |
| # Resize to 256x256 | |
| img_resized = resize( | |
| img, (IMG_SIZE, IMG_SIZE), | |
| preserve_range=True, anti_aliasing=True | |
| ).astype(np.float32) | |
| # Build uint8 display (contrast-stretched) | |
| lo, hi = img_resized.min(), img_resized.max() | |
| if hi > lo: | |
| display = ((img_resized - lo) / (hi - lo) * 255).astype(np.uint8) | |
| else: | |
| display = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8) | |
| # Normalize to [0, 1] | |
| img_norm = img_resized.copy() | |
| if img_norm.max() > 1.0: | |
| img_norm = img_norm / 255.0 | |
| img_norm = np.clip(img_norm, 0.0, 1.0) | |
| model_input = img_norm.reshape(1, IMG_SIZE, IMG_SIZE, 1) | |
| return model_input, display | |
| def preprocess_from_hu(hu_slice, window=None): | |
| """ | |
| Preprocess a Hounsfield Unit slice directly (from DICOM or NIfTI load). | |
| Applies lung window first, then runs standard preprocessing. | |
| Args: | |
| hu_slice : np.ndarray [H, W], float32, Hounsfield Units | |
| window : dict with 'center' and 'width', or None for LUNG_WINDOW | |
| Returns: | |
| model_input : np.ndarray [1, 256, 256, 1], float32 | |
| display : np.ndarray [256, 256], uint8 | |
| """ | |
| if window is None: | |
| window = LUNG_WINDOW | |
| windowed = apply_window(hu_slice, **window) | |
| return preprocess_image(windowed) | |
| # =========================================================================== | |
| # INFERENCE | |
| # =========================================================================== | |
| def validate_ct_image(img_array): | |
| """ | |
| Basic sanity check that the uploaded image looks like a CT scan. | |
| Returns (is_valid: bool, message: str) | |
| Only rejects clearly wrong inputs — blank images or pure white photos. | |
| Real CT images from the web may have white backgrounds, borders, or | |
| different windowing, so we keep this check very permissive. | |
| """ | |
| arr = img_array[0, :, :, 0] # [256, 256] | |
| std_val = float(arr.std()) | |
| bright_frac = float((arr > 0.95).mean()) # fraction of near-white pixels | |
| # Reject completely blank/uniform images (e.g., solid black or solid grey) | |
| if std_val < 0.01: | |
| return False, "Image appears completely blank. Please upload a chest CT scan." | |
| # Reject images that are almost entirely white (e.g., accidental screenshot of browser) | |
| if bright_frac > 0.85: | |
| return False, "Image appears to be mostly white/blank. Please upload a grayscale chest CT scan." | |
| return True, "OK" | |
| def predict_mask(model, img_array, threshold=0.5): | |
| """ | |
| Run segmentation model and return binary + probability lung masks. | |
| Args: | |
| model : loaded Keras segmentation model | |
| img_array : np.ndarray [1, 256, 256, 1], float32 | |
| threshold : float, sigmoid threshold for lung vs background | |
| Returns: | |
| binary_mask : np.ndarray [256, 256], float32, values {0.0, 1.0} | |
| prob_mask : np.ndarray [256, 256], float32, probabilities [0, 1] | |
| """ | |
| raw_pred = model.predict(img_array, verbose=0) | |
| prob_mask = raw_pred[0, :, :, 0] | |
| binary_mask = (prob_mask >= threshold).astype(np.float32) | |
| return binary_mask, prob_mask | |
| # =========================================================================== | |
| # NODULE DETECTION [NEW] | |
| # =========================================================================== | |
| def detect_nodules( | |
| seg_model, | |
| detector_model, | |
| img_array, | |
| binary_mask, | |
| patch_size=32, | |
| stride=16, | |
| confidence_threshold=0.5, | |
| ): | |
| """ | |
| Sliding-window nodule detector that only scans inside the lung mask. | |
| Pipeline: | |
| 1. Use binary_mask to identify the lung region | |
| 2. Slide a patch_size x patch_size window across the lung region only | |
| (skipping background pixels entirely — saves compute) | |
| 3. For each patch, run detector_model to get nodule probability | |
| 4. Apply Non-Maximum Suppression (NMS) to remove duplicate detections | |
| 5. Return a list of candidate nodule locations with confidence scores | |
| This is the standard approach used in clinical CAD systems (e.g., LUNA16 | |
| challenge pipelines: sliding window -> false positive reduction). | |
| Args: | |
| seg_model : Keras segmentation model (used for context — not rerun) | |
| detector_model : Keras nodule classification model | |
| Input: [1, patch_size, patch_size, 1] | |
| Output: [1, 1] float, nodule probability | |
| img_array : np.ndarray [1, 256, 256, 1], preprocessed CT slice | |
| binary_mask : np.ndarray [256, 256], float32, lung mask from seg_model | |
| patch_size : int, size of the sliding window in pixels | |
| stride : int, step size between patches (smaller = more sensitive) | |
| confidence_threshold : float, minimum score to report as a detection | |
| Returns: | |
| detections : list of dicts, each containing: | |
| { | |
| "x" : int, top-left x of bounding box (in 256x256 coords) | |
| "y" : int, top-left y of bounding box | |
| "w" : int, patch width (= patch_size) | |
| "h" : int, patch height (= patch_size) | |
| "cx" : int, center x | |
| "cy" : int, center y | |
| "confidence" : float, nodule probability [0, 1] | |
| "size_mm" : str, estimated diameter category | |
| } | |
| """ | |
| img_2d = img_array[0, :, :, 0] # [256, 256] | |
| H, W = img_2d.shape | |
| candidates = [] | |
| # Scan only within the lung mask bounding box for efficiency | |
| lung_rows = np.where(binary_mask.sum(axis=1) > 0)[0] | |
| lung_cols = np.where(binary_mask.sum(axis=0) > 0)[0] | |
| # ── Fallback: if seg model found no lung mask, scan the whole image ────── | |
| # This happens with synthetic images or scans the 2D model wasn't trained on. | |
| if len(lung_rows) == 0 or len(lung_cols) == 0: | |
| margin = patch_size | |
| row_start, row_end = margin, H - patch_size - margin | |
| col_start, col_end = margin, W - patch_size - margin | |
| use_mask_filter = False # skip the mask-coverage check below | |
| else: | |
| row_start = max(0, int(lung_rows[0])) | |
| row_end = min(H - patch_size, int(lung_rows[-1])) | |
| col_start = max(0, int(lung_cols[0])) | |
| col_end = min(W - patch_size, int(lung_cols[-1])) | |
| use_mask_filter = True | |
| for r in range(row_start, row_end + 1, stride): | |
| for c in range(col_start, col_end + 1, stride): | |
| # Skip patches that are mostly outside the lung mask | |
| if use_mask_filter: | |
| patch_mask = binary_mask[r:r+patch_size, c:c+patch_size] | |
| if patch_mask.mean() < 0.3: # less than 30% lung coverage | |
| continue | |
| # Extract and normalize patch | |
| patch = img_2d[r:r+patch_size, c:c+patch_size].copy() | |
| if patch.shape != (patch_size, patch_size): | |
| continue | |
| patch_input = patch.reshape(1, patch_size, patch_size, 1) | |
| # Run detector | |
| score = float(detector_model.predict(patch_input, verbose=0)[0, 0]) | |
| if score >= confidence_threshold: | |
| # Estimate nodule size from patch area and confidence | |
| if score >= 0.85: | |
| size_cat = "Large (>10mm)" | |
| elif score >= 0.70: | |
| size_cat = "Medium (6-10mm)" | |
| else: | |
| size_cat = "Small (3-6mm)" | |
| candidates.append({ | |
| "x" : c, | |
| "y" : r, | |
| "w" : patch_size, | |
| "h" : patch_size, | |
| "cx" : c + patch_size // 2, | |
| "cy" : r + patch_size // 2, | |
| "confidence" : score, | |
| "size_mm" : size_cat, | |
| }) | |
| # Non-Maximum Suppression: remove overlapping boxes, keep highest-confidence | |
| detections = _nms(candidates, iou_threshold=0.3) | |
| # Sort by confidence descending | |
| detections.sort(key=lambda d: d["confidence"], reverse=True) | |
| return detections | |
| def _nms(candidates, iou_threshold=0.3): | |
| """ | |
| Non-Maximum Suppression for bounding box deduplication. | |
| Without NMS, the sliding window generates many overlapping boxes for the | |
| same nodule. NMS keeps only the highest-confidence box from each cluster. | |
| Algorithm: | |
| 1. Sort by confidence (highest first) | |
| 2. For each box, remove all remaining boxes with IoU > threshold | |
| 3. Repeat until all boxes are processed | |
| Args: | |
| candidates : list of detection dicts (must contain x, y, w, h, confidence) | |
| iou_threshold : float, boxes with IoU > this are considered duplicates | |
| Returns: | |
| kept : list of surviving detection dicts | |
| """ | |
| if not candidates: | |
| return [] | |
| candidates = sorted(candidates, key=lambda d: d["confidence"], reverse=True) | |
| kept = [] | |
| while candidates: | |
| best = candidates.pop(0) | |
| kept.append(best) | |
| remaining = [] | |
| for cand in candidates: | |
| if _box_iou(best, cand) < iou_threshold: | |
| remaining.append(cand) | |
| candidates = remaining | |
| return kept | |
| def _box_iou(a, b): | |
| """Compute IoU between two bounding box dicts (x, y, w, h).""" | |
| ax1, ay1 = a["x"], a["y"] | |
| ax2, ay2 = ax1 + a["w"], ay1 + a["h"] | |
| bx1, by1 = b["x"], b["y"] | |
| bx2, by2 = bx1 + b["w"], by1 + b["h"] | |
| inter_x1 = max(ax1, bx1) | |
| inter_y1 = max(ay1, by1) | |
| inter_x2 = min(ax2, bx2) | |
| inter_y2 = min(ay2, by2) | |
| inter_w = max(0, inter_x2 - inter_x1) | |
| inter_h = max(0, inter_y2 - inter_y1) | |
| inter = inter_w * inter_h | |
| area_a = a["w"] * a["h"] | |
| area_b = b["w"] * b["h"] | |
| union = area_a + area_b - inter | |
| return inter / union if union > 0 else 0.0 | |
| def draw_detections(rgb_image, detections, scale=1.0): | |
| """ | |
| Draw bounding boxes and confidence labels on an RGB image. | |
| Color coding: | |
| - Red (confidence >= 0.85): High-confidence, likely pathological | |
| - Orange(confidence >= 0.70): Medium-confidence, needs follow-up | |
| - Yellow(confidence < 0.70): Low-confidence, possible false positive | |
| Args: | |
| rgb_image : np.ndarray [256, 256, 3], uint8 base image | |
| detections : list of detection dicts from detect_nodules() | |
| scale : float, scale factor if displaying at larger than 256px | |
| Returns: | |
| annotated : np.ndarray [256, 256, 3], uint8 with boxes drawn | |
| """ | |
| annotated = rgb_image.copy().astype(np.uint8) | |
| for det in detections: | |
| conf = det["confidence"] | |
| x, y, w, h = det["x"], det["y"], det["w"], det["h"] | |
| # Color by confidence | |
| if conf >= 0.85: | |
| color = (220, 50, 50) # Red — high confidence | |
| elif conf >= 0.70: | |
| color = (255, 140, 0) # Orange — medium | |
| else: | |
| color = (255, 220, 0) # Yellow — low | |
| # Draw rectangle (2px border) | |
| for thick in range(2): | |
| t = thick | |
| # Top edge | |
| annotated[max(0,y-t):y-t+1, x:x+w, :] = color | |
| # Bottom edge | |
| annotated[y+h+t:y+h+t+1, x:x+w, :] = color | |
| # Left edge | |
| annotated[y:y+h, max(0,x-t):x-t+1, :] = color | |
| # Right edge | |
| annotated[y:y+h, x+w+t:x+w+t+1, :] = color | |
| # Draw small center dot | |
| cx, cy = det["cx"], det["cy"] | |
| for dr in range(-2, 3): | |
| for dc in range(-2, 3): | |
| nr, nc = cy + dr, cx + dc | |
| if 0 <= nr < 256 and 0 <= nc < 256: | |
| annotated[nr, nc, :] = color | |
| return annotated | |
| # =========================================================================== | |
| # OVERLAY GENERATION | |
| # =========================================================================== | |
| def create_overlay(original_gray, binary_mask, alpha=0.5, color=(50, 220, 100)): | |
| """Blend coloured lung mask onto grayscale CT scan.""" | |
| base = np.stack([original_gray]*3, axis=-1).astype(np.float32) | |
| mask_rgb = np.zeros((IMG_SIZE, IMG_SIZE, 3), dtype=np.float32) | |
| mask_rgb[:,:,0] = binary_mask * color[0] | |
| mask_rgb[:,:,1] = binary_mask * color[1] | |
| mask_rgb[:,:,2] = binary_mask * color[2] | |
| mask_3ch = np.stack([binary_mask]*3, axis=-1) | |
| overlay = np.where(mask_3ch > 0, (1-alpha)*base + alpha*mask_rgb, base) | |
| return np.clip(overlay, 0, 255).astype(np.uint8) | |
| def create_heatmap_overlay(original_gray, prob_mask, alpha=0.6): | |
| """Create probability heatmap overlay (green=low, red=high probability).""" | |
| import matplotlib.cm as cm | |
| colormap = cm.get_cmap("RdYlGn_r") | |
| heatmap_rgb = (colormap(prob_mask)[:,:,:3] * 255).astype(np.float32) | |
| base = np.stack([original_gray]*3, axis=-1).astype(np.float32) | |
| overlay = (1-alpha)*base + alpha*heatmap_rgb | |
| return np.clip(overlay, 0, 255).astype(np.uint8) | |
| # =========================================================================== | |
| # METRICS | |
| # =========================================================================== | |
| def calc_dice(y_true, y_pred, smooth=1.0): | |
| """Dice Similarity Coefficient. DSC = 2|A∩B| / (|A|+|B|). Range [0,1].""" | |
| y_true_f = y_true.flatten().astype(np.float32) | |
| y_pred_f = y_pred.flatten().astype(np.float32) | |
| intersection = np.sum(y_true_f * y_pred_f) | |
| return float((2.0*intersection + smooth) / (np.sum(y_true_f) + np.sum(y_pred_f) + smooth)) | |
| def calc_iou(y_true, y_pred, smooth=1.0): | |
| """Intersection over Union (Jaccard). IoU = |A∩B| / |A∪B|. Range [0,1].""" | |
| y_true_f = y_true.flatten().astype(np.float32) | |
| y_pred_f = y_pred.flatten().astype(np.float32) | |
| intersection = np.sum(y_true_f * y_pred_f) | |
| union = np.sum(y_true_f) + np.sum(y_pred_f) - intersection | |
| return float((intersection + smooth) / (union + smooth)) | |
| def calc_precision(y_true, y_pred, smooth=1.0): | |
| """Precision = TP / (TP + FP)""" | |
| y_true_f = y_true.flatten().astype(np.float32) | |
| y_pred_f = y_pred.flatten().astype(np.float32) | |
| tp = np.sum(y_true_f * y_pred_f) | |
| fp = np.sum((1 - y_true_f) * y_pred_f) | |
| return float((tp + smooth) / (tp + fp + smooth)) | |
| def calc_recall(y_true, y_pred, smooth=1.0): | |
| """Recall (Sensitivity) = TP / (TP + FN)""" | |
| y_true_f = y_true.flatten().astype(np.float32) | |
| y_pred_f = y_pred.flatten().astype(np.float32) | |
| tp = np.sum(y_true_f * y_pred_f) | |
| fn = np.sum(y_true_f * (1 - y_pred_f)) | |
| return float((tp + smooth) / (tp + fn + smooth)) | |