change func to return images

s23_openai_clip.py

@@ -436,15 +436,8 @@ def find_matches(model, image_embeddings, query, image_filenames, n=9):
 
     values, indices = torch.topk(dot_similarity.squeeze(0), n * 5)
     matches = [image_filenames[idx] for idx in indices[::5]]
+    return matches
 
-    _, axes = plt.subplots(3, 3, figsize=(10, 10))
-    for match, ax in zip(matches, axes.flatten()):
-        image = cv2.imread(f"{CFG.image_path}/{match}")
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-        ax.imshow(image)
-        ax.axis("off")
-
-    plt.show()
 
 """This is how we use this function. Aaaannnndddd the results:
 (The results in the blog post and the one at the beginning of the notebook were achieved with training on the 30k version)
@@ -455,7 +448,7 @@ def find_matches(model, image_embeddings, query, image_filenames, n=9):
 def inference_CLIP(query_text):
   _, valid_df = make_train_valid_dfs()
   model, image_embeddings = get_image_embeddings(valid_df, "best.pt")
-  find_matches(model,
+  return find_matches(model,
              image_embeddings,
              query=query_text,
              # query="dogs on the grass",