create a project using streamlit,clip model,pinecone db

.gitignore

@@ -0,0 +1,3 @@
+config.yaml
+.venv
+logs

requirements.txt

@@ -0,0 +1,7 @@
+streamlit
+requests
+Pillow
+pandas
+fastapi[standard]
+pinecone
+transformers 

src/app/app.py

@@ -0,0 +1,16 @@
+import homepage
+
+search_option = ['Select an option','Search by text', 'Search by image']
+
+homepage.setup_page()
+
+choosen_option = homepage.get_user_selection(search_option)
+if choosen_option.lower() == 'search by text':
+    user_query = homepage.get_search_text_input()
+    if user_query:
+        homepage.get_images_by_text(user_query)
+elif choosen_option.lower() == 'search by image':
+    image_input = homepage.get_search_image_input()
+    if image_input:
+        homepage.get_images_by_image(image_input)
+

src/app/homepage.py

@@ -0,0 +1,61 @@
+import os
+import sys
+src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
+sys.path.append(src_directory)
+import streamlit as st
+from utils import logger
+from database_pinecone import querry_database
+from model.clip_model import ClipModel
+
+clip_model = ClipModel()
+logger = logger.get_logger()
+
+PAGE_TITLE = "Look A Like - Image Finder"
+PAGE_LAYOUT = "centered"
+SIDEBAR_TITLE = "Find Similar Images"
+
+def setup_page():
+    if 'is_page_configured' not in st.session_state:
+        st.set_page_config(page_title=PAGE_TITLE, layout=PAGE_LAYOUT)
+        st.title(PAGE_TITLE)
+        st.sidebar.title(SIDEBAR_TITLE)
+        logger.info(f"Page configured with title '{PAGE_TITLE}', layout '{PAGE_LAYOUT}', and sidebar title '{SIDEBAR_TITLE}'")
+        st.session_state.is_page_configured = True
+    else:
+        logger.info("Page configuration already completed. Skipping setup.")
+
+def get_user_selection(options):
+    selected_option = st.sidebar.selectbox("Select the option", options)
+    return selected_option
+
+def get_search_image_input():
+    uploaded_image = st.sidebar.file_uploader("Upload the image to get similar images", type=['png', 'jpeg'])
+    return uploaded_image
+
+def get_search_text_input():
+    user_search = st.sidebar.text_input("Enter the text to search")
+    return user_search
+
+def display_images(response):
+    if response:
+        cols = st.columns(2)
+        for i, result in enumerate(response.matches):
+            with cols[i % 2]:
+                st.image(result.metadata["url"])
+
+def write_message(message):
+    st.write(message)
+
+def get_images_by_text(query):
+    embedding = clip_model.get_text_embedding(query)
+    response = querry_database.fetch_data(embedding)
+    message = f"Showing search results for {query}"
+    write_message(message)
+    images = display_images(response)
+
+def get_images_by_image(query):
+    embedding = clip_model.get_uploaded_image_embedding(query)
+    response = querry_database.fetch_data(embedding)
+    message = f"Showing search results of relevant images"
+    write_message(message)
+    images = display_images(response)

src/config/config.py

@@ -0,0 +1,14 @@
+import yaml
+from utils import logger
+
+logger = logger.get_logger()
+
+def load_config():
+    try:
+        with open('config.yaml', 'r') as file:
+            config_data = yaml.load(file, Loader=yaml.FullLoader)
+            logger.info("Successfully loaded the config.")
+            return config_data
+    except Exception as e:
+        logger.error(f"Unexpected error occurred while loading the config: {e}")
+        raise Exception(f"Error loading configuration: {e}")

src/data/data_set.py

@@ -0,0 +1,26 @@
+import pandas as pd
+from utils import logger
+
+logger = logger.get_logger()
+
+file_name = 'src/data/image_dataset.csv'
+
+tsv_file = 'src/data/photos.tsv000'
+
+def convert_tsc_to_csv(tsv_file):
+    df = pd.read_csv(tsv_file, sep='\t', header=0)
+    dataset = df.to_csv(file_name)
+    return dataset
+
+def get_df(start_index,end_index):
+    try:
+        logger.info("Loading the dataframe")
+        image_df = pd.read_csv(file_name)
+        final_df = image_df[['photo_id','photo_image_url']]
+        df = final_df[start_index:end_index]
+        logger.info("Successfully loaded the data frame") 
+        return df
+    except Exception as e:
+        logger.error(f"Unable to load the dataframe {e}")
+        raise
+    

