Merge pull request #38 from andreped/gui

Fixed Model3D camera position; Added inference info to UI; Minor refactoring

demo/src/gui.py

@@ -55,8 +55,10 @@ class WebUI:
         self.volume_renderer = gr.Model3D(
             clear_color=[0.0, 0.0, 0.0, 0.0],
             label="3D Model",
+            show_label=True,
             visible=True,
             elem_id="model-3d",
+            camera_position=[90, 180, 768],
         ).style(height=512)
 
     def set_class_name(self, value):
@@ -79,7 +81,7 @@ class WebUI:
             task=self.class_names[self.class_name],
             name=self.result_names[self.class_name],
         )
-        LOGGER.info("Converting prediction NIfTI to GLB...")
+        LOGGER.info("Converting prediction NIfTI to OBJ...")
         nifti_to_glb("prediction.nii.gz")
 
         LOGGER.info("Loading CT to numpy...")
@@ -123,6 +125,7 @@ class WebUI:
                         show_copy_button=True,
                         scroll_to_output=False,
                         container=True,
+                        line_breaks=True,
                     )
                     demo.load(read_logs, None, logs, every=1)
 
@@ -132,7 +135,8 @@ class WebUI:
                             sidebar_state = gr.State(True)
 
                             btn_toggle_sidebar = gr.Button(
-                                "Toggle Sidebar", elem_id="toggle-button"
+                                "Toggle Sidebar",
+                                elem_id="toggle-button",
                             )
                             btn_toggle_sidebar.click(
                                 self.toggle_sidebar,
@@ -155,7 +159,7 @@ class WebUI:
                         model_selector = gr.Dropdown(
                             list(self.class_names.keys()),
                             label="Task",
-                            info="Which task to perform",
+                            info="Which structure to segment.",
                             multiselect=False,
                             size="sm",
                         )
@@ -191,6 +195,13 @@ class WebUI:
                             cache_examples=True,
                         )
 
+                        gr.Markdown(
+                            """
+                            **NOTE:** Inference might take several minutes (Airways: ~8 minutes), see logs to the left. \\
+                            The segmentation will be available in the 2D and 3D viewers below when finished.
+                            """
+                        )
+
                     with gr.Row():
                         with gr.Box():
                             with gr.Column():

demo/src/utils.py

@@ -46,12 +46,19 @@ def nifti_to_glb(path, output="prediction.obj"):
     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: