| import numpy as np | |
| import gradio as gr | |
| from cxas import CXAS | |
| import torch | |
| import torch.nn.functional as F | |
| from cxas.visualize import visualize_mask | |
| from cxas.label_mapper import label_mapper | |
| model = CXAS( | |
| model_name = 'UNet_ResNet50_default', | |
| gpus = 'cpu' | |
| ) | |
| def predict(image, class_names): | |
| with torch.no_grad(): | |
| if type(class_names) == str: | |
| class_names = [class_names] | |
| normalized_image = {'data':F.interpolate(model.fileloader.normalize(torch.tensor(np.transpose(image, [2,0,1])/255)).unsqueeze(0), 512).float()} | |
| out = model(normalized_image)['segmentation_preds'][0].cpu().numpy() | |
| return np.array(visualize_mask(class_names = class_names, | |
| mask = out, | |
| image = np.transpose(image, [2,0,1]), | |
| img_size = 512, | |
| cat = True, | |
| axis = 1, | |
| )) | |
| dropdown_option_class_names = list(label_mapper.keys()) | |
| with gr.Blocks() as demo: | |
| with gr.Tab("Anatomy Segmentation"): | |
| with gr.Row(): | |
| gr.Markdown('This is a demo for Chest X-Ray Anatomy Segmentation with 159 classes. \n As it is running on CPU, it is not that fast. \n The demo only supports .jpg and .png files.') | |
| with gr.Row(): | |
| gr.Markdown('Note: for lateral views, only patients facing to the right side of the image are supported.') | |
| with gr.Row(): | |
| gr.Markdown('Not intended for clinical use!') | |
| with gr.Row(): | |
| gr.Markdown('To use it locally and to segment other data types like dicom, check out \'pip install cxas\'.') | |
| with gr.Row(): | |
| class_options = gr.Dropdown(dropdown_option_class_names, value = ['thoracic spine'], label="class_names", multiselect=True) | |
| with gr.Row(): | |
| image_input = gr.Image() | |
| with gr.Row(): | |
| image_output = gr.Image() | |
| image_button = gr.Button("Segment Chest X-Ray Anatomy") | |
| image_button.click(predict, inputs=[image_input,class_options], outputs=image_output) | |
| demo.launch() | |