Abhilashvj commited on
Commit
98e4681
Β·
1 Parent(s): dca463b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -4
app.py CHANGED
@@ -7,6 +7,8 @@ import sys
7
  import traceback
8
  import uuid
9
  from typing import List, Optional
 
 
10
 
11
  import cv2
12
  import numpy as np
@@ -30,6 +32,10 @@ IMAGE_SIMILARITY_PINECONE_DEMO = "/find-similar-image-pinecone/"
30
  INDEX_NAME = "imagesearch-demo"
31
  INDEX_DIMENSION = 512
32
  TMP_DIR = "tmp"
 
 
 
 
33
 
34
  image_sim_model = SentenceTransformer("clip-ViT-B-32")
35
 
@@ -178,6 +184,12 @@ async def image_search_local(
178
  raise HTTPException(status_code=500, detail=str(e))
179
 
180
 
 
 
 
 
 
 
181
  @app.post(IMAGE_SIMILARITY_PINECONE_DEMO)
182
  async def image_search_pinecone(
183
  images_to_search: Optional[List[UploadFile]] = File(None),
@@ -204,7 +216,10 @@ async def image_search_pinecone(
204
  query_image = Image.open(io.BytesIO(contents))
205
  print("Indexing query image...")
206
  query_image = enhance_image(query_image)
207
- prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
 
 
 
208
  if INDEX_NAME not in pinecone.list_indexes():
209
  return {"similar_images": [], "status": "No index found for images"}
210
 
@@ -281,6 +296,7 @@ async def image_search_pinecone(
281
  # read image contain
282
  contents = await image.read()
283
  pil_image = Image.open(io.BytesIO(contents))
 
284
  tmp_file = f"{TMP_DIR}/{image.filename}"
285
  pil_image.save(tmp_file)
286
  meta_datas.append({"file_path": tmp_file})
@@ -288,9 +304,12 @@ async def image_search_pinecone(
288
  ids.append(str(uuid.uuid1()).replace("-",""))
289
 
290
  print("Encoding images to vectors...")
291
- corpus_embeddings = image_sim_model.encode(
292
- search_images, convert_to_tensor=True, show_progress_bar=True
293
- ).tolist()
 
 
 
294
  print(f"Indexing images to pinecone Index : {INDEX_NAME}")
295
  index.upsert(
296
  vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace
 
7
  import traceback
8
  import uuid
9
  from typing import List, Optional
10
+ from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
11
+ from transformers import CLIPProcessor, CLIPModel
12
 
13
  import cv2
14
  import numpy as np
 
32
  INDEX_NAME = "imagesearch-demo"
33
  INDEX_DIMENSION = 512
34
  TMP_DIR = "tmp"
35
+ model_type = "huggingface"
36
+ # model_type = "sentence_transformers"
37
+ processor = AutoProcessor.from_pretrained("patrickjohncyh/fashion-clip")
38
+ model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")
39
 
40
  image_sim_model = SentenceTransformer("clip-ViT-B-32")
41
 
 
184
  raise HTTPException(status_code=500, detail=str(e))
185
 
186
 
187
+ def get_clip_vectors(image):
188
+ img_inputs = processor(images=image, return_tensors="pt")
189
+ img_emb = model.get_image_features(**img_inputs)
190
+
191
+ return img_emb
192
+
193
  @app.post(IMAGE_SIMILARITY_PINECONE_DEMO)
194
  async def image_search_pinecone(
195
  images_to_search: Optional[List[UploadFile]] = File(None),
 
216
  query_image = Image.open(io.BytesIO(contents))
217
  print("Indexing query image...")
218
  query_image = enhance_image(query_image)
219
+ if model_type =="huggingface":
220
+ prompt_embedding = get_clip_vectors(query_image)
221
+ else:
222
+ prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
223
  if INDEX_NAME not in pinecone.list_indexes():
224
  return {"similar_images": [], "status": "No index found for images"}
225
 
 
296
  # read image contain
297
  contents = await image.read()
298
  pil_image = Image.open(io.BytesIO(contents))
299
+ pil_image = enhance_image(pil_image)
300
  tmp_file = f"{TMP_DIR}/{image.filename}"
301
  pil_image.save(tmp_file)
302
  meta_datas.append({"file_path": tmp_file})
 
304
  ids.append(str(uuid.uuid1()).replace("-",""))
305
 
306
  print("Encoding images to vectors...")
307
+ if model_type =="huggingface":
308
+ corpus_embeddings = get_clip_vectors(query_image).tolist()
309
+ else:
310
+ corpus_embeddings = image_sim_model.encode(
311
+ search_images, convert_to_tensor=True, show_progress_bar=True
312
+ ).tolist()
313
  print(f"Indexing images to pinecone Index : {INDEX_NAME}")
314
  index.upsert(
315
  vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace