Update app.py

app.py

@@ -14,13 +14,12 @@ nos = [1,2,3,4,5,6,7,8,9,10]
 index = np.load("indexes/knn_10752_65.npy")
 ds = load_dataset("SDBiaseval/identities", split="train")
 
-def get_index(gender, ethnicity, model, no):
-    df = ds.remove_columns(["image","image_path"]).to_pandas()
-    return df.loc[(df['ethnicity'] == ethnicity)  & (df['gender'] == gender) & (df['no'] == no) & (df['model'] == model)].index[0]
-
 def get_nearest_64(gender, ethnicity, model, no):
-    ix = get_index(gender, ethnicity, model, no)
-    return ds.select([index[ix][0]])["image"], ds.select(index[ix][1:5])["image"]
+    df = ds.remove_columns(["image","image_path"]).to_pandas()
+    ix = df.loc[(df['ethnicity'] == ethnicity)  & (df['gender'] == gender) & (df['no'] == no) & (df['model'] == model)].index[0]
+    image = ds.select([index[ix][0]])["image"]
+    neighbor_list = ds.select(index[ix][1:])["image"]
+    return image, neighbor_list 
 
 with gr.Blocks() as demo:
     gender = gr.Radio(gender_labels, label="Gender label")
@@ -29,7 +28,5 @@ with gr.Blocks() as demo:
     no = gr.Radio(nos, label="Image number")
 
     button = gr.Button(value="Get nearest neighbors")
-    button.click(get_nearest_64, inputs=[gender, ethnicity, model, no], outputs=[gr.Image(), gr.Gallery()])
+    button.click(get_nearest_64, inputs=[gender, ethnicity, model, no], outputs=[gr.Image(), gr.Gallery().style(grid=[8], height="auto")])
 demo.launch()
-# demo = gr.Interface(fn=query_db, inputs="text", outputs=[gr.Image(), gr.Gallery().style(grid=[8])])
-# demo.launch(debug=True)