Merge branch 'main' of https://huggingface.co/spaces/clip-italian/clip-italian-demo into main

image2text.py

@@ -1,4 +1,72 @@
 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():
-    pass
+    st.title("From Image to Text")
+    st.markdown(
+        """
+
+        ### 👋 Ciao!
+
+        Here you can find the captions that are most related to a given image.
+
+        🤌 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)

introduction.md

@@ -54,6 +54,8 @@ a dataset with 700K translated captions.
 
 ## Better Augmentations
 
+We knew that without a good augmentation strategy we could never get competitive results to a model trained on 400 million images. Therefor we implemented heavy augmentations to make the training more data efficient. We made sure to keep hue augmentations limited however to still give the model the ability to learn color definitions. While we would have liked to have augmentations for the captions as well after some experimentation we settled with random sampling from the five captions available in MSCOCO and leaving the rest of the captions unmodified. 
+
 ## Better Training
 
 After different trials, we realized that the usual way of training this model was
@@ -62,17 +64,15 @@ training pipeline: the optimizer and the training with frozen components.
 
 ### Optimizer
 
-The standard AdamW didn't seem enough to train the model and thus we opted for a different optimization strategy. We eventually used AdaBelief with AGC and Cosine Annealing.
+While the initial code used AdamW as an optimizer we soon noticed that it introduced some bad properties into the training. The model strated to overfit relatively quickly and the weight decay made this effect worse. We eventually decided to an optimization strategy that had worked well for us in similar cases and used AdaBelief with Adaptive Gradient Clipping (AGC) and a Cosine Annealing Schedule. Together with slightly tuning the learning rate this helped us to reduce the validation loss by 25%.
 Our implementation is available online [here](https://github.com/clip-italian/clip-italian/blob/master/hybrid_clip/run_hybrid_clip.py#L667).
 
 ### Backbone Freezing
 
 The ViT used by OpenAI was already trained on 400million images and it is the element in our architecture that probably required less training.
-The same is true for the BERT model we use. Thus, we decided to do a first training with the backbone of our architecture completely frozen, to allow
-the deeper layer to adapt to the new setting. Eventually, we run a new training, by fine-tuning al the components. This technique allowed us to
-reach a much better validation loss.
+The same is true for the BERT model we use. To allow the randomly initialized Re-projection Layers to warm up without messing with the tuned weights of the backbones we decided to do a first training with the backbones of our architecture completely frozen. Only after these layers converged did we unfreeze the rest of the model to fine-tune all the components. This technique allowed us to reach a much better validation loss.
 
-<img src="https://huggingface.co/spaces/clip-italian/clip-italian-demo/raw/main/static/img/clip-italian.png" alt="drawing" width="600"/>
+<img src="https://huggingface.co/spaces/clip-italian/clip-italian-demo/raw/main/static/img/clip-italian.png" alt="drawing" width="50%"/>
 
 # Scientific Validity
 

requirements.txt

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

static/features/{cc_features.npy → CC_val_embeddings.npy}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:63f185e851ff9cd0a19c5b1877087d860ca53ec5fc9e6a7d608249b9aacb77df
-size 2050773120
+oid sha256:775803a42011b09e8f5d19fcbdd67123cc3447154e1f8e5990cae1bce4581662
+size 27369600

text2image.py

@@ -22,9 +22,15 @@ def get_model():
     return FlaxHybridCLIP.from_pretrained("clip-italian/clip-italian")
 
 
-@st.cache(hash_funcs={transformers.models.bert.tokenization_bert_fast.BertTokenizerFast: lambda _: None})
+@st.cache(
+    hash_funcs={
+        transformers.models.bert.tokenization_bert_fast.BertTokenizerFast: lambda _: None
+    }
+)
 def get_tokenizer():
-    return AutoTokenizer.from_pretrained("dbmdz/bert-base-italian-xxl-uncased", cache_dir="./", use_fast=True)
+    return AutoTokenizer.from_pretrained(
+        "dbmdz/bert-base-italian-xxl-uncased", cache_dir="./", use_fast=True
+    )
 
 
 @st.cache(suppress_st_warning=True)
@@ -37,10 +43,14 @@ def download_images():
         photo_filename = "unsplash-25k-photos.zip"
         if not os.path.exists(photo_filename):  # Download dataset if does not exist
             print(f"Downloading {photo_filename}...")
-            response = requests.get(f"http://sbert.net/datasets/{photo_filename}", stream=True)
-            total_size_in_bytes = int(response.headers.get('content-length', 0))
+            response = requests.get(
+                f"http://sbert.net/datasets/{photo_filename}", stream=True
+            )
+            total_size_in_bytes = int(response.headers.get("content-length", 0))
             block_size = 1024  # 1 Kb
-            progress_bar = stqdm(total=total_size_in_bytes)  # , unit='iB', unit_scale=True
+            progress_bar = stqdm(
+                total=total_size_in_bytes
+            )  # , unit='iB', unit_scale=True
             content = io.BytesIO()
             for data in response.iter_content(block_size):
                 progress_bar.update(len(data))
@@ -54,53 +64,106 @@ def download_images():
 
 
 @st.cache()
-def get_image_features():
-    return jnp.load("static/features/features.npy")
-
-def app():
+def get_image_features(dataset_name):
+    if dataset_name == "Unsplash":
+        return jnp.load("static/features/features.npy")
+    else:
+        return jnp.load("static/features/CC_val_embeddings.npy")
 
-    """
 
-    # 👋 Ciao!
-
-    # CLIP Italian Demo
-    ## HF-Flax Community Week
+@st.cache()
+def load_urls(dataset_name):
+    if dataset_name == "CC":
+        with open("static/CC_val_urls.txt") as fp:
+            urls = [l.strip() for l in fp.readlines()]
+        return urls
+    else:
+        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),
+            ),
+        ]
+    )
 
-    In this demo you can search for images in the Unsplash 25k Photos dataset.
 
-    🤌 Italian mode on! 🤌
+def app():
 
-    """
+    st.title("From Text to Image")
+    st.markdown(
+        """
+    
+        ### 👋 Ciao!
+
+        Here you can search for images in the Unsplash 25k Photos dataset.
+        
+        🤌 Italian mode on! 🤌
+
+        """
+    )
+
+    if "suggestion" not in st.session_state:
+        st.session_state.suggestion = ""
+
+    def update_query(value=""):
+        st.session_state.suggestion = value
+
+    col1, col2, col3, col4 = st.beta_columns(4)
+    with col1:
+        st.button("Un gatto", on_click=update_query, kwargs=dict(value="Un gatto"))
+    with col2:
+        st.button("Due gatti", on_click=update_query, kwargs=dict(value="Due gatti"))
+    with col3:
+        st.button(
+            "Un fiore giallo",
+            on_click=update_query,
+            kwargs=dict(value="Un fiore giallo"),
+        )
+    with col4:
+        st.button(
+            "Un fiore blu", on_click=update_query, kwargs=dict(value="Un fiore blu")
+        )
+
+    col1, col2 = st.beta_columns([3, 1])
+    with col1:
+        query = st.text_input(
+            "Insert an italian query text here...", st.session_state.suggestion
+        )
+    with col2:
+        dataset_name = st.selectbox("IR dataset", ["Unsplash", "CC"])
 
-    query = st.text_input("Insert an italian query text here...")
     if query:
-        with st.spinner("Computing in progress..."):
+        with st.spinner("Computing..."):
+
             model = get_model()
-            download_images()
 
-            image_features = get_image_features()
+            if dataset_name == "Unsplash":
+                download_images()
 
+            image_features = get_image_features(dataset_name)
             model = get_model()
             tokenizer = get_tokenizer()
 
-            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=val_preprocess)
+            if dataset_name == "Unsplash":
+                image_size = model.config.vision_config.image_size
+                dataset = utils.CustomDataSet(
+                    "photos/", transform=get_image_transform(image_size)
+                )
+            elif dataset_name == "CC":
+                dataset = load_urls(dataset_name)
+            else:
+                raise ValueError()
 
             image_paths = utils.find_image(
-                query, model, dataset, tokenizer, image_features, n=2
+                query, model, dataset, tokenizer, image_features, 2, dataset_name
             )
 
         st.image(image_paths)

utils.py

@@ -41,24 +41,36 @@ 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)):
         images = images.permute(0, 2, 3, 1).numpy()
-        features = model.get_image_features(
-            images,
-        )
+        features = model.get_image_features(images,)
         features /= jnp.linalg.norm(features, axis=-1, keepdims=True)
         image_features.extend(features)
     return jnp.array(image_features)
 
 
-def find_image(text_query, model, dataset, tokenizer, image_features, n=1):
+def find_image(text_query, model, dataset, tokenizer, image_features, n, dataset_name):
     zeroshot_weights = text_encoder(text_query, model, tokenizer)
     zeroshot_weights /= jnp.linalg.norm(zeroshot_weights)
     distances = jnp.dot(image_features, zeroshot_weights.reshape(-1, 1))
     file_paths = []
     for i in range(1, n + 1):
         idx = jnp.argsort(distances, axis=0)[-i, 0]
-        file_paths.append("photos/" + dataset.get_image_name(idx))
+
+        if dataset_name == "Unsplash":
+            file_paths.append("photos/" + dataset.get_image_name(idx))
+        elif dataset_name == "CC":
+            file_paths.append(dataset[idx])
+        else:
+            raise ValueError(f"{dataset_name} not supported here")
     return file_paths