Justin Donaldson commited on
Commit
e3d923c
1 Parent(s): 9dda517

fixing gallery

Browse files
Files changed (2) hide show
  1. .envrc +2 -0
  2. app.py +40 -35
.envrc ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ eval "$(conda shell.zsh hook)"
2
+ conda activate vibesearch
app.py CHANGED
@@ -1,10 +1,11 @@
 
 
 
1
  import gradio as gr
2
- from transformers import CLIPProcessor, CLIPModel, CLIPTokenizer
3
  import sentence_transformers
4
- from sentence_transformers import SentenceTransformer, util
5
- import pickle
6
  from PIL import Image
7
- import os
 
8
 
9
  # (Pdb) query_emb.shape
10
  # torch.Size([1, 512])
@@ -12,25 +13,24 @@ import os
12
  # (24996, 512)
13
 
14
 
15
-
16
  ## Define model
17
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
18
  processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
19
  tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
20
 
21
- #Open the precomputed embeddings
22
- emb_filename = 'lv-handbags.pkl'
23
  # emb_filename = 'unsplash-25k-photos-embeddings.pkl'
24
 
25
- with open(emb_filename, 'rb') as fIn:
26
- img_names, img_emb = pickle.load(fIn)
27
- #print(f'img_emb: {print(img_emb)}')
28
- #print(f'img_names: {print(img_names)}')
29
 
30
 
31
  def search_text(query, top_k=4):
32
- """" Search an image based on the text query.
33
-
34
  Args:
35
  query ([string]): [query you want search for]
36
  top_k (int, optional): [Amount of images o return]. Defaults to 1.
@@ -39,44 +39,49 @@ def search_text(query, top_k=4):
39
  [list]: [list of images that are related to the query.]
40
  """
41
  # First, we encode the query.
42
- inputs = tokenizer([query], padding=True, return_tensors="pt")
43
  query_emb = model.get_text_features(**inputs)
44
 
45
  # import pdb; pdb.set_trace()
46
-
47
  # Then, we use the util.semantic_search function, which computes the cosine-similarity
48
  # between the query embedding and all image embeddings.
49
  # It then returns the top_k highest ranked images, which we output
50
  hits = util.semantic_search(query_emb, img_emb, top_k=top_k)[0]
51
-
52
- image=[]
53
  for hit in hits:
54
- #print(img_names[hit['corpus_id']])
55
  # object = Image.open(os.path.join("photos/", img_names[hit['corpus_id']]))
56
- object = Image.open(os.path.join("lvphotos/", img_names[hit['corpus_id']]))
57
  image.append(object)
58
- #print(f'array length is: {len(image)}')
59
 
60
  return image
61
-
62
 
63
  iface = gr.Interface(
64
- title = "Hushh Vibe Search Model on Louis Vuitton API",
65
- description = "Quick demo of using text to perform vector search on an image collection",
66
- article = "TBD",
67
- fn=search_text,
68
- inputs=[gr.Textbox(lines=4,
69
- label="Write what you are looking for in an image...",
70
- placeholder="Text Here...")],
71
- outputs=[gr.Gallery(
72
- label="Generated images", show_label=False, elem_id="gallery"
73
- ).style(grid=[2], height="auto")]
74
-
75
- ,examples=[
 
 
 
 
 
76
  [("Vacation Star")],
77
  [("Rock Star")],
78
  [("Barbie")],
79
  [("Small Purse")],
80
  [("Big Bag")],
81
- ]
82
- ).launch(debug=True)
 
1
+ import os
2
+ import pickle
3
+
4
  import gradio as gr
 
5
  import sentence_transformers
 
 
6
  from PIL import Image
7
+ from sentence_transformers import SentenceTransformer, util
8
+ from transformers import CLIPModel, CLIPProcessor, CLIPTokenizer
9
 
10
  # (Pdb) query_emb.shape
11
  # torch.Size([1, 512])
 
13
  # (24996, 512)
14
 
15
 
 
16
  ## Define model
17
  model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
18
  processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
19
  tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
20
 
21
+ # Open the precomputed embeddings
22
+ emb_filename = "lv-handbags.pkl"
23
  # emb_filename = 'unsplash-25k-photos-embeddings.pkl'
24
 
25
+ with open(emb_filename, "rb") as fIn:
26
+ img_names, img_emb = pickle.load(fIn)
27
+ # print(f'img_emb: {print(img_emb)}')
28
+ # print(f'img_names: {print(img_names)}')
29
 
30
 
31
  def search_text(query, top_k=4):
32
+ """ " Search an image based on the text query.
33
+
34
  Args:
35
  query ([string]): [query you want search for]
36
  top_k (int, optional): [Amount of images o return]. Defaults to 1.
 
39
  [list]: [list of images that are related to the query.]
40
  """
41
  # First, we encode the query.
42
+ inputs = tokenizer([query], padding=True, return_tensors="pt")
43
  query_emb = model.get_text_features(**inputs)
44
 
45
  # import pdb; pdb.set_trace()
46
+
47
  # Then, we use the util.semantic_search function, which computes the cosine-similarity
48
  # between the query embedding and all image embeddings.
49
  # It then returns the top_k highest ranked images, which we output
50
  hits = util.semantic_search(query_emb, img_emb, top_k=top_k)[0]
51
+
52
+ image = []
53
  for hit in hits:
54
+ # print(img_names[hit['corpus_id']])
55
  # object = Image.open(os.path.join("photos/", img_names[hit['corpus_id']]))
56
+ object = Image.open(os.path.join("lvphotos/", img_names[hit["corpus_id"]]))
57
  image.append(object)
58
+ # print(f'array length is: {len(image)}')
59
 
60
  return image
61
+
62
 
63
  iface = gr.Interface(
64
+ title="Hushh Vibe Search Model on Louis Vuitton API",
65
+ description="Quick demo of using text to perform vector search on an image collection",
66
+ article="TBD",
67
+ fn=search_text,
68
+ inputs=[
69
+ gr.Textbox(
70
+ lines=4,
71
+ label="Write what you are looking for in an image...",
72
+ placeholder="Text Here...",
73
+ )
74
+ ],
75
+ outputs=[
76
+ gr.Gallery(
77
+ label="Generated images", show_label=False, elem_id="gallery", columns=2
78
+ )
79
+ ],
80
+ examples=[
81
  [("Vacation Star")],
82
  [("Rock Star")],
83
  [("Barbie")],
84
  [("Small Purse")],
85
  [("Big Bag")],
86
+ ],
87
+ ).launch(debug=True)