Spaces:
Sleeping
Sleeping
| 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 = 300 | |
| 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 + "test_thorax_CT.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=False) | |