Spaces:
Sleeping
Sleeping
File size: 3,391 Bytes
d1df841 2d90d0f d1df841 2d90d0f d1df841 2d90d0f d1df841 2ee277d |
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 |
import gradio as gr
import numpy as np
from PIL import Image
from engine.search import ImageSearchModule
from engine.utils import unzip_file
import os
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parent
def check_dirs():
dirs = {
"Data": (PROJECT_ROOT / "data"),
"Images": (PROJECT_ROOT / "data" / "images"),
"Features": (PROJECT_ROOT / "data" / "features")
}
if not dirs['Data'].exists():
unzip_file(zip_file_path=str(PROJECT_ROOT / "data.zip"), destination_folder=str(PROJECT_ROOT ))
for dir_name, dir_path in dirs.items():
if not dir_path.exists():
raise FileNotFoundError(f"{dir_name} directory not found: {dir_path}")
print("All data directories exist ✅")
check_dirs()
# Initialize the ImageSearchModule
search = ImageSearchModule(
image_embeddings_dir=str(PROJECT_ROOT / "data/features"),
original_images_dir=str(PROJECT_ROOT / "data/images"),
)
print("Add image embeddings and caption embeddings to vector database")
search.add_images()
def search_images(input_data, search_type):
if search_type == "image" and input_data is not None:
# Fix: Get the file path directly from the input data
results = search.search_by_image(input_data, top_k=10, similarity_threshold=0)
elif search_type == "text" and input_data.strip():
results = search.search_by_text(input_data, top_k=10, similarity_threshold=0)
else:
return [(Image.new("RGB", (100, 100), color="gray"), "No results")] * 10
images_with_captions = []
for image_name, similarity in results:
image_path = os.path.join(search.original_images_dir, f"resized_{image_name}")
matching_files = [
f
for f in os.listdir(search.original_images_dir)
if f.startswith(f"resized_{image_name}")
]
if matching_files:
img = Image.open(
os.path.join(search.original_images_dir, matching_files[0])
)
images_with_captions.append((img, f"Similarity: {similarity:.2f}"))
else:
images_with_captions.append(
(Image.new("RGB", (100, 100), color="gray"), "Image not found")
)
# Pad the results if less than 10 images are found
while len(images_with_captions) < 10:
images_with_captions.append(
(Image.new("RGB", (100, 100), color="gray"), "No result")
)
return images_with_captions
with gr.Blocks() as demo:
gr.Markdown("# Image Search App")
with gr.Tab("Image Search"):
# Fix: Change input type to 'filepath'
image_input = gr.Image(type="filepath", label="Upload an image")
image_button = gr.Button("Search by Image")
with gr.Tab("Text Search"):
text_input = gr.Textbox(label="Enter text query")
text_button = gr.Button("Search by Text")
gallery = gr.Gallery(
label="Search Results",
show_label=False,
elem_id="gallery",
columns=2,
height="auto",
)
image_button.click(
fn=search_images,
inputs=[image_input, gr.Textbox(value="image", visible=False)],
outputs=[gallery],
)
text_button.click(
fn=search_images,
inputs=[text_input, gr.Textbox(value="text", visible=False)],
outputs=[gallery],
)
demo.queue().launch()
|