RobotJelly commited on
Commit
1fbfc7c
1 Parent(s): 2619d84
Files changed (1) hide show
  1. app.py +7 -3
app.py CHANGED
@@ -71,11 +71,15 @@ model = SentenceTransformer('clip-ViT-B-32')
71
 
72
  # pre-computed embeddings
73
  emb_filename = 'unsplash-25k-photos-embeddings.pkl'
74
- with open(emb_filename, 'rb') as fIn:
75
- img_names, img_emb = pickle.load(fIn)
76
 
77
  def display_matches(similarity):
78
- best_matched_images = [Image.open(IMAGES_DIR / img_names[top_k_best_image]) for top_k_best_image in torch.topk(similarity, 4, 0).indices]
 
 
 
 
79
  return best_matched_images
80
 
81
  def image_search(search_text, search_image, option):
71
 
72
  # pre-computed embeddings
73
  emb_filename = 'unsplash-25k-photos-embeddings.pkl'
74
+ with open(emb_filename, 'rb') as emb:
75
+ img_names, img_emb = pickle.load(emb)
76
 
77
  def display_matches(similarity):
78
+ best_matched_images = []
79
+ top_k_indices = torch.topk(similarity, 4, 0).indices
80
+ for matched_image in top_k_indices:
81
+ img = Image.open(IMAGES_DIR / img_names[matched_image])
82
+ best_matched_images.append(img)
83
  return best_matched_images
84
 
85
  def image_search(search_text, search_image, option):