Spaces:
Sleeping
Sleeping
File size: 3,156 Bytes
39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 07d2580 39bee95 |
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 |
import json
import os
import time
import streamlit as st
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from transformers import AutoProcessor, CLIPModel
from PIL import Image
import torch
global processor, model
processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
load_dotenv()
# Connect to Pinecone
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))
# Create an index if it does not exist
index_name = "index-search"
unsplash_index = None
if not pc.has_index(index_name):
pc.create_index(name=index_name, metric="cosine",
dimension=512,
spec=ServerlessSpec(cloud="aws", region="us-east-1"))
# Wait for the index to be ready
while True:
index = pc.describe_index(index_name)
if index.status.get("ready", False):
unsplash_index = pc.Index(index_name)
break
print("Waiting for index to be ready...")
time.sleep(1)
else:
unsplash_index = pc.Index(index_name)
# Streamlit UI
st.title("🔍 CLIP-Powered Image Search")
st.markdown("Search images using **text** or **image**!")
# Search type selection
search_type = st.radio("Select Search Type", ["Text Search", "Image Search"], horizontal=True)
def get_text_embedding(query):
inputs = processor(text=query, return_tensors="pt")
text_features = model.get_text_features(**inputs)
return text_features.detach().numpy().flatten().tolist()
def get_image_embedding(image):
image = image.convert("RGB").resize((224, 224))
inputs = processor(images=image, return_tensors="pt")
image_features = model.get_image_features(**inputs)
return image_features.detach().numpy().flatten().tolist()
if search_type == "Text Search":
search_query = st.text_input("Enter a search query (min 3 characters)")
if len(search_query) >= 3:
with st.spinner("Searching images..."):
text_embedding = get_text_embedding(search_query)
response = unsplash_index.query(top_k=10, vector=text_embedding, namespace="image-search-dataset", include_metadata=True)
# Display images in two columns
cols = st.columns(2)
for i, result in enumerate(response.matches):
with cols[i % 2]:
st.image(result.metadata["url"], caption=f"Match {i+1}")
elif search_type == "Image Search":
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "png", "jpeg"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption="Uploaded Image")
with st.spinner("Searching for similar images..."):
image_embedding = get_image_embedding(image)
response = unsplash_index.query(top_k=10, vector=image_embedding, namespace="image-search-dataset", include_metadata=True)
# Display images in two columns
cols = st.columns(2)
for i, result in enumerate(response.matches):
with cols[i % 2]:
st.image(result.metadata["url"], caption=f"Match {i+1}")
|