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
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
# setup logging
LOGGER = setup_logger()
class WebUI:
def __init__(
self,
model_name: str = None,
cwd: str = "/home/user/app/",
share: int = 1,
):
# global states
self.images = []
self.pred_images = []
# @TODO: This should be dynamically set based on chosen volume size
self.nb_slider_items = 820
self.model_name = model_name
self.cwd = cwd
self.share = share
self.filename = None
self.extension = None
self.class_name = "airways" # default
self.class_names = {
"airways": "CT_Airways",
"lungs": "CT_Lungs",
}
self.result_names = {
"airways": "Airways",
"lungs": "Lungs",
}
# define widgets not to be rendered immediantly, but later on
self.slider = gr.Slider(
minimum=1,
maximum=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",
show_label=True,
visible=True,
elem_id="model-3d",
camera_position=[90, 180, 768],
height=512,
)
def set_class_name(self, value):
LOGGER.info(f"Changed task to: {value}")
self.class_name = value
def combine_ct_and_seg(self, img, pred):
return (img, [(pred, self.class_name)])
def upload_file(self, file):
out = file.name
LOGGER.info(f"File uploaded: {out}")
return out
def process(self, mesh_file_name):
path = mesh_file_name.name
curr = path.split("/")[-1]
self.extension = ".".join(curr.split(".")[1:])
self.filename = (
curr.split(".")[0] + "-" + self.class_names[self.class_name]
)
run_model(
path,
model_path=os.path.join(self.cwd, "resources/models/"),
task=self.class_names[self.class_name],
name=self.result_names[self.class_name],
output_filename=self.filename + "." + self.extension,
)
LOGGER.info("Converting prediction NIfTI to OBJ...")
nifti_to_obj(path=self.filename + "." + self.extension)
LOGGER.info("Loading CT to numpy...")
self.images = load_ct_to_numpy(path)
LOGGER.info("Loading prediction volume to numpy..")
self.pred_images = load_pred_volume_to_numpy(
self.filename + "." + self.extension
)
return "./prediction.obj"
def download_prediction(self):
if (self.filename is None) or (self.extension is None):
LOGGER.error(
"The prediction is not available or ready to download. Wait until the result is available in the 3D viewer."
)
raise ValueError("Run inference before downloading!")
return self.filename + "." + self.extension
def get_img_pred_pair(self, k):
k = int(k)
out = gr.AnnotatedImage(
self.combine_ct_and_seg(self.images[k], self.pred_images[k]),
visible=True,
elem_id="model-2d",
color_map={self.class_name: "#ffae00"},
height=512,
width=512,
)
return out
def toggle_sidebar(self, state):
state = not state
return gr.update(visible=state), state
def run(self):
with gr.Blocks(css=css) as demo:
with gr.Row():
with gr.Column(visible=True, scale=0.2) as sidebar_left:
logs = gr.Textbox(
placeholder="\n" * 16,
label="Logs",
info="Verbose from inference will be displayed below.",
lines=36,
max_lines=36,
autoscroll=True,
elem_id="logs",
show_copy_button=True,
container=True,
)
demo.load(read_logs, None, logs, every=1)
with gr.Column():
with gr.Row():
with gr.Column(scale=1, min_width=150):
sidebar_state = gr.State(True)
btn_toggle_sidebar = gr.Button(
"Toggle Sidebar",
elem_id="toggle-button",
)
btn_toggle_sidebar.click(
self.toggle_sidebar,
[sidebar_state],
[sidebar_left, sidebar_state],
)
btn_clear_logs = gr.Button(
"Clear logs", elem_id="logs-button"
)
btn_clear_logs.click(flush_logs, [], [])
file_output = gr.File(
file_count="single",
elem_id="upload",
scale=3,
)
file_output.upload(
self.upload_file, file_output, file_output
)
model_selector = gr.Dropdown(
list(self.class_names.keys()),
label="Task",
info="Which structure to segment.",
multiselect=False,
scale=1,
)
model_selector.input(
fn=lambda x: self.set_class_name(x),
inputs=model_selector,
outputs=None,
)
with gr.Column(scale=1, min_width=150):
run_btn = gr.Button(
"Run analysis",
variant="primary",
elem_id="run-button",
)
run_btn.click(
fn=lambda x: self.process(x),
inputs=file_output,
outputs=self.volume_renderer,
)
download_btn = gr.DownloadButton(
"Download prediction",
visible=True,
variant="secondary",
elem_id="download",
)
download_btn.click(
fn=self.download_prediction,
inputs=None,
outputs=download_btn,
)
with gr.Row():
gr.Examples(
examples=[
os.path.join(self.cwd, "test_thorax_CT.nii.gz"),
],
inputs=file_output,
outputs=file_output,
fn=self.upload_file,
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.Group():
with gr.Column():
# create dummy image to be replaced by loaded images
t = gr.AnnotatedImage(
visible=True,
elem_id="model-2d",
color_map={self.class_name: "#ffae00"},
# height=512,
# width=512,
)
self.slider.input(
self.get_img_pred_pair,
self.slider,
t,
)
self.slider.render()
with gr.Group(): # gr.Box():
self.volume_renderer.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=self.share
)