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 )