[image2text] add initial version

image2text.py

@@ -1,4 +1,10 @@
 import streamlit as st
+from text2image import get_model, get_tokenizer, get_image_transform
+from utils import text_encoder, image_encoder
+from PIL import Image
+from jax import numpy as jnp
+import pandas as pd
+
 
 def app():
     st.title("From Image to Text")
@@ -12,4 +18,55 @@ def app():
         🤌 Italian mode on! 🤌
 
         """
-    )
+    )
+
+    filename = st.file_uploader(
+        "Choose an image from your computer", type=["jpg", "jpeg", "png"]
+    )
+
+    MAX_CAP = 4
+
+    col1, col2 = st.beta_columns([3, 1])
+
+    with col2:
+        captions_count = st.selectbox(
+            "Number of captions", options=range(1, MAX_CAP + 1)
+        )
+        compute = st.button("Compute")
+
+    with col1:
+        captions = list()
+        for idx in range(min(MAX_CAP, captions_count)):
+            captions.append(st.text_input(f"Insert Caption {idx+1}"))
+
+    if compute:
+        captions = [c for c in captions if c != ""]
+
+        if not captions or not filename:
+            st.error("Please choose one image and at least one caption")
+        else:
+            with st.spinner("Computing..."):
+                model = get_model()
+                tokenizer = get_tokenizer()
+
+                text_embeds = list()
+                for i, c in enumerate(captions):
+                    text_embeds.extend(text_encoder(c, model, tokenizer))
+
+                text_embeds = jnp.array(text_embeds)
+
+                image = Image.open(filename).convert("RGB")
+                transform = get_image_transform(model.config.vision_config.image_size)
+                image_embed = image_encoder(transform(image), model)
+
+                # we could have a softmax here
+                cos_similarities = jnp.matmul(image_embed, text_embeds.T)
+
+                chart_data = pd.Series(cos_similarities[0], index=captions)
+
+                col1, col2 = st.beta_columns(2)
+                with col1:
+                    st.bar_chart(chart_data)
+
+                with col2:
+                    st.image(image)

requirements.txt

@@ -4,4 +4,5 @@ transformers
 torch
 torchvision
 natsort
-stqdm
+stqdm
+pandas

text2image.py

@@ -81,6 +81,20 @@ def load_urls(dataset_name):
         ValueError(f"{dataset_name} not supported here")
 
 
+def get_image_transform(image_size):
+    return Compose(
+        [
+            Resize([image_size], interpolation=InterpolationMode.BICUBIC),
+            CenterCrop(image_size),
+            ToTensor(),
+            Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
+        ]
+    )
+
+
 def app():
 
     st.title("From Text to Image")
@@ -140,18 +154,9 @@ def app():
 
             if dataset_name == "Unsplash":
                 image_size = model.config.vision_config.image_size
-                val_preprocess = Compose(
-                    [
-                        Resize([image_size], interpolation=InterpolationMode.BICUBIC),
-                        CenterCrop(image_size),
-                        ToTensor(),
-                        Normalize(
-                            (0.48145466, 0.4578275, 0.40821073),
-                            (0.26862954, 0.26130258, 0.27577711),
-                        ),
-                    ]
+                dataset = utils.CustomDataSet(
+                    "photos/", transform=get_image_transform(image_size)
                 )
-                dataset = utils.CustomDataSet("photos/", transform=val_preprocess)
             elif dataset_name == "CC":
                 dataset = load_urls(dataset_name)
             else:

utils.py

@@ -41,6 +41,14 @@ def text_encoder(text, model, tokenizer):
     return jnp.expand_dims(embedding, axis=0)
 
 
+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)
+    return features
+
+
 def precompute_image_features(model, loader):
     image_features = []
     for i, (images) in enumerate(tqdm(loader)):