import io import os import requests import zipfile import natsort import gc from PIL import Image from PIL import UnidentifiedImageError os.environ["TOKENIZERS_PARALLELISM"] = "false" from stqdm import stqdm import streamlit as st from jax import numpy as jnp import transformers from transformers import AutoTokenizer from torchvision.transforms import Compose, CenterCrop, Normalize, Resize, ToTensor from torchvision.transforms.functional import InterpolationMode from modeling_hybrid_clip import FlaxHybridCLIP import utils @st.cache(hash_funcs={FlaxHybridCLIP: lambda _: None}) def get_model(): return FlaxHybridCLIP.from_pretrained("clip-italian/clip-italian") @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 ) @st.cache(suppress_st_warning=True) def download_images(): # from sentence_transformers import SentenceTransformer, util img_folder = "photos/" if not os.path.exists(img_folder) or len(os.listdir(img_folder)) == 0: os.makedirs(img_folder, exist_ok=True) 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)) block_size = 1024 # 1 Kb 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)) content.write(data) progress_bar.close() z = zipfile.ZipFile(content) # content.close() print("Extracting the dataset...") z.extractall(path=img_folder) print("Done.") @st.cache() def get_image_features(dataset_name): if dataset_name == "Unsplash": return jnp.load("static/features/features.npy") else: return jnp.load("static/features/CC_embeddings.npy") @st.cache() def load_urls(dataset_name): if dataset_name == "CC": with open("static/CC_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), ), ] ) headers = { #'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/71.0.3578.98 Safari/537.36', 'User-Agent': 'Googlebot-Image/1.0', # Pretend to be googlebot 'X-Forwarded-For': '64.18.15.200' } def app(): #st.title("From Text to Image") st.markdown("

Image Retrieval

", unsafe_allow_html=True) st.markdown("

Text to Image

", unsafe_allow_html=True) st.markdown( """ 👋 Ciao! Here you can type Italian query and search from ~150k images in the Conceptual Captions (CC) dataset or 25k Photos in the Unsplash dataset. Though these images were not used for training the model, you will see most queries make sense. Rare errors might be due to 2 possibilities: the model is answering in a wrong way or the image you are looking for are not in the dataset & the model is giving you the best answer it can get. You can choose from one of the following examples... """ ) suggestions = [ "Un gatto", "Due gatti", "Un fiore giallo", "Un fiore blu", "Una coppia in montagna", "Una coppia al tramonto" ] sugg_idx = -1 col1, col2, col3, col4, col5, col6 = st.beta_columns([1, 1, 1.2, 1.2, 1.4, 1.4]) with col1: if st.button(suggestions[0]): sugg_idx = 0 with col2: if st.button(suggestions[1]): sugg_idx = 1 with col3: if st.button(suggestions[2]): sugg_idx = 2 with col4: if st.button(suggestions[3]): sugg_idx = 3 with col5: if st.button(suggestions[4]): sugg_idx = 4 with col6: if st.button(suggestions[5]): sugg_idx = 5 col1, col2 = st.beta_columns([3, 1]) with col1: query = st.text_input("... or insert an Italian query text") with col2: dataset_name = st.selectbox("IR dataset", ["CC", "Unsplash"]) query = suggestions[sugg_idx] if sugg_idx > -1 else query if query else "" if query: with st.spinner("Computing..."): if dataset_name == "Unsplash": download_images() image_features = get_image_features(dataset_name) model = get_model() tokenizer = get_tokenizer() 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() N = 3 image_paths = utils.find_image( query, model, dataset, tokenizer, image_features, N, dataset_name ) for i, image_url in enumerate(image_paths): try: if dataset_name == "Unsplash": st.image(image_url) elif dataset_name == "CC": image_raw = requests.get(image_url, stream=True, allow_redirects=True, headers=headers).raw image = Image.open(image_raw).convert("RGB") st.image(image, use_column_width=True) break except (UnidentifiedImageError) as e: if i == N - 1: st.text(f'Tried to show {N} different image URLS but none of them were reachabele.\ Maybe try a different query?') gc.collect() sugg_idx = -1