Update app.py

app.py

@@ -1,179 +1,183 @@
 import gradio as gr
-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 fastai.learner import load_learner
-from fastai.basics import load_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
-
-# 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:
-        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)
-    
-    mask_data = inference(learn, reorder=reorder, resample=resample,
-                          org_img=org_img, input_img=input_img,
-                          org_size=org_size).data
-    
-    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]
-    
-    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:
-        variables = load_pickle(models_path / variables_fn)
-        if isinstance(variables, dict):
-            reorder = variables.get('reorder')
-            resample = variables.get('resample')
-            if reorder is None or resample is None:
-                raise ValueError("'reorder' or 'resample' not found in vars.pkl. Using default values.")
-                reorder = True  # Set a default value
-                resample = True  # Set a default value
-        else:
-            raise ValueError("vars.pkl does not contain a dictionary. Using default values for reorder and resample.")
-            reorder = True  # Set a default value
-            resample = True  # Set a default value
-    except Exception as e:
-        raise ValueError(f"Error loading variables: {str(e)}")
-        reorder = True  # Set a default value
-        resample = True  # Set a default value
-
-    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)")
-
-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
-demo.launch()
+   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 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
+
+   # 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:
+           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)
+       
+       mask_data = inference(learn, reorder=reorder, resample=resample,
+                             org_img=org_img, input_img=input_img,
+                             org_size=org_size).data
+       
+       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]
+       
+       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)")
+
+   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
+   demo.launch()