KalbeDigitalLab/PathologyNucleiClassification

1import torch2from monai.bundle import ConfigParser3import gradio as gr4 5from utils import page_utils6 7parser = ConfigParser() #  load configuration files that specify various parameters for running the MONAI workflow.8parser.read_config(f="configs/inference.json") # read the config from specified JSON file9parser.read_meta(f="configs/metadata.json") # read the metadata from specified JSON file10 11inference = parser.get_parsed_content("inferer")12network = parser.get_parsed_content("network_def")13preprocess = parser.get_parsed_content("preprocessing")14state_dict = torch.load("models/model.pt", map_location=torch.device('cpu'))15network.load_state_dict(state_dict, strict=True) #  Loads a model’s parameter dictionary16 17class_names = {18    0: "Other",19    1: "Inflammatory",20    2: "Epithelial",21    3: "Spindle-Shaped",22}23 24def classify_image(image_file, label_file):25    if image_file is None:26        raise gr.Error("Need a histology image")27    if label_file is None:28        raise gr.Error("Need a label image")29    data = {"image":image_file, "label":label_file}30    batch = preprocess(data)31    batch['image'] = batch['image']32    network.eval()33    with torch.no_grad():34        pred = inference(batch['image'].unsqueeze(dim=0), network) # expect 4 channels input  (3 RGB, 1 Label mask)35    prob = pred.softmax(-1).detach().cpu().numpy()[0]36    confidences = {class_names[i]: float(prob[i]) for i in range(len(class_names))}37    return confidences38 39example_files1 = [40    ['sample_data/Images/test_11_2_0628.png',41    'sample_data/Labels/test_11_2_0628.png'],42    ['sample_data/Images/test_9_4_0149.png',43    'sample_data/Labels/test_9_4_0149.png'],44    ['sample_data/Images/test_12_3_0292.png',45    'sample_data/Labels/test_12_3_0292.png'],46    ['sample_data/Images/test_9_4_0019.png',47    'sample_data/Labels/test_9_4_0019.png']48]49 50example_files2 = [51    ['sample_data/Images/test_14_3_0433.png',52    'sample_data/Labels/test_14_3_0433.png'],53    ['sample_data/Images/test_14_4_0544.png',54    'sample_data/Labels/test_14_4_0544.png'],55    ['sample_data/Images/train_1_1_0095.png',56    'sample_data/Labels/train_1_1_0095.png'],57    ['sample_data/Images/train_1_3_0020.png',58    'sample_data/Labels/train_1_3_0020.png'],59]60 61with open('index.html', encoding='utf-8') as file:62   html_content = file.read()63 64with gr.Blocks(theme=gr.themes.Default(primary_hue=page_utils.KALBE_THEME_COLOR, secondary_hue=page_utils.KALBE_THEME_COLOR).set(65        button_primary_background_fill='*primary_600',66        button_primary_background_fill_hover='*primary_500',67        button_primary_text_color='white',68    )) as app:69    gr.HTML(html_content)70    with gr.Row():71        with gr.Column():72            with gr.Row():73                inp_img = gr.Image(type="filepath", image_mode="RGB", label="Histology Image", show_label=True)74                label_img = gr.Image(type="filepath", image_mode="L", label="Label Image", show_label=True)75            with gr.Row():76                clear_btn = gr.Button(value="Clear")77                process_btn = gr.Button(value="Process", variant="primary")78        out_txt = gr.Label(label="Probabilities", num_top_classes=4)79 80    process_btn.click(fn=classify_image, inputs=[inp_img, label_img], outputs=out_txt)81    clear_btn.click(lambda:(82        gr.update(value=None),83        gr.update(value=None),84        gr.update(value=None)85        ),86        inputs=None,87        outputs=[inp_img, label_img,out_txt]88        )89 90    gr.Markdown("## Image Examples")91    with gr.Row():92        for file in example_files1:93            gr.Examples(94                [file], inputs=[inp_img, label_img]95            )96    with gr.Row():97        for file in example_files2:98            gr.Examples(99                [file], inputs=[inp_img, label_img]100            )101app.launch()