Demo improvements

demo/src/gui.py

@@ -11,7 +11,7 @@ class WebUI:
         self.pred_images = []
 
         # @TODO: This should be dynamically set based on chosen volume size
-        self.nb_slider_items = 100
+        self.nb_slider_items = 160
 
         self.class_name = class_name
         self.cwd = cwd
@@ -34,9 +34,9 @@ class WebUI:
     def load_mesh(self, mesh_file_name):
         path = mesh_file_name.name
         run_model(path)
-        nifti_to_glb("prediction-livermask.nii")
+        nifti_to_glb("./prediction.nii.gz")
         self.images = load_ct_to_numpy(path)
-        self.pred_images = load_pred_volume_to_numpy("./prediction-livermask.nii")
+        self.pred_images = load_pred_volume_to_numpy("./prediction.nii.gz")
         self.slider = self.slider.update(value=2)
         return "./prediction.obj"
     
@@ -74,7 +74,7 @@ class WebUI:
             
             with gr.Row():
                 gr.Examples(
-                    examples=[self.cwd + "test-volume.nii"],
+                    examples=[self.cwd + "lung_001.nii.gz"],
                     inputs=file_output,
                     outputs=file_output,
                     fn=self.upload_file,

demo/src/utils.py

@@ -11,8 +11,8 @@ def load_ct_to_numpy(data_path):
 
     data = np.rot90(data, k=1, axes=(0, 1))
 
-    data[data < -150] = -150
-    data[data > 250] = 250
+    data[data < -1024] = 1024
+    data[data > 1024] = 1024
 
     data = data - np.amin(data)
     data = data / np.amax(data) * 255