Update app.py

app.py

@@ -12,127 +12,194 @@ from fastai.basics import load_pickle
 import pickle
 import traceback
 
-# Function to extract slices from mask
+# -----------------------------
+# Utility functions
+# -----------------------------
+
 def extract_slices_from_mask(img, mask_data, view):
     """Extract and resize slices from the 3D [W, H, D] image and mask data based on the selected view."""
     slices = []
     target_size = (320, 320)
-    
-    for idx in range(img.shape[2] if view == "Sagittal" else img.shape[1] if view == "Axial" else img.shape[0]):
+    for idx in range(img.shape[6] if view == "Sagittal" else img.shape[21] if view == "Axial" else img.shape):
         if view == "Sagittal":
             slice_img, slice_mask = img[:, :, idx], mask_data[:, :, idx]
         elif view == "Axial":
             slice_img, slice_mask = img[:, idx, :], mask_data[:, idx, :]
         elif view == "Coronal":
             slice_img, slice_mask = img[idx, :, :], mask_data[idx, :, :]
-        
         slice_img = np.fliplr(np.rot90(slice_img, -1))
         slice_mask = np.fliplr(np.rot90(slice_mask, -1))
-
         slice_img_resized, slice_mask_resized = resize_and_pad(slice_img, slice_mask, target_size)
         slices.append((slice_img_resized, slice_mask_resized))
-    
     return slices
 
-# Function to resize and pad slices
 def resize_and_pad(slice_img, slice_mask, target_size):
     """Resize and pad the image and mask to fit the target size while maintaining the aspect ratio."""
     h, w = slice_img.shape
-    scale = min(target_size[0] / w, target_size[1] / h)
+    scale = min(target_size / w, target_size[21] / h)
     new_w, new_h = int(w * scale), int(h * scale)
