Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from Clustering import ClusteringData | |
| import numpy as np | |
| from PIL import Image | |
| import requests | |
| import tempfile | |
| import os | |
| import logging | |
| import json | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| cd = ClusteringData() | |
| cd.load_model_data() | |
| logger.info("Clustering data loaded") | |
| def search_images(text_query, uploaded_image, search_mode, top_k): | |
| preview = None | |
| results = [] | |
| if search_mode == "Text" and text_query.strip(): | |
| response = requests.get( | |
| f"https://ashish-001-clip-api.hf.space/embedding?text={text_query.strip()}") | |
| if response.status_code == 200: | |
| logger.info("Embedding returned successfully by text API") | |
| data = json.loads(response.content) | |
| embedding = data["embedding"] | |
| results = cd.find_similar_records(embedding, k=top_k) | |
| else: | |
| logger.info(f"{response.status_code} returned by the text API") | |
| results = [] | |
| results = [os.path.join("coco", "val2017", "val2017", fname) | |
| for i, fname in enumerate(results)] | |
| return None, results | |
| elif search_mode == "Image": | |
| if uploaded_image is not None: | |
| preview = uploaded_image | |
| tmp_path = uploaded_image | |
| # with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file: | |
| # uploaded_image.save(tmp_file.name) | |
| # tmp_path = tmp_file.name | |
| else: | |
| preview = 'Image.jpg' | |
| tmp_path = 'Image.jpg' | |
| url = "https://ashish-001-clip-api.hf.space/clip/process" | |
| files = {"file": open(tmp_path, "rb")} | |
| response = requests.post(url, files=files) | |
| if response.status_code == 200: | |
| embedding = np.array(response.json()['embedding']).squeeze() | |
| logger.info("Embedding returned successfully by image API") | |
| results = cd.find_similar_records(embedding, k=top_k) | |
| else: | |
| logger.info( | |
| f"{response.status_code} returned by the image API") | |
| results = [] | |
| results = [os.path.join("coco", "val2017", "val2017", fname) | |
| for i, fname in enumerate(results)] | |
| return preview, results | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## Multimodal Image Search with CLIP") | |
| gr.Markdown("Search images using **text** or **image upload**.") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Inputs | |
| search_mode = gr.Radio( | |
| ["Text", "Image"], label="Search Mode", value="Text") | |
| text_input = gr.Textbox( | |
| label="Enter text query", placeholder="Type something...", visible=True, value='Empty street') | |
| file_input = gr.Image( | |
| type="filepath", | |
| label="Upload image", | |
| value="Image.jpg", | |
| visible=False | |
| ) | |
| top_k = gr.Slider(1, 20, value=6, step=1, | |
| label="Number of results") | |
| submit_btn = gr.Button("Search") | |
| with gr.Column(scale=2): | |
| preview_img = gr.Image(label="Uploaded / Default Image") | |
| result_gallery = gr.Gallery( | |
| label="Results", columns=3, height="auto") | |
| def toggle_inputs(mode): | |
| if mode == "Text": | |
| return ( | |
| gr.update(visible=True), | |
| gr.update(visible=False, value=None), | |
| [], | |
| None | |
| ) | |
| else: | |
| return ( | |
| gr.update(visible=False), | |
| gr.update(visible=True, value=None), | |
| [], | |
| "Image.jpg" | |
| ) | |
| search_mode.change(toggle_inputs, inputs=search_mode, | |
| outputs=[text_input, file_input, result_gallery, preview_img]) | |
| submit_btn.click(fn=search_images, | |
| inputs=[text_input, | |
| file_input, search_mode, top_k], | |
| outputs=[preview_img, result_gallery,]) | |
| demo.launch() | |