Spaces:
Runtime error
Runtime error
File size: 5,523 Bytes
ac1c6ae 25ab7f8 fe476eb ac1c6ae 230d2d3 ac1c6ae 1f0f3e5 6ee3bc1 e64c54e ac1c6ae 0ab69ad ddfac1f ac1c6ae 8e7c132 ac1c6ae ddfac1f ac1c6ae 8e7c132 ac1c6ae a329e3c cfa8379 fe476eb b5ddef2 fe476eb 25ab7f8 cfa8379 26c2408 ac1c6ae 25ab7f8 d90092b ac1c6ae cfa8379 a5c02ba ac1c6ae 26c2408 cb03df9 f365545 26c2408 ac1c6ae cb03df9 ac1c6ae 9a5647a ac1c6ae f9494ca f365545 f9494ca d8a45b1 cb03df9 d8a45b1 cb03df9 ac1c6ae 9a5647a ac1c6ae 9a5647a 25ab7f8 c4eb027 25ab7f8 ac1c6ae cbf56b1 676881e ac1c6ae 0ab69ad 676881e 0ab69ad 676881e 5a5b371 676881e 9a5647a f9494ca 676881e 9a5647a 676881e f9494ca 0ab69ad |
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 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
from colordescriptor import ColorDescriptor
from CLIP import CLIPImageEncoder
from LBP import LBPImageEncoder
from helper import chi2_distance, euclidean_distance, merge_features
import gradio as gr
import os
import cv2
import numpy as np
from datasets import *
dataset = load_dataset("huggan/CelebA-faces", download_mode='force_redownload')
dataset.cleanup_cache_files()
candidate_subset = dataset["train"].select(range(10)) # This is a small CBIR app! :D
def index_dataset(dataset):
print(dataset)
print("LBP Embeddings")
lbp_model = LBPImageEncoder(8,2)
dataset_with_embeddings = dataset.map(lambda row: {'lbp_embeddings': lbp_model.describe(row["image"])})
print("Color Embeddings")
cd = ColorDescriptor((8, 12, 3))
dataset_with_embeddings = dataset_with_embeddings.map(lambda row: {'color_embeddings': cd.describe(row["image"])})
print("CLIP Embeddings")
clip_model = CLIPImageEncoder()
dataset_with_embeddings = dataset_with_embeddings.map(clip_model.encode_images, batched=True, batch_size=16)
print("LBP and Color")
dataset_with_embeddings = dataset_with_embeddings.map(lambda row: {'lbp_color_embeddings': merge_features(row['lbp_embeddings'], row['color_embeddings'])})
# Add index
dataset_with_embeddings.add_faiss_index(column='color_embeddings')
dataset_with_embeddings.save_faiss_index('color_embeddings', 'color_index.faiss')
dataset_with_embeddings.add_faiss_index(column='clip_embeddings')
dataset_with_embeddings.add_faiss_index(column='lbp_embeddings')
dataset_with_embeddings.save_faiss_index('clip_embeddings', 'clip_index.faiss')
print(dataset_with_embeddings)
return dataset_with_embeddings
def check_index(ds):
index_path = "my_index.faiss"
if os.path.isfile('color_index.faiss') and os.path.isfile('clip_index.faiss'):
ds.load_faiss_index('color_embeddings', 'color_index.faiss')
return ds.load_faiss_index('clip_embeddings', 'clip_index.faiss')
else:
return index_dataset(ds)
dataset_with_embeddings = check_index(candidate_subset)
# Main function, to find similar images
# TODO: implement different distance measures
def get_neighbors(query_image, selected_descriptor, selected_distance, top_k=5):
"""Returns the top k nearest examples to the query image.
Args:
query_image: A PIL object representing the query image.
top_k: An integer representing the number of nearest examples to return.
Returns:
A list of the top_k most similar images as PIL objects.
"""
if "Color Descriptor" == selected_descriptor:
cd = ColorDescriptor((8, 12, 3))
qi_embedding = cd.describe(query_image)
qi_np = np.array(qi_embedding)
if selected_distance == "FAISS":
scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
'color_embeddings', qi_np, k=top_k)
elif selected_distance == "Chi-squared":
tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': chi2_distance(histA=qi_embedding, histB=row['color_embeddings'])})
retrieved_examples = tmp_dataset.sort("distance")[:5]
else:
tmp_dataset = dataset_with_embeddings.map(lambda row: {'distance': euclidean_distance(histA=qi_embedding, histB=row['color_embeddings'])})
retrieved_examples = tmp_dataset.sort("distance")[:5]
images = retrieved_examples['image'] #retrieved images is a dict, with images and embeddings
return images
if "CLIP" == selected_descriptor:
clip_model = CLIPImageEncoder()
qi_embedding = clip_model.encode_image(query_image)
scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
'clip_embeddings', qi_embedding, k=top_k)
images = retrieved_examples['image']
return images
if "LBP" == selected_descriptor:
lbp_model = LBPImageEncoder(8,2)
qi_embedding = lbp_model.describe(query_image)
scores, retrieved_examples = dataset_with_embeddings.get_nearest_examples(
'lbp_embeddings', qi_embedding, k=top_k)
images = retrieved_examples['image']
return images
else:
print("This descriptor is not yet supported :(")
return []
def load_cbir_dataset(datasetname, size=1000):
pass
# Define the Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("""
# Welcome to this CBIR app
This is a CBIR app focused on the retrieval of similar faces.
## Find similar images
Here you can upload an image, that is compared with existing image in our dataset.
""")
with gr.Row():
image_input = gr.Image(type="pil", label="Please upload your image")
gallery_output = gr.Gallery()
btn = gr.Button(value="Submit")
gr.Markdown("""
## Settings
Here you can adjust how the images are found
""")
with gr.Row():
descr_dropdown = gr.Dropdown(["Color Descriptor", "LBP", "CLIP"], value="LBP", label="Please choose an descriptor")
dist_dropdown = gr.Dropdown(["FAISS", "Chi-squared", "Euclid"], value="FAISS", label="Please choose a distance measure")
dataset_dropdown = gr.Dropdown(
["huggan/CelebA-faces", "EIT/cbir-eit"],
value="huggan/CelebA-faces",
label="Please select a dataset"
)
btn.click(get_neighbors, inputs=[image_input, descr_dropdown, dist_dropdown], outputs=[gallery_output])
if __name__ == "__main__":
demo.launch()
|