Update app.py

app.py

@@ -7,6 +7,8 @@ import sys
 import traceback
 import uuid
 from typing import List, Optional
+from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
+from transformers import CLIPProcessor, CLIPModel
 
 import cv2
 import numpy as np
@@ -30,6 +32,10 @@ IMAGE_SIMILARITY_PINECONE_DEMO = "/find-similar-image-pinecone/"
 INDEX_NAME = "imagesearch-demo"
 INDEX_DIMENSION = 512
 TMP_DIR = "tmp"
+model_type = "huggingface"
+# model_type = "sentence_transformers"
+processor = AutoProcessor.from_pretrained("patrickjohncyh/fashion-clip")
+model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")
 
 image_sim_model = SentenceTransformer("clip-ViT-B-32")
 
@@ -178,6 +184,12 @@ async def image_search_local(
         raise HTTPException(status_code=500, detail=str(e))
 
 
+def get_clip_vectors(image):
+    img_inputs = processor(images=image, return_tensors="pt")
+    img_emb = model.get_image_features(**img_inputs)
+    
+    return img_emb
+    
 @app.post(IMAGE_SIMILARITY_PINECONE_DEMO)
 async def image_search_pinecone(
     images_to_search: Optional[List[UploadFile]] = File(None),
@@ -204,7 +216,10 @@ async def image_search_pinecone(
             query_image = Image.open(io.BytesIO(contents))
             print("Indexing query image...")
             query_image = enhance_image(query_image)
-            prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
+            if model_type =="huggingface":
+                prompt_embedding = get_clip_vectors(query_image)
+            else:
+                prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
             if INDEX_NAME not in pinecone.list_indexes():
                 return {"similar_images": [], "status": "No index found for images"}
 
@@ -281,6 +296,7 @@ async def image_search_pinecone(
                 # read image contain
                 contents = await image.read()
                 pil_image = Image.open(io.BytesIO(contents))
+                pil_image = enhance_image(pil_image)
                 tmp_file = f"{TMP_DIR}/{image.filename}"
                 pil_image.save(tmp_file)
                 meta_datas.append({"file_path": tmp_file})
@@ -288,9 +304,12 @@ async def image_search_pinecone(
                 ids.append(str(uuid.uuid1()).replace("-",""))
 
             print("Encoding images to vectors...")
-            corpus_embeddings = image_sim_model.encode(
-                search_images, convert_to_tensor=True, show_progress_bar=True
-            ).tolist()
+            if model_type =="huggingface":
+                corpus_embeddings = get_clip_vectors(query_image).tolist()
+            else:
+                corpus_embeddings = image_sim_model.encode(
+                    search_images, convert_to_tensor=True, show_progress_bar=True
+                ).tolist()
             print(f"Indexing images to pinecone Index : {INDEX_NAME}")
             index.upsert(
                 vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace