update

app.py

@@ -1,89 +1,82 @@
 import gradio as gr
-from fastMONAI.vision_all import *
-from huggingface_hub import snapshot_download
-from pathlib import Path
 import torch
 import cv2
+from huggingface_hub import snapshot_download
+from fastMONAI.vision_all import *
 
 def initialize_system():
     """Initial setup of model paths and other constants."""
-    models_path = Path(snapshot_download(repo_id="skaliy/endometrial_cancer_segmentation", cache_dir='models', revision='main'))
+    models_path = Path(snapshot_download(repo_id="skaliy/endometrial_cancer_segmentation",
+                                         cache_dir='models',
+                                         revision='main'))
     save_dir =  Path.cwd() / 'ec_pred'    
     save_dir.mkdir(parents=True, exist_ok=True)
     download_example_endometrial_cancer_data(path=save_dir, multi_channel=False)
     
     return models_path, save_dir
 
-def load_system_resources(models_path):
-    """Load necessary resources like learner and variables."""
+def extract_slice_from_mask(img, mask_data):
+    """Extract a slice from the 3D [W, H, D] image and mask data based on mask data."""
     
-    learner = load_learner(models_path / 'vibe-learner.pkl', cpu=True)  # TODO: add an option to run on GPU
-    vars_fn = models_path / 'vars.pkl'
-    _, reorder, resample = load_variables(pkl_fn=vars_fn)
+    sums = mask_data.sum(axis=(0, 1))
+    idx = np.argmax(sums)
+    img, mask_data = img[:, :, idx], mask_data[:, :, idx]
 
-    return learner, reorder, resample
-
-def get_mid_slice(img, mask_data): 
-    """Extract the middle slice of the mask in a 3D array."""
-    
-    sums = mask_data.sum(axis=(0,1))
-    mid_idx = np.argmax(sums)
-    img, mask_data = img[:, :, mid_idx], mask_data[:, :, mid_idx]
-    
     return np.fliplr(np.rot90(img, -1)), np.fliplr(np.rot90(mask_data, -1))
 
-
+#| export
 def get_fused_image(img, pred_mask, alpha=0.8):
-    """Overlay the mask on the image."""
+    """Fuse a grayscale image with a mask overlay."""
+    
     gray_img_colored = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
     mask_color = np.array([0, 0, 255])
     colored_mask = (pred_mask[..., None] * mask_color).astype(np.uint8)
-    
+
     return cv2.addWeighted(gray_img_colored, alpha, colored_mask, 1 - alpha, 0)
 
-def compute_tumor_volume(mask_data):
-    """Compute the volume of the tumor in milliliters (ml)."""
-    
-    dx, dy, dz = mask_data.spacing
-    voxel_volume_ml = dx * dy * dz / 1000  
-    return np.sum(mask_data) * voxel_volume_ml
 
-def predict(fileobj, learner, reorder, resample, save_dir):
+def gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir):
     """Predict function using the learner and other resources."""
     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)
+    org_img, input_img, org_size = med_img_reader(img_path, 
+                                                  reorder=reorder,
+                                                  resample=resample,
+                                                  only_tensor=False)
     
-    mask_data = inference(learner, reorder=reorder, resample=resample, org_img=org_img, input_img=input_img, org_size=org_size).data 
+    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 #TEMP
-    
+    img = org_img.data
     org_img.set_data(mask_data)
     org_img.save(save_path)
 
-    img, pred_mask = get_mid_slice(img[0], mask_data[0])
+    img, pred_mask = extract_slice_from_mask(img[0], mask_data[0])
     img = ((img - img.min()) / (img.max() - img.min()) * 255).astype(np.uint8) #normalize
-    volume = compute_tumor_volume(org_img)
+    
+    volume = compute_binary_tumor_volume(org_img)
    
     return get_fused_image(img, pred_mask), round(volume, 2)
 
 
 models_path, save_dir = initialize_system()
-learner, reorder, resample = load_system_resources(models_path)
+learn, reorder, resample = load_system_resources(models_path=models_path,
+                                                 learner_fn='vibe-learner.pkl',
+                                                 variables_fn='vars.pkl')
 output_text = gr.Textbox(label="Volume of the predicted tumor:")
 
 demo = gr.Interface(
-    fn=lambda fileobj: predict(fileobj, learner, reorder, resample, save_dir),
+    fn=lambda fileobj: gradio_image_segmentation(fileobj, learn, reorder, resample, save_dir),
     inputs=["file"],
     outputs=["image", output_text],
-    examples=[[save_dir/"vibe.nii.gz"]]
-)
+    examples=[[save_dir/"vibe.nii.gz"]])
 
 demo.launch()