Reworked app + 2D + 3D viewer

Dockerfile

@@ -22,6 +22,8 @@ WORKDIR /code
 RUN apt-get update -y
 RUN apt install git --fix-missing -y
 
+RUN ls -la
+
 # install dependencies
 COPY ./demo/requirements.txt /code/demo/requirements.txt
 RUN python3.7 -m pip install --no-cache-dir --upgrade -r /code/demo/requirements.txt
@@ -32,8 +34,6 @@ RUN python3.7 -m pip install --force-reinstall typing_extensions==4.0.0
 # Install wget
 RUN apt install wget -y
 
-RUN ls -la
-
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user
 

demo/app.py

@@ -1,178 +1,16 @@
-import gradio as gr
-import subprocess as sp
-from skimage.measure import marching_cubes
-import nibabel as nib
-from nibabel.processing import resample_to_output
-import numpy as np
-import random
+from src.gui import WebUI
 
 
-def nifti_to_glb(path):
-    # 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")
-
-    # extract surface
-    verts, faces, normals, values = marching_cubes(data, 0)
-    faces += 1
-
-    with open('prediction.obj', 'w') as thefile:
-        for item in verts:
-            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]))
-
-
-def run_model(input_path):
-    from livermask.utils.run import run_analysis
-    
-    run_analysis(cpu=True, extension='.nii', path=input_path, output='prediction', verbose=True, vessels=False, name="/home/user/app/model.h5", mp_enabled=False)
-
-
-def load_mesh(mesh_file_name):
-    path = mesh_file_name.name
-    run_model(path)
-    nifti_to_glb("prediction-livermask.nii")
-    return "./prediction.obj"
-
-
-def setup_gallery(data_path, pred_path):
-    image = nib.load(data_path)
-    resampled = resample_to_output(image, [1, 1, 1], order=1)
-    data = resampled.get_fdata().astype("uint8")
-
-    image = nib.load(pred_path)
-    resampled = resample_to_output(image, [1, 1, 1], order=0)
-    pred = resampled.get_fdata().astype("uint8")
-
-
-def load_ct_to_numpy(data_path):
-    if type(data_path) != str:
-        data_path = data_path.name
-
-    image = nib.load(data_path)
-    data = image.get_fdata()
-
-    data = np.rot90(data, k=1, axes=(0, 1))
-
-    data[data < -150] = -150
-    data[data > 250] = 250
-
-    data = data - np.amin(data)
-    data = data / np.amax(data) * 255
-    data = data.astype("uint8")
-
-    print(data.shape)
-    return [data[..., i] for i in range(data.shape[-1])]
-
-
-def upload_file(file):
-    return file.name
-
-#def select_section(evt: gr.SelectData):
-#    return section_labels[evt.index]
-
-
-if __name__ == "__main__":
+def main():
     print("Launching demo...")
-    with gr.Blocks() as demo:
-        """
-        with gr.Blocks() as demo:
-            with gr.Row():
-                text1 = gr.Textbox(label="t1")
-                slider2 = gr.Textbox(label="slide")
-                drop3 = gr.Dropdown(["a", "b", "c"], label="d3")
-            with gr.Row():
-                with gr.Column(scale=1, min_width=600):
-                    text1 = gr.Textbox(label="prompt 1")
-                    text2 = gr.Textbox(label="prompt 2")
-                    inbtw = gr.Button("Between")
-                    text4 = gr.Textbox(label="prompt 1")
-                    text5 = gr.Textbox(label="prompt 2")
-                with gr.Column(scale=2, min_width=600):
-                    img1 = gr.Image("images/cheetah.jpg")
-                    btn = gr.Button("Go").style(full_width=True)
-        
-        greeter_1 = gr.Interface(lambda name: f"Hello {name}!", inputs="textbox", outputs=gr.Textbox(label="Greeter 1"))
-        greeter_2 = gr.Interface(lambda name: f"Greetings {name}!", inputs="textbox", outputs=gr.Textbox(label="Greeter 2"))
-        demo = gr.Parallel(greeter_1, greeter_2)
-
-        volume_renderer = gr.Interface(
-            fn=load_mesh,
-            inputs=gr.UploadButton(label="Click to Upload a File", file_types=[".nii", ".nii.nz"], file_count="single"),
-            outputs=gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], label="3D Model"),
-            title="livermask: Automatic Liver Parenchyma segmentation in CT",
-            description="Using pretrained deep learning model trained on the LiTS17 dataset",
-        )
-        """
-
-        with gr.Row():
-            # file_output = gr.File()
-            upload_button = gr.UploadButton(label="Click to Upload a File", file_types=[".nii", ".nii.nz"], file_count="single")
-            # upload_button.upload(upload_file, upload_button, file_output)
-
-            #select_btn = gr.Button("Run analysis")
-            #select_btn.click(fn=upload_file, inputs=upload_button, outputs=output, api_name="Analysis")
-        
-            #upload_button.click(section, [img_input, num_boxes, num_segments], img_output)
-        
-        #print("file output:", file_output)
 
