Update app.py

app.py

@@ -1,222 +1,115 @@
-import gradio as gr
+from pathlib import Path
 import torch
 import cv2
 import numpy as np
-from pathlib import Path
-from huggingface_hub import snapshot_download
-from fastMONAI.vision_all import *
-from git import Repo
-import os
+from fastMONAI.vision_all import med_img_reader   # keep the import
 from fastai.learner import load_learner
-from fastai.basics import load_pickle
 import pickle
 
-# 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[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[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[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:  # 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)
-    
-    # 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
-
-# Define the inference function
+def read_image_as_tensor(path: str, reorder: bool, resample: list) -> tuple:
+    """Read a medical image and always return a torch.Tensor as the second element."""
+    org_img, raw, *rest = med_img_reader(
+        path,
+        reorder=reorder,
+        resample=resample,
+        only_tensor=False,
+        dtype=torch.Tensor,
+    )
+    # raw may be a dict or a tensor
+    if isinstance(raw, dict):
+        # fastMONAI convention: the actual tensor lives under the key "tensor"
+        tensor = raw.get("tensor")
+        if tensor is None:
+            # fallback: first torch.Tensor found in the dict
+            tensor = next(v for v in raw.values() if isinstance(v, torch.Tensor))
+    else:
+        tensor = raw
+
+    # original size (used later for volume calculation)
+    org_size = rest[0] if rest else org_img.shape[1:]
+
+    return org_img, tensor, org_size
+
+
 def inference(learn, reorder, resample, org_img, input_img, org_size):
-    """Perform segmentation using the loaded model."""
-    # Ensure input_img is a torch.Tensor
+    """Run the learner on a single 3‑D volume."""
     if not isinstance(input_img, torch.Tensor):
-        raise ValueError(f"Expected input_img to be a torch.Tensor, but got {type(input_img)}")
-    
-    # Perform the segmentation
+        raise TypeError(f"input_img must be a torch.Tensor, got {type(input_img)}")
+
+    # Ensure batch dimension exists (fastai expects N×C×D×H×W)
+    if input_img.dim() == 4:               # (C, D, H, W) → (1, C, D, H, W)
+        input_img = input_img.unsqueeze(0)
+
     with torch.no_grad():
         pred = learn.predict(input_img)
-    
-    # Process the prediction if necessary
-    mask_data = pred[0]  # Assuming the first element of the prediction is the mask
-    
-    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."""
+    # fastai returns (tensor, ...) – we only need the mask tensor
+    mask = pred[0] if isinstance(pred, (list, tuple)) else pred
+    return mask
 
-    if view is None:
-        view = 'Sagittal'
-    
-    img_path = Path(fileobj.name)
 
-    # Convert PosixPath to string
-    img_path_str = str(img_path)
+def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir, view):
+    """Main Gradio callback – reads the file, runs inference, returns visualisation."""
+    view = view or "Sagittal"
 
-    save_fn = 'pred_' + img_path.stem
+    img_path = Path(fileobj.name)
+    save_fn = f"pred_{img_path.stem}.nii.gz"
     save_path = save_dir / save_fn
-    
-    # Ensure only_tensor is set to False to get all values
-    try:
-        org_img, input_img, org_size = med_img_reader(img_path_str, 
-                                                      reorder=reorder,
-                                                      resample=resample,
-                                                      only_tensor=False,
-                                                      dtype=torch.Tensor)
-    except ValueError:
-        # Handle the case where med_img_reader returns only two values
-        org_img, input_img = med_img_reader(img_path_str, 
-                                           reorder=reorder,
-                                           resample=resample,
-                                           only_tensor=False,
-                                           dtype=torch.Tensor)
-        # Infer org_size from org_img
-        org_size = org_img.shape[1:]  # Assuming org_img has a shape attribute
-
-    # 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)}")
-    
-    mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=input_img,
-                          org_size=org_size)
-    
-    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]
-
-    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)
-    fused_images = [(get_fused_image(
-                        normalize_image(slice_img),  # Normalize safely
-                        slice_mask, view))
-                    for slice_img, slice_mask in slices]
-    
+
+    # ------------------------------------------------------------------
+    # 1️⃣  Read the image and guarantee a torch.Tensor for the model
+    # ------------------------------------------------------------------
+    org_img, input_tensor, org_size = read_image_as_tensor(
+        str(img_path),
+        reorder=reorder,
+        resample=resample,
+    )
+
+    # ------------------------------------------------------------------
+    # 2️⃣  Run the model
+    # ------------------------------------------------------------------
+    mask_data = inference(
+        learn,
+        reorder=reorder,
+        resample=resample,
+        org_img=org_img,
+        input_img=input_tensor,
+        org_size=org_size,
+    )
+
+    # ------------------------------------------------------------------
+    # 3️⃣  Post‑process orientation (keep your original logic)
+    # ------------------------------------------------------------------
+    if "".join(org_img.orientation) == "LSA":
+        # Adjust axes to match the original orientation
+        mask_data = mask_data.permute(0, 1, 3, 2)          # (B, C, H, W, D) → (B, C, H, D, W)
+        mask_data = torch.flip(mask_data[0], dims=[1])   # remove batch, flip dim‑1
+        mask_data = mask_data.unsqueeze(0)               # add batch back
+
+    # ------------------------------------------------------------------
+    # 4️⃣  Save the mask as a NIfTI file (optional)
+    # ------------------------------------------------------------------
+    img = org_img.data                     # original image data (torch.Tensor)
+    org_img.set_data(mask_data)            # replace image data with mask
+    org_img.save(save_path)                # writes a .nii.gz file
+
+    # ------------------------------------------------------------------
+    # 5️⃣  Build gallery of fused slices
+    # ------------------------------------------------------------------
+    slices = extract_slices_from_mask(img[0].cpu().numpy(), mask_data[0].cpu().numpy(), view)
+
+    fused_images = [
+        get_fused_image(
+            normalize_image(slice_img),   # safe 0‑255 uint8
+            slice_mask,
+            view,
+        )
+        for slice_img, slice_mask in slices
+    ]
+
+    # ------------------------------------------------------------------
+    # 6️⃣  Compute volume (your helper expects a FastMRIImage with mask inside)
+    # ------------------------------------------------------------------
     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)
-    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
-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 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
-demo.launch()
+    return fused_images, round(volume, 2)