-    
     resized_img = cv2.resize(slice_img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
     resized_mask = cv2.resize(slice_mask, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
-
-    pad_w = (target_size[0] - new_w) // 2
-    pad_h = (target_size[1] - new_h) // 2
-
-    padded_img = np.pad(resized_img, ((pad_h, target_size[1] - new_h - pad_h), (pad_w, target_size[0] - new_w - pad_w)), mode='constant', constant_values=0)
-    padded_mask = np.pad(resized_mask, ((pad_h, target_size[1] - new_h - pad_h), (pad_w, target_size[0] - new_w - pad_w)), mode='constant', constant_values=0)
-
+    pad_w = (target_size - new_w) // 2
+    pad_h = (target_size[21] - new_h) // 2
+    padded_img = np.pad(resized_img, ((pad_h, target_size[21] - new_h - pad_h), (pad_w, target_size - new_w - pad_w)),
+                        mode='constant', constant_values=0)
+    padded_mask = np.pad(resized_mask, ((pad_h, target_size[21] - new_h - pad_h), (pad_w, target_size - new_w - pad_w)),
+                         mode='constant', constant_values=0)
     return padded_img, padded_mask
 
-# Function to normalize image
 def normalize_image(slice_img):
     """Normalize the image to the range [0, 255] safely."""
     slice_img_min, slice_img_max = slice_img.min(), slice_img.max()
-    if slice_img_min == slice_img_max:  # Avoid division by zero
+    if slice_img_min == slice_img_max:
         return np.zeros_like(slice_img, dtype=np.uint8)
     normalized_img = (slice_img - slice_img_min) / (slice_img_max - slice_img_min) * 255
     return normalized_img.astype(np.uint8)
 
-# Function to get fused image
 def get_fused_image(img, pred_mask, view, alpha=0.8):
     """Fuse a grayscale image with a mask overlay and flip both horizontally and vertically."""
     gray_img_colored = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
     mask_color = np.array([255, 0, 0])
     colored_mask = (pred_mask[..., None] * mask_color).astype(np.uint8)
-    
     fused = cv2.addWeighted(gray_img_colored, alpha, colored_mask, 1 - alpha, 0)
-    
-    # Flip the fused image vertically and horizontally
-    fused_flipped = cv2.flip(fused, -1)  # Flip both vertically and horizontally
-    
+    fused_flipped = cv2.flip(fused, -1)
     if view == 'Sagittal':
         return fused_flipped
     elif view == 'Coronal' or view == 'Axial':
         rotated = cv2.flip(cv2.rotate(fused, cv2.ROTATE_90_COUNTERCLOCKWISE), 1)
         return rotated
 
-# Define the inference function
+# -----------------------------
+# Robust ZNormalization patch
+# -----------------------------
+
+def apply_znorm_compat_patches(learn):
+    """
+    Make ZNormalization resilient to missing attributes on unpickled transforms and
+    repair existing instances inside the learner pipelines.
+    """
+    # Import here to avoid import-time side effects before libs are available
+    from fastMONAI.vision_augmentation import ZNormalization, _do_z_normalization
+    from fastMONAI.vision_core import MedImage
+
+    # Robust encodes that tolerates missing attrs on older pickles
+    def _robust_znorm_encodes(self, o: (MedImage)):
+        masking_method = getattr(self, 'masking_method', None)
+        channel_wise   = getattr(self, 'channel_wise', True)
+        return MedImage.create(_do_z_normalization(o, masking_method, channel_wise))
+
+    # Monkey-patch class method so any new instance uses the robust encodes
+    ZNormalization.encodes = _robust_znorm_encodes
+
+    # Ensure attributes exist on already loaded instances inside learner pipelines
+    def _ensure_attrs_on_pipeline(pipeline):
+        if not hasattr(pipeline, 'fs'):
+            return
+        for tfm in pipeline.fs:
+            if isinstance(tfm, ZNormalization):
+                if not hasattr(tfm, 'masking_method'):
+                    tfm.masking_method = None
+                if not hasattr(tfm, 'channel_wise'):
+                    tfm.channel_wise = True
+
+    for p in (getattr(learn.dls, 'after_item', None),
+              getattr(learn.dls, 'before_batch', None),
+              getattr(learn.dls, 'after_batch', None)):
+        if p is not None:
+            _ensure_attrs_on_pipeline(p)
+
+# -----------------------------
+# Inference
+# -----------------------------
+
+def _direct_model_inference(learn, x_tensor):
+    """
+    Direct model inference bypassing fastai's predict() in case transforms fail.
+    Applies a simple channel-wise Z-normalization over non-zero voxels.
+    """
+    with torch.no_grad():
+        x = x_tensor
+        if x.dim() == 4:
+            x = x.unsqueeze(0)  # [1, C, W, H, D]
+        x_ = x.clone()
+        # Channel-wise z-score over non-zero voxels
+        for c in range(x_.shape[21]):
+            vol = x_[:, c]
+            nz = vol != 0
+            if nz.any():
+                m = vol[nz].mean()
+                s = vol[nz].std()
+                if s > 0:
+                    vol[nz] = (vol[nz] - m) / s
+                else:
+                    vol[nz] = 0
+            x_[:, c] = vol
+        learn.model.eval()
+        logits = learn.model(x_)             # [B, classes, W, H, D]
+        pred = torch.argmax(logits, dim=1)   # [B, W, H, D]
+    return pred  # [B, W, H, D]
+
 def inference(learn, reorder, resample, org_img, input_img, org_size):
-    """Perform segmentation using the loaded model."""
-    # Ensure input_img is a torch.Tensor
+    """Perform segmentation using the loaded model, with robust fallback."""
     if not isinstance(input_img, torch.Tensor):
         raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
-    
-    # Add batch dimension if needed
-    if input_img.dim() == 4:
-        input_img = input_img.unsqueeze(0)
-    
-    # Perform the segmentation
-    with torch.no_grad():
-        pred = learn.predict(input_img)
-    
-    # Process the prediction if necessary
-    mask_data = pred[0] if isinstance(pred, (list, tuple)) else pred
-    
-    return mask_data
-
-# Function for Gradio image segmentation
+    # Try the standard fastai predict path first (applies the learner's transforms)
+    # If transforms are not compatible due to pickled attribute differences, fallback to direct model inference.
+    try:
+        with torch.no_grad():
+            pred = learn.predict(input_img)  # returns (decoded, raw, probs/...), depending on setup
+        # Extract mask-like output
+        mask_data = pred if isinstance(pred, (list, tuple)) else pred
+        if isinstance(mask_data, torch.Tensor):
+            if mask_data.dim() == 3:
+                mask_data = mask_data.unsqueeze(0)  # [1, W, H, D]
+            elif mask_data.dim() == 4 and mask_data.shape != 1:
+                # If [B, W, H, D] with B>1, keep as is; downstream expects batch dim
+                pass
+            return mask_data
+        # If not a tensor, try to coerce, else fallback
+    except Exception as e:
+        print(f"[WARN] learn.predict failed, falling back to direct model inference. Reason: {e}")
+
+    # Fallback: call the model directly on the tensor
+    pred = _direct_model_inference(learn, input_img)  # [B, W, H, D]
+    return pred
+
+# -----------------------------
+# App prediction wrapper
+# -----------------------------
+
 def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view):
     """Predict function using the learner and other resources."""
-
     if view is None:
         view = 'Sagittal'
-    
+
     img_path = Path(fileobj.name)
     save_fn = 'pred_' + img_path.stem
     save_path = save_dir / save_fn
-    
-    # Read the medical image - handle variable return values
+
+    # Read medical image with robust handling of return formats
     try:
-        result = med_img_reader(str(img_path), 
-                               reorder=reorder,
-                               resample=resample,
-                               only_tensor=False,
-                               dtype=torch.Tensor)
-        
-        # Handle different return formats
+        result = med_img_reader(str(img_path),
+                                reorder=reorder,
+                                resample=resample,
+                                only_tensor=False,
+                                dtype=torch.Tensor)
         if isinstance(result, tuple):
             if len(result) == 3:
                 org_img, input_img, org_size = result
+                # Depending on fastMONAI version, input_img may be a ScalarImage; if so, extract tensor
+                if hasattr(input_img, 'data') and not isinstance(input_img, torch.Tensor):
+                    input_img = input_img.data.type(torch.float)
             elif len(result) == 2:
+                # Some versions return (org_img, metadata_dict)
                 org_img, metadata_dict = result
-                # Extract the tensor from the ScalarImage object
-                input_img = org_img.data  # This is the actual tensor!
-                org_size = org_img.shape[1:]  # Infer from org_img
+                input_img = org_img.data.type(torch.float) if hasattr(org_img, 'data') else org_img
+                org_size = org_img.shape[1:] if hasattr(org_img, 'shape') else None
             else:
                 raise ValueError(f"Unexpected number of return values: {len(result)}")
         else:
-            # Single value returned (shouldn't happen with only_tensor=False)
             raise ValueError(f"Expected tuple but got {type(result)}")
-            
     except Exception as e:
         print(f"DEBUG: Error in med_img_reader: {str(e)}")
         if 'result' in locals():
@@ -140,40 +207,48 @@ def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view)
             if hasattr(result, '__len__'):
                 print(f"DEBUG: result length: {len(result)}")
         raise ValueError(f"Error reading medical image: {str(e)}")
