ArtMatch / app.py
LeBuH's picture
Upload app.py
3bda56e verified
raw
history blame
3.66 kB
import gradio as gr
from PIL import Image
import torch
import numpy as np
import faiss
from transformers import (
BlipProcessor,
BlipForConditionalGeneration,
CLIPProcessor,
CLIPModel
)
from datasets import load_dataset
wikiart_dataset = load_dataset("huggan/wikiart", split="train", streaming=True)
def get_item_streaming(dataset, idx):
for i, item in enumerate(dataset):
if i == idx:
return item
raise IndexError("Index out of range")
device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device).eval()
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval()
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
image_index = faiss.read_index("image_index.faiss")
text_index = faiss.read_index("text_index.faiss")
def generate_caption(image: Image.Image):
inputs = blip_processor(image, return_tensors="pt").to(device)
with torch.no_grad():
caption_ids = blip_model.generate(**inputs)
caption = blip_processor.decode(caption_ids[0], skip_special_tokens=True)
return caption
def get_clip_text_embedding(text):
inputs = clip_processor(text=[text], return_tensors="pt", padding=True).to(device)
with torch.no_grad():
features = clip_model.get_text_features(**inputs)
features = features.cpu().numpy().astype("float32")
faiss.normalize_L2(features)
return features
def get_clip_image_embedding(image):
inputs = clip_processor(images=image, return_tensors="pt").to(device)
with torch.no_grad():
features = clip_model.get_image_features(**inputs)
features = features.cpu().numpy().astype("float32")
faiss.normalize_L2(features)
return features
def get_results_with_images(embedding, index, top_k=2):
D, I = index.search(embedding, top_k)
results = []
for idx in I[0]:
try:
item = get_item_streaming(wikiart_dataset, int(idx))
img = item["image"]
title = item.get("title", "Untitled")
artist = item.get("artist", "Unknown")
caption = f"ID: {idx}\n{title}{artist}"
results.append((img, caption))
except IndexError:
continue
return results
# Основная функция поиска
def search_similar_images(image: Image.Image):
caption = generate_caption(image)
text_emb = get_clip_text_embedding(caption)
image_emb = get_clip_image_embedding(image)
text_results = get_results_with_images(text_emb, text_index)
image_results = get_results_with_images(image_emb, image_index)
return caption, text_results, image_results
demo = gr.Interface(
fn=search_similar_images,
inputs=gr.Image(label="Загрузите изображение", type="pil"),
outputs=[
gr.Textbox(label="📜 Сгенерированное описание"),
gr.Gallery(label="🔍 Похожие по описанию (CLIP)", height="auto", columns=2),
gr.Gallery(label="🎨 Похожие по изображению (CLIP)", height="auto", columns=2)
],
title="🎨 Semantic WikiArt Search (BLIP + CLIP)",
description="Загрузите изображение. Модель BLIP сгенерирует описание, а CLIP найдёт похожие картины по тексту и изображению."
)
demo.launch()