-        images = load_ct_to_numpy("./test-volume.nii")
+    model_name = "/home/user/app/model.h5"  # "/Users/andreped/workspace/livermask/model.h5"
+    class_name = "parenchyma"
 
-        def variable_outputs(k):
-            k = int(k) - 1
-            out = [gr.AnnotatedImage.update(visible=False)] * len(images)
-            out[k] = gr.AnnotatedImage.update(visible=True)
-            return out
-        
-        def section(img, num_segments):
-            sections = []
-            for b in range(num_segments):
-                x = random.randint(0, img.shape[1])
-                y = random.randint(0, img.shape[0])
-                r = random.randint(0, min(x, y, img.shape[1] - x, img.shape[0] - y))
-                mask = np.zeros(img.shape[:2])
-                for i in range(img.shape[0]):
-                    for j in range(img.shape[1]):
-                        dist_square = (i - y) ** 2 + (j - x) ** 2
-                        if dist_square < r**2:
-                            mask[i, j] = round((r**2 - dist_square) / r**2 * 4) / 4
-                sections.append((mask, "parenchyma"))
-            return (img, sections)
-        
-        with gr.Row():
-            s = gr.Slider(1, len(images), value=1, step=1, label="Which 2D slice to show")
-        
-        with gr.Row():
-            with gr.Box():
-                images_boxes = []
-                for i, image in enumerate(images):
-                    visibility = True if i == 1 else False  # only first slide visible - change slide through slider
-                    t = gr.AnnotatedImage(value=section(image, 1), visible=visibility).style(color_map={"parenchyma": "#ffae00"}, width=image.shape[1])
-                    images_boxes.append(t)
+    # initialize and run app
+    app = WebUI(model_name=model_name, class_name=class_name)
+    app.run()
 
-                s.change(variable_outputs, s, images_boxes)
 
-        
-        #upload_button.upload(upload_file, upload_button, file_output)
-        
-        #section_btn.click(section, [images[40], num_boxes, num_segments], img_output)
-        #ct_images.upload(section, [images[40], num_boxes, num_segments], img_output)
-
-        #demo = gr.Interface(
-        #    fn=load_ct_to_numpy,
-        #    inputs=gr.UploadButton(label="Click to Upload a File", file_types=[".nii", ".nii.nz"], file_count="single"),
-        #    outputs=gr.Gallery(label="CT slices").style(columns=[4], rows=[4], object_fit="contain", height="auto"),
-        #    title="livermask: Automatic Liver Parenchyma segmentation in CT",
-        #    description="Using pretrained deep learning model trained on the LiTS17 dataset",
-        #)
-
-    # sharing app publicly -> share=True: https://gradio.app/sharing-your-app/
-    # inference times > 60 seconds -> need queue(): https://github.com/tloen/alpaca-lora/issues/60#issuecomment-1510006062
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)
+if __name__ == "__main__":
+    main()

demo/src/compute.py

@@ -0,0 +1,6 @@
+
+
+def run_model(input_path, model_name="/home/user/app/model.h5"):
+    from livermask.utils.run import run_analysis
+    run_analysis(cpu=True, extension='.nii', path=input_path, output='prediction', verbose=True, vessels=False, name=model_name, mp_enabled=False)
+

demo/src/convert.py

@@ -0,0 +1,24 @@
+import nibabel as nib
+from nibabel.processing import resample_to_output
+from skimage.measure import marching_cubes
+
+
+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")
+
+    # extract surface
+    verts, faces, normals, values = marching_cubes(data, 0)
+    faces += 1
+
+    with open(output, 'w') as thefile:
+        for item in verts:
+            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

