Spaces:
Runtime error
Runtime error
| import os | |
| import gradio as gr | |
| import torch | |
| from monai import bundle | |
| from monai.transforms import ( | |
| Compose, | |
| LoadImaged, | |
| EnsureChannelFirstd, | |
| Orientationd, | |
| NormalizeIntensityd, | |
| Activationsd, | |
| AsDiscreted, | |
| ScaleIntensityd, | |
| ) | |
| # Define the bundle name and path for downloading | |
| BUNDLE_NAME = 'spleen_ct_segmentation_v0.1.0' | |
| BUNDLE_PATH = os.path.join(torch.hub.get_dir(), 'bundle', BUNDLE_NAME) | |
| # Title and description | |
| title = '<h1 style="text-align: center;">Segment Brain Tumors with MONAI! 🧠 </h1>' | |
| description = """ | |
| ## 🚀 To run | |
| Upload a brain MRI image file, or try out one of the examples below! | |
| If you want to see a different slice, update the slider. | |
| More details on the model can be found [here!](https://huggingface.co/katielink/brats_mri_segmentation_v0.1.0) | |
| ## ⚠️ Disclaimer | |
| This is an example, not to be used for diagnostic purposes. | |
| """ | |
| references = """ | |
| ## 👀 References | |
| 1. Myronenko, Andriy. "3D MRI brain tumor segmentation using autoencoder regularization." International MICCAI Brainlesion Workshop. Springer, Cham, 2018. https://arxiv.org/abs/1810.11654. | |
| 2. Menze BH, et al. "The Multimodal Brain Tumor Image Segmentation Benchmark (BRATS)", IEEE Transactions on Medical Imaging 34(10), 1993-2024 (2015) DOI: 10.1109/TMI.2014.2377694 | |
| 3. Bakas S, et al. "Advancing The Cancer Genome Atlas glioma MRI collections with expert segmentation labels and radiomic features", Nature Scientific Data, 4:170117 (2017) DOI:10.1038/sdata.2017.117 | |
| """ | |
| examples = [ | |
| ['examples/BRATS_485.nii.gz', 65], | |
| ['examples/BRATS_486.nii.gz', 80] | |
| ] | |
| # Load the MONAI pretrained model from Hugging Face Hub | |
| model, _, _ = bundle.load( | |
| name = BUNDLE_NAME, | |
| source = 'huggingface_hub', | |
| repo = 'katielink/brats_mri_segmentation_v0.1.0', | |
| load_ts_module=True, | |
| ) | |
| # Use GPU if available | |
| device = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| # Load the parser from the MONAI bundle's inference config | |
| parser = bundle.load_bundle_config(BUNDLE_PATH, 'inference.json') | |
| # Compose the preprocessing transforms | |
| preproc_transforms = Compose( | |
| [ | |
| LoadImaged(keys=["image"]), | |
| EnsureChannelFirstd(keys="image"), | |
| Orientationd(keys=["image"], axcodes="RAS"), | |
| NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True), | |
| ] | |
| ) | |
| # Get the inferer from the bundle's inference config | |
| inferer = parser.get_parsed_content( | |
| 'inferer', | |
| lazy=True, eval_expr=True, instantiate=True | |
| ) | |
| # Compose the postprocessing transforms | |
| post_transforms = Compose( | |
| [ | |
| Activationsd(keys='pred', sigmoid=True), | |
| AsDiscreted(keys='pred', threshold=0.5), | |
| ScaleIntensityd(keys='image', minv=0., maxv=1.) | |
| ] | |
| ) | |
| # Define the predict function for the demo | |
| def predict(input_file, z_axis, model=model, device=device): | |
| # Load and process data in MONAI format | |
| data = {'image': [input_file.name]} | |
| data = preproc_transforms(data) | |
| # Run inference and post-process predicted labels | |
| model.to(device) | |
| model.eval() | |
| with torch.no_grad(): | |
| inputs = data['image'].to(device) | |
| data['pred'] = inferer(inputs=inputs[None,...], network=model) | |
| data = post_transforms(data) | |
| # Convert tensors back to numpy arrays | |
| data['image'] = data['image'].numpy() | |
| data['pred'] = data['pred'].cpu().detach().numpy() | |
| # Magnetic resonance imaging sequences | |
| t1c = data['image'][0, :, :, z_axis] # T1-weighted, post contrast | |
| t1 = data['image'][1, :, :, z_axis] # T1-weighted, pre contrast | |
| t2 = data['image'][2, :, :, z_axis] # T2-weighted | |
| flair = data['image'][3, :, :, z_axis] # FLAIR | |
| # BraTS labels | |
| tc = data['pred'][0, 0, :, :, z_axis] # Tumor core | |
| wt = data['pred'][0, 1, :, :, z_axis] # Whole tumor | |
| et = data['pred'][0, 2, :, :, z_axis] # Enhancing tumor | |
| return [t1c, t1, t2, flair], [tc, wt, et] | |
| # Use blocks to set up a more complex demo | |
| with gr.Blocks() as demo: | |
| with gr.Row(): | |
| # Get the input file and slice slider as inputs | |
| input_file = gr.File(label='input file') | |
| z_axis = gr.Slider(0, 200, label='slice', value=50) | |
| with gr.Row(): | |
| # Show the button with custom label | |
| button = gr.Button("Segment Tumor!") | |
| with gr.Row(): | |
| with gr.Column(): | |
| # Show the input image with different MR sequences | |
| input_image = gr.Gallery(label='input MRI sequences (T1+, T1, T2, FLAIR)') | |
| with gr.Column(): | |
| # Show the segmentation labels | |
| output_segmentation = gr.Gallery(label='output segmentations (TC, WT, ET)') | |
| # Run prediction on button click | |
| button.click( | |
| predict, | |
| inputs=[input_file, z_axis], | |
| outputs=[input_image, output_segmentation] | |
| ) | |
| # Have some example for the user to try out | |
| examples = gr.Examples( | |
| examples=examples, | |
| inputs=[input_file, z_axis], | |
| outputs=[input_image, output_segmentation], | |
| fn=predict, | |
| cache_examples=False | |
| ) | |
| # Launch the demo | |
| demo.launch() |