""" 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))