HarshithReddy01/srmamamba-liver-segmentation
1
1import os2import numpy as np3import torch4import nibabel as nib5from monai import transforms6from scipy import ndimage7from scipy.ndimage import binary_closing, binary_opening, binary_fill_holes, median_filter8from skimage.morphology import ball9from datetime import datetime10 11def load_nifti_ras(file_path):12 """13 Load NIfTI and reorient to RAS so orientation matches preprocess_nifti output.14 Returns (data_ras, affine_ras, spacing) for use as original_data and saving.15 """16 ras_transform = transforms.Compose([17 transforms.LoadImaged(keys=["image"]),18 transforms.Orientationd(keys=["image"], axcodes="RAS"),19 ])20 data = {"image": file_path}21 out = ras_transform(data)22 img = out["image"]23 if hasattr(img, "numpy"):24 data_ras = img.numpy().astype(np.float32)25 elif hasattr(img, "cpu"):26 data_ras = img.cpu().numpy().astype(np.float32)27 else:28 data_ras = np.asarray(img, dtype=np.float32)29 affine_ras = None30 spacing = (1.0, 1.0, 1.0)31 if hasattr(img, "meta") and img.meta is not None:32 affine_ras = img.meta.get("affine")33 if affine_ras is not None:34 affine_ras = np.array(affine_ras)35 pixdim = img.meta.get("pixdim")36 if pixdim is not None and len(pixdim) >= 4:37 spacing = (float(pixdim[1]), float(pixdim[2]), float(pixdim[3]))38 if affine_ras is None:39 nifti_img = nib.load(file_path)40 ras_nifti = nib.as_closest_canonical(nifti_img)41 affine_ras = np.array(ras_nifti.affine)42 spacing = ras_nifti.header.get_zooms()[:3]43 if len(spacing) < 3:44 spacing = (1.0, 1.0, 1.0)45 # Ensure 3D (D, H, W) to match model output spatial shape46 if data_ras.ndim == 4:47 data_ras = data_ras.squeeze(0)48 return data_ras, affine_ras, spacing49 50 51def validate_nifti(nifti_img):52 shape = nifti_img.shape53 if len(shape) < 3:54 raise ValueError(f"Invalid NIfTI shape: {shape}. Expected at least 3 dimensions.")55 if any(s <= 0 for s in shape):56 raise ValueError(f"Invalid NIfTI shape: {shape}. All dimensions must be positive.")57 if any(s > 2000 for s in shape):58 raise ValueError(f"Volume too large: {shape}. Maximum dimension size is 2000.")59 60 voxel_spacing = nifti_img.header.get_zooms()[:3] if len(nifti_img.header.get_zooms()) >= 3 else (1.0, 1.0, 1.0)61 if any(sp <= 0 for sp in voxel_spacing):62 raise ValueError(f"Invalid voxel spacing: {voxel_spacing}. All values must be positive.")63 64 raw_data = nifti_img.get_fdata()65 if np.isnan(raw_data).any():66 raise ValueError("NIfTI contains NaN values")67 if np.isinf(raw_data).any():68 raise ValueError("NIfTI contains Inf values")69 70 return True71 72def preprocess_nifti(file_path, device=None):73 try:74 print(f"Preprocessing file: {file_path}")75 if not os.path.exists(file_path):76 raise FileNotFoundError(f"File not found: {file_path}")77 78 file_size = os.path.getsize(file_path) / (1024**2)79 file_size_kb = os.path.getsize(file_path) / 102480 if file_size == 0:81 raise ValueError("NIfTI file is empty")82 if file_size > 2000:83 raise ValueError(f"NIfTI file too large: {file_size:.1f} MB. Maximum processing size is 2 GB. For larger files, consider compression or resampling.")84 85 if file_size_kb < 100:86 print(f" ⚠ WARNING: File size is very small ({file_size_kb:.1f} KB). This may indicate:")87 print(f" - Low resolution/compressed data (may lose texture and boundary cues)")88 print(f" - Single slice upload (incomplete anatomy)")89 print(f" - Data compression artifacts (may distort intensity gradients)")90 91 print(f"Loading NIfTI file with nibabel...")92 if file_size > 100:93 nifti_img = nib.load(file_path, mmap=True)94 else:95 nifti_img = nib.load(file_path)96 print(f"NIfTI shape: {nifti_img.shape}, dtype: {nifti_img.get_fdata().dtype}")97 98 if len(nifti_img.shape) == 3:99 if any(s < 10 for s in nifti_img.shape):100 print(f" ⚠ WARNING: Very small dimension detected ({nifti_img.shape}). May be a single slice or cropped volume.")101 if nifti_img.shape[2] < 20:102 print(f" ⚠ WARNING: Only {nifti_img.shape[2]} slices detected. Model expects full 3D volumes for best results.")103 104 validate_nifti(nifti_img)105 106 voxel_spacing = nifti_img.header.get_zooms()[:3] if len(nifti_img.header.get_zooms()) >= 3 else (1.0, 1.0, 1.0)107 if voxel_spacing == (1.0, 1.0, 1.0):108 print(f" ⚠ WARNING: Voxel spacing is (1.0, 1.0, 1.0) - metadata may be missing or lost during conversion.")109 print(f" This can cause incorrect volume calculations and scaling issues.")110 111 affine = nifti_img.affine112 affine_det = np.linalg.det(affine[:3, :3])113 print(f" → Voxel spacing: {voxel_spacing}")114 print(f" → Affine determinant: {affine_det:.6f}")115 116 if abs(affine_det) < 0.1 or abs(affine_det) > 100:117 print(f" ⚠ WARNING: Unusual affine determinant ({affine_det:.6f}). Spatial metadata may be corrupted.")118 119 raw_data_dtype = nifti_img.get_fdata().dtype120 if raw_data_dtype == np.uint8 or raw_data_dtype == np.uint16:121 print(f" ⚠ WARNING: Input data type is {raw_data_dtype} (integer). Model expects float32.")122 print(f" Integer data may indicate compression or conversion artifacts.")123 print(f" Converting to float32, but quality may be reduced.")124 125 raw_data = nifti_img.get_fdata(dtype=np.float32)126 print(f" → Raw data stats: min={raw_data.min():.4f}, max={raw_data.max():.4f}, mean={raw_data.mean():.4f}, std={raw_data.std():.4f}")127 128 if raw_data.max() - raw_data.min() < 1e-6:129 raise ValueError(f"Input NIfTI file contains constant values (min=max={raw_data.min():.4f}). Cannot process.")130 131 if raw_data.std() < 1e-3:132 print(f" ⚠ WARNING: Very low data variance (std={raw_data.std():.4f}). Data may be corrupted or over-compressed.")133 134 if raw_data.max() > 10000 or raw_data.min() < -1000:135 print(f" ⚠ WARNING: Extreme intensity values detected (range: [{raw_data.min():.1f}, {raw_data.max():.1f}]).")136 print(f" Data may not be properly normalized. Model expects normalized float32 tensors.")137 138 nonzero_mask = raw_data > 1e-6139 nonzero_count = nonzero_mask.sum()140 total_count = raw_data.size141 nonzero_ratio = nonzero_count / total_count if total_count > 0 else 0.0142 143 print(f" → Non-zero voxels: {nonzero_count:,} / {total_count:,} ({100*nonzero_ratio:.2f}%)")144 145 is_prenormalized = (raw_data.max() <= 1.0 and raw_data.min() >= 0.0)146 if is_prenormalized:147 print(f" → Detected pre-normalized data (range [0, 1]). Still applying training-matched NormalizeIntensityd.")148 149 use_enhanced_preprocessing = os.environ.get("USE_ENHANCED_PREPROCESSING", "false").lower() == "true"150 151 if use_enhanced_preprocessing:152 try:153 transform = transforms.Compose([154 transforms.LoadImaged(keys=["image"]),155 transforms.Orientationd(keys=["image"], axcodes="RAS"),156 transforms.Spacingd(keys=["image"], pixdim=(1.5, 1.5, 3.0), mode="bilinear"),157 transforms.EnsureChannelFirstD(keys=["image"], channel_dim="no_channel"),158 transforms.ScaleIntensityRangePercentilesd(keys="image", lower=2, upper=98, b_min=0.0, b_max=1.0, clip=True),159 transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),160 transforms.ToTensord(keys=["image"])161 ])162 print(" → Using enhanced preprocessing (orientation + spacing + percentile scaling)")163 except Exception as e:164 print(f" ⚠ Warning: Could not create enhanced transform pipeline: {e}. Falling back to training-matched preprocessing...")165 use_enhanced_preprocessing = False166 167 if not use_enhanced_preprocessing:168 # Training-matched: LoadImaged, Orientationd(RAS), EnsureChannelFirstD, NormalizeIntensityd, ToTensord169 # Same as SRMA-Mamba test.py but with RAS for consistent orientation with mask saving.170 transform = transforms.Compose([171 transforms.LoadImaged(keys=["image"]),172 transforms.Orientationd(keys=["image"], axcodes="RAS"),173 transforms.EnsureChannelFirstD(keys=["image"], channel_dim="no_channel"),174 transforms.NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),175 transforms.ToTensord(keys=["image"])176 ])177 print(" → Using training-matched preprocessing (RAS + NormalizeIntensityd)")178 179 data = {"image": file_path}180 print("Applying transforms...")181 182 try:183 augmented = transform(data)184 image_data = augmented["image"]185 except Exception as e:186 print(f" ⚠ Transform failed: {e}. Trying fallback preprocessing...")187 try:188 raw_data_norm = (raw_data - raw_data.min()) / (raw_data.max() - raw_data.min() + 1e-8)189 if raw_data_norm.std() < 1e-6:190 raise ValueError("Normalized data is still constant")191 image_data = torch.from_numpy(raw_data_norm).float()192 image_data = image_data.unsqueeze(0)193 print(" → Used fallback normalization (min-max scaling)")194 except Exception as e2:195 raise ValueError(f"Both standard and fallback preprocessing failed: {e2}")196 197 if not isinstance(image_data, torch.Tensor):198 image_data = torch.from_numpy(np.array(image_data))199 200 if image_data.dtype != torch.float32:201 image_data = image_data.float()202 203 img_np = image_data.numpy() if not hasattr(image_data, 'device') or image_data.device.type == 'cpu' else image_data.cpu().numpy()204 vmin, vmax = float(img_np.min()), float(img_np.max())205 206 if vmax - vmin < 1e-6:207 print(f" ⚠ WARNING: Preprocessing produced near-constant image (min={vmin:.6f}, max={vmax:.6f}). Trying alternative preprocessing...")208 try:209 if nonzero_ratio > 0.01:210 nonzero_mean = raw_data[nonzero_mask].mean()211 nonzero_std = raw_data[nonzero_mask].std() + 1e-8212 raw_data_norm = np.zeros_like(raw_data)213 raw_data_norm[nonzero_mask] = (raw_data[nonzero_mask] - nonzero_mean) / nonzero_std214 raw_data_norm = (raw_data_norm - raw_data_norm.min()) / (raw_data_norm.max() - raw_data_norm.min() + 1e-8)215 else:216 raw_data_norm = (raw_data - raw_data.min()) / (raw_data.max() - raw_data.min() + 1e-8)217 218 if raw_data_norm.std() < 1e-6:219 raise ValueError("Alternative normalization also produced constant data")220 221 image_data = torch.from_numpy(raw_data_norm).float()222 image_data = image_data.unsqueeze(0)223 img_np = image_data.numpy()224 vmin, vmax = float(img_np.min()), float(img_np.max())225 print(f" → Alternative preprocessing successful: min={vmin:.4f}, max={vmax:.4f}, mean={img_np.mean():.4f}, std={img_np.std():.4f}")226 except Exception as e3:227 raise ValueError(f"Preprocessing produced near-constant image: min={vmin:.6f}, max={vmax:.6f}. Alternative preprocessing also failed: {e3}")228 229 print(f" → After transforms: min={vmin:.4f}, max={vmax:.4f}, mean={img_np.mean():.4f}, std={img_np.std():.4f}")230 231 if device is not None and device.type == 'cuda':232 if image_data.is_pinned():233 image_data = image_data.to(device, non_blocking=True)234 else:235 image_data = image_data.pin_memory().to(device, non_blocking=True)236 237 if len(image_data.shape) >= 4:238 try:239 if hasattr(torch, "channels_last_3d"):240 image_data = image_data.contiguous(memory_format=torch.channels_last_3d)241 if image_data.is_contiguous(memory_format=torch.channels_last_3d):242 print(f" → Using channels-last 3D memory layout (optimized for GPU)")243 except:244 pass245 246 print(f"Preprocessed shape: {image_data.shape}, dtype: {image_data.dtype}, device: {image_data.device if hasattr(image_data, 'device') else 'CPU'}")247 if image_data.numel() == 0:248 raise ValueError("Preprocessed image is empty")249 return image_data250 except Exception as e:251 error_msg = f"Preprocessing error: {e}"252 print(f"✗ {error_msg}")253 import traceback254 traceback.print_exc()255 raise ValueError(f"Failed to preprocess NIfTI file: {e}") from e256 257def refine_liver_mask_enhanced(mask, voxel_spacing, pred_probabilities, threshold, modality):258 259 original_shape = mask.shape260 original_sum = mask.sum()261 262 was_4d = len(mask.shape) == 4263 was_5d = len(mask.shape) == 5264 265 if was_5d:266 mask_3d = mask[0, 0, 0] if mask.shape[0] == 1 and mask.shape[1] == 1 and mask.shape[2] == 1 else mask[0, 0]267 elif was_4d:268 mask_3d = mask[0, 0] if mask.shape[0] == 1 and mask.shape[1] == 1 else mask[0]269 else:270 mask_3d = mask.copy()271 272 if mask_3d.dtype != np.uint8:273 mask_3d = (mask_3d > 0.5).astype(np.uint8)274 275 if mask_3d.sum() == 0:276 return np.zeros(original_shape, dtype=np.uint8), {277 "original_voxels": 0, "refined_voxels": 0, "removed_voxels": 0,278 "connected_components_before": 0, "connected_components_after": 0,279 "volume_change_ml": 0.0, "volume_change_percent": 0.0,280 "guards_ok": False281 }, 0.0282 283 H, W, D = mask_3d.shape284 guards_ok = True285 286 print(f" NOTE: Spatial priors assume RAS orientation (Right-Anterior-Superior).")287 print(f" Input should be reoriented to RAS using nib.as_closest_canonical() before processing.")288 print(f" If orientation is unknown, spatial priors may remove valid liver tissue.")289 290 top_remove = max(1, int(0.15 * D))291 mask_3d[:, :, :top_remove] = 0292 if top_remove > 0:293 print(f" Spatial prior: Removed top {top_remove} slices (15% - diaphragm protection, assumes Superior axis)")294 295 right_trim = max(0, int(0.30 * W))296 mask_3d[:, W-right_trim:, :] = 0297 if right_trim > 0:298 print(f" Spatial prior: Removed right {right_trim} pixels (30% - stomach protection, assumes Right axis)")299 300 left_trim = max(0, int(0.15 * W))301 mask_3d[:, :left_trim, :] = 0302 if left_trim > 0:303 print(f" Spatial prior: Removed left {left_trim} pixels (15% - spleen protection, assumes Left axis)")304 305 bottom_remove = max(1, int(0.10 * D))306 mask_3d[:, :, -bottom_remove:] = 0307 if bottom_remove > 0:308 print(f" Spatial prior: Removed bottom {bottom_remove} slices (10% - lower abdomen protection, assumes Inferior axis)")309 310 if D > 2:311 bottom_slices = mask_3d[:, :, -2:]312 if bottom_slices.sum() > 0:313 mask_3d[:, :, -2:] = 0314 print(f" Bottom-cap trim: Removed bottom 2 slices (diaphragm protection)")315 guards_ok = False316 317 labels_before, num_components_before = ndimage.label(mask_3d)318 319 if num_components_before == 0:320 print(f" QC FAIL: No components after spatial priors. Attempting auto-rethreshold...")321 guards_ok = False322 if hasattr(pred_probabilities, 'shape') and len(pred_probabilities.shape) >= 3:323 if len(pred_probabilities.shape) == 4:324 pred_3d = pred_probabilities[0, 0]325 elif len(pred_probabilities.shape) == 5:326 pred_3d = pred_probabilities[0, 0, 0]327 else:328 pred_3d = pred_probabilities329 330 top_remove = max(1, int(0.15 * D))331 right_trim = max(0, int(0.30 * W))332 left_trim = max(0, int(0.15 * W))333 bottom_remove = max(1, int(0.10 * D))334 335 for retry_threshold in [0.70, 0.65, 0.60, 0.55, 0.50]:336 mask_retry = (pred_3d > retry_threshold).astype(np.uint8)337 mask_retry[:, :, :top_remove] = 0338 mask_retry[:, W-right_trim:, :] = 0339 mask_retry[:, :left_trim, :] = 0340 mask_retry[:, :, -bottom_remove:] = 0341 if mask_retry.sum() > 1000:342 mask_3d = mask_retry343 print(f" Auto-rethreshold: Found mask at threshold {retry_threshold:.3f}")344 break345 else:346 return np.zeros(original_shape, dtype=np.uint8), {347 "original_voxels": original_sum, "refined_voxels": 0, "removed_voxels": int(original_sum),348 "connected_components_before": 0, "connected_components_after": 0,349 "volume_change_ml": 0.0, "volume_change_percent": -100.0,350 "guards_ok": False351 }, 0.0352 353 labels_before, num_components_before = ndimage.label(mask_3d)354 component_sizes = ndimage.sum(mask_3d, labels_before, range(1, num_components_before + 1))355 largest_label = component_sizes.argmax() + 1356 mask_3d = (labels_before == largest_label).astype(np.uint8)357 print(f" Kept largest connected component ({component_sizes.max():,} voxels)")358 359 coords = np.where(mask_3d > 0)360 if len(coords[0]) > 0:361 z_span = (coords[2].max() - coords[2].min() + 1) / D if D > 0 else 0362 363 if z_span < 0.25:364 print(f" QC FAIL: Z-span only {z_span*100:.1f}% (<25%). Attempting iterative rethreshold...")365 guards_ok = False366 367 if hasattr(pred_probabilities, 'shape') and len(pred_probabilities.shape) >= 3:368 if len(pred_probabilities.shape) == 4:369 pred_3d = pred_probabilities[0, 0]370 elif len(pred_probabilities.shape) == 5:371 pred_3d = pred_probabilities[0, 0, 0]372 else:373 pred_3d = pred_probabilities374 375 best_mask = mask_3d376 best_z_span = z_span377 378 top_remove = max(1, int(0.12 * D))379 right_trim = max(0, int(0.25 * W))380 left_trim = max(0, int(0.10 * W))381 382 for retry_threshold in [0.65, 0.60, 0.55, 0.50, 0.45, 0.40, 0.35]:383 mask_retry = (pred_3d > retry_threshold).astype(np.uint8)384 mask_retry[:, :, :top_remove] = 0385 mask_retry[:, W-right_trim:, :] = 0386 mask_retry[:, :left_trim, :] = 0387 388 if mask_retry.sum() < 1000:389 continue390 391 labels_retry, _ = ndimage.label(mask_retry)392 if labels_retry.max() > 0:393 comp_sizes_retry = ndimage.sum(mask_retry, labels_retry, range(1, labels_retry.max() + 1))394 largest_retry = comp_sizes_retry.argmax() + 1395 mask_retry = (labels_retry == largest_retry).astype(np.uint8)396 397 coords_retry = np.where(mask_retry > 0)398 if len(coords_retry[0]) > 0:399 z_span_retry = (coords_retry[2].max() - coords_retry[2].min() + 1) / D400 401 if z_span_retry >= 0.25:402 mask_3d = mask_retry403 print(f" Auto-rethreshold SUCCESS: threshold={retry_threshold:.3f}, z-span={z_span_retry*100:.1f}%")404 break405 elif z_span_retry > best_z_span:406 best_mask = mask_retry407 best_z_span = z_span_retry408 else:409 if best_z_span > z_span:410 mask_3d = best_mask411 print(f" Auto-rethreshold: Using best z-span={best_z_span*100:.1f}% (still <25%)")412 else:413 print(f" Auto-rethreshold FAILED: No threshold yielded z-span >= 25%")414 415 labels_before_morph, _ = ndimage.label(mask_3d)416 if labels_before_morph.max() > 0:417 component_sizes_before_morph = ndimage.sum(mask_3d, labels_before_morph, range(1, labels_before_morph.max() + 1))418 if len(component_sizes_before_morph) > 0:419 largest_label_before_morph = component_sizes_before_morph.argmax() + 1420 mask_3d = (labels_before_morph == largest_label_before_morph).astype(np.uint8)421 print(f" Kept largest component before morphology")422 423 try:424 mask_3d = mask_3d.astype(bool)425 structure = ball(2)426 mask_3d = binary_closing(mask_3d, structure=structure)427 mask_3d = mask_3d.astype(np.uint8)428 print(f" Applied binary closing (ball radius=2)")429 except Exception as e:430 print(f" Binary closing failed: {e}")431 432 try:433 mask_3d = mask_3d.astype(bool)434 mask_3d = binary_fill_holes(mask_3d)435 mask_3d = mask_3d.astype(np.uint8)436 print(f" Filled holes")437 except Exception as e:438 print(f" Hole filling failed: {e}")439 440 try:441 mask_3d = median_filter(mask_3d, size=3)442 print(f" Applied 3D median filter (size=3)")443 except Exception as e:444 print(f" Median filter failed: {e}")445 446 labels_after_morph, _ = ndimage.label(mask_3d)447 if labels_after_morph.max() > 0:448 component_sizes_morph = ndimage.sum(mask_3d, labels_after_morph, range(1, labels_after_morph.max() + 1))449 if len(component_sizes_morph) > 0:450 largest_label_morph = component_sizes_morph.argmax() + 1451 mask_3d = (labels_after_morph == largest_label_morph).astype(np.uint8)452 print(f" Re-kept largest component after morphology")453 454 labels_after, num_components_after = ndimage.label(mask_3d)455 456 refined_sum = mask_3d.sum()457 removed_voxels = int(np.int64(original_sum) - np.int64(refined_sum))458 459 voxel_volume = voxel_spacing[0] * voxel_spacing[1] * voxel_spacing[2]460 volume_change_ml = (removed_voxels * voxel_volume) / 1000.0461 volume_change_percent = (removed_voxels / float(original_sum) * 100.0) if original_sum > 0 else 0.0462 463 volume_ml = (refined_sum * voxel_volume) / 1000.0464 465 coords_final = np.where(mask_3d > 0)466 if len(coords_final[0]) > 0:467 z_span_final = (coords_final[2].max() - coords_final[2].min() + 1) / D if D > 0 else 0468 x_centroid = np.mean(coords_final[1]) if len(coords_final) > 1 else W / 2469 y_centroid = np.mean(coords_final[0]) if len(coords_final) > 0 else H / 2470 471 if volume_ml < 800 or volume_ml > 2500:472 print(f" QC FAIL: Volume {volume_ml:.1f}ml outside normal range [800-2500ml]")473 guards_ok = False474 475 if z_span_final < 0.20:476 print(f" QC FAIL: Z-span {z_span_final*100:.1f}% too small (<20%)")477 guards_ok = False478 479 liver_x_min = 0.15 * W480 liver_x_max = 0.55 * W481 if x_centroid < liver_x_min or x_centroid > liver_x_max:482 print(f" QC FAIL: x-centroid {x_centroid:.1f} outside expected liver band [15%-55% of width]")483 guards_ok = False484 485 liver_y_min = 0.25 * H486 liver_y_max = 0.75 * H487 if y_centroid < liver_y_min or y_centroid > liver_y_max:488 print(f" QC FAIL: y-centroid {y_centroid:.1f} outside expected liver band [25%-75% of height]")489 guards_ok = False490 491 if volume_ml < 800:492 print(f" QC WARNING: Volume {volume_ml:.1f}ml suspiciously low - may be wrong organ")493 guards_ok = False494 495 if volume_change_percent > 80:496 print(f" QC FAIL: Refinement removed {volume_change_percent:.1f}% - too aggressive")497 guards_ok = False498 499 if was_5d:500 if original_shape[0] == 1 and original_shape[1] == 1 and original_shape[2] == 1:501 mask_3d = mask_3d[np.newaxis, np.newaxis, np.newaxis, :, :, :]502 else:503 mask_3d = mask_3d[np.newaxis, np.newaxis, :, :, :]504 elif was_4d:505 if original_shape[0] == 1 and original_shape[1] == 1:506 mask_3d = mask_3d[np.newaxis, np.newaxis, :, :, :]507 else:508 mask_3d = mask_3d[np.newaxis, :, :, :]509 510 mask_3d = mask_3d.astype(np.uint8)511 512 if mask_3d.shape != original_shape:513 if len(original_shape) == 3:514 while mask_3d.ndim > 3:515 mask_3d = mask_3d.squeeze(0)516 elif len(original_shape) == 4:517 while mask_3d.ndim < 4:518 mask_3d = mask_3d[np.newaxis, ...]519 while mask_3d.ndim > 4:520 mask_3d = mask_3d.squeeze(0)521 elif len(original_shape) == 5:522 while mask_3d.ndim < 5:523 mask_3d = mask_3d[np.newaxis, ...]524 while mask_3d.ndim > 5:525 mask_3d = mask_3d.squeeze(0)526 527 print(f" Refinement complete: {original_sum:,} -> {refined_sum:,} voxels ({removed_voxels:,} removed, {volume_change_percent:.2f}%)")528 print(f" Connected components: {num_components_before} -> {num_components_after}")529 530 confidence_score = calculate_confidence_score(mask_3d, pred_probabilities, threshold, num_components_after, volume_change_percent, guards_ok, voxel_spacing)531 532 metrics = {533 "original_voxels": int(original_sum),534 "refined_voxels": int(refined_sum),535 "removed_voxels": removed_voxels,536 "connected_components_before": int(num_components_before),537 "connected_components_after": int(num_components_after),538 "volume_change_ml": float(volume_change_ml),539 "volume_change_percent": float(volume_change_percent),540 "guards_ok": guards_ok541 }542 543 return mask_3d, metrics, confidence_score544 545def calculate_confidence_score(mask, pred_probabilities, threshold, num_components, volume_change_percent, guards_ok=True, voxel_spacing=(1.0, 1.0, 1.0)):546 if mask.sum() == 0:547 return 0.0548 549 if len(mask.shape) == 4:550 mask_3d = mask[0, 0]551 elif len(mask.shape) == 5:552 mask_3d = mask[0, 0, 0]553 else:554 mask_3d = mask555 556 if len(pred_probabilities.shape) == 4:557 pred_3d = pred_probabilities[0, 0]558 elif len(pred_probabilities.shape) == 5:559 pred_3d = pred_probabilities[0, 0, 0]560 else:561 pred_3d = pred_probabilities562 563 mask_indices = mask_3d > 0564 if mask_indices.sum() == 0:565 return 0.0566 567 avg_p = float(np.clip(pred_3d[mask_indices].mean(), 0.0, 1.0))568 comp_pen = 1.0 if num_components == 1 else max(0.5, 1.0 - 0.1 * (num_components - 1))569 vol_pen = 1.0 if abs(volume_change_percent) < 50 else 0.7570 571 if not guards_ok:572 guard_pen = 0.5573 else:574 guard_pen = 1.0575 576 volume_ml = (mask_3d.sum() * (voxel_spacing[0] * voxel_spacing[1] * voxel_spacing[2])) / 1000.0577 if volume_ml < 800:578 volume_penalty = 0.5579 elif volume_ml < 1000:580 volume_penalty = 0.7581 elif volume_ml < 1200:582 volume_penalty = 0.9583 else:584 volume_penalty = 1.0585 586 confidence = 100 * avg_p * comp_pen * vol_pen * guard_pen * volume_penalty587 confidence = float(np.clip(confidence, 0, 100))588 589 return confidence590 591def refine_liver_mask(mask, voxel_spacing=(1.0, 1.0, 1.0), enable_smoothing=True, min_component_size=None):592 """593 Refine liver segmentation mask to remove fragmentation, smooth boundaries, and ensure single connected component.594 595 Args:596 mask: 3D or 4D numpy array (H, W, D) or (1, 1, H, W, D) with binary values (0 or 1)597 voxel_spacing: Tuple of (z, y, x) voxel spacing in mm598 enable_smoothing: Whether to apply median filter smoothing (default: True)599 min_component_size: Minimum size for connected components to keep (None = keep only largest)600 601 Returns:602 refined_mask: Refined binary mask (same shape as input)603 metrics: Dictionary with refinement statistics604 """605 original_shape = mask.shape606 original_sum = mask.sum()607 608 was_4d = len(mask.shape) == 4609 was_5d = len(mask.shape) == 5610 611 if was_5d:612 mask = mask[0, 0, 0] if mask.shape[0] == 1 and mask.shape[1] == 1 and mask.shape[2] == 1 else mask[0, 0]613 elif was_4d:614 mask = mask[0, 0] if mask.shape[0] == 1 and mask.shape[1] == 1 else mask[0]615 616 if mask.dtype != np.uint8:617 mask = (mask > 0.5).astype(np.uint8)618 619 if mask.sum() == 0:620 print(" ⚠ Empty mask - no refinement possible")621 return np.zeros(original_shape, dtype=np.uint8), {622 "original_voxels": 0,623 "refined_voxels": 0,624 "removed_voxels": 0,625 "connected_components_before": 0,626 "connected_components_after": 0,627 "volume_change_ml": 0.0,628 "volume_change_percent": 0.0629 }630 631 labels_before, num_components_before = ndimage.label(mask)632 633 if num_components_before == 0:634 print(" ⚠ No connected components found")635 return np.zeros(original_shape, dtype=np.uint8), {636 "original_voxels": original_sum,637 "refined_voxels": 0,638 "removed_voxels": int(original_sum),639 "connected_components_before": 0,640 "connected_components_after": 0,641 "volume_change_ml": 0.0,642 "volume_change_percent": -100.0643 }644 645 component_sizes = ndimage.sum(mask, labels_before, range(1, num_components_before + 1))646 647 if min_component_size is None:648 largest_label = component_sizes.argmax() + 1649 mask = (labels_before == largest_label).astype(np.uint8)650 print(f" → Kept largest connected component ({component_sizes.max():,} voxels)")651 else:652 valid_labels = np.where(component_sizes >= min_component_size)[0] + 1653 if len(valid_labels) == 0:654 largest_label = component_sizes.argmax() + 1655 mask = (labels_before == largest_label).astype(np.uint8)656 print(f" → No components >= {min_component_size} voxels, kept largest ({component_sizes.max():,} voxels)")657 else:658 mask = np.isin(labels_before, valid_labels).astype(np.uint8)659 print(f" → Kept {len(valid_labels)} component(s) >= {min_component_size} voxels")660 661 after_cc = mask.sum()662 663 try:664 structure = ball(3)665 mask = binary_closing(mask, structure=structure)666 print(f" → Applied binary closing (ball radius=3)")667 except Exception as e:668 print(f" ⚠ Binary closing failed: {e}")669 670 try:671 mask = binary_fill_holes(mask)672 print(f" → Filled holes")673 except Exception as e:674 print(f" ⚠ Hole filling failed: {e}")675 676 try:677 structure = ball(2)678 mask = binary_opening(mask, structure=structure)679 print(f" → Applied binary opening (ball radius=2)")680 except Exception as e:681 print(f" ⚠ Binary opening failed: {e}")682 683 if enable_smoothing:684 try:685 mask = median_filter(mask, size=3)686 print(f" → Applied 3D median filter (size=3)")687 except Exception as e:688 print(f" ⚠ Median filter failed: {e}")689 690 labels_after, num_components_after = ndimage.label(mask)691 692 refined_sum = mask.sum()693 removed_voxels = int(original_sum - refined_sum)694 695 voxel_volume = voxel_spacing[0] * voxel_spacing[1] * voxel_spacing[2]696 volume_change_ml = (removed_voxels * voxel_volume) / 1000.0697 volume_change_percent = (removed_voxels / original_sum * 100.0) if original_sum > 0 else 0.0698 699 if was_5d:700 if original_shape[0] == 1 and original_shape[1] == 1 and original_shape[2] == 1:701 mask = mask[np.newaxis, np.newaxis, np.newaxis, :, :, :]702 else:703 mask = mask[np.newaxis, np.newaxis, :, :, :]704 elif was_4d:705 if original_shape[0] == 1 and original_shape[1] == 1:706 mask = mask[np.newaxis, np.newaxis, :, :, :]707 else:708 mask = mask[np.newaxis, :, :, :]709 710 mask = mask.astype(np.uint8)711 712 if mask.shape != original_shape:713 print(f" ⚠ Shape mismatch: {mask.shape} vs {original_shape}, fixing...")714 if len(original_shape) == 3:715 while mask.ndim > 3:716 mask = mask.squeeze(0)717 elif len(original_shape) == 4:718 while mask.ndim < 4:719 mask = mask[np.newaxis, ...]720 while mask.ndim > 4:721 mask = mask.squeeze(0)722 elif len(original_shape) == 5:723 while mask.ndim < 5:724 mask = mask[np.newaxis, ...]725 while mask.ndim > 5:726 mask = mask.squeeze(0)727 728 print(f" ✓ Refinement complete: {original_sum:,} → {refined_sum:,} voxels ({removed_voxels:,} removed, {volume_change_percent:.2f}%)")729 print(f" → Connected components: {num_components_before} → {num_components_after}")730 731 metrics = {732 "original_voxels": int(original_sum),733 "refined_voxels": int(refined_sum),734 "removed_voxels": removed_voxels,735 "connected_components_before": int(num_components_before),736 "connected_components_after": int(num_components_after),737 "volume_change_ml": float(volume_change_ml),738 "volume_change_percent": float(volume_change_percent)739 }740 741 return mask, metrics742 743def calculate_liver_volume(pred_binary, voxel_spacing=(1.0, 1.0, 1.0)):744 voxel_volume = voxel_spacing[0] * voxel_spacing[1] * voxel_spacing[2]745 liver_voxels = pred_binary.sum()746 volume_ml = liver_voxels * voxel_volume / 1000.0747 return volume_ml748 749def analyze_liver_morphology(pred_binary):750 if len(pred_binary.shape) == 4:751 mask_3d = pred_binary[0]752 elif len(pred_binary.shape) == 5:753 mask_3d = pred_binary[0, 0]754 else:755 mask_3d = pred_binary756 757 labeled_mask, num_features = ndimage.label(mask_3d)758 if num_features == 0:759 return {"connected_components": 0, "largest_component_ratio": 0.0, "fragmentation": "high"}760 761 component_sizes = [np.sum(labeled_mask == i) for i in range(1, num_features + 1)]762 largest_component = max(component_sizes)763 total_liver = pred_binary.sum()764 largest_ratio = largest_component / total_liver if total_liver > 0 else 0.0765 766 if largest_ratio > 0.95:767 fragmentation = "low"768 elif largest_ratio > 0.80:769 fragmentation = "moderate"770 else:771 fragmentation = "high"772 773 return {774 "connected_components": int(num_features),775 "largest_component_ratio": float(largest_ratio),776 "fragmentation": fragmentation777 }778 779def check_volume_sanity(volume_ml):780 normal_range = (float(os.getenv("LIVER_VOL_LOW", "1200")), float(os.getenv("LIVER_VOL_HIGH", "1800")))781 if volume_ml < normal_range[0] * 0.5:782 return "CRITICAL", f"Volume ({volume_ml:.1f} ml) is extremely low (<50% of normal). Please visually inspect overlay for segmentation errors."783 elif volume_ml < normal_range[0]:784 return "WARNING", f"Volume ({volume_ml:.1f} ml) is below normal range. Please visually inspect overlay."785 elif volume_ml > normal_range[1] * 1.5:786 return "CRITICAL", f"Volume ({volume_ml:.1f} ml) is extremely high (>150% of normal). Please visually inspect overlay for segmentation errors."787 elif volume_ml > normal_range[1]:788 return "WARNING", f"Volume ({volume_ml:.1f} ml) is above normal range. Please visually inspect overlay."789 return "OK", None790 791def generate_medical_report(statistics, volume_ml, morphology, modality, confidence_score=0.0):792 liver_percentage = statistics["liver_percentage"]793 volume_shape = statistics["volume_shape"]794 liver_voxels = statistics.get("liver_voxels", 0)795 total_voxels = statistics.get("total_voxels", 0)796 797 normal_liver_volume_range = (1200, 1800)798 normal_liver_percentage_range = (2.0, 3.5)799 800 findings = []801 recommendations = []802 clinical_notes = []803 quality_assessment = []804 805 if liver_voxels == 0:806 severity = "failure"807 status = "FAILURE"808 findings.append("**SEGMENTATION FAILURE:** No liver tissue detected (0 voxels segmented).")809 recommendations.append("**CRITICAL:** Segmentation failed completely. Possible causes:")810 recommendations.append(" • Input quality issues (low resolution, compression, missing metadata)")811 recommendations.append(" • Threshold too high for prediction distribution")812 recommendations.append(" • Model mismatch with input modality or preprocessing")813 recommendations.append(" • Please check input file quality and try again, or contact support.")814 clinical_notes.append("The automated segmentation system failed to identify any liver tissue. This indicates a technical failure rather than an anatomical finding.")815 quality_assessment.append("**Segmentation Failure:** No voxels were segmented. Manual review and re-processing required.")816 impression_parts = ["Automated liver segmentation FAILED. No liver tissue was detected."]817 impression_parts.append("This is a technical failure requiring investigation of input quality and model compatibility.")818 else:819 severity = "normal"820 status = "NORMAL"821 822 num_components = morphology.get("connected_components", 1)823 largest_ratio = morphology.get("largest_component_ratio", 1.0)824 825 if num_components > 1 and largest_ratio < 0.9:826 severity = "critical"827 status = "CRITICAL"828 findings.append(f"**CRITICAL: Fragmented Segmentation:** {num_components} disconnected components detected. Largest component is only {largest_ratio*100:.1f}% of total volume.")829 recommendations.append("**URGENT:** Segmentation shows severe fragmentation. Manual correction required.")830 clinical_notes.append("The segmentation contains multiple disconnected regions, indicating possible segmentation artifacts or severe anatomical abnormalities.")831 elif volume_ml < normal_liver_volume_range[0] * 0.5 or volume_ml > normal_liver_volume_range[1] * 1.5:832 if confidence_score < 50:833 severity = "critical"834 status = "CRITICAL"835 else:836 severity = "moderate"837 status = "WARNING"838 elif volume_ml < normal_liver_volume_range[0] or volume_ml > normal_liver_volume_range[1]:839 severity = "moderate"840 status = "WARNING"841 842 volume_sanity_status, volume_sanity_msg = check_volume_sanity(volume_ml)843 if volume_sanity_status == "CRITICAL":844 if severity != "critical":845 severity = "critical"846 status = "CRITICAL"847 findings.append(f"**CRITICAL FINDING:** {volume_sanity_msg}")848 recommendations.append("**URGENT:** Visual inspection and manual review required. Segmentation may contain significant errors that could affect clinical interpretation.")849 clinical_notes.append("The automated segmentation has produced results that fall outside expected physiological ranges. This may indicate technical issues with the segmentation algorithm or unusual patient anatomy.")850 elif volume_sanity_status == "WARNING":851 if severity == "normal":852 severity = "moderate"853 status = "WARNING"854 findings.append(f"**WARNING:** {volume_sanity_msg}")855 recommendations.append("Visual inspection recommended to verify segmentation accuracy and ensure clinical validity.")856 clinical_notes.append("The segmentation results are outside the typical range but may still be clinically valid depending on patient-specific factors.")857 clinical_notes.append("Note: Normal liver volume range (1200-1800 ml) is for average adult body size. Pediatric patients or extreme body sizes may have different normal ranges.")858 859 if volume_ml < normal_liver_volume_range[0]:860 findings.append(f"**Liver Volume Assessment:** Measured liver volume is **{volume_ml:.1f} ml**, which is below the normal reference range of {normal_liver_volume_range[0]}-{normal_liver_volume_range[1]} ml.")861 clinical_notes.append(f"This represents approximately **{((normal_liver_volume_range[0] - volume_ml) / normal_liver_volume_range[0] * 100):.1f}% reduction** compared to the lower limit of normal. Possible etiologies include:")862 clinical_notes.append(" • Chronic liver disease with parenchymal loss")863 clinical_notes.append(" • Post-surgical resection")864 clinical_notes.append(" • Cirrhosis with volume loss")865 clinical_notes.append(" • Age-related atrophy")866 recommendations.append("Consider follow-up imaging to monitor liver volume changes over time. Correlation with clinical history and liver function tests is recommended.")867 if severity == "normal":868 severity = "mild" if volume_ml > normal_liver_volume_range[0] * 0.7 else "moderate"869 elif volume_ml > normal_liver_volume_range[1]:870 findings.append(f"**Liver Volume Assessment:** Measured liver volume is **{volume_ml:.1f} ml**, which exceeds the normal reference range of {normal_liver_volume_range[0]}-{normal_liver_volume_range[1]} ml.")871 clinical_notes.append(f"This represents approximately **{((volume_ml - normal_liver_volume_range[1]) / normal_liver_volume_range[1] * 100):.1f}% increase** compared to the upper limit of normal, consistent with hepatomegaly. Potential causes include:")872 clinical_notes.append(" • Fatty liver disease (steatosis)")873 clinical_notes.append(" • Congestive hepatopathy")874 clinical_notes.append(" • Inflammatory conditions")875 clinical_notes.append(" • Storage diseases")876 clinical_notes.append(" • Neoplastic processes")877 recommendations.append("Further clinical evaluation recommended to identify underlying etiology. Consider correlation with laboratory findings, clinical history, and additional imaging studies.")878 if severity == "normal":879 severity = "mild" if volume_ml < normal_liver_volume_range[1] * 1.3 else "moderate"880 else:881 findings.append(f"**Liver Volume Assessment:** Measured liver volume is **{volume_ml:.1f} ml**, which falls within the normal reference range of {normal_liver_volume_range[0]}-{normal_liver_volume_range[1]} ml.")882 clinical_notes.append("The liver volume is within expected physiological parameters for an adult patient.")883 884 if morphology["connected_components"] > 1:885 if morphology["largest_component_ratio"] < 0.9:886 if severity != "critical":887 severity = "critical"888 status = "CRITICAL"889 findings.append(f"**CRITICAL: Fragmented Segmentation:** The liver segmentation identified **{morphology['connected_components']} separate connected components**. The largest component represents only **{morphology['largest_component_ratio']*100:.1f}%** of the total segmented volume.")890 quality_assessment.append("**Severe Fragmentation Detected:** Multiple disconnected regions suggest possible segmentation artifacts or severe anatomical variations.")891 recommendations.append("**URGENT:** Manual review and correction required. Fragmentation indicates potential segmentation errors.")892 elif morphology["largest_component_ratio"] < 0.95:893 if severity == "normal":894 severity = "moderate"895 status = "WARNING"896 findings.append(f"**Segmentation Quality:** The liver segmentation identified **{morphology['connected_components']} separate connected components**. The largest component represents **{morphology['largest_component_ratio']*100:.1f}%** of the total segmented volume.")897 quality_assessment.append("**Moderate Fragmentation Detected:** Multiple disconnected regions suggest possible segmentation artifacts or anatomical variations.")898 quality_assessment.append("Post-processing filters (largest-component selection, hole-filling, morphological operations) have been applied to optimize the segmentation.")899 recommendations.append("Review the segmentation overlay carefully. The presence of multiple components may indicate:")900 recommendations.append(" • Segmentation artifacts requiring manual correction")901 recommendations.append(" • Anatomical variants (e.g., accessory liver lobes)")902 recommendations.append(" • Pathological processes causing liver fragmentation")903 else:904 findings.append(f"**Segmentation Quality:** The liver segmentation shows **{morphology['connected_components']} components**, with the largest component comprising **{morphology['largest_component_ratio']*100:.1f}%** of the total volume, indicating good segmentation continuity.")905 quality_assessment.append("The segmentation demonstrates good connectivity with a dominant main component.")906 else:907 quality_assessment.append("**Excellent Segmentation Quality:** Single connected component indicates robust segmentation with good anatomical continuity.")908 909 if morphology["fragmentation"] == "high":910 findings.append(f"**High Fragmentation Detected:** The liver segmentation demonstrates high morphological fragmentation, which may reflect irregular liver surface or segmentation challenges.")911 quality_assessment.append("High fragmentation suggests the liver may have irregular borders or that the segmentation encountered challenging anatomical features.")912 recommendations.append("Manual review and potential refinement of the segmentation may be beneficial for optimal clinical interpretation.")913 if severity == "normal":914 severity = "mild"915 elif morphology["fragmentation"] == "moderate":916 quality_assessment.append("Moderate fragmentation observed, which is acceptable for clinical use but may benefit from review.")917 else:918 quality_assessment.append("Low fragmentation indicates smooth, well-defined liver boundaries.")919 920 if liver_percentage < normal_liver_percentage_range[0]:921 findings.append(f"**Spatial Distribution:** The liver occupies **{liver_percentage:.2f}%** of the total scan volume, which is below the typical range of {normal_liver_percentage_range[0]}-{normal_liver_percentage_range[1]}%.")922 clinical_notes.append("This may reflect a smaller liver relative to the field of view, or indicate that the scan includes a larger portion of the abdomen.")923 elif liver_percentage > normal_liver_percentage_range[1]:924 findings.append(f"**Spatial Distribution:** The liver occupies **{liver_percentage:.2f}%** of the total scan volume, which is above the typical range.")925 clinical_notes.append("This may indicate an enlarged liver or a scan field of view focused on the upper abdomen.")926 else:927 findings.append(f"**Spatial Distribution:** The liver occupies **{liver_percentage:.2f}%** of the scan volume, within the expected range.")928 929 if total_voxels > 0:930 voxel_density = liver_voxels / total_voxels * 100931 quality_assessment.append(f"**Segmentation Coverage:** {liver_voxels:,} voxels segmented out of {total_voxels:,} total voxels ({voxel_density:.2f}% coverage).")932 933 if volume_shape:934 quality_assessment.append(f"**Image Dimensions:** {volume_shape[0]} × {volume_shape[1]} × {volume_shape[2]} voxels")935 936 impression_parts = []937 if severity == "normal":938 impression_parts.append("Automated liver segmentation completed successfully using the SRMA-Mamba deep learning model.")939 impression_parts.append("The segmentation demonstrates good quality with measurements within expected physiological ranges.")940 elif severity == "mild":941 impression_parts.append("Automated liver segmentation completed with minor findings.")942 impression_parts.append("The segmentation is generally acceptable but requires clinical correlation and visual review.")943 elif severity == "moderate":944 impression_parts.append("Automated liver segmentation completed with notable findings requiring attention.")945 impression_parts.append("Visual inspection and clinical correlation are recommended to ensure accuracy.")946 elif severity == "critical":947 impression_parts.append("Automated liver segmentation completed with critical findings.")948 impression_parts.append("Immediate visual inspection and manual review are strongly recommended.")949 950 impression_parts.append(f"**{len(findings)} key finding(s)** identified during automated analysis.")951 952 report = {953 "patient_id": "N/A",954 "study_date": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),955 "modality": modality.upper(),956 "status": status,957 "findings": findings,958 "clinical_notes": clinical_notes,959 "quality_assessment": quality_assessment,960 "measurements": {961 "liver_volume_ml": round(volume_ml, 2),962 "liver_volume_liters": round(volume_ml / 1000.0, 3),963 "liver_percentage": round(liver_percentage, 2),964 "liver_voxels": int(liver_voxels),965 "total_voxels": int(total_voxels),966 "volume_shape": volume_shape,967 "morphology": morphology,968 "confidence_score": round(confidence_score, 1)969 },970 "impression": " ".join(impression_parts) if liver_voxels > 0 else impression_parts[0] if impression_parts else "Segmentation failed.",971 "recommendations": recommendations,972 "severity": severity,973 "methodology": "SRMA-Mamba: State Space Model for Medical Image Segmentation using Mamba architecture with sliding window inference",974 "disclaimer": "**IMPORTANT:** This is an automated analysis generated by artificial intelligence. Results should be reviewed and validated by a qualified radiologist or physician. This report is not intended for diagnostic use without appropriate clinical correlation and professional medical interpretation."975 }976 977 return report978 979