Renamed to nifti_to_obj and moved to convert

demo/src/convert.py

@@ -0,0 +1,35 @@
+import nibabel as nib
+from nibabel.processing import resample_to_output
+from skimage.measure import marching_cubes
+
+
+def nifti_to_obj(path, output="prediction.obj"):
+    # load NIFTI into numpy array
+    image = nib.load(path)
+    resampled = resample_to_output(image, [1, 1, 1], order=1)
+    data = resampled.get_fdata().astype("uint8")
+
+    # Create a material with a red diffuse color (RGB value)
+    red_material = "newmtl RedMaterial\nKd 1 0 0"  # Red diffuse color (RGB)
+
+    # extract surface
+    verts, faces, normals, values = marching_cubes(data, 0)
+    faces += 1
+
+    with open(output, "w") as thefile:
+        # Write the material definition to the OBJ file
+        thefile.write(red_material + "\n")
+
+        for item in verts:
+            # thefile.write('usemtl RedMaterial\n')
+            thefile.write("v {0} {1} {2}\n".format(item[0], item[1], item[2]))
+
+        for item in normals:
+            thefile.write("vn {0} {1} {2}\n".format(item[0], item[1], item[2]))
+
+        for item in faces:
+            thefile.write(
+                "f {0}//{0} {1}//{1} {2}//{2}\n".format(
+                    item[0], item[1], item[2]
+                )
+            )

demo/src/gui.py

@@ -2,6 +2,7 @@ import os
 
 import gradio as gr
 
+from .convert import nifti_to_obj
 from .css_style import css
 from .inference import run_model
 from .logger import flush_logs
@@ -9,7 +10,7 @@ from .logger import read_logs
 from .logger import setup_logger
 from .utils import load_ct_to_numpy
 from .utils import load_pred_volume_to_numpy
-from .utils import nifti_to_glb
+
 
 # setup logging
 LOGGER = setup_logger()
@@ -82,7 +83,7 @@ class WebUI:
             name=self.result_names[self.class_name],
         )
         LOGGER.info("Converting prediction NIfTI to OBJ...")
-        nifti_to_glb("prediction.nii.gz")
+        nifti_to_obj("prediction.nii.gz")
 
         LOGGER.info("Loading CT to numpy...")
         self.images = load_ct_to_numpy(path)
@@ -113,7 +114,6 @@ class WebUI:
         with gr.Blocks(css=css) as demo:
             with gr.Row():
                 with gr.Column(visible=True, scale=0.2) as sidebar_left:
-                    # gr.Markdown("SideBar Left")
                     logs = gr.Textbox(
                         placeholder="\n" * 16,
                         label="Logs",

demo/src/utils.py

@@ -1,7 +1,5 @@
 import nibabel as nib
 import numpy as np
-from nibabel.processing import resample_to_output
-from skimage.measure import marching_cubes
 
 
 def load_ct_to_numpy(data_path):
@@ -38,35 +36,3 @@ def load_pred_volume_to_numpy(data_path):
 
     print(data.shape)
     return [data[..., i] for i in range(data.shape[-1])]
-
-
-def nifti_to_glb(path, output="prediction.obj"):
-    # load NIFTI into numpy array
-    image = nib.load(path)
-    resampled = resample_to_output(image, [1, 1, 1], order=1)
-    data = resampled.get_fdata().astype("uint8")
-
-    # Create a material with a red diffuse color (RGB value)
-    red_material = "newmtl RedMaterial\nKd 1 0 0"  # Red diffuse color (RGB)
-
-    # extract surface
-    verts, faces, normals, values = marching_cubes(data, 0)
-    faces += 1
-
-    with open(output, "w") as thefile:
-        # Write the material definition to the OBJ file
-        thefile.write(red_material + "\n")
-
-        for item in verts:
-            # thefile.write('usemtl RedMaterial\n')
-            thefile.write("v {0} {1} {2}\n".format(item[0], item[1], item[2]))
-
-        for item in normals:
-            thefile.write("vn {0} {1} {2}\n".format(item[0], item[1], item[2]))
-
-        for item in faces:
-            thefile.write(
-                "f {0}//{0} {1}//{1} {2}//{2}\n".format(
-                    item[0], item[1], item[2]
-                )
-            )