File size: 2,776 Bytes
36107ce | 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 and precomputed clips
import pickle
precomputed_filename = 'precomputed_clips'
def load_precomputed(precomputed_filename):
with open(precomputed_filename + '.pickle', 'rb') as f:
return pickle.load(f)
precomputed_dict = load_precomputed(precomputed_filename)
# embeddings and similar pictures
import torch
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import os
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
def get_clip_embeddings(input_data, input_type='text'):
# Load the CLIP model and processor
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# Prepare the input based on the type
if input_type == 'text':
inputs = processor(text=input_data, return_tensors="pt", padding=True, truncation=True)
elif input_type == 'image':
if isinstance(input_data, str):
image = Image.open(input_data)
elif isinstance(input_data, Image.Image):
image = input_data
else:
raise ValueError("For image input, provide either a file path or a PIL Image object")
inputs = processor(images=image, return_tensors="pt")
else:
raise ValueError("Invalid input_type. Choose 'text' or 'image'")
# Get the embeddings
with torch.no_grad():
if input_type == 'text':
embeddings = model.get_text_features(**inputs)
else:
embeddings = model.get_image_features(**inputs)
return embeddings.numpy()
def find_similar_images(text_input, image_embeddings, all_images, take_best = 4):
# Získání embeddingu pro text
text_embedding = get_clip_embeddings(text_input, input_type='text')
# Výpočet kosinové podobnosti mezi textem a obrázky
similarities = cosine_similarity(text_embedding, image_embeddings)
# Seřazení podle podobnosti
best_indices = np.argsort(similarities[0])[::-1][:take_best]
# Výběr nejlepších 4 obrázků
best_images = [all_images[i] for i in best_indices]
return [Image.open(img) for img in best_images]
# find the most similar pictures compared to text inserted
def find_most_similar(text_input):
return find_similar_images(text_input, precomputed_dict['image_clips'], precomputed_dict['image_paths'])
# gradio run
import gradio as gr # Importing Gradio for creating the web interface
# vytvoření Gradio rozhraní
interface = gr.Interface(
fn=find_most_similar,
inputs="text",
outputs=gr.Gallery(label="Most Similar Images"),
title="Find Similar Images with CLIP",
description="Enter a text prompt to find the most similar images."
)
# app launch
interface.launch()
|