Update app.py

app.py

@@ -10,7 +10,6 @@ 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):
@@ -84,18 +83,20 @@ 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):
-        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)
+        # 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
     
     # Perform the segmentation
     with torch.no_grad():
-        pred = learn.predict(input_img)
+        pred = learn.predict(input_tensor)
     
     # Process the prediction if necessary
-    mask_data = pred[0] if isinstance(pred, (list, tuple)) else pred
+    mask_data = pred[0]  # Assuming the first element of the prediction is the mask
     
     return mask_data
 
@@ -107,47 +108,122 @@ 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
     
-    # Read the medical image
+    # Debug: Let's check what med_img_reader actually returns
     try:
-        org_img, input_img, org_size = med_img_reader(str(img_path), 
-                                                     reorder=reorder,
-                                                     resample=resample,
-                                                     only_tensor=False,
-                                                     dtype=torch.Tensor)
+        # 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)
+        
+        # Debug print to understand the structure
+        print(f"DEBUG: med_img_reader returned type: {type(result)}")
+        
+        # Handle different return types
+        if isinstance(result, tuple):
+            if len(result) == 3:
+                org_img, input_img, org_size = result
+            elif len(result) == 2:
+                org_img, input_img = result
+                # Infer org_size from org_img
+                if hasattr(org_img, 'shape'):
+                    org_size = org_img.shape[1:]
+                else:
+                    org_size = None
+            else:
+                raise ValueError(f"Unexpected tuple length from med_img_reader: {len(result)}")
+        elif isinstance(result, dict):
+            # If it's a dict, try to extract the needed values
+            # This might happen if the function returns something wrapped
+            print(f"DEBUG: Dictionary keys: {result.keys() if isinstance(result, dict) else 'N/A'}")
+            # Try to extract values - adjust based on actual keys
+            org_img = result.get('org_img', result.get('original', None))
+            input_img = result.get('input_img', result.get('input', result.get('data', None)))
+            org_size = result.get('org_size', result.get('size', None))
+            
+            if org_img is None or input_img is None:
+                # If we can't extract from dict, try getting tensor directly
+                input_tensor = med_img_reader(img_path_str, 
+                                            reorder=reorder,
+                                            resample=resample,
+                                            only_tensor=True,
+                                            dtype=torch.Tensor)
+                # Create dummy org_img
+                from torchio import ScalarImage
+                org_img = ScalarImage(img_path_str)
+                input_img = input_tensor
+                org_size = input_tensor.shape[1:] if hasattr(input_tensor, 'shape') else None
+        else:
+            # If it's neither tuple nor dict, it might be the tensor directly
+            # Try to get the full data with only_tensor=True
+            input_tensor = med_img_reader(img_path_str, 
+                                        reorder=reorder,
+                                        resample=resample,
+                                        only_tensor=True,
+                                        dtype=torch.Tensor)
+            # Create dummy org_img
+            from torchio import ScalarImage
+            org_img = ScalarImage(img_path_str)
+            input_img = input_tensor
+            org_size = input_tensor.shape[1:] if hasattr(input_tensor, 'shape') else None
+            
     except Exception as e:
-        raise ValueError(f"Error reading medical image: {str(e)}")
+        print(f"DEBUG: Error in med_img_reader: {str(e)}")
+        # Fallback: try with only_tensor=True
+        try:
+            input_tensor = med_img_reader(img_path_str, 
+                                        reorder=reorder,
+                                        resample=resample,
+                                        only_tensor=True,
+                                        dtype=torch.Tensor)
+            # Create dummy org_img
+            from torchio import ScalarImage
+            org_img = ScalarImage(img_path_str)
+            input_img = input_tensor
+            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)}")
     
-    # 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)}")
+    # Ensure input_img is proper format for inference
+    if hasattr(input_img, 'data') and isinstance(input_img.data, torch.Tensor):
+        input_tensor = input_img.data
+    elif isinstance(input_img, torch.Tensor):
+        input_tensor = input_img
+    else:
+        raise ValueError(f"Cannot extract tensor from input_img of type {type(input_img)}")
     
-    # Perform inference
     mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=input_img,
+                          org_img=org_img, input_img=input_tensor,
                           org_size=org_size)
     
-    # Handle orientation if needed
-    if hasattr(org_img, 'orientation') and "".join(org_img.orientation) == "LSA":        
+    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]
 
-    # Save the prediction
-    img = org_img.data
+    # Extract data from org_img properly
+    if hasattr(org_img, 'data'):
+        img = org_img.data
+    else:
+        img = org_img
+        
     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)
@@ -157,7 +233,6 @@ 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)}")
 
@@ -185,62 +260,38 @@ def load_system_resources(models_path, learner_fn='heart_model.pkl', variables_f
     return learn, reorder, resample
 
 # Initialize the system
-try:
-    print("🚀 Initializing application...")
-    
-    clone_dir = Path.cwd() / 'clone_dir'
-    URI = os.getenv('PAT_Token_URI')
+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")
+if not URI:
+    raise ValueError("PAT_Token_URI environment variable is not set")
 
-    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}")
-
-    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,
-        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)
+if os.path.exists(clone_dir):
+    pass
+else:
+    Repo.clone_from(URI, clone_dir)
+
+models_path = clone_dir
+save_dir = Path.cwd() / 'hs_pred'
+save_dir.mkdir(parents=True, exist_ok=True)
+
+# Load the model and other required resources
+learn, reorder, resample = load_system_resources(models_path=models_path)
+
+# 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")
+
+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]],
+    allow_flagging='never')
 
 # Launch the Gradio interface
-print("🌐 Launching Gradio interface...")
 demo.launch()