src/data/request_images.py

@@ -0,0 +1,24 @@
+import requests
+from PIL import Image
+from utils import logger
+
+logger = logger.get_logger()
+
+def get_image_url(url):
+    try:
+        logger.info("Loading image from url to embed")
+        res = requests.get(url,stream = True).raw
+        img = Image.open(res)
+        logger.info("Loaded the image to embed successfully")
+        return img
+    except Exception as e:
+        logger.error(f"Unable to load the image to embed {e}")
+
+def convert_image_to_embedding_format(query_image):
+    try:
+        logger.info("Loading the image to embed")
+        image = Image.open(query_image)
+        logger.info("Loaded the image to embed successfully")
+        return image
+    except Exception as e:
+        logger.error(f"Unable to load the image to embed {e}")

src/database_pinecone/create_database.py

@@ -0,0 +1,94 @@
+import os
+import sys
+src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
+sys.path.append(src_directory)
+from pinecone import Pinecone, ServerlessSpec
+import time
+from model.clip_model import ClipModel
+from data import request_images
+from data import data_set
+from config import config
+from utils import logger
+
+config = config.load_config()
+logger = logger.get_logger()
+
+clip_model = ClipModel()
+
+def create_index(pinecone, index_name):
+    pinecone.create_index(
+    name=index_name,
+    dimension=512,
+    metric="cosine",
+    spec=ServerlessSpec(
+    cloud="aws",
+    region="us-east-1"
+        ) 
+    )
+
+def wait_till_index_loaded(pinecone, index_name):
+    while True:
+        index = pinecone.describe_index(index_name)
+        if index.status.get("ready", False):
+            index = pinecone.Index(index_name)
+            logger.info(f"Index '{index_name}' is ready and is now accessible.")
+            return index
+        else:
+            logger.debug(f"Index '{index_name}' is not ready yet. Checking again in 1 second.")
+            time.sleep(1)
+
+def get_index():
+    try:
+        pincone_api_key = config['pinecone_db']['pincone_api_key']
+        pc = Pinecone(api_key=pincone_api_key)
+        index = None
+        index_name = "imagesearch"
+        logger.info(f"Checking if the index '{index_name}' exists...")
+        if not pc.has_index(index_name):
+            logger.info(f"Index '{index_name}' does not exist. Creating a new index...")
+            create_index(pc,index_name)
+            logger.info(f"Index '{index_name}' creation initiated. Waiting for it to be ready...")
+            index = wait_till_index_loaded(index_name,pc)
+        else:
+            index = pc.Index(index_name)
+            logger.info(f"Index '{index_name}' already exists. Returning the existing index.")
+        return index
+    except Exception as e:
+        logger.info(f"Error occurred while getting or creating the Pinecone index: {str(e)}", exc_info=True)
+        return index
+    
+def upsert_data(index,embeddings,id,url):
+    try :
+        logger.info("Started to upsert the data")
+        index.upsert(
+            vectors=[{
+                "id": id,
+                "values": embeddings,
+                "metadata": {
+                "url": url,
+                "photo_id": id
+                }
+            }],
+            namespace="image-search-dataset",
+        )
+        logger.info(f"Successfully upserted the data in database")
+    except Exception as e:
+        logger.info(f"Unable to upsert the data {e}")
+        raise
+    
+def add_data_to_database(df):
+    try:
+        index = get_index()
+        logger.info("Starting to add the embeddings to the database")
+        for _, data in df.iterrows():
+            url = data['photo_image_url']
+            id = data['photo_id']
+            embeddings = clip_model.get_image_embedding(url)
+            upsert_data(index,embeddings,id,url)
+        logger.info("Added embeddings to the database successfully")
+    except Exception as e:
+        logger.info("Unable to add the data. Error : {e}")
+
+
+# df = data_set.get_df(8000,8500)     
+# add_data_to_database(df)

src/database_pinecone/querry_database.py

@@ -0,0 +1,25 @@
+import os
+import sys
+src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
+sys.path.append(src_directory)
+from utils import logger
+import streamlit as st
+from model.clip_model import ClipModel
+from database_pinecone import create_database
+
+clip_model = ClipModel()
+logger = logger.get_logger()
+
+index = create_database.get_index()
+namespace = 'image-search-dataset'
+
+def fetch_data(embedding):
+    try:
+        response = index.query(
+            top_k=10,
+            vector=embedding,
+            namespace=namespace,
+            include_metadata=True)
+        return response
+    except Exception as e:
+        raise

