Update app.py

app.py

@@ -7,13 +7,21 @@ from huggingface_hub import snapshot_download
 from fastMONAI.vision_all import *
 from git import Repo
 import os
-from fastai.learner import load_learner
-from fastai.basics import load_pickle
-import pickle
-import traceback
-import torchio as tio
 
-# Function to extract slices from mask
+#Additional support for local execution:-
+#import pathlib
+#temp = pathlib.PosixPath
+#pathlib.PosixPath = pathlib.WindowsPath
+#pathlib.PosixPath = temp
+
+clone_dir = Path.cwd() / 'clone_dir'
+URI = os.getenv('PAT_Token_URI')
+
+if os.path.exists(clone_dir):
+    pass
+else:
+    Repo.clone_from(URI, clone_dir)
+
 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 = []
@@ -35,7 +43,6 @@ def extract_slices_from_mask(img, mask_data, view):
     
     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
@@ -53,7 +60,6 @@ def resize_and_pad(slice_img, slice_mask, target_size):
 
     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()
@@ -62,7 +68,6 @@ def normalize_image(slice_img):
     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)
@@ -74,369 +79,71 @@ def get_fused_image(img, pred_mask, view, alpha=0.8):
     # Flip the fused image vertically and horizontally
     fused_flipped = cv2.flip(fused, -1)  # Flip both vertically and horizontally
     
-    if view == 'Sagittal':
+    if view=='Sagittal':
         return fused_flipped
-    elif view == 'Coronal' or view == 'Axial':
+    elif view=='Coronal' or 'Axial':
         rotated = cv2.flip(cv2.rotate(fused, cv2.ROTATE_90_COUNTERCLOCKWISE), 1)
         return rotated
 
-# Define the inference function
-def inference(learn, reorder, resample, org_img, input_img, org_size):
-    """Perform segmentation using the loaded model."""
-    print("\n=== Inference Debug ===")
-    print(f"Initial input shape: {input_img.shape if isinstance(input_img, torch.Tensor) else 'Not a tensor'}")
-    
-    # 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)}")
-    
-    # Add batch dimension if needed
-    if input_img.dim() == 4:
-        input_img = input_img.unsqueeze(0)
-        print(f"Added batch dim - shape: {input_img.shape}")
-    
-    # Verify we now have the correct dimensions
-    if input_img.dim() != 5:
-        raise ValueError(f"Input tensor must be 5D after preprocessing, got {input_img.dim()}D tensor with shape {input_img.shape}")
-    
-    print(f"Input tensor shape before prediction: {input_img.shape}")
-    
-    # Perform the segmentation
-    with torch.no_grad():
-        try:
-            print(f"\nModel input shape before processing: {input_img.shape}")
-            
-            # Convert to plain tensor if wrapped in MedImage
-            if hasattr(input_img, 'data'):
-                input_img = input_img.data
-                print(f"Unwrapped MedImage to tensor with shape: {input_img.shape}")
-            
-            # Ensure we have exactly 5 dimensions [1,1,320,320,140]
-            if input_img.dim() == 6:
-                input_img = input_img.squeeze(0)  # Remove extra batch dim
-            elif input_img.dim() == 4:
-                input_img = input_img.unsqueeze(0)  # Add batch dim
-            
-            # Verify final shape
-            if input_img.shape != torch.Size([1,1,320,320,140]):
-                input_img = input_img.view(1, 1, 320, 320, 140)
-                print(f"Reshaped to required dimensions: {input_img.shape}")
-            
-            print(f"Final input shape to model: {input_img.shape}")
-            pred = learn.predict(input_img)
-            print(f"Prediction output shape: {pred[0].shape if isinstance(pred, (list, tuple)) else pred.shape}")
-        except Exception as e:
-            print(f"Error during prediction: {str(e)}")
-            print(f"Current input shape: {input_img.shape}")
-            print(f"Input tensor contents: {input_img}")
-            raise
-    
-    # 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:
+    if view == None:
         view = 'Sagittal'
     
     img_path = Path(fileobj.name)
+
     save_fn = 'pred_' + img_path.stem
     save_path = save_dir / save_fn
+    org_img, input_img, org_size = med_img_reader(img_path, 
+                                                  reorder=reorder,
+                                                  resample=resample,
+                                                  only_tensor=False)
     
