Update app.py

app.py

@@ -11,195 +11,129 @@ from fastai.learner import load_learner
 from fastai.basics import load_pickle
 import pickle
 import traceback
+import torchio as tio  # Added for the patch
 
-# -----------------------------
-# Utility functions
-# -----------------------------
-
+# Function to extract slices from mask
 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[6] if view == "Sagittal" else img.shape[21] if view == "Axial" else img.shape):
+    
+    for idx in range(img.shape[2] if view == "Sagittal" else img.shape[1] if view == "Axial" else img.shape[0]):
         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 / w, target_size[21] / h)
+    scale = min(target_size[0] / w, target_size[1] / 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 - 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)
+
+    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)
+
     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:
+    if slice_img_min == slice_img_max:  # Avoid division by zero
         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)
-    fused_flipped = cv2.flip(fused, -1)
+    
+    # Flip the fused image vertically and horizontally
+    fused_flipped = cv2.flip(fused, -1)  # Flip both vertically and horizontally
+    
     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
 
-# -----------------------------
-# 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]
-
+# Define the inference function
 def inference(learn, reorder, resample, org_img, input_img, org_size):
-    """Perform segmentation using the loaded model, with robust fallback."""
+    """Perform segmentation using the loaded model."""
+    # Ensure input_img is a torch.Tensor
     if not isinstance(input_img, torch.Tensor):
         raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
-    # 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
-# -----------------------------
-
+    
+    # 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
 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 medical image with robust handling of return formats
+    
+    # Read the medical image - handle variable return values
     try:
-        result = med_img_reader(str(img_path),
-                                reorder=reorder,
-                                resample=resample,
-                                only_tensor=False,
-                                dtype=torch.Tensor)
+        result = med_img_reader(str(img_path), 
+                               reorder=reorder,
+                               resample=resample,
+                               only_tensor=False,
+                               dtype=torch.Tensor)
+        
+        # Handle different return formats
         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
-                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
+                # 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
             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():
@@ -207,48 +141,40 @@ 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 tensor with appropriate dims [C, W, H, D] or [1, C, W, H, D]
+    
+    # Ensure input_img is a torch.Tensor
     if not isinstance(input_img, torch.Tensor):
         raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
-    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)
+    
+    # Perform inference
     mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=x_in,
+                          org_img=org_img, input_img=input_img,
                           org_size=org_size)
-
-    # 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])
+    
+    # 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])
         mask_data = torch.Tensor(mask_data)[None]
 
-    # Save the prediction as a TorchIO image using original image object
-    img = org_img.data if hasattr(org_img, 'data') else None
+    # Save the prediction
+    img = org_img.data
     org_img.set_data(mask_data)
     org_img.save(save_path)
 
-    # 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
-    ]
-
+    # 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]
+    
     # Compute volume
     volume = compute_binary_tumor_volume(org_img)
-    return fused_images, round(volume, 2)
 
-# -----------------------------
-# Resources
-# -----------------------------
+    return fused_images, round(volume, 2)
 
+# Function to load system resources
 def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_fn='vars.pkl'):
     """Load the model and other required resources."""
     try:
@@ -260,44 +186,33 @@ 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)
-
-        # Accept dict or list/tuple; be defensive against shape changes
-        shape = None
-        if isinstance(variables, dict):
-            reorder = bool(variables.get('reorder', False))
-            resample = variables.get('resample', [1.25, 1.25, 1.37])
-        elif isinstance(variables, (list, tuple)):
-            if len(variables) >= 3:
-                shape, reorder, resample = variables, variables[19], variables[20]
-            elif len(variables) == 2:
-                reorder, resample = variables, variables[19]
-            else:
-                raise ValueError(f"vars.pkl has unexpected length: {len(variables)}")
-        else:
-            raise ValueError(f"vars.pkl has unexpected type: {type(variables)}")
-
-        # Validate types strictly
+           
+        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]
+           
         if not isinstance(reorder, bool):
-            raise ValueError(f"vars.pkl 'reorder' must be bool, got {type(reorder)}")
-        if not (isinstance(resample, (list, tuple)) and len(resample) == 3):
-            raise ValueError(f"vars.pkl 'resample' must be a list/tuple of len 3, got {resample}")
-        resample = list(resample)
-
+            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
 
-
-# -----------------------------
-# Main
-# -----------------------------
-
+# Initialize the system
 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")
 
@@ -307,26 +222,72 @@ try:
         print(f"📥 Cloning repository from {URI}")
         Repo.clone_from(URI, clone_dir)
         print(f"✅ Repository cloned to {clone_dir}")
-
+    
+    # ===== MONKEY PATCH FOR ZNORMALIZATION ERROR =====
+    # This patch fixes the 'ZNormalization' object has no attribute 'masking_method' error
+    # Import necessary modules
+    from fastMONAI.vision_augmentation import ZNormalization
+    from fastMONAI.vision_core import MedImage
+    
+    # Create a patched version of the encodes method
+    def patched_encodes(self, o):
+        """Patched version of ZNormalization.encodes to handle attribute mismatch."""
+        # Handle the attribute mismatch - ensure we have what we need
+        if not hasattr(self, 'z_normalization'):
+            # Get masking_method with fallback
+            masking_method = getattr(self, 'masking_method', None)
+            self.z_normalization = tio.ZNormalization(masking_method=masking_method)
+        
+        # Ensure channel_wise exists
+        if not hasattr(self, 'channel_wise'):
+            self.channel_wise = True
+        
+        # Use the current implementation logic
+        try:
+            if self.channel_wise:
+                o = torch.stack([self.z_normalization(c[None])[0] for c in o])
+            else: 
+                o = self.z_normalization(o)
+        except RuntimeError as e:
+            if "Standard deviation is 0" in str(e):
+                mean = float(o.mean())
+                error_msg = (
+                    f"Standard deviation is 0 for image (mean={mean:.3f}).\n"
+                    f"This indicates uniform pixel values.\n\n"
+                    f"Possible causes:\n"
+                    f"• Corrupted or blank image\n"
+                    f"• Oversaturated regions\n"
+                    f"• Background-only regions\n"
+                    f"• All-zero mask being processed as image\n\n"
+                    f"Suggested solutions:\n"
+                    f"• Check image quality and acquisition\n"
+                    f"• Verify image vs mask data loading"
+                )
+                raise RuntimeError(error_msg) from e
+
+        return MedImage.create(o)
+    
+    # Replace the problematic method
+    ZNormalization.encodes = patched_encodes
+    print("✅ Applied ZNormalization patch to fix attribute error")
+    # ===== END OF MONKEY PATCH =====
+    
     models_path = clone_dir
     save_dir = Path.cwd() / 'hs_pred'
     save_dir.mkdir(parents=True, exist_ok=True)
     print(f"✅ Save directory created at {save_dir}")
 
-    # Load model and variables
+    # Load the model and other required resources
     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}")
 
-    # Apply robust ZNormalization compatibility patches BEFORE any predict is called
-    apply_znorm_compat_patches(learn)
-    print("✅ Applied ZNormalization compatibility patches")
-
-    # Gradio UI
+    # Gradio interface setup
     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}")
@@ -340,17 +301,18 @@ 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,
-        allow_flagging='never',
-        cache_examples=False  # avoid eager caching running predict at startup
-    )
-
+        cache_examples=False,  # Disable example caching to prevent errors during startup
+        allow_flagging='never')
+    
     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()