Update app.py

app.py

@@ -148,12 +148,12 @@ def run_model_pipeline(model, original_image, comparison_image, threshold=0.3):
         col3, col4 = st.columns(2)
         
         with col3:
-            x1 = rearrange(img1_tensor.float().squeeze(0), 'b h w c -> b c h w')
+            x1 = rearrange(img1_tensor.float(), 'b h w c -> b c h w')
             fig1 = visualize_heatmap(model, x1)
             st.pyplot(fig1)
             
         with col4:
-            x2 = rearrange(img2_tensor.float().squeeze(0), 'b h w c -> b c h w')
+            x2 = rearrange(img2_tensor.float(), 'b h w c -> b c h w')
             fig2 = visualize_heatmap(model, x2)
             st.pyplot(fig2)