update localization

localization.py

@@ -127,22 +127,22 @@ def get_heatmap(image_url, text, pixel_size=10, iterations=3):
         # if full_sim is None:
         #     full_sim = sim
         # sim = jax.nn.relu(sim - full_sim)
-        emb = jnp.expand_dims(e, axis=0) * n
+        emb = jnp.expand_dims(e, axis=0) #* n
         
         if v:
             vm = jnp.any(m, axis=0)
-            vertical_scores = vertical_scores + (emb * vm) #/ jnp.mean(vm)
-            vertical_masks = vertical_masks + vm #/ jnp.mean(vm)
+            vertical_scores = vertical_scores + (emb * vm) / jnp.mean(vm)
+            vertical_masks = vertical_masks + vm / jnp.mean(vm)
         if h:
             hm = jnp.any(m, axis=1)
-            horizontal_scores = horizontal_scores + (emb * hm) #/ jnp.mean(hm)
-            horizontal_masks = horizontal_masks + hm #/ jnp.mean(hm)
+            horizontal_scores = horizontal_scores + (emb * hm) / jnp.mean(hm)
+            horizontal_masks = horizontal_masks + hm / jnp.mean(hm)
                 
     
     embs_1 = jnp.expand_dims((vertical_scores), axis=0) * jnp.expand_dims(jnp.abs(horizontal_scores), axis=1)
     embs_2 = jnp.expand_dims(jnp.abs(vertical_scores), axis=0) * jnp.expand_dims((horizontal_scores), axis=1)
     full_embs = jnp.minimum(embs_1, embs_2)
-    mask_sum = jnp.expand_dims(vertical_masks, axis=0) * jnp.expand_dims(horizontal_masks, axis=1)
+    mask_sum = jnp.expand_dims(vertical_masks + 1, axis=0) * jnp.expand_dims(horizontal_masks + 1, axis=1)
     full_embs = (full_embs / mask_sum)
     
     orig_shape = full_embs.shape
@@ -196,9 +196,9 @@ def app():
     col1, col2 = st.columns([0.75, 0.25])
 
     with col2:
-        pixel_size = st.selectbox("Pixel Size", options=range(10, 26, 5), index=2)
+        pixel_size = st.selectbox("Pixel Size", options=range(5, 26, 5), index=3)
 
-        iterations = st.selectbox("Refinement Steps", options=range(1, 6, 1), index=0)
+        iterations = st.selectbox("Refinement Steps", options=range(0, 6, 1), index=0)
 
         compute = st.button("LOCATE")