add missing utils file

utils/vector_database.py

@@ -0,0 +1,91 @@
+from PIL import Image
+from pymilvus import connections, Collection
+from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
+
+def load_collection(name):
+    collection = Collection(name)
+    return collection
+
+def create_collection(name, description):
+    fields = [
+        FieldSchema(name="text_embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
+        FieldSchema(name="image_embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
+        FieldSchema(name="avg_embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
+        FieldSchema(name="weighted_avg_embedding", dtype=DataType.FLOAT_VECTOR, dim=512),
+        FieldSchema(name="image_id", dtype=DataType.INT64, is_primary=True),
+        FieldSchema(name="metadata", dtype=DataType.VARCHAR, max_length=5000)
+    ]
+    
+    schema = CollectionSchema(fields, description=description)
+    collection = Collection(name=name, schema=schema)
+    
+    return collection
+
+def create_hnsw_index(collection):
+    index_params = {
+        "index_type": "HNSW",
+        "metric_type": "IP",  # IP for cosine similarity
+        "params": {"M": 32, "efConstruction": 200}
+    }
+    
+    collection.create_index(field_name="text_embedding", index_params=index_params)
+    collection.create_index(field_name="image_embedding", index_params=index_params)
+    collection.create_index(field_name="avg_embedding", index_params=index_params)
+    collection.create_index(field_name="weighted_avg_embedding", index_params=index_params)
+
+def insert_data(collection, catalog, column, text_embeds, image_embeds, avg_embeds, w_avg_embeds):
+    image_ids = catalog['Id'].tolist()
+    metadata = catalog[column].tolist()
+    
+    collection.insert([
+        text_embeds,
+        image_embeds,
+        avg_embeds,
+        w_avg_embeds,
+        image_ids,
+        metadata
+    ])
+
+def search_in_milvus(collection, search_field, query_embedding, top_k=6):
+
+    # Step 1: Perform search in Milvus
+    search_params = {"metric_type": "IP", "params": {"ef": 128}}
+    results = collection.search(
+        query_embedding.tolist(), # Query vector
+        search_field, # Field to search in
+        param=search_params,
+        limit=top_k, # Top k results
+        output_fields=["image_id", "metadata"]
+    )
+
+    # Step 2: Extract the relevant information from the search results
+    search_results = []
+    for result in results[0]:  # The first element of 'results' contains the search results
+        image_id = result.entity.get("image_id")  # Retrieve the image ID
+        metadata = result.entity.get("metadata")  # Retrieve metadata (such as description, brand, etc.)
+        similarity_score = result.distance  # Retrieve similarity score (distance)
+
+        # Load the image (you can use PIL to load the image)
+        #image_path = "/content/drive/MyDrive/images/" + str(image_id) + ".jpg"
+        image_path = "/images/" + str(image_id) + ".jpg"
+        image = Image.open(image_path)
+
+        # Append the image, metadata, and score to the search results
+        search_results.append({
+            "image": image,
+            "metadata": metadata,
+            "similarity_score": similarity_score
+        })
+
+    # Step 3: Return the search results
+    return search_results
+
+conn = None 
+
+conn = connections.connect("default",
+                    uri='https://in03-6efb78578dde7a3.serverless.gcp-us-west1.cloud.zilliz.com',
+                    token='78a82c19d7a02c531dab34d97ffde11caba0aa18b58ad02c46ee98df99d912291043835a002e427d89d5ddbb65b342191c36c1ae'
+                    )
+
+fashionclip_collection = load_collection("fashionclip")
+fashionsiglip_collection = load_collection("fashionsiglip")