Update app.py

app.py

@@ -10,6 +10,7 @@ import os
 from fastai.learner import load_learner
 from fastai.basics import load_pickle
 import pickle
+import traceback
 
 # Function to extract slices from mask
 def extract_slices_from_mask(img, mask_data, view):
@@ -83,61 +84,21 @@ def inference(learn, reorder, resample, org_img, input_img, org_size):
     """Perform segmentation using the loaded model."""
     # Ensure input_img is a torch.Tensor
     if not isinstance(input_img, torch.Tensor):
-        # If input_img is not a tensor, try to extract the tensor data
-        if hasattr(input_img, 'data'):
-            input_tensor = input_img.data
-        else:
-            raise ValueError(f"Expected input_img to be a torch.Tensor or have a 'data' attribute, but got {type(input_img)}")
-    else:
-        input_tensor = input_img
+        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_tensor)
+        pred = learn.predict(input_img)
     
     # Process the prediction if necessary
-    mask_data = pred[0]  # Assuming the first element of the prediction is the mask
+    mask_data = pred[0] if isinstance(pred, (list, tuple)) else pred
     
     return mask_data
 
-# Function to extract tensor from different data types
-def extract_tensor_from_data(data, fallback_path=None):
-    """Extract tensor from various data types returned by med_img_reader."""
-    print(f"DEBUG: Extracting tensor from data type: {type(data)}")
-    
-    if isinstance(data, torch.Tensor):
-        return data
-    elif isinstance(data, dict):
-        print(f"DEBUG: Dictionary keys: {list(data.keys())}")
-        # Try common key names for tensor data
-        possible_keys = ['data', 'tensor', 'image', 'img', 'input', 'volume']
-        for key in possible_keys:
-            if key in data:
-                tensor_data = data[key]
-                if isinstance(tensor_data, torch.Tensor):
-                    print(f"DEBUG: Found tensor in key '{key}'")
-                    return tensor_data
-                elif hasattr(tensor_data, 'data') and isinstance(tensor_data.data, torch.Tensor):
-                    print(f"DEBUG: Found tensor in key '{key}'.data")
-                    return tensor_data.data
-        
-        # If no tensor found in dict, try to load directly
-        if fallback_path:
-            print("DEBUG: No tensor found in dict, trying direct load")
-            try:
-                direct_tensor = med_img_reader(fallback_path, only_tensor=True, dtype=torch.Tensor)
-                return direct_tensor
-            except Exception as e:
-                print(f"DEBUG: Direct load failed: {e}")
-    elif hasattr(data, 'data') and isinstance(data.data, torch.Tensor):
-        print("DEBUG: Found tensor in data.data attribute")
-        return data.data
-    elif hasattr(data, 'tensor') and isinstance(data.tensor, torch.Tensor):
-        print("DEBUG: Found tensor in tensor attribute")
-        return data.tensor
-    
-    raise ValueError(f"Cannot extract tensor from data of type {type(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."""
@@ -146,159 +107,68 @@ def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view)
         view = 'Sagittal'
     
     img_path = Path(fileobj.name)
-
-    # Convert PosixPath to string
-    img_path_str = str(img_path)
-
     save_fn = 'pred_' + img_path.stem
     save_path = save_dir / save_fn
     
-    # Try different approaches to load the image
-    org_img = None
-    input_tensor = None
-    org_size = None
-    
+    # Read the medical image - handle variable return values
     try:
-        # First try with only_tensor=False to get all values
-        result = med_img_reader(img_path_str, 
-                              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)
         
-        print(f"DEBUG: med_img_reader returned type: {type(result)}")
-        
-        # Handle different return types
+        # Handle different return formats
         if isinstance(result, tuple):
-            print(f"DEBUG: Tuple length: {len(result)}")
-            if len(result) >= 2:
-                org_img = result[0]
-                input_img = result[1]
-                org_size = result[2] if len(result) > 2 else None
-                
-                # Extract tensor from input_img
-                input_tensor = extract_tensor_from_data(input_img, img_path_str)
-                
-                # Handle org_size if not provided
-                if org_size is None and hasattr(input_tensor, 'shape'):
-                    org_size = input_tensor.shape[1:]
+            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 tuple length from med_img_reader: {len(result)}")
-                
-        elif isinstance(result, dict):
-            print(f"DEBUG: Dictionary keys: {list(result.keys())}")
-            
-            # Try to extract org_img
-            org_img_keys = ['org_img', 'original', 'original_image', 'source']
-            for key in org_img_keys:
-                if key in result:
-                    org_img = result[key]
-                    break
-            
-            # Extract input tensor
-            input_tensor = extract_tensor_from_data(result, img_path_str)
-            
-            # Extract org_size
-            if 'org_size' in result:
-                org_size = result['org_size']
-            elif hasattr(input_tensor, 'shape'):
-                org_size = input_tensor.shape[1:]
-                
+                raise ValueError(f"Unexpected number of return values: {len(result)}")
         else:
