fix streamlit issues and update localization

localization.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from text2image import get_model, get_tokenizer, get_image_transform
 from utils import text_encoder
-from torchvision import transforms
+from transformers import AutoProcessor
 from PIL import Image
 from jax import numpy as jnp
 import pandas as pd
@@ -13,30 +13,34 @@ import jax
 import gc
 
 
-preprocess = transforms.Compose(
-    [
-        transforms.ToTensor(),
-        transforms.Normalize(
-            (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
-        ),
-    ]
-)
+preprocess = AutoProcessor.from_pretrained("clip-italian/clip-italian")
 
 
-def pad_to_square(image, size=224):
-    ratio = float(size) / max(image.size)
-    new_size = tuple([int(x * ratio) for x in image.size])
+def resize_longer(image, longer_size=224):
+    old_size = image.size
+    ratio = float(longer_size) / max(old_size)
+    new_size = tuple([int(x * ratio) for x in old_size])
     image = image.resize(new_size, Image.ANTIALIAS)
-    new_image = Image.new("RGB", size=(size, size), color=(128, 128, 128))
-    new_image.paste(image, ((size - new_size[0]) // 2, (size - new_size[1]) // 2))
-    return new_image
+    return image
+
+
+def pad_to_square(image):
+    (a,b)=image.shape[:2]
+    if a<b:
+        ah = (b - a) // 2
+        padding=((ah,b - a -ah), (0,0), (0,0))
+    else:
+        bh = (a - b) // 2
+        padding=((0,0), (bh,a-b-bh), (0,0))
+    return np.pad(image, padding,mode='constant',constant_values=127)
 
 
 def image_encoder(image, model):
     image = np.transpose(image, (0, 2, 3, 1))
     features = model.get_image_features(image)
-    features /= jnp.linalg.norm(features, keepdims=True)
-    return features
+    feature_norms = jnp.linalg.norm(features, axis=-1, keepdims=True)
+    features = features / feature_norms
+    return features, feature_norms
 
 
 def gen_image_batch(image_url, image_size=224, pixel_size=10):
@@ -44,64 +48,130 @@ def gen_image_batch(image_url, image_size=224, pixel_size=10):
 
     image_batch = []
     masks = []
+    is_vertical = []
+    is_horizontal = []
+    
     image_raw = requests.get(image_url, stream=True).raw
     image = Image.open(image_raw).convert("RGB")
-    image = pad_to_square(image, size=image_size)
-    gray = np.ones_like(image) * 128
-    mask = np.ones_like(image)
+    image = np.array(resize_longer(image, longer_size=image_size))
+    gray = np.ones_like(image) * 127
+    mask = np.ones_like(image[:,:,:1])
 
     image_batch.append(image)
     masks.append(mask)
+    is_vertical.append(True)
+    is_horizontal.append(True)
+
 
-    for i in range(0, n_pixels):
-        for j in range(i + 1, n_pixels):
+    for i in range(0, image.shape[0] // pixel_size + 1):
+        for j in range(i+1, image.shape[0] // pixel_size + 2):
             m = mask.copy()
-            m[: min(i * pixel_size, image_size) + 1, :] = 0
-            m[min(j * pixel_size, image_size) + 1 :, :] = 0
+            m[:min(i*pixel_size, image_size), :] = 0
+            m[min(j*pixel_size, image_size):, :] = 0
             neg_m = 1 - m
-            image_batch.append(image * m + gray * neg_m)
+            image_batch.append(image.copy() * m + gray * neg_m)
             masks.append(m)
+            is_vertical.append(False)
+            is_horizontal.append(True)
 
-    for i in range(0, n_pixels + 1):
-        for j in range(i + 1, n_pixels + 1):
+    for i in range(0, image.shape[1] // pixel_size + 1):
+        for j in range(i+1, image.shape[1] // pixel_size + 2):
             m = mask.copy()
-            m[:, : min(i * pixel_size + 1, image_size)] = 0
-            m[:, min(j * pixel_size + 1, image_size) :] = 0
+            m[:, :min(i*pixel_size, image_size)] = 0
+            m[:, min(j*pixel_size, image_size):] = 0
             neg_m = 1 - m
-            image_batch.append(image * m + gray * neg_m)
+            image_batch.append(image.copy() * m + gray * neg_m)
             masks.append(m)
+            is_vertical.append(True)
+            is_horizontal.append(False)
 
-    return image_batch, masks
+    return image_batch, masks, is_vertical, is_horizontal
 
 
 def get_heatmap(image_url, text, pixel_size=10, iterations=3):
-    tokenizer = get_tokenizer()
+    # tokenizer = get_tokenizer()
     model = get_model()
     image_size = model.config.vision_config.image_size
-    text_embedding = text_encoder(text, model, tokenizer)
-    images, masks = gen_image_batch(
-        image_url, image_size=image_size, pixel_size=pixel_size
-    )
 
+    images, masks, vertical, horizontal = gen_image_batch(image_url, pixel_size=pixel_size)
     input_image = images[0].copy()
-    images = np.stack([preprocess(image) for image in images], axis=0)
-    image_embeddings = jnp.asarray(image_encoder(images, model))
-
-    sims = []
-    scores = []
-    mask_val = jnp.zeros_like(masks[0])
-
-    for e, m in zip(image_embeddings, masks):
-        sim = jnp.matmul(e, text_embedding.T)
-        sims.append(sim)
-        if len(sims) > 1:
-            scores.append(sim * m)
-            mask_val += 1 - m
 
-    score = jnp.mean(jnp.clip(jnp.array(scores) - sims[0], 0, jnp.inf), axis=0)
+    inputs = preprocess(text=[text], images=images, return_tensors="np")
+        
+    image_embeddings, embedding_norms = image_encoder(inputs['pixel_values'], model)
+    text_embedding = model.get_text_features(inputs["input_ids"], inputs["attention_mask"])[0]
+    text_embedding = text_embedding / jnp.linalg.norm(text_embedding, axis=-1, keepdims=True)
+    
+    vertical_scores = jnp.zeros((masks[0].shape[1], 512))
+    vertical_masks = jnp.zeros((masks[0].shape[1], 1))
+    horizontal_scores = jnp.zeros((masks[0].shape[0], 512))
+    horizontal_masks = jnp.zeros((masks[0].shape[0], 1))
+
+    for e, n, m, v, h in zip(image_embeddings, embedding_norms, masks, vertical, horizontal):
+#         sim = (jnp.matmul(e, text_embedding.T)) #  + 1) / 2
+        
+#         sim = jax.nn.relu(sim)
+        
+        # if full_sim is None:
+        #     full_sim = sim
+        # sim = jax.nn.relu(sim - full_sim)
+        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)
+        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)
+                
+    
+    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)
+    
+    print(full_embs.shape)
+            
+    #full_embs = full_embs / jnp.linalg.norm(full_embs, axis=-1, keepdims=True)
+    full_embs = (full_embs / mask_sum)
+    
+    orig_shape = full_embs.shape
+    sims = jnp.matmul(jnp.reshape(full_embs, (-1, 512)), text_embedding.T)
+    sims = jnp.reshape(sims, (*orig_shape[:2], 1))
+    #sims = jax.nn.relu(sims)
+    
+    
+    
+    
+    
+            
+    # mean_vertical_scores = vertical_scores / vertical_masks
+    # mean_horizontal_scores = horizontal_scores / horizontal_masks
+    
+    # print(mean_vertical_score)
+    # print(mean_horizontal_score)
+    
+    # score = jnp.matmul(mean_vertical_scores, mean_horizontal_scores.T)
+    
+    #mask = jnp.matmul(vertical_masks, horizontal_scores.T)
+    #score = score / mask
+    
+    score = sims   # jnp.expand_dims(score.T, axis=-1)
+    #score = jax.nn.relu(score) / jnp.max(jnp.abs(score))
+
+    #score = jax.nn.relu(score - sims[0])
+    
+    # score = jnp.square(score)
+    
     for i in range(iterations):
         score = jnp.clip(score - jnp.mean(score), 0, jnp.inf)
+        
     score = (score - jnp.min(score)) / (jnp.max(score) - jnp.min(score))
+    
+    print(jnp.min(score), jnp.max(score))
+    
     return np.asarray(score), input_image
 
 
@@ -144,7 +214,7 @@ def app():
     with col2:
         pixel_size = st.selectbox("Pixel Size", options=range(10, 26, 5), index=2)
 
-        iterations = st.selectbox("Refinement Steps", options=range(3, 30, 3), index=0)
+        iterations = st.selectbox("Refinement Steps", options=range(1, 6, 1), index=0)
 
         compute = st.button("LOCATE")
 

modeling_hybrid_clip.py

@@ -136,8 +136,10 @@ class FlaxHybridCLIP(FlaxPreTrainedModel):
     ):
         if input_shape is None:
             input_shape = ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))
+            
+        print(kwargs)
 
-        module = self.module_class(config=config, dtype=dtype, **kwargs)
+        module = self.module_class(config=config, dtype=dtype)   # , **kwargs)
         super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
 
     def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:

requirements.txt

@@ -8,4 +8,4 @@ stqdm
 pandas
 requests
 psutil
-streamlit==1.2.0
+streamlit