-    # Read the medical image - handle variable return values
-    try:
-        print(f"\nReading medical image: {img_path}")
-        result = med_img_reader(str(img_path), 
-                               reorder=reorder,
-                               resample=resample,
-                               only_tensor=False,
-                               dtype=torch.Tensor)
-        print(f"med_img_reader returned type: {type(result)}")
-        if isinstance(result, tuple):
-            print(f"Tuple length: {len(result)}")
-            for i, item in enumerate(result):
-                print(f"Item {i} type: {type(item)}")
-                if hasattr(item, 'shape'):
-                    print(f"Item {i} shape: {item.shape}")
-        
-        # Handle different return formats
-        if isinstance(result, tuple):
-            if len(result) == 3:
-                org_img, input_img, org_size = result
-            elif len(result) == 2:
-                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
-            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():
-            print(f"DEBUG: med_img_reader returned: {type(result)}")
-            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
-    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
-    print(f"\nInput tensor shape before inference: {input_img.shape}")
-    
-    # Convert to plain tensor if wrapped in MedImage
-    if hasattr(input_img, 'data'):
-        input_img = input_img.data
-        print(f"Unwrapped MedImage to tensor with shape: {input_img.shape}")
-    
-    # Directly reshape to required 5D format [1,1,320,320,140]
-    try:
-        input_img = input_img.reshape(1, 1, 320, 320, 140)
-        print(f"Reshaped to required dimensions: {input_img.shape}")
-    except Exception as e:
-        raise ValueError(f"Could not reshape tensor from {input_img.shape} to [1,1,320,320,140]: {str(e)}")
     mask_data = inference(learn, reorder=reorder, resample=resample,
                           org_img=org_img, input_img=input_img,
-                          org_size=org_size)
-    print(f"Raw mask data shape: {mask_data.shape}")
+                          org_size=org_size).data
     
-    # Handle orientation if needed
-    if hasattr(org_img, 'orientation') and "".join(org_img.orientation) == "LSA":        
-        print("Applying LSA orientation correction")
+    if "".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]
-        print(f"Orientation corrected shape: {mask_data.shape}")
 
-    # Save the prediction
     img = org_img.data
     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)