-            # If it's neither tuple nor dict, try to extract tensor directly
-            input_tensor = extract_tensor_from_data(result, img_path_str)
-            
-        # If org_img is still None, try to create it
-        if org_img is None:
-            try:
-                from torchio import ScalarImage
-                org_img = ScalarImage(img_path_str)
-            except Exception as e:
-                print(f"DEBUG: Failed to create ScalarImage: {e}")
-                # Create a minimal org_img object
-                class MinimalImage:
-                    def __init__(self, tensor, path):
-                        self.data = tensor
-                        self.path = path
-                        self.orientation = "RAS"  # Default orientation
-                    
-                    def set_data(self, data):
-                        self.data = data
-                    
-                    def save(self, path):
-                        # Minimal save implementation
-                        pass
-                
-                org_img = MinimalImage(input_tensor, img_path_str)
-        
-        # Ensure org_size is set
-        if org_size is None and hasattr(input_tensor, 'shape'):
-            org_size = input_tensor.shape[1:]
+            # 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 handling: {str(e)}")
-        # Final fallback: try to load tensor directly
-        try:
-            input_tensor = med_img_reader(img_path_str, 
-                                        reorder=reorder,
-                                        resample=resample,
-                                        only_tensor=True,
-                                        dtype=torch.Tensor)
-            
-            # Create minimal org_img
-            class MinimalImage:
-                def __init__(self, tensor, path):
-                    self.data = tensor
-                    self.path = path
-                    self.orientation = "RAS"
-                
-                def set_data(self, data):
-                    self.data = data
-                
-                def save(self, path):
-                    pass
-            
-            org_img = MinimalImage(input_tensor, img_path_str)
-            org_size = input_tensor.shape[1:] if hasattr(input_tensor, 'shape') else None
-            
-        except Exception as e2:
-            raise ValueError(f"Failed to load image: {str(e2)}")
+        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 we have a valid input tensor
-    if input_tensor is None:
-        raise ValueError("Could not extract input tensor from loaded image data")
-    
-    print(f"DEBUG: Final input_tensor type: {type(input_tensor)}, shape: {getattr(input_tensor, 'shape', 'No shape')}")
+    # 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
     mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=input_tensor,
+                          org_img=org_img, input_img=input_img,
                           org_size=org_size)
     
-    # Handle orientation if org_img has orientation attribute
-    if hasattr(org_img, 'orientation') and "".join(getattr(org_img, 'orientation', '')) == "LSA":        
+    # 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]
 
-    # Extract data from org_img properly
-    if hasattr(org_img, 'data'):
-        img = org_img.data
-    else:
-        img = input_tensor  # Use input_tensor as fallback
-    
-    # Set data and save
-    if hasattr(org_img, 'set_data'):
-        org_img.set_data(mask_data)
-    if hasattr(org_img, 'save'):
-        org_img.save(save_path)
+    # Save the prediction
+    img = org_img.data
+    org_img.set_data(mask_data)
+    org_img.save(save_path)
 
-    slices = extract_slices_from_mask(img[0], mask_data[0], view)
+    # 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),  # Normalize safely
+                        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)
@@ -308,6 +178,7 @@ def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_f
     """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)}")
 
@@ -335,38 +206,62 @@ def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_f
     return learn, reorder, resample
 
 # Initialize the system
-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")
+try:
+    print("🚀 Initializing application...")
+    
+    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)
+    if not URI:
+        raise ValueError("PAT_Token_URI environment variable is not set")
 
-models_path = clone_dir
-save_dir = Path.cwd() / 'hs_pred'
-save_dir.mkdir(parents=True, exist_ok=True)
+    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}")
 
-# Load the model and other required resources
-learn, reorder, resample = load_system_resources(models_path=models_path)
+    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}")
 
-# Gradio interface setup
-output_text = gr.Textbox(label="Volume of the Left Atrium (mL):")
+    # 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}")
 
-view_selector = gr.Radio(choices=["Axial", "Coronal", "Sagittal"], value='Sagittal', label="Select View (Sagittal by default)")
+    # 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")
+    # 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=[[example_path, 'Sagittal']],  # Added view parameter to example
-    allow_flagging='never')
+    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,
+        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()