Spaces:
Build error
Build error
| 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() |