Update app.py

app.py

@@ -88,11 +88,12 @@ def generate_activation_maps(image, patch_size=16):
     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("w, h:", w, h)
+    # min_size = min(w, h)
     print("Image shape:", img.shape)
     preprocess = transforms.Compose([
             transforms.Resize((img.shape[0], img.shape[1])),
-            transforms.CenterCrop(min_size),
+            transforms.CenterCrop((w, h)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize the tensors
         ])
@@ -105,27 +106,29 @@ def generate_activation_maps(image, patch_size=16):
 
 # Streamlit UI
 st.title("PathDino - Compact ViT for Histopathology Image Analysis")
-st.write("Upload a histology image with 512x512 dimension of 20X magnification to view the activation maps.")
+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=400)
+    # columns = st.columns(3)
+    st.image(uploaded_image, caption="Uploaded Image", width=500)
 
     # 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(2)
     columns2 = st.columns(2)
     columns3 = st.columns(2)
+    # for index in range(6):
     for index, col in enumerate(columns):
         # Create a plot
-        plt.plot(512, 512)
+        plt.plot(600, 600)
 
         # Remove x and y axis labels
         plt.xticks([])  # Hide x-axis ticks and labels
@@ -134,16 +137,21 @@ if uploaded_image is not None:
         # 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)
+        
+        print(type(attention_list[index]))
+        print(attention_list[index].shape)
 
         plt.imshow(attention_list[index])
         col.pyplot(plt)
+        # col
+        # st.image(plt, caption=f"Head-{index+1}", width=display_w)
         plt.close()
         
     for index, col in enumerate(columns2):
         
         index = index + 2
         # Create a plot
-        plt.plot(512, 512)
+        plt.plot(600, 600)
 
         # Remove x and y axis labels
         plt.xticks([])  # Hide x-axis ticks and labels
@@ -161,7 +169,7 @@ if uploaded_image is not None:
         
         index = index + 4
         # Create a plot
-        plt.plot(512, 512)
+        plt.plot(600, 600)
 
         # Remove x and y axis labels
         plt.xticks([])  # Hide x-axis ticks and labels