Upload 5009 files

app.py

@@ -0,0 +1,145 @@
+import chromadb
+from chromadb.config import Settings
+import torchvision.models as models
+import torch
+from torchvision import transforms
+from PIL import Image
+import logging
+import streamlit as st
+import requests
+import json
+import uuid
+import os
+
+try:
+
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+
+    @st.cache_resource
+    def load_mobilenet_model():
+        device = 'cpu'
+        model = models.mobilenet_v3_small(pretrained=False)
+        model.classifier[3] = torch.nn.Linear(1024, 768)
+        model.load_state_dict(torch.load(
+            'mobilenet_v3_small_distilled_new_state_dict.pth', map_location=device))
+        model.eval().to(device)
+        return model
+
+    @st.cache_resource
+    def load_chromadb():
+        chroma_client = chromadb.PersistentClient(
+            path='data', settings=Settings(anonymized_telemetry=False))
+        collection = chroma_client.get_collection(name='images')
+        return collection
+
+    model = load_mobilenet_model()
+    logger.info("MobileNet loaded")
+    collection = load_chromadb()
+    logger.info("ChromaDB loaded")
+    logger.info(
+        f"Connected to ChromaDB collection images with {collection.count()} items")
+
+    preprocess = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
+                             0.229, 0.224, 0.225])
+    ])
+
+    def get_image_embedding(image):
+        if isinstance(image, str):
+            img = Image.open(image).convert('RGB')
+        else:
+            img = Image.open(image).convert('RGB')
+        input_tensor = preprocess(img).unsqueeze(0).to('cpu')
+        with torch.no_grad():
+            student_embedding = model(input_tensor)
+
+        return torch.nn.functional.normalize(student_embedding, p=2, dim=1).squeeze(0).tolist()
+
+    def save_image(image_file):
+        unique_filename = f"{image_file.name}"
+        save_path = os.path.join('images', unique_filename)
+        with open(save_path, "wb") as f:
+            f.write(image_file.getbuffer())
+        return save_path
+
+    def resize_image(image_path, size=(224, 224)):
+        if isinstance(image_path, str):
+            img = Image.open(image_path).convert("RGB")
+        else:
+            # Handle uploaded file
+            img = Image.open(image_path).convert("RGB")
+        img_resized = img.resize(size, Image.LANCZOS)  # High-quality resizing
+        return img_resized
+
+    st.sidebar.header("Upload Images")
+    image_files = st.sidebar.file_uploader(
+        "Upload images", type=["png", "jpg", "jpeg"], accept_multiple_files=True)
+    num_images = st.sidebar.slider(
+        "Number of results to return", min_value=1, max_value=10, value=3)
+
+    if image_files:
+        st.sidebar.subheader(
+            "Add Images to collection")
+        if st.sidebar.button("Add uploaded images"):
+            for idx, image_file in enumerate(image_files):
+                image_embedding = get_image_embedding(image_file)
+                saved_path = save_image(image_file)
+                unique_id = str(uuid.uuid4())
+                metadata = {
+                    'path': f'images/{image_file.name}', "type": "photo"
+                }
+                collection.add(
+                    embeddings=[image_embedding],
+                    ids=[unique_id],
+                    metadatas=[metadata]
+                )
+                st.sidebar.success(
+                    f"Image {image_file.name} added to the collection")
+
+    st.title('Image Search Using  Text')
+    st.write(
+        "The images stored in this database are sourced from the [COCO 2017 Validation Dataset](https://cocodataset.org/#download).")
+    st.write('Enter the text to search for images with matching description')
+    text_input = st.text_input("Description", "Playground")
+    if st.button("Search"):
+        if text_input.strip():
+            params = {'text': text_input}
+            response = requests.get(
+                'https://ashish-001-text-embedding-api.hf.space/embedding', params=params)
+            if response.status_code == 200:
+                logger.info("Embedding returned by API successfully")
+                data = json.loads(response.content)
+                embedding = data['embedding']
+                results = collection.query(
+                    query_embeddings=[embedding],
+                    n_results=num_images
+                )
+                images = [results['metadatas'][0][i]['path']
+                          for i in range(len(results['metadatas'][0]))]
+                distances = [results['distances'][0][i]
+                             for i in range(len(results['metadatas'][0]))]
+                if images:
+                    cols_per_row = 3
+                    rows = (len(images)+cols_per_row-1)//cols_per_row
+                    for row in range(rows):
+                        cols = st.columns(cols_per_row)
+                        for col_idx, col in enumerate(cols):
+                            img_idx = row*cols_per_row+col_idx
+                            if img_idx < len(images):
+                                resized_img = resize_image(
+                                    images[img_idx], size=(224, 224))
+                                col.image(resized_img,
+                                          caption=f"Image {img_idx+1}\ndistance {distances[img_idx]}", use_container_width=True)
+                else:
+                    st.write("No image found")
+            else:
+                st.write("Please try again later")
+                logger.info(f"status code {response.status_code} returned")
+        else:
+            st.write("Please enter text in the text area")
+
+except Exception as e:
+    logger.info(f"Exception occured: {e}")

data/2b9adfc7-6437-464a-9a6e-68d5e8f953a8/data_level0.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2615c0e69a7f28c18f6e0617b4bcea2e66e69d562e34b1b87636a78dcf368b7
+size 16060000

data/2b9adfc7-6437-464a-9a6e-68d5e8f953a8/header.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aa80aba5314c8ce52cf6f6e7f1d45e6b8021453f9c7f7b292ee6a5dd1a776372
+size 100

data/2b9adfc7-6437-464a-9a6e-68d5e8f953a8/index_metadata.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dc389628547534447cc61f658675b83226ca3ac7d507fc74e3700d004400bfab
+size 187994

data/2b9adfc7-6437-464a-9a6e-68d5e8f953a8/length.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e3e5a4c01928bba72fd04ee278bf472b3293b60b979c0104fdc71105a9cbc587
+size 20000

data/2b9adfc7-6437-464a-9a6e-68d5e8f953a8/link_lists.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1de5334c70f1165524f9ccc0309d5c67d6b37783b27179109edd39233e529914
+size 42780

data/chroma.sqlite3

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cc86101247d3b61131d449cb045e687c8b808f69cc94ee0a48c9da0d52df6bc
+size 22306816