Update app.py

app.py

@@ -1,183 +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
-   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)
+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, :, :]
        
-       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))
+       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)
        
-       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)
+       if not isinstance(variables, list) or len(variables) != 3:
+           raise ValueError(f"vars.pkl does not contain the expected list format. Found: {variables}")
        
-       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)
+       # Assuming the format is [shape, reorder, resample]
+       shape = variables[0]
+       reorder = variables[1]
+       resample = variables[2]
        
-       fused = cv2.addWeighted(gray_img_colored, alpha, colored_mask, 1 - alpha, 0)
+       if not isinstance(reorder, bool):
+           raise ValueError(f"vars.pkl does not contain a valid 'reorder' value. Found: {reorder}")
        
-       # Flip the fused image vertically and horizontally
-       fused_flipped = cv2.flip(fused, -1)  # Flip both vertically and horizontally
+       if not isinstance(resample, list) or len(resample) != 3:
+           raise ValueError(f"vars.pkl does not contain a valid 'resample' value. Found: {resample}")
        
-       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()
+   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()