Create app.py

app.py

@@ -0,0 +1,120 @@
+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()