Spaces:
Sleeping
Sleeping
File size: 3,948 Bytes
51e6591 05efdb7 51e6591 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #Importing libraries
import chromadb
import time
from transformers import CLIPModel, CLIPProcessor
import gradio as gr
from sklearn.metrics.pairwise import cosine_similarity
import torch
import numpy
from PIL import Image
#create the client
client = chromadb.Client()
collection = client.create_collection('image_collection')
#define the model and the proccessor
model = CLIPModel.from_pretrained('openai/clip-vit-base-patch32')
processor = CLIPProcessor.from_pretrained ('openai/clip-vit-base-patch32')
image_paths = [
'image-01.jpg',
'image-02.jpg',
'image-03.jpeg'
]
images = [Image.open(image_path) for image_path in image_paths]
inputs = processor(images = images, return_tensors='pt', padding=True)
#Generate the embeddings
start_time = time.time()
with torch.no_grad():
image_embeddings = model.get_image_features(**inputs).numpy()
image_embeddings = [embedding.tolist() for embedding in image_embeddings]
end_time = time.time()
ingestion_time = end_time - start_time
collection.add(
embeddings = image_embeddings,
metadatas = [{'image':image} for image in image_paths],
ids = [str(i) for i in range(len(image_paths))]
)
print(f'image ingestion time {ingestion_time:.4f} seconds')
def calculate_accuracy(image_embedding, query_embedding):
similarity = cosine_similarity([image_embedding], [query_embedding])[0][0]
return similarity
def search_image(query):
if not query.strip():
return None, 'Ooopsy, you forgot to input something?, please try again ☠👻'
print(f'\nQuery: {query}')
start_query = time.time()
inputs = processor(text = query, return_tensors='pt', padding=True)
with torch.no_grad():
query_embeding = model.get_text_features(**inputs).numpy()
query_embeding = query_embeding.tolist()
end_query = time.time()
query_time = end_query - start_query
result = collection.query(query_embeddings = query_embeding, n_results=1)
result_image_path = result['metadatas'][0][0]['image']
result_image_index = int(result['ids'][0][0])
matched_image_embedding = image_embeddings[result_image_index]
accuracy_score = calculate_accuracy(matched_image_embedding, query_embeding[0])
result_image = Image.open(result_image_path)
file_name = result_image_path.split('/')[-1]
return result_image, f'Accuracy score: {accuracy_score:.4f}\nQuery Time: {query_time:.4f} seconds', file_name
queries = [
'A group of polar bears',
'A famous landmark in paris',
'A hot pizza fresh from the oven',
'food',
'A place',
'A structure in Europe',
'Animals'
]
def populate_queries(suggested_querries):
return suggested_querries
#defining a gradio interface
with gr.Blocks() as gr_interface:
gr.Markdown('# This is an image retrieval application Made by Allan Munene🤗😁🤑')
with gr.Row():
with gr.Column():
gr.Markdown(f'***Image Ingestion Time***: {ingestion_time:.4f} seconds')
gr.Markdown('## Input Panel')
custom_query = gr.Textbox(placeholder = 'Input your query here', label='What are you looking for?')
with gr.Row():
submit_button = gr.Button('Submit Query')
cancel_button = gr.Button('Cancel')
with gr.Row(elem_id='button-container'):
for query in queries:
gr.Button(query).click(fn=lambda q=query: q, outputs=custom_query)
with gr.Column():
gr.Markdown('Output Image')
Output_image = gr.Image(type='pil', label='Result_image')
accuracy = gr.Textbox(label='Performance')
f_name = gr.Textbox(label='File Name')
submit_button.click(fn=search_image, inputs=custom_query, outputs=[Output_image, accuracy, f_name])
cancel_button.click(fn=lambda: (None, ''), outputs=[Output_image, accuracy, f_name])
gr_interface.launch(share=True)
|