Spaces:
Build error
Build error
File size: 3,592 Bytes
087fe06 |
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 |
# Import Libraries
from pathlib import Path
import pandas as pd
import numpy as np
import torch
import clip
from PIL import Image
from io import BytesIO
import requests
import gradio as gr
# Load the openAI's CLIP model
model, preprocess = clip.load("ViT-B/32", jit=False)
#display output photo
def show_output_image(matched_images) :
image=[]
for photo_id in matched_images:
photo_image_url = f"https://unsplash.com/photos/{photo_id}/download?w=280"
#photo_image_url = f"https://unsplash.com/photos/{photo_id}?w=280"
response = requests.get(photo_image_url)
img = Image.open(BytesIO(response.content))
#return img
image.append(img)
return image
# Encode and normalize the search query using CLIP
def encode_search_query(search_query, model, device):
with torch.no_grad():
text_encoded = model.encode_text(clip.tokenize(search_query).to(device))
text_encoded /= text_encoded.norm(dim=-1, keepdim=True)
# Retrieve the feature vector from the GPU and convert it to a numpy array
return text_encoded.cpu().numpy()
# Find all matched photos
def find_matches(text_features, photo_features, photo_ids, results_count=4):
# Compute the similarity between the search query and each photo using the Cosine similarity
similarities = (photo_features @ text_features.T).squeeze(1)
# Sort the photos by their similarity score
best_photo_idx = (-similarities).argsort()
# Return the photo IDs of the best matches
return [photo_ids[i] for i in best_photo_idx[:results_count]]
def image_search(search_text, search_image, option):
# taking photo IDs
photo_ids = pd.read_csv("./photo_ids.csv")
photo_ids = list(photo_ids['photo_id'])
# taking features vectors
photo_features = np.load("./features.npy")
# check if CUDA available
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the openAI's CLIP model
#model, preprocess = clip.load("ViT-B/32", device=device, jit=False)
model = model.to(device)
# Input Text Query
#search_query = "The feeling when your program finally works"
if option == "Text-To-Image" :
# Extracting text features
text_features = encode_search_query(search_text, model, device)
# Find the matched Images
matched_images = find_matches(text_features, photo_features, photo_ids, 4)
# ---- debug purpose ------#
print(matched_images[0])
id = matched_images[0]
photo_image_url = f"https://unsplash.com/photos/{id}/download?w=280"
print(photo_image_url)
#--------------------------#
return show_output_image(matched_images)
elif option == "Image-To-Image":
# Input Image for Search
with torch.no_grad():
image_feature = model.encode_image(preprocess(search_image).unsqueeze(0).to(device))
image_feature = (image_feature / image_feature.norm(dim=-1, keepdim=True)).cpu().numpy()
# Find the matched Images
matched_images = find_matches(image_feature, photo_features, photo_ids, 4)
#is_input_image = True
images = show_output_image(matched_images)
return images
gr.Interface(fn=image_search,
inputs=[gr.inputs.Textbox(lines=7, label="Input Text"),
gr.inputs.Image(type="pil", optional=True),
gr.inputs.Dropdown(["Text-To-Image", "Image-To-Image"])
],
outputs=gr.outputs.Carousel([gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil"), gr.outputs.Image(type="pil")]),
enable_queue=True
).launch(debug=True) |