src/model/clip_model.py

@@ -0,0 +1,80 @@
+import os
+import sys
+from transformers import AutoProcessor, CLIPModel, AutoTokenizer
+src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
+sys.path.append(src_directory)
+from data import request_images
+from utils import logger
+
+logger = logger.get_logger()
+
+class ClipModel:
+    _models = {}
+
+    def __init__(self, model_name: str = "openai/clip-vit-base-patch32", tokenizer_name: str = "openai/clip-vit-large-patch14"):
+        self.model_name = model_name
+        self.tokenizer_name = tokenizer_name
+
+        if model_name not in ClipModel._models:
+            ClipModel._models[model_name] = self.load_models()
+
+    def load_models(self):
+        try:
+            logger.info(f"Loading the models: {self.model_name}")
+            model = CLIPModel.from_pretrained(self.model_name)
+            processor = AutoProcessor.from_pretrained(self.model_name)
+            tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_name)
+            return {
+                'model': model,
+                'processor': processor,
+                'tokenizer': tokenizer
+            }
+        except Exception as e:
+            logger.error(f"Unable to load the model {e}")
+            raise
+
+    def get_text_embedding(self, text: str):
+        try:
+            logger.info(f"Getting embedding for the text: {text}")
+            inputs = self._models[self.model_name]['tokenizer']([text], padding=True, return_tensors="pt")
+            text_features = self._models[self.model_name]['model'].get_text_features(**inputs)
+            text_embedding = text_features.detach().numpy().flatten().tolist()
+            logger.info("Text embedding successfully retrieved.")
+            return text_embedding
+        except Exception as e:
+            logger.error(f"Error while getting embedding for text: {e}")
+            raise
+
+    def get_image_embedding(self, image):
+        try:
+            logger.info(f"Getting embedding for the image")
+            image = request_images.get_image_url(image)
+            inputs = self._models[self.model_name]['processor'](images=image, return_tensors="pt")
+            image_features = self._models[self.model_name]['model'].get_image_features(**inputs)
+            embeddings = image_features.detach().cpu().numpy().flatten().tolist()
+            logger.info("Image embedding successfully retrieved.")
+            return embeddings
+        except Exception as e:
+            logger.error(f"Error while getting embedding for image: {e}")
+            raise
+        
+    def get_uploaded_image_embedding(self, image):
+        try:
+            logger.info(f"Getting embedding for the image")
+            image = request_images.convert_image_to_embedding_format(image)
+            inputs = self._models[self.model_name]['processor'](images=image, return_tensors="pt")
+            image_features = self._models[self.model_name]['model'].get_image_features(**inputs)
+            embeddings = image_features.detach().cpu().numpy().flatten().tolist()
+            logger.info("Image embedding successfully retrieved.")
+            return embeddings
+        except Exception as e:
+            logger.error(f"Error while getting embedding for image: {e}")
+            raise
+
+if __name__ == "__main__":
+    try:
+        logger.info("Starting the initialization of the ClipModel class...")
+        clip_model = ClipModel()
+        logger.info("ClipModel class initialized successfully.")
+    except Exception as e:
+        logger.error(f"Error during ClipModel initialization: {str(e)}")

src/utils/logger.py

@@ -0,0 +1,35 @@
+import logging
+from logging.handlers import RotatingFileHandler
+import os
+
+log_file = 'image_search.log'
+log_dir = 'src/logs'
+log_level=logging.INFO
+
+def get_logger( ):
+
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+
+    log_file_path = os.path.join(log_dir, log_file)
+
+    logger = logging.getLogger(__name__)
+
+    if not logger.hasHandlers():
+        logger.setLevel(log_level)
+
+        console_handler = logging.StreamHandler()
+        console_handler.setLevel(logging.DEBUG) 
+
+        file_handler = RotatingFileHandler(log_file_path, maxBytes=5*1024*1024, backupCount=3)
+        file_handler.setLevel(logging.INFO)
+
+        log_format = '%(asctime)s - %(levelname)s - %(message)s'
+        formatter = logging.Formatter(log_format, datefmt='%Y-%m-%d %H:%M')
+        console_handler.setFormatter(formatter)
+        file_handler.setFormatter(formatter)
+
+        logger.addHandler(console_handler)
+        logger.addHandler(file_handler)
+
+    return logger