+    slices = extract_slices_from_mask(img[0], mask_data[0], view)
     fused_images = [(get_fused_image(
-                        normalize_image(slice_img),
+                        normalize_image(slice_img),  # Normalize safely
                         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
-def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_fn='vars.pkl'):
-    """Load the model and other required resources."""
-    try:
-        learn = load_learner(models_path / learner_fn)
-        print(f"✅ Model loaded from {models_path / learner_fn}")
-    except Exception as e:
-        raise ValueError(f"Error loading the model: {str(e)}")
-
-    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]
-           
-        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
-try:
-    print("🚀 Initializing application...")
-    
-    clone_dir = Path.cwd() / 'clone_dir'
-    URI = os.getenv('PAT_Token_URI')
+models_path = Path.cwd() / 'clone_dir'
+save_dir = Path.cwd() / 'hs_pred'
+save_dir.mkdir(parents=True, exist_ok=True)
 
-    if not URI:
-        raise ValueError("PAT_Token_URI environment variable is not set")
+# Load the model and other required resources
+learn, reorder, resample = load_system_resources(models_path=models_path,
+                                                 learner_fn='heart_model.pkl',
+                                                 variables_fn='vars.pkl')
 
-    if os.path.exists(clone_dir):
-        print(f"✅ Repository already exists at {clone_dir}")
-    else:
-        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 is the ONLY patch needed - it directly fixes the root cause
-    from fastMONAI.vision_augmentation import ZNormalization
-    from fastMONAI.vision_core import MedImage
-    
-    # Save the original encodes method
-    original_encodes = ZNormalization.encodes
-
-    def patched_encodes(self, o):
-        """Patched version of ZNormalization.encodes with custom normalization."""
-        print("\n=== Custom ZNormalization ===")
-        print(f"Input type: {type(o)}")
-        
-        try:
-            # Get the input tensor data
-            if hasattr(o, 'data') and isinstance(o.data, torch.Tensor):
-                tensor_data = o.data
-            elif isinstance(o, torch.Tensor):
-                tensor_data = o
-            else:
-                print("Input is not a tensor - converting")
-                tensor_data = torch.tensor(o)
-            
-            print(f"Input shape: {tensor_data.shape}")
-            print(f"Input dtype: {tensor_data.dtype}")
-            
-            # Ensure we have the right dimensions
-            if tensor_data.dim() == 6:
-                tensor_data = tensor_data.squeeze(0)
-                print(f"Removed extra dim - new shape: {tensor_data.shape}")
-            
-            # Calculate mean and std
-            mean = tensor_data.mean()
-            std = tensor_data.std()
-            print(f"Calculated mean: {mean:.4f}, std: {std:.4f}")
-            
-            # Avoid division by zero
-            if std == 0:
-                std = 1.0
-                print("Warning: std is 0, setting to 1.0")
-            
-            # Apply normalization
-            normalized = (tensor_data - mean) / std
-            print(f"Normalized tensor shape: {normalized.shape}")
-            
-            # Return as plain tensor to avoid MedImage wrapping issues
-            return normalized
-            
-        except Exception as e:
-            print(f"Error in custom normalization: {str(e)}")
-            print("Fallback - returning input as-is")
-            return MedImage.create(o)
-        
-        print("=== End Custom ZNormalization ===\n")
+# Gradio interface setup
+output_text = gr.Textbox(label="Volume of the Left Atrium (mL):")
 
-    # Add PadOrCrop debug prints
-    from fastMONAI.vision_augmentation import PadOrCrop
-    original_padorcrop_encodes = PadOrCrop.encodes
+view_selector = gr.Radio(choices=["Axial", "Coronal", "Sagittal"], value='Sagittal', label="Select View (Sagittal by default)")
 
-    def patched_padorcrop_encodes(self, o):
-        """Patched version of PadOrCrop.encodes that ensures tensor output."""
-        print("\n=== PadOrCrop Debug ===")
-        print(f"Input type: {type(o)}")
-        if hasattr(o, 'shape'): print(f"Input shape: {o.shape}")
-        
-        # Safely get attributes with fallbacks
-        target_shape = getattr(self, 'size', o.shape[1:])  # Default to input shape
-        padding_mode = getattr(self, 'padding_mode', 'constant')
-        mask_name = getattr(self, 'mask_name', None)
-        
-        print(f"Target shape: {target_shape}")
-        print(f"Padding mode: {padding_mode}")
-        print(f"Mask name: {mask_name}")
-        
-        try:
-            # Get the transformation result
-            result = original_padorcrop_encodes(self, o)
-            print(f"Initial output type: {type(result)}")
-            
-            # Extract tensor data from result
-            if hasattr(result, 'data'):
-                print("Extracting .data attribute")
-                tensor_data = result.data
-            elif isinstance(result, torch.Tensor):
-                print("Using tensor directly")
-                tensor_data = result
-            else:
-                print("Unknown output type - using input")
-                tensor_data = o.data if hasattr(o, 'data') else o
-            
-            # Ensure we have a tensor
-            if not isinstance(tensor_data, torch.Tensor):
-                print("Converting to tensor")
-                tensor_data = torch.tensor(tensor_data)
-            
-            print(f"Final output shape: {tensor_data.shape}")
-            return MedImage.create(tensor_data)
-                
-        except Exception as e:
-            print(f"Error in PadOrCrop: {str(e)}")
-            print("Fallback - returning input as MedImage")
-            return MedImage.create(o.data if hasattr(o, 'data') else o)
-        
-        print("=== End PadOrCrop Debug ===\n")
-
-    PadOrCrop.encodes = patched_padorcrop_encodes
-    print("✅ Applied PadOrCrop patch with debug prints")
-
-    # Replace the encodes method
-    ZNormalization.encodes = patched_encodes
-    print("✅ Applied ZNormalization patch to handle missing attributes in encodes method")
-    # ===== 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 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}")
-
-    # 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}")
-        examples = None
-    else:
-        print(f"✅ Example file found: {example_path}")
-        examples = [[example_path]]
-
-    demo = gr.Interface(
-        fn=lambda fileobj, view='Sagittal': gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view),
-        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,  # 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)
+demo = gr.Interface(
+    fn=lambda fileobj, view='Sagittal': gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view),
+    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=[[str(Path.cwd() /"sample.nii.gz")]],
+    allow_flagging='never')
 
 # Launch the Gradio interface
-print("🌐 Launching Gradio interface...")
-demo.launch()
+demo.launch()