Update app.py

app.py

@@ -104,7 +104,7 @@ comparison_image = st.file_uploader(
     "Upload the signature to compare", type=["png", "jpg", "jpeg"]
 )
 
-def run_model_pipeline(model, original_image, comparison_image, threshold=0.5):
+def run_model_pipeline(model, original_image, comparison_image, threshold=0.3):
     if original_image is not None and comparison_image is not None: # ensure both images are uploaded
 
         # Preprocess images
@@ -146,11 +146,15 @@ def run_model_pipeline(model, original_image, comparison_image, threshold=0.5):
         # Visualize heatmaps from extracted model features
         st.subheader("Feature Heatmaps")
         col3, col4 = st.columns(2)
+        
         with col3:
-            fig1 = visualize_heatmap(model, img1_tensor.float().squeeze(0))
+            x1 = rearrange(img1_tensor.float().squeeze(0), 'b h w c -> b c h w')
+            fig1 = visualize_heatmap(model, x1)
             st.pyplot(fig1)
+            
         with col4:
-            fig2 = visualize_heatmap(model, img2_tensor.float().squeeze(0))
+            x2 = rearrange(img2_tensor.float().squeeze(0), 'b h w c -> b c h w')
+            fig2 = visualize_heatmap(model, x2)
             st.pyplot(fig2)
 
     else: