veb-101/UWMGI_Medical_Image_Segmentation

1import os2import numpy as np3import gradio as gr4from glob import glob5from functools import partial6from dataclasses import dataclass7 8import torch9import torch.nn.functional as F10import torchvision.transforms as TF11from transformers import SegformerForSemanticSegmentation12 13 14@dataclass15class Configs:16    NUM_CLASSES: int = 4  # including background.17    CLASSES: tuple = ("Large bowel", "Small bowel", "Stomach")18    IMAGE_SIZE: tuple[int, int] = (288, 288)  # W, H19    MEAN: tuple = (0.485, 0.456, 0.406)20    STD: tuple = (0.229, 0.224, 0.225)21    MODEL_PATH: str = os.path.join(os.getcwd(), "segformer_trained_weights")22 23 24def get_model(*, model_path, num_classes):25    model = SegformerForSemanticSegmentation.from_pretrained(model_path, num_labels=num_classes, ignore_mismatched_sizes=True)26    return model27 28 29@torch.inference_mode()30def predict(input_image, model=None, preprocess_fn=None, device="cpu"):31    shape_H_W = input_image.size[::-1]32    input_tensor = preprocess_fn(input_image)33    input_tensor = input_tensor.unsqueeze(0).to(device)34 35    # Generate predictions36    outputs = model(pixel_values=input_tensor.to(device), return_dict=True)37    predictions = F.interpolate(outputs["logits"], size=shape_H_W, mode="bilinear", align_corners=False)38 39    preds_argmax = predictions.argmax(dim=1).cpu().squeeze().numpy()40 41    seg_info = [(preds_argmax == idx, class_name) for idx, class_name in enumerate(Configs.CLASSES, 1)]42 43    return (input_image, seg_info)44 45 46if __name__ == "__main__":47    class2hexcolor = {"Stomach": "#007fff", "Small bowel": "#009A17", "Large bowel": "#FF0000"}48 49    DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")50 51    model = get_model(model_path=Configs.MODEL_PATH, num_classes=Configs.NUM_CLASSES)52    model.to(DEVICE)53    model.eval()54    _ = model(torch.randn(1, 3, *Configs.IMAGE_SIZE[::-1], device=DEVICE))55 56    preprocess = TF.Compose(57        [58            TF.Resize(size=Configs.IMAGE_SIZE[::-1]),59            TF.ToTensor(),60            TF.Normalize(Configs.MEAN, Configs.STD, inplace=True),61        ]62    )63 64    with gr.Blocks(title="Medical Image Segmentation") as demo:65        gr.Markdown("""<h1><center>Medical Image Segmentation with UW-Madison GI Tract Dataset</center></h1>""")66        with gr.Row():67            img_input = gr.Image(type="pil", height=360, width=360, label="Input image")68            img_output = gr.AnnotatedImage(label="Predictions", height=360, width=360, color_map=class2hexcolor)69 70        section_btn = gr.Button("Generate Predictions")71        section_btn.click(partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE), img_input, img_output)72 73        images_dir = glob(os.path.join(os.getcwd(), "samples") + os.sep + "*.png")74        examples = [i for i in np.random.choice(images_dir, size=10, replace=False)]75        gr.Examples(examples=examples, inputs=img_input, outputs=img_output)76 77    demo.launch()78