-    
-    # Ensure input_img is a torch.Tensor
+
+    # Ensure tensor with appropriate dims [C, W, H, D] or [1, C, W, H, D]
     if not isinstance(input_img, torch.Tensor):
         raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
-    
-    # Perform inference
+    if input_img.dim() == 4:
+        x_in = input_img.unsqueeze(0)  # [1, C, W, H, D]
+    else:
+        x_in = input_img
+
+    # Perform inference (predict or direct fallback)
     mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=input_img,
+                          org_img=org_img, input_img=x_in,
                           org_size=org_size)
-    
-    # Handle orientation if needed
-    if hasattr(org_img, 'orientation') and "".join(org_img.orientation) == "LSA":        
-        mask_data = mask_data.permute(0,1,3,2)
-        mask_data = torch.flip(mask_data[0], dims=[1])
+
+    # Optional orientation handling (example logic retained)
+    if hasattr(org_img, 'orientation') and "".join(org_img.orientation) == "LSA":
+        mask_data = mask_data.permute(0, 1, 3, 2)
+        mask_data = torch.flip(mask_data, dims=[21])
         mask_data = torch.Tensor(mask_data)[None]
 
-    # Save the prediction
-    img = org_img.data
+    # Save the prediction as a TorchIO image using original image object
+    img = org_img.data if hasattr(org_img, 'data') else None
     org_img.set_data(mask_data)
     org_img.save(save_path)
 
