Update app.py

app.py

@@ -13,152 +13,155 @@ 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
+    """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)
+    """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
+    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)
+    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
+    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)
+    """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
+    """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)
+    """Predict function using the learner and other resources."""
+
+    if view is None:
+        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
+    org_img, input_img, org_size = med_img_reader(img_path_str, 
+                                                  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
+    """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")
+    raise ValueError("PAT_Token_URI environment variable is not set")
 
 if os.path.exists(clone_dir):
-   pass
+    pass
 else:
-   Repo.clone_from(URI, clone_dir)
+    Repo.clone_from(URI, clone_dir)
 
 models_path = clone_dir
 save_dir = Path.cwd() / 'hs_pred'
@@ -172,12 +175,15 @@ 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=[[str(Path.cwd() / "sample.nii.gz")]],
-   allow_flagging='never')
+    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()