Fixed n_searched_images=1 case

app.py

@@ -23,10 +23,12 @@ def text_2_image(model, img_emb, img_names, img_urls, n_top_k_images):
         if st.button("Convert"):
             st.write("The image with the most similar embedding is:")
             cosine_sim = get_match(model, text, img_emb)
-            logger.info(cosine_sim.shape)
-            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices
-            logger.info(top_k_images_indices.squeeze().tolist())
-            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices.squeeze().tolist()]
+            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices.squeeze()
+            if top_k_images_indices.nelement() == 1:
+                top_k_images_indices = [top_k_images_indices.tolist()]
+            else:
+                top_k_images_indices = top_k_images_indices.tolist()
+            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices]
             cols = st.columns(n_top_k_images)
             for i, image_found in enumerate(images_found):
                 logger.success(f"Image match found: {image_found}")
@@ -51,11 +53,12 @@ def image_2_image(model, img_emb, img_names, img_urls,n_top_k_images):
         if st.button("Convert"):
             st.write("The image with the most similar embedding is:")
             cosine_sim = get_match(model, image.convert("RGB"), img_emb)
-            logger.info(cosine_sim.shape)
-            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices
-            logger.info(top_k_images_indices.squeeze().tolist())
-            # The problem is here
-            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices.squeeze().tolist()]
+            top_k_images_indices = torch.topk(cosine_sim, n_top_k_images, 1).indices.squeeze()
+            if top_k_images_indices.nelement() == 1:
+                top_k_images_indices = [top_k_images_indices.tolist()]
+            else:
+                top_k_images_indices = top_k_images_indices.tolist()
+            images_found = [img_names[top_k_best_image] for top_k_best_image in top_k_images_indices]
             cols = st.columns(n_top_k_images)
             for i, image_found in enumerate(images_found):
                 logger.success(f"Image match found: {image_found}")