Update clip_model.py

clip_model.py

@@ -306,7 +306,7 @@ 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]]
     
-    _, axes = plt.subplots(3, 3, figsize=(10, 10))
+    _, axes = plt.subplots(4, 4, figsize=(10, 10))
 
     results = []
     for match, ax in zip(matches, axes.flatten()):
@@ -321,11 +321,11 @@ def find_matches(model, image_embeddings, query, image_filenames, n=9):
 def clip_image_search(model,image_embeddings,
              query,
              image_filenames,
-             n=9 ):
+             n=16):
   _, valid_df = make_train_valid_dfs()
   model, image_embeddings = get_image_embeddings(valid_df, "best.pt")
   return find_matches(model,
              image_embeddings,
              query,
              image_filenames = valid_df['image'].values,
-             n=9)
+             n)