Update app.py

app.py

@@ -4,20 +4,20 @@ from utils import get_patch_embeddings, compute_patch_similarity, overlay_simila
 
 selected_patch = {"row": 0, "col": 0}  
 
-def init_states(img):
+def init_states(img, alpha):
     if img is None:
         return gr.update(value=None), None
     patch_embs, patch_embs_norm, rows, cols = get_patch_embeddings(img, ps=16, device=device)
     
     sim_map = compute_patch_similarity(patch_embs, patch_embs_norm, 0, 0)
-    result_img = overlay_similarity(img, sim_map, alpha=0.6, cmap="hot")
+    result_img = overlay_similarity(img, sim_map, alpha=alpha, cmap="hot")
 
     state = {
         "img": img,
         "patch_embs": patch_embs,
         "patch_embs_norm": patch_embs_norm,
         "grid_size": rows,
-        "alpha": 0.6,
+        "alpha": alpha,
         "overlay_img":result_img,
     }
 
@@ -49,7 +49,7 @@ def store_patch(evt, state):
     return state
 
 
-def reload_overlay(evt: gr.SelectData,state):
+def reload_overlay(evt: gr.SelectData,state,alpha):
     if state is None:
         return None
     store_patch(evt, state)
@@ -57,7 +57,6 @@ def reload_overlay(evt: gr.SelectData,state):
     img = state["img"]
     patch_embs = state["patch_embs"]
     patch_embs_norm = state["patch_embs_norm"]
-    alpha = state["alpha"]
     sim_map = compute_patch_similarity(patch_embs, patch_embs_norm, row, col)
     result_img = overlay_similarity(img, sim_map, alpha=alpha, cmap="hot")
     return result_img
@@ -125,14 +124,15 @@ with gr.Blocks() as demo:
 
     img_input.change(
         fn=init_states,
-        inputs=[img_input],
+        inputs=[img_input, alpha_slider],
         outputs=[state_store, output_img]
     )
 
     output_img.select(
         fn=reload_overlay,
-        inputs=[state_store],
+        inputs=[state_store, alpha_slider],
         outputs=[output_img]
     )
 
+
 demo.launch()