Merge pull request #30 from andreped/huggingface

Developed gradio demo + HF Space integration and continuous deployment

.github/workflows/deploy.yml

@@ -0,0 +1,20 @@
+name: Deploy
+on:
+  push:
+    branches: [ main ]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          fetch-depth: 0
+          lfs: true
+      - name: Push to hub
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: git push https://andreped:[email protected]/spaces/andreped/LungTumorMask main

.github/workflows/filesize.yml

@@ -0,0 +1,16 @@
+name: Check file size
+on:               # or directly `on: [push]` to run the action on every push on any branch
+  pull_request:
+    branches: [ main ]
+
+  # to run this workflow manually from the Actions tab
+  workflow_dispatch:
+
+jobs:
+  sync-to-hub:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check large files
+        uses: ActionsDesk/[email protected]
+        with:
+          filesizelimit: 10485760 # this is 10MB so we can sync to HF Spaces

README.md

@@ -1,8 +1,21 @@
+---
+title: 'LungTumorMask: Automatic lung tumor segmentation in CT'
+colorFrom: indigo
+colorTo: indigo
+sdk: docker
+app_port: 7860
+emoji: 🔎
+pinned: false
+license: mit
+app_file: demo/app.py
+---
+
 # Automatic lung tumor segmentation in CT
 
 [![license](https://img.shields.io/github/license/DAVFoundation/captain-n3m0.svg?style=flat-square)](https://github.com/DAVFoundation/captain-n3m0/blob/master/LICENSE)
 [![Build Actions Status](https://github.com/VemundFredriksen/LungTumorMask/workflows/Build/badge.svg)](https://github.com/VemundFredriksen/LungTumorMask/actions)
 [![Paper](https://zenodo.org/badge/DOI/10.1371/journal.pone.0266147.svg)](https://doi.org/10.1371/journal.pone.0266147)
+<a target="_blank" href="https://huggingface.co/spaces/andreped/LungTumorMask"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-yellow.svg"></a>
 
 This is the official repository for the paper [_"Teacher-student approach for lung tumor segmentation from mixed-supervised datasets"_](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0266147), published in PLOS ONE.
 
@@ -14,7 +27,10 @@ The source code used to train the model and conduct the experiments presented in
 
 <img src="https://github.com/VemundFredriksen/LungTumorMask/releases/download/0.0.1/sample_renders.png" width="70%">
 
-&nbsp;
+## [Demo](https://github.com/VemundFredriksen/LungTumorMask#demo) <a target="_blank" href="https://huggingface.co/spaces/andreped/LungTumorMask"><img src="https://img.shields.io/badge/🤗%20Hugging%20Face-Spaces-yellow.svg"></a>
+
+An online version of the tool has been made openly available at Hugging Face spaces, to enable researchers to easily test the software on their own data without downloading it. To access it, click on the badge above.
+
 
 ## [Installation](https://github.com/VemundFredriksen/LungTumorMask#installation)
 

demo/app.py

@@ -0,0 +1,18 @@
+from src.gui import WebUI
+
+
+def main():
+    print("Launching demo...")
+
+    # cwd = "/Users/andreped/workspace/LungTumorMask/"  # local testing -> macOS
+    cwd = "/home/user/app/"  # production -> docker
+
+    class_name = "tumor"
+
+    # initialize and run app
+    app = WebUI(model_name=model_name, class_name=class_name, cwd=cwd)
+    app.run()
+
+
+if __name__ == "__main__":
+    main()

demo/requirements.txt

@@ -0,0 +1,2 @@
+lungtumormask @ git+https://github.com/vemundfredriksen/LungTumorMask.git
+gradio==3.32.0

demo/src/compute.py

@@ -0,0 +1,3 @@
+def run_model(input_path):
+    from lungtumormask import mask
+    mask.mask(input_path, "./output.nii.gz", lung_filter=True, threshold=0.5, radius=1, batch_size=1)

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,103 @@
+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, class_name:str = None, cwd:str = None):
+        # global states
+        self.images = []
+        self.pred_images = []
+
+        # @TODO: This should be dynamically set based on chosen volume size
+        self.nb_slider_items = 160
+
+        self.class_name = class_name
+        self.cwd = cwd
+
+        # 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,
+            elem_id="model-3d",
+        ).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):
+        path = mesh_file_name.name
+        run_model(path)
+        nifti_to_glb("./prediction.nii.gz")
+        self.images = load_ct_to_numpy(path)
+        self.pred_images = load_pred_volume_to_numpy("./prediction.nii.gz")
+        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):
+        css="""
+        #model-3d {
+        height: 512px;
+        }
+        #model-2d {
+        height: 512px;
+        margin: auto;
+        }
+        """
+        with gr.Blocks(css=css) as demo:
+
+            with gr.Row():
+                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),
+                    inputs=file_output,
+                    outputs=self.volume_renderer
+                )
+            
+            with gr.Row():
+                gr.Examples(
+                    examples=[self.cwd + "lung_001.nii.gz"],
+                    inputs=file_output,
+                    outputs=file_output,
+                    fn=self.upload_file,
+                    cache_examples=True,
+                )
+            
+            with gr.Row():
+                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, elem_id="model-2d")\
+                            .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 < -1024] = 1024
+    data[data > 1024] = 1024
+
+    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])]