Update app.py

app.py

@@ -227,29 +227,78 @@ try:
         Repo.clone_from(URI, clone_dir)
         print(f"✅ Repository cloned to {clone_dir}")
     
-    # ===== MONKEY PATCH FOR PADORCROP ERROR =====
-    # This patch fixes the 'PadOrCrop' object has no attribute 'size' error
-    from fastMONAI.vision_augmentation import PadOrCrop
-    
-    def patched_pad_or_crop_encodes(self, o):
-        """Patched version of PadOrCrop.encodes to handle attribute mismatch."""
-        # Access the attributes through self.pad_or_crop instead of directly
-        return do_pad_or_crop(
-            o,
-            target_shape=self.pad_or_crop.target_shape,
-            padding_mode=self.pad_or_crop.padding_mode,
-            mask_name=self.pad_or_crop.mask_name,
-            dtype=type(o)
-        )
+    # ===== MONKEY PATCH FOR ZNORMALIZATION ERROR =====
+    # This patch fixes the 'ZNormalization' object has no attribute 'masking_method' error
+    from fastMONAI.vision_augmentation import ZNormalization
+    from fastMONAI.vision_core import MedImage
     
+    def patched_znormalization_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 with proper dimension handling
+        try:
+            if self.channel_wise:
+                if o.dim() == 5:  # (batch, channels, x, y, z)
+                    results = []
+                    for batch_item in o:
+                        channel_results = []
+                        for channel in batch_item:
+                            # channel should be 3D (x, y, z)
+                            normalized = self.z_normalization(channel[None])[0]
+                            channel_results.append(normalized)
+                        results.append(torch.stack(channel_results))
+                    o = torch.stack(results)
+                elif o.dim() == 4:  # (channels, x, y, z)
+                    channel_results = []
+                    for channel in o:
+                        # channel should be 3D (x, y, z)
+                        normalized = self.z_normalization(channel[None])[0]
+                        channel_results.append(normalized)
+                    o = torch.stack(channel_results)
+                else:  # (x, y, z) - single channel image
+                    o = self.z_normalization(o[None])[0]
+            else:
+                if o.dim() == 5:  # (batch, channels, x, y, z)
+                    results = []
+                    for batch_item in o:
+                        normalized = self.z_normalization(batch_item)
+                        results.append(normalized)
+                    o = torch.stack(results)
+                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
-    PadOrCrop.encodes = patched_pad_or_crop_encodes
-    print("✅ Applied PadOrCrop patch to fix attribute errors")
+    ZNormalization.encodes = patched_znormalization_encodes
+    print("✅ Applied ZNormalization patch to fix attribute errors")
     # ===== END OF MONKEY PATCH =====
     
-    # Note: We don't need the ZNormalization patch anymore
-    # The error is specifically with PadOrCrop
-    
     models_path = clone_dir
     save_dir = Path.cwd() / 'hs_pred'
     save_dir.mkdir(parents=True, exist_ok=True)