import io import os import requests import zipfile import natsort 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_val_embeddings.npy") @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), ), ] ) 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"]) if query: with st.spinner("Computing..."): model = get_model() 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() image_paths = utils.find_image( query, model, dataset, tokenizer, image_features, 2, dataset_name ) st.image(image_paths)