Working version

app.py

@@ -67,21 +67,22 @@ def process_dataset(dataset_name, dataset_split, dataset_column_image):
 def compute(dataset_name, dataset_split, dataset_column_image, img):
     img_dataset = process_dataset(dataset_name, dataset_split, dataset_column_image)
     closest_idx, distances = find_closest_images(img, img_dataset)
-    return [img_dataset[i] for i in closest_idx]
+    return [img_dataset[i][dataset_column_image] for i in closest_idx]
 
 
 with gr.Blocks() as demo:
     gr.Markdown("# Find if your images are in a public dataset!")
     with gr.Row():
         with gr.Column(scale=1, min_width=600):
-            dataset_name = gr.Textbox(label="Enter the name of a dataset containing images")
-            dataset_split = gr.Textbox(label="Enter the split of this dataset to consider")
-            dataset_column_image = gr.Textbox(label="Enter the name of the column of this dataset that contains images")
+            dataset_name = gr.Textbox(label="Enter the name of a dataset containing images", value="huggan/few-shot-pokemon")
+            dataset_split = gr.Textbox(label="Enter the split of this dataset to consider", value="train")
+            dataset_column_image = gr.Textbox(label="Enter the name of the column of this dataset that contains images", value="image")
             img = gr.Image(label="Input your image that will be compared against images of the dataset", type="pil")
             btn = gr.Button("Find").style(full_width=True)
 
         with gr.Column(scale=2, min_width=600):
             gallery_similar = gr.Gallery(label="similar images")
+            gallery_similar.style(grid=3)
         
     event = btn.click(compute, [dataset_name, dataset_split, dataset_column_image, img], gallery_similar)