Update utils.py

utils.py

@@ -4,16 +4,13 @@ from pathlib import Path
 
 import matplotlib as mpl
 import matplotlib.pyplot as plt
-import nibabel as nib
 import numpy as np
+import SimpleITK as sitk
 import torch
-from monai.transforms import LoadImage
 from mrsegmentator import inference
 from mrsegmentator.utils import add_postfix
-from PIL import Image, ImageColor, ImageDraw, ImageEnhance
+from PIL import Image
 from scipy import ndimage
-from monai.transforms import LoadImage, Orientation, Spacing
-import SimpleITK as sitk
 
 import streamlit as st
 
@@ -63,7 +60,7 @@ def run(tmpdirname):
     if st.session_state.option is not None:
         image = Path(__file__).parent / str(st.session_state.option)
 
-        inference.infer([image], tmpdirname, [0], split_level=2, cpu_only=True)
+        inference.infer([image], tmpdirname, st.session_state.folds, split_level=1)
         seg_name = add_postfix(image.name, "seg")
         preds_path = tmpdirname + "/" + seg_name
         st.session_state.preds_3D = read_image(preds_path)
@@ -78,9 +75,7 @@ def reflect_box_into_model(box_3d):
     x2_prompt = int(x2 * 256.0 / 325.0)
     y2_prompt = int(y2 * 256.0 / 325.0)
     z2_prompt = int(z2 * 32.0 / 325.0)
-    return torch.tensor(
-        np.array([z1_prompt, y1_prompt, x1_prompt, z2_prompt, y2_prompt, x2_prompt])
-    )
+    return torch.tensor(np.array([z1_prompt, y1_prompt, x1_prompt, z2_prompt, y2_prompt, x2_prompt]))
 
 
 def reflect_json_data_to_3D_box(json_data, view):
@@ -88,19 +83,17 @@ def reflect_json_data_to_3D_box(json_data, view):
         st.session_state.rectangle_3Dbox[1] = json_data["objects"][0]["top"]
         st.session_state.rectangle_3Dbox[2] = json_data["objects"][0]["left"]
         st.session_state.rectangle_3Dbox[4] = (
-            json_data["objects"][0]["top"]
-            + json_data["objects"][0]["height"] * json_data["objects"][0]["scaleY"]
+            json_data["objects"][0]["top"] + json_data["objects"][0]["height"] * json_data["objects"][0]["scaleY"]
         )
         st.session_state.rectangle_3Dbox[5] = (
-            json_data["objects"][0]["left"]
-            + json_data["objects"][0]["width"] * json_data["objects"][0]["scaleX"]
+            json_data["objects"][0]["left"] + json_data["objects"][0]["width"] * json_data["objects"][0]["scaleX"]
         )
     print(st.session_state.rectangle_3Dbox)
 
 
-def make_fig(image, preds, px_range = (10, 400), transparency=0.5):
+def make_fig(image, preds, px_range=(10, 400), transparency=0.5):
 
-    fig, ax = plt.subplots(1, 1, figsize=(4,4))
+    fig, ax = plt.subplots(1, 1, figsize=(4, 4))
     image_slice = image.clip(*px_range)
 
     ax.imshow(
@@ -150,7 +143,7 @@ def make_fig(image, preds, px_range = (10, 400), transparency=0.5):
     plt.axis("off")
     ax.set_xticks([])
     ax.set_yticks([])
-    
+
     fig.canvas.draw()
 
     # transform to image
@@ -160,8 +153,8 @@ def make_fig(image, preds, px_range = (10, 400), transparency=0.5):
 #######################################
 
 
-def make_isotropic(image, interpolator = sitk.sitkLinear, spacing = None):
-    '''
+def make_isotropic(image, interpolator=sitk.sitkLinear, spacing=None):
+    """
     Many file formats (e.g. jpg, png,...) expect the pixels to be isotropic, same
     spacing for all axes. Saving non-isotropic data in these formats will result in
     distorted images. This function makes an image isotropic via resampling, if needed.
@@ -175,7 +168,7 @@ def make_isotropic(image, interpolator = sitk.sitkLinear, spacing = None):
     Returns:
         SimpleITK.Image with isotropic spacing which occupies the same region in space as
         the input image.
-    '''
+    """
     original_spacing = image.GetSpacing()
     # Image is already isotropic, just return a copy.
     if all(spc == original_spacing[0] for spc in original_spacing):
@@ -184,18 +177,26 @@ def make_isotropic(image, interpolator = sitk.sitkLinear, spacing = None):
     original_size = image.GetSize()
     if spacing is None:
         spacing = min(original_spacing)
-    new_spacing = [spacing]*image.GetDimension()
-    new_size = [int(round(osz*ospc/spacing)) for osz, ospc in zip(original_size, original_spacing)]
-    return sitk.Resample(image, new_size, sitk.Transform(), interpolator,
-                         image.GetOrigin(), new_spacing, image.GetDirection(), 0, # default pixel value
-                         image.GetPixelID())
+    new_spacing = [spacing] * image.GetDimension()
+    new_size = [int(round(osz * ospc / spacing)) for osz, ospc in zip(original_size, original_spacing)]
+    return sitk.Resample(
+        image,
+        new_size,
+        sitk.Transform(),
+        interpolator,
+        image.GetOrigin(),
+        new_spacing,
+        image.GetDirection(),
+        0,  # default pixel value
+        image.GetPixelID(),
+    )
 
 
 def read_image(path):
-    
+
     img = sitk.ReadImage(path)
     img = sitk.DICOMOrient(img, "LPS")
     img = make_isotropic(img)
     img = sitk.GetArrayFromImage(img)
-    
+
     return img