Spaces:
Runtime error
Runtime error
File size: 3,912 Bytes
66f8989 |
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 88 89 90 91 92 93 94 95 96 97 98 |
from sentence_transformers import SentenceTransformer, util
from PIL import Image
import pickle
import os
import gradio as gr
import zipfile
# Load CLIP model
text_model = SentenceTransformer("clip-ViT-B-32-multilingual-v1")
image_model = SentenceTransformer("clip-ViT-B-32")
image_model.parallel_tokenization = False
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):
util.http_get("http://sbert.net/datasets/" +
photo_filename, photo_filename)
# Extract all images
with zipfile.ZipFile(photo_filename, "r") as zf:
for member in zf.infolist():
zf.extract(member, img_folder)
emb_filename = ".\\unsplash-25k-photos-embeddings.pkl"
if not os.path.exists(emb_filename):
util.http_get(
"http://sbert.net/datasets/unsplash-25k-photos-embeddings.pkl", emb_filename
)
with open(emb_filename, "rb") as fIn:
img_names, img_emb = pickle.load(fIn)
img_folder = ".\\photos\\"
duplicates = util.paraphrase_mining_embeddings(img_emb)
def search_text(query, top_k=1):
""" " Search an image based on the text query.
Args:
query ([string]): [query you want search for]
top_k (int, optional): [Amount of images o return]. Defaults to 1.
Returns:
[list]: [list of images that are related to the query.]
"""
# First, we encode the query.
query_emb = text_model.encode([query])
# Then, we use the util.semantic_search function, which computes the cosine-similarity
# between the query embedding and all image embeddings.
# It then returns the top_k highest ranked images, which we output
hits = util.semantic_search(query_emb, img_emb, top_k=top_k)[0]
image = []
for hit in hits:
object = Image.open(os.path.join(
".\\photos\\", img_names[hit["corpus_id"]]))
image.append((object, img_names[hit["corpus_id"]]))
return image
iface_search = gr.Interface(
title="Семантический поиск по картинке - d8a.ai",
description="""Демо-версия семантического поиска изображений, использующая
современные алгоритмы искусственного интеллекта для получения высокоточных
результатов поиска. Пользователи могут искать изображения с помощью запросов
на естественном языке, просматривать и предварительно просматривать результаты.
Это приложение идеально подходит для создателей контента, маркетологов и менеджеров
социальных сетей и обеспечивает более интеллектуальный и интуитивно понятный
способ поиска и управления визуальным контентом.""",
fn=search_text,
allow_flagging="never",
inputs=[
gr.inputs.Textbox(
lines=4,
label="Поисковый текст",
placeholder="Что вы хотите найти?",
default="Горы Кыргызстана",
),
gr.inputs.Slider(minimum=0, maximum=9, default=5,
step=1, label="Количество"),
],
outputs=gr.Gallery(
label="Найденные изображения", show_label=False, elem_id="gallery"
).style(grid=[5], height="auto"),
examples=[[("Горы Кыргызстана"), 5], [("Бишкек"), 5],
[("A boy with a ball"), 5]],
)
iface_search.launch()
|