Update app.py

app.py

@@ -4,6 +4,7 @@ import matplotlib.pyplot as plt
 from torchvision import transforms
 from PIL import Image
 import torch.nn as nn
+import numpy as np
 
 from PathDino import get_pathDino_model
 
@@ -82,13 +83,19 @@ def visualize_attention_ViT(model, img, patch_size=16):
         return attention_list
 
 # Define the function to generate activation maps
-def generate_activation_maps(image):
+def generate_activation_maps(image, patch_size=16):
+    # Convert the image to a NumPy array
+    img = np.array(image)
+    # make the image divisible by the patch size
+    w, h = img.shape[1] - img.shape[0] % patch_size, img.shape[1] - img.shape[1] % patch_size
+    min_size = min(w, h)
+    print("Image shape:", img.shape)
     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
-    ])
+            transforms.Resize((img.shape[0], img.shape[1])),
+            transforms.CenterCrop(min_size),
+            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
@@ -106,15 +113,16 @@ 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)
+    columns[1].image(uploaded_image, caption="Uploaded Image", width=400)
 
     # 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)
+    columns = st.columns(2)
+    columns2 = st.columns(2)
+    columns3 = st.columns(2)
     for index, col in enumerate(columns):
         # Create a plot
         plt.plot(512, 512)
@@ -133,7 +141,25 @@ if uploaded_image is not None:
         
     for index, col in enumerate(columns2):
         
-        index = index + len(attention_list)//2
+        index = index + 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()
+        
+    for index, col in enumerate(columns3):
+        
+        index = index + 4
         # Create a plot
         plt.plot(512, 512)