Spaces:
Runtime error
Runtime error
File size: 2,706 Bytes
4f12085 6b7167b 4f58d6c 4f12085 90ecddb 4f12085 2c174d9 4f12085 28dd995 2c174d9 4f12085 16b83cd 4f12085 e3b861a 4f12085 260e0f1 16b83cd 9492b8f a7746f5 4f12085 e3b861a 4f12085 77c44c9 4f12085 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import streamlit as st
import pandas as pd
import numpy as np
from html import escape
import os
import torch
import transformers
from transformers import RobertaModel, AutoTokenizer
#@st.cache(show_spinner=False)
#def load():
# text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text')
# image_embeddings = torch.load('embedding.pt')
# links = np.load('data.npy', allow_pickle=True)
# return text_encoder, links, image_embeddings
tokenizer = AutoTokenizer.from_pretrained('SajjadAyoubi/clip-fa-text')
text_encoder = RobertaModel.from_pretrained('SajjadAyoubi/clip-fa-text').eval()
image_embeddings = torch.load('embedding.pt')
links = np.load('data.npy', allow_pickle=True)
def get_html(url_list, height=180):
html = "<div style='margin-top: 20px; max-width: 1200px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>"
for url in url_list:
html2 = f"<img style='height: {height}px; margin: 5px' src='{escape(url)}'>"
html = html + html2
html += "</div>"
return html
def image_search(query, top_k=10):
with torch.no_grad():
text_embedding = text_encoder(**tokenizer(query, return_tensors='pt')).pooler_output
values, indices = torch.cosine_similarity(text_embedding, image_embeddings).sort(descending=True)
return [links[i] for i in indices[:top_k]]
description = '''
# Semantic image search :)
'''
def main():
st.markdown('''
<style>
.block-container{
max-width: 1200px;
}
div.row-widget.stRadio > div{
flex-direction:row;
display: flex;
justify-content: center;
}
div.row-widget.stRadio > div > label{
margin-left: 5px;
margin-right: 5px;
}
section.main>div:first-child {
padding-top: 0px;
}
section:not(.main)>div:first-child {
padding-top: 30px;
}
div.reportview-container > section:first-child{
max-width: 320px;
}
#MainMenu {
visibility: hidden;
}
footer {
visibility: hidden;
}
</style>''',
unsafe_allow_html=True)
st.sidebar.markdown(description)
_, c, _ = st.columns((1, 3, 1))
query = c.text_input('Search text', value='مرغ دریای')
if len(query) > 0:
results = image_search(query)
st.markdown(get_html(results), unsafe_allow_html=True)
if __name__ == '__main__':
main() |