Update app.py

app.py

@@ -1,145 +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}")
+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", "Road")
+    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}")