@@ -0,0 +1,76 @@
+import gradio as gr
+from .utils import load_ct_to_numpy, load_pred_volume_to_numpy
+from .compute import run_model
+from .convert import nifti_to_glb
+
+
+class WebUI:
+    def __init__(self, model_name, class_name):
+        # global states
+        self.images = []
+        self.pred_images = []
+
+        self.nb_slider_items = 100
+
+        self.model_name = model_name
+        self.class_name = class_name
+
+        # define widgets not to be rendered immediantly, but later on
+        self.slider = gr.Slider(1, self.nb_slider_items, value=1, step=1, label="Which 2D slice to show")
+        self.volume_renderer = gr.Model3D(
+            clear_color=[0.0, 0.0, 0.0, 0.0],
+            label="3D Model",
+            visible=True
+        ).style(height=512)
+
+    def combine_ct_and_seg(self, img, pred):
+        return (img, [(pred, self.class_name)])
+    
+    def upload_file(self, file):
+        return file.name
+    
+    def load_mesh(self, mesh_file_name, model_name="/home/user/app/model.h5"):
+        path = mesh_file_name.name
+        run_model(path, model_name)
+        nifti_to_glb("prediction-livermask.nii")
+        self.images = load_ct_to_numpy("./files/test_ct.nii")
+        self.pred_images = load_pred_volume_to_numpy("./prediction-livermask.nii")
+        self.slider = self.slider.update(value=2)
+        return "./prediction.obj"
+    
+    def get_img_pred_pair(self, k):
+        k = int(k) - 1
+        out = [gr.AnnotatedImage.update(visible=False)] * self.nb_slider_items
+        out[k] = gr.AnnotatedImage.update(self.combine_ct_and_seg(self.images[k], self.pred_images[k]), visible=True)
+        return out
+
+    def run(self):
+        with gr.Blocks() as demo:
+
+            with gr.Row().style(equal_height=True):
+                file_output = gr.File(file_types=[".nii", ".nii.nz"], file_count="single").style(full_width=False, size="sm")
+                file_output.upload(self.upload_file, file_output, file_output)
+
+                run_btn = gr.Button("Run analysis").style(full_width=False, size="sm")
+                run_btn.click(fn=lambda x: self.load_mesh(x, model_name=self.model_name), inputs=file_output, outputs=self.volume_renderer)
+            
+            with gr.Row().style(equal_height=True):
+                with gr.Box():
+                    image_boxes = []
+                    for i in range(self.nb_slider_items):
+                        visibility = True if i == 1 else False
+                        t = gr.AnnotatedImage(visible=visibility)\
+                            .style(color_map={self.class_name: "#ffae00"}, height=512, width=512)
+                        image_boxes.append(t)
+
+                    self.slider.change(self.get_img_pred_pair, self.slider, image_boxes)
+                
+                with gr.Box():
+                    self.volume_renderer.render()
+            
+            with gr.Row():
+                self.slider.render()
+
+        # sharing app publicly -> share=True: https://gradio.app/sharing-your-app/
+        # inference times > 60 seconds -> need queue(): https://github.com/tloen/alpaca-lora/issues/60#issuecomment-1510006062
+        demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)

demo/src/utils.py

@@ -0,0 +1,38 @@
+import nibabel as nib
+import numpy as np
+
+
+def load_ct_to_numpy(data_path):
+    if type(data_path) != str:
+        data_path = data_path.name
+
+    image = nib.load(data_path)
+    data = image.get_fdata()
+
+    data = np.rot90(data, k=1, axes=(0, 1))
+
+    data[data < -150] = -150
+    data[data > 250] = 250
+
+    data = data - np.amin(data)
+    data = data / np.amax(data) * 255
+    data = data.astype("uint8")
+
+    print(data.shape)
+    return [data[..., i] for i in range(data.shape[-1])]
+
+
+def load_pred_volume_to_numpy(data_path):
+    if type(data_path) != str:
+        data_path = data_path.name
+
+    image = nib.load(data_path)
+    data = image.get_fdata()
+
+    data = np.rot90(data, k=1, axes=(0, 1))
+
+    data[data > 0] = 1
+    data = data.astype("uint8")
+
+    print(data.shape)
+    return [data[..., i] for i in range(data.shape[-1])]