Spaces:
Build error
Build error
File size: 2,855 Bytes
632e4ae 1f13a56 632e4ae fe1a29a 632e4ae 1f13a56 632e4ae 85cbea1 632e4ae 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 fe1a29a 1f13a56 |
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 |
import os
import sys
src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
sys.path.append(src_directory)
from transformers import AutoProcessor, CLIPModel
import streamlit as st
from utils import logger
from database import pinecone_index
from PIL import Image
logger = logger.get_logger()
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = AutoProcessor.from_pretrained("openai/clip-vit-base-patch32")
PINECONE_INDEX = pinecone_index.create_index()
def search_by_text(query_text, index):
inputs = processor(text=query_text, return_tensors="pt")
text_features = model.get_text_features(**inputs)
query_vector = text_features.detach().cpu().numpy().flatten().tolist()
results = index.query(vector=query_vector, top_k=10, include_metadata=True, namespace="image-search-dataset")
return results
def search_by_image(image, index):
inputs = processor(images=image, return_tensors="pt")
image_features = model.get_image_features(**inputs)
query_vector = image_features.detach().cpu().numpy().flatten().tolist()
results = index.query(vector=query_vector, top_k=5, include_metadata=True, namespace="image-search-dataset")
return results
def main():
st.set_page_config(page_title="Clip Search", layout="wide")
st.title("📸Image Search with Pinecone and CLIP")
option = st.selectbox("Choose Input Type", ["Text", "Image Upload"])
if option == "Text":
user_text = st.text_input("Enter your search text", placeholder = "for eg: dogs or cat etc..")
if st.button("Search"):
results = search_by_text(user_text, PINECONE_INDEX)
columns = st.columns(2)
for idx, match in enumerate(results['matches']):
with columns[idx % 2]:
st.image(
match['metadata']['url'],
caption=f"Match: {match['metadata']['photo_id']}",
width=500
)
elif option == "Image Upload":
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")
if st.button("Search by Image"):
results = search_by_image(image, PINECONE_INDEX)
columns = st.columns(2)
for idx, match in enumerate(results['matches']):
with columns[idx % 2]:
st.image(
match['metadata']['url'],
caption=f"Match: {match['metadata']['photo_id']}",
width=500
)
if __name__ == "__main__":
main() |