RobotJelly commited on
Commit
207f77e
1 Parent(s): 4640c21
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -31,7 +31,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
31
  # taking features vectors
32
  #photo_features = np.load("./features.npy")
33
 
34
- #IMAGES_DIR = './photos'
35
 
36
 
37
 
@@ -79,8 +79,9 @@ with open(emb_filename, 'rb') as fIn:
79
 
80
  def display_matches(similarity):
81
  best_matched_images = []
82
- for best_img in torch.topk(similarity, 4).indices:
83
- img = Image.open(os.path.join('./photos', img_names[best_img]))
 
84
  best_matched_images.append(img)
85
  return best_matched_images
86
 
 
31
  # taking features vectors
32
  #photo_features = np.load("./features.npy")
33
 
34
+ IMAGES_DIR = './photos/'
35
 
36
 
37
 
 
79
 
80
  def display_matches(similarity):
81
  best_matched_images = []
82
+ indices = torch.topk(input=similarity, k=4, dim=0)
83
+ for best_img_idx in indices:
84
+ img = Image.open(IMAGES_DIR / img_names[best_img_idx])
85
  best_matched_images.append(img)
86
  return best_matched_images
87