Spaces:
Runtime error
Runtime error
import torch | |
import matplotlib.pyplot as plt | |
from monai.networks.nets import SegResNet | |
from monai.inferers import sliding_window_inference | |
from monai.transforms import ( | |
Activations, | |
AsDiscrete, | |
Compose, | |
) | |
model = SegResNet( | |
blocks_down=[1, 2, 2, 4], | |
blocks_up=[1, 1, 1], | |
init_filters=16, | |
in_channels=4, | |
out_channels=3, | |
dropout_prob=0.2, | |
) | |
model.load_state_dict( | |
torch.load("weights/model.pt", map_location=torch.device('cpu')) | |
) | |
# define inference method | |
VAL_AMP = True | |
def inference(input): | |
def _compute(input): | |
return sliding_window_inference( | |
inputs=input, | |
roi_size=(240, 240, 160), | |
sw_batch_size=1, | |
predictor=model, | |
overlap=0.5, | |
) | |
if VAL_AMP: | |
with torch.cuda.amp.autocast(): | |
return _compute(input) | |
else: | |
return _compute(input) | |
post_trans = Compose( | |
[Activations(sigmoid=True), AsDiscrete(threshold=0.5)] | |
) | |
import gradio as gr | |
def load_sample1(): | |
return load_sample(1) | |
def load_sample2(): | |
return load_sample(2) | |
def load_sample3(): | |
return load_sample(3) | |
def load_sample4(): | |
return load_sample(4) | |
def load_sample5(): | |
return load_sample(5) | |
def load_sample6(): | |
return load_sample(6) | |
def load_sample7(): | |
return load_sample(7) | |
def load_sample8(): | |
return load_sample(8) | |
import torchvision | |
def load_sample(index): | |
#sample_index = index | |
image_filenames = [] | |
for i in range(4): | |
image_filenames.append(f"thumbnails/image{index-1}_{i}.png") | |
label_filenames = [] | |
for i in range(3): | |
label_filenames.append(f"thumbnails_label/label{index-1}_{i}.png") | |
return [index, image_filenames[0], image_filenames[1], image_filenames[2], image_filenames[3], | |
label_filenames[0], label_filenames[1], label_filenames[2]] | |
def predict(sample_index): | |
sample = torch.load(f"samples/val{sample_index-1}.pt") | |
model.eval() | |
with torch.no_grad(): | |
# select one image to evaluate and visualize the model output | |
val_input = sample["image"].unsqueeze(0) | |
roi_size = (128, 128, 64) | |
sw_batch_size = 4 | |
val_output = inference(val_input) | |
val_output = post_trans(val_output[0]) | |
imgs_output = [] | |
for i in range(3): | |
imgs_output.append(val_output[i, :, :, 70]) | |
pil_images_output = [] | |
for i in range(3): | |
pil_images_output.append(torchvision.transforms.functional.to_pil_image(imgs_output[i])) | |
return [pil_images_output[0], pil_images_output[1], pil_images_output[2]] | |
with gr.Blocks(title="Brain tumor 3D segmentation with MONAI - ClassCat", | |
css=".gradio-container {background:azure;}" | |
) as demo: | |
sample_index = gr.State([]) | |
gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">Brain tumor 3D segmentation with MONAI</div>""") | |
gr.HTML("""<h4 style="color:navy;">1. Select an example, which includes input images and label images, by clicking "Example x" button.</h4>""") | |
with gr.Row(): | |
input_image0 = gr.Image(label="image channel 0", type="filepath", shape=(240, 240)) | |
input_image1 = gr.Image(label="image channel 1", type="filepath", shape=(240, 240)) | |
input_image2 = gr.Image(label="image channel 2", type="filepath", shape=(240, 240)) | |
input_image3 = gr.Image(label="image channel 3", type="filepath", shape=(240, 240)) | |
with gr.Row(): | |
label_image0 = gr.Image(label="label channel 0", type="filepath", shape=(240, 240)) | |
label_image1 = gr.Image(label="label channel 1", type="filepath", shape=(240, 240)) | |
label_image2 = gr.Image(label="label channel 2", type="filepath", shape=(240, 240)) | |
with gr.Row(): | |
example1_btn = gr.Button("Example 1") | |
example2_btn = gr.Button("Example 2") | |
example3_btn = gr.Button("Example 3") | |
example4_btn = gr.Button("Example 4") | |
example5_btn = gr.Button("Example 5") | |
example6_btn = gr.Button("Example 6") | |
example7_btn = gr.Button("Example 7") | |
example8_btn = gr.Button("Example 8") | |
example1_btn.click(fn=load_sample1, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example2_btn.click(fn=load_sample2, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example3_btn.click(fn=load_sample3, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example4_btn.click(fn=load_sample4, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example5_btn.click(fn=load_sample5, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example6_btn.click(fn=load_sample6, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example7_btn.click(fn=load_sample7, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
example8_btn.click(fn=load_sample8, inputs=None, | |
outputs=[sample_index, input_image0, input_image1, input_image2, input_image3, | |
label_image0, label_image1, label_image2]) | |
gr.HTML("""<br/>""") | |
gr.HTML("""<h4 style="color:navy;">2. Then, click "Infer" button to predict segmentation images. It will take about 30 seconds (on cpu)</h4>""") | |
with gr.Row(): | |
output_image0 = gr.Image(label="output channel 0", type="pil") | |
output_image1 = gr.Image(label="output channel 1", type="pil") | |
output_image2 = gr.Image(label="output channel 2", type="pil") | |
send_btn = gr.Button("Infer") | |
send_btn.click(fn=predict, inputs=[sample_index], outputs=[output_image0, output_image1, output_image2]) | |
gr.HTML("""<br/>""") | |
gr.HTML("""<h4 style="color:navy;">Reference</h4>""") | |
gr.HTML("""<ul>""") | |
gr.HTML("""<li><a href="https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/brats_segmentation_3d.ipynb" target="_blank">Brain tumor 3D segmentation with MONAI</a></li>""") | |
gr.HTML("""</ul>""") | |
#demo.queue() | |
demo.launch(debug=True) | |
### EOF ### | |