-    # Extract slices and create fused images
-    slices = extract_slices_from_mask(img[0].numpy(), mask_data[0].numpy(), view)
-    fused_images = [(get_fused_image(
-                        normalize_image(slice_img),
-                        slice_mask, view))
-                    for slice_img, slice_mask in slices]
-    
+    # Build fused gallery
+    base_img_np = img.numpy() if img is not None else input_img.numpy()
+    pred_np = mask_data.detach().cpu().numpy().astype(np.uint8)
+    slices = extract_slices_from_mask(base_img_np, pred_np, view)
+    fused_images = [
+        get_fused_image(normalize_image(slice_img), slice_mask, view)
+        for slice_img, slice_mask in slices
+    ]
+
     # Compute volume
     volume = compute_binary_tumor_volume(org_img)
-
     return fused_images, round(volume, 2)
 
-# Function to load system resources
+# -----------------------------
+# Resources
+# -----------------------------
+
 def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_fn='vars.pkl'):
     """Load the model and other required resources."""
     try:
@@ -185,33 +260,29 @@ def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_f
     try:
         with open(models_path / variables_fn, 'rb') as f:
             variables = pickle.load(f)
-           
         if not isinstance(variables, list) or len(variables) != 3:
             raise ValueError(f"vars.pkl does not contain the expected list format. Found: {variables}")
-           
-        # Assuming the format is [shape, reorder, resample]
-        shape = variables[0]
-        reorder = variables[1]
-        resample = variables[2]
-           
+        shape = variables
+        reorder = variables[21]
+        resample = variables[6]
         if not isinstance(reorder, bool):
             raise ValueError(f"vars.pkl does not contain a valid 'reorder' value. Found: {reorder}")
-           
         if not isinstance(resample, list) or len(resample) != 3:
             raise ValueError(f"vars.pkl does not contain a valid 'resample' value. Found: {resample}")
-           
     except Exception as e:
         raise ValueError(f"Error loading variables: {str(e)}")
 
     return learn, reorder, resample
 
-# Initialize the system
+# -----------------------------
+# Main
+# -----------------------------
+
 try:
     print("🚀 Initializing application...")
-    
+
     clone_dir = Path.cwd() / 'clone_dir'
     URI = os.getenv('PAT_Token_URI')
-
     if not URI:
         raise ValueError("PAT_Token_URI environment variable is not set")
 
@@ -227,17 +298,20 @@ try:
     save_dir.mkdir(parents=True, exist_ok=True)
     print(f"✅ Save directory created at {save_dir}")
 
-    # Load the model and other required resources
+    # Load model and variables
     print("🔍 Loading model and resources...")
     learn, reorder, resample = load_system_resources(models_path=models_path)
     print(f"✅ Model loaded successfully")
     print(f"✅ Reorder: {reorder}, Resample: {resample}")
 
-    # Gradio interface setup
+    # Apply robust ZNormalization compatibility patches BEFORE any predict is called
+    apply_znorm_compat_patches(learn)
+    print("✅ Applied ZNormalization compatibility patches")
+
+    # Gradio UI
     output_text = gr.Textbox(label="Volume of the Left Atrium (mL):")
     view_selector = gr.Radio(choices=["Axial", "Coronal", "Sagittal"], value='Sagittal', label="Select View (Sagittal by default)")
 
-    # Ensure the example file path is correct
     example_path = str(clone_dir / "sample.nii.gz")
     if not os.path.exists(example_path):
         print(f"⚠️  Example file not found: {example_path}")
@@ -251,18 +325,17 @@ try:
         inputs=["file", view_selector],
         outputs=[gr.Gallery(label="Click an Image, and use Arrow Keys to scroll slices", columns=3, height=450), output_text],
         examples=examples,
-        cache_examples=False,  # ADD THIS PARAMETER
-        allow_flagging='never')
-    
+        allow_flagging='never',
+        cache_examples=False  # avoid eager caching running predict at startup
+    )
+
     print("✅ Gradio interface initialized successfully")
-    
+
 except Exception as e:
     print(f"❌ Error during initialization: {str(e)}")
     print(f"Error type: {type(e).__name__}")
     traceback.print_exc()
-    # Exit with error code
     exit(1)
 
-# Launch the Gradio interface
 print("🌐 Launching Gradio interface...")
-demo.launch()
+demo.launch()