fix encoder loading

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 transformers import AutoProcessor
+from torchvision import transforms
 from PIL import Image
 from jax import numpy as jnp
 import pandas as pd
@@ -13,7 +13,16 @@ import jax
 import gc
 
 
-preprocess = AutoProcessor.from_pretrained("clip-italian/clip-italian")
+preprocess = transforms.Compose(
+    [
+        transforms.ToTensor(),
+        transforms.Resize(224),
+        transforms.Normalize(
+            (0.48145466, 0.4578275, 0.40821073), 
+            (0.26862954, 0.26130258, 0.27577711)
+        ),
+    ]
+)
 
 
 def resize_longer(image, longer_size=224):
@@ -89,18 +98,16 @@ def gen_image_batch(image_url, image_size=224, pixel_size=10):
 
 
 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
 
     images, masks, vertical, horizontal = gen_image_batch(image_url, pixel_size=pixel_size)
     input_image = images[0].copy()
 
-    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)
+    images = np.stack([preprocess(pad_to_square(image)) for image in images], axis=0)
+    image_embeddings, embedding_norms = image_encoder(images, model)
+    text_embeddings, _ = text_encoder(text, model, tokenizer)
     
     vertical_scores = jnp.zeros((masks[0].shape[1], 512))
     vertical_masks = jnp.zeros((masks[0].shape[1], 1))
@@ -131,39 +138,11 @@ def get_heatmap(image_url, text, pixel_size=10, iterations=3):
     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)
+    sims = jnp.matmul(jnp.reshape(full_embs, (-1, 512)), text_embeddings.T)
+    score = jnp.reshape(sims, (*orig_shape[:2], 1))
     
     for i in range(iterations):
         score = jnp.clip(score - jnp.mean(score), 0, jnp.inf)

utils.py

@@ -34,18 +34,20 @@ def text_encoder(text, model, tokenizer):
         padding="max_length",
         return_tensors="np",
     )
-    embedding = model.get_text_features(inputs["input_ids"], inputs["attention_mask"])[
-        0
-    ]
-    embedding /= jnp.linalg.norm(embedding)
-    return jnp.expand_dims(embedding, axis=0)
+    embedding = model.get_text_features(
+        inputs["input_ids"], 
+        inputs["attention_mask"])[0]
+    norms = jnp.linalg.norm(embedding, axis=-1, keepdims=True)
+    embedding = embedding / norms
+    return jnp.expand_dims(embedding, axis=0), norms
 
 
 def image_encoder(image, model):
     image = image.permute(1, 2, 0).numpy()
     image = jnp.expand_dims(image, axis=0)  #  add batch size
     features = model.get_image_features(image,)
-    features /= jnp.linalg.norm(features, axis=-1, keepdims=True)
+    norms = jnp.linalg.norm(features, axis=-1, keepdims=True)
+    features = features / norms
     return features