import streamlit as st import torch import matplotlib.pyplot as plt from torchvision import transforms from PIL import Image import torch.nn as nn from PathDino import get_pathDino_model import os os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE' device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Load PathDino model and image transforms model, image_transforms = get_pathDino_model("PathDino512.pth") st.sidebar.markdown("### PathDino") st.sidebar.markdown( "PathDino is a lightweight histology transformer consisting of just five small vision transformer blocks. " "PathDino is a customized ViT architecture, finely tuned to the nuances of histological images. It not only exhibits " "superior performance but also effectively reduces susceptibility to overfitting, a common challenge in histology " "image analysis.\n\n" ) default_image_url_compare = "images/HistRotate.png" st.sidebar.image(default_image_url_compare, caption='A 360 rotation augmentation for training models on histopathology images. Unlike training on natural images where the rotation may change the context of the visual data, rotating a histopathology patch does not change the context and it improves the learning process for better reliable embedding learning.', width=500) default_image_url_compare = "images/FigPathDino_parameters_FLOPs_compare.png" st.sidebar.image(default_image_url_compare, caption='PathDino Vs its counterparts. Number of Parameters (Millions) vs the patch-level retrieval with macro avg F-score of majority vote (MV@5) on CAMELYON16 dataset. The bubble size represents the FLOPs.', width=500) default_image_url_compare = "images/ActivationMap.png" st.sidebar.image(default_image_url_compare, caption='Attention Visualization. When visualizing attention patterns, our PathDino transformer outperforms HIPT-small and DinoSSLPath, despite being trained on a smaller dataset of 6 million TCGA patches. In contrast, DinoSSLPath and HIPT were trained on much larger datasets, with 19 million and 104 million TCGA patches, respectively.', width=500) def visualize_attention_ViT(model, img, patch_size=16): attention_list = [] device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") w_featmap = img.shape[-2] // patch_size h_featmap = img.shape[-1] // patch_size attentions = model.get_last_selfattention(img.to(device)) nh = attentions.shape[1] # number of head # we keep only the output patch attention attentions = attentions[0, :, 0, 1:].reshape(nh, -1) attentions = attentions.reshape(nh, w_featmap, h_featmap) attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=patch_size, mode="nearest")[0].detach().numpy() for j in range(nh): attention_list.append(attentions[j]) return attention_list # Define the function to generate activation maps def generate_activation_maps(image): preprocess = transforms.Compose([ transforms.Resize((512, 512)), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize the tensors ]) image_tensor = preprocess(image) img = image_tensor.unsqueeze(0).to(device) # Generate activation maps with torch.no_grad(): attention_list = visualize_attention_ViT(model=model, img=img, patch_size=16) return attention_list # Streamlit UI st.title("PathDino - Compact ViT for Histolopathology Image Analysis") st.write("Upload a histology image to view the activation maps.") # uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"]) uploaded_image = "images/HistRotate.png" uploaded_image = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"]) if uploaded_image is not None: columns = st.columns(3) columns[1].image(uploaded_image, caption="Uploaded Image", width=300) # Load the image and apply preprocessing uploaded_image = Image.open(uploaded_image).convert('RGB') attention_list = generate_activation_maps(uploaded_image) print(len(attention_list)) st.subheader(f"Attention Maps of the input image") columns = st.columns(len(attention_list)//2) columns2 = st.columns(len(attention_list)//2) for index, col in enumerate(columns): # Create a plot plt.plot(512, 512) # Remove x and y axis labels plt.xticks([]) # Hide x-axis ticks and labels plt.yticks([]) # Hide y-axis ticks and labels # Alternatively, if you only want to hide the labels and keep the ticks: plt.gca().axes.get_xaxis().set_visible(False) plt.gca().axes.get_yaxis().set_visible(False) plt.imshow(attention_list[index]) col.pyplot(plt) plt.close() for index, col in enumerate(columns2): index = index + len(attention_list)//2 # Create a plot plt.plot(512, 512) # Remove x and y axis labels plt.xticks([]) # Hide x-axis ticks and labels plt.yticks([]) # Hide y-axis ticks and labels # Alternatively, if you only want to hide the labels and keep the ticks: plt.gca().axes.get_xaxis().set_visible(False) plt.gca().axes.get_yaxis().set_visible(False) plt.imshow(attention_list[index]) col.pyplot(plt) plt.close()