Add static features

.gitattributes

@@ -14,3 +14,4 @@
 *.pb filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text

app.py

@@ -16,6 +16,8 @@ from torchvision.transforms.functional import InterpolationMode
 from tqdm import tqdm
 from modeling_hybrid_clip import FlaxHybridCLIP
 
+import utils
+
 
 @st.cache
 def get_model():
@@ -39,92 +41,9 @@ def download_images():
     print("Done.")
 
 
-@st.cache(allow_output_mutation=True)
-def get_image_features(model, image_dir):
-    image_size = model.config.vision_config.image_size
-
-    val_preprocess = transforms.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 = CustomDataSet(image_dir, transform=val_preprocess)
-
-    loader = torch.utils.data.DataLoader(
-        dataset,
-        batch_size=16,
-        shuffle=False,
-        num_workers=4,
-        drop_last=False,
-    )
-
-    return precompute_image_features(loader), dataset
-
-
-class CustomDataSet(torch.utils.data.Dataset):
-    def __init__(self, main_dir, transform):
-        self.main_dir = main_dir
-        self.transform = transform
-        all_imgs = os.listdir(main_dir)
-        self.total_imgs = natsort.natsorted(all_imgs)
-
-    def __len__(self):
-        return len(self.total_imgs)
-
-    def get_image_name(self, idx):
-        return self.total_imgs[idx]
-
-    def __getitem__(self, idx):
-        img_loc = os.path.join(self.main_dir, self.total_imgs[idx])
-        image = PilImage.open(img_loc).convert("RGB")
-        tensor_image = self.transform(image)
-        return tensor_image
-
-
-def text_encoder(text, tokenizer):
-    inputs = tokenizer(
-        [text],
-        max_length=96,
-        truncation=True,
-        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)
-
-
-@st.cache
-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 /= jnp.linalg.norm(features, axis=-1, keepdims=True)
-        image_features.extend(features)
-    return jnp.array(image_features)
-
-
-def find_image(text_query, dataset, tokenizer, image_features, n=1):
-    zeroshot_weights = text_encoder(text_query, 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))
-    return file_paths
+@st.cache()
+def get_image_features():
+    return jnp.load("static/features/features.npy")
 
 
 """
@@ -142,6 +61,9 @@ if query:
         model = get_model()
         download_images()
 
+        image_features = get_image_features()
+
+        model = get_model()
         tokenizer = AutoTokenizer.from_pretrained(
             "dbmdz/bert-base-italian-xxl-uncased", cache_dir=None, use_fast=True
         )
@@ -160,18 +82,10 @@ if query:
             ]
         )
 
-        dataset = CustomDataSet("photos/", transform=val_preprocess)
+        dataset = utils.CustomDataSet("photos/", transform=val_preprocess)
 
-        loader = torch.utils.data.DataLoader(
-            dataset,
-            batch_size=16,
-            shuffle=False,
-            num_workers=2,
-            drop_last=False,
+        image_paths = utils.find_image(
+            query, model, dataset, tokenizer, image_features, n=2
         )
 
-        image_features = precompute_image_features(model, loader)
-
-        image_paths = find_image(query, dataset, tokenizer, image_features, n=2)
-
     st.image(image_paths)

static/features/features.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53a956386a27089b0bfe84bc311fbee885983815f5a6e9d9e58ec5c3a52015e9
+size 51191936

utils.py

@@ -0,0 +1,69 @@
+import torch
+import os
+from transformers import AutoTokenizer
+from jax import numpy as jnp
+import json
+import requests
+import zipfile
+import io
+import natsort
+from PIL import Image as PilImage
+from tqdm import tqdm
+
+
+class CustomDataSet(torch.utils.data.Dataset):
+    def __init__(self, main_dir, transform):
+        self.main_dir = main_dir
+        self.transform = transform
+        all_imgs = os.listdir(main_dir)
+        self.total_imgs = natsort.natsorted(all_imgs)
+
+    def __len__(self):
+        return len(self.total_imgs)
+
+    def get_image_name(self, idx):
+        return self.total_imgs[idx]
+
+    def __getitem__(self, idx):
+        img_loc = os.path.join(self.main_dir, self.total_imgs[idx])
+        image = PilImage.open(img_loc).convert("RGB")
+        tensor_image = self.transform(image)
+        return tensor_image
+
+
+def text_encoder(text, model, tokenizer):
+    inputs = tokenizer(
+        [text],
+        max_length=96,
+        truncation=True,
+        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)
+
+
+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 /= 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):
+    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))
+    return file_paths