Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import json
 import numpy as np
 import pandas as pd
+from datasets import load_from_disk
+from itertools import chain
 import operator
 
 pd.options.plotting.backend = "plotly"
@@ -9,6 +11,12 @@ pd.options.plotting.backend = "plotly"
 
 TITLE = "Diffusion Professions Cluster Explorer"
 
+professions = load_from_disk("professions")
+professions_df = professions.to_pandas()
+
+def get_image(model, fname):
+    return professions.select(professions_df[(professions_df["image_path"]==fname) & (professions_df["model"]==model)].index)["image"][0]
+
 clusters_dicts = dict(
     (num_cl, json.load(open(f"clusters/professions_to_clusters_{num_cl}.json")))
     for num_cl in [12, 24, 48]
@@ -142,7 +150,8 @@ def show_examplars(num_clusters, prof_name, mod_name, cl_id):
     examplars_dict = clusters_dicts[num_clusters][df_models[mod_name]][prof_name][
         "cluster_examplars"
     ][str(cl_id)]
-    return json.dumps(examplars_dict)
+    l = list(chain(*[examplars_dict[k] for k in examplars_dict]))
+    return [get_image(model,fname) for _,model,fname in l]
 
 
 with gr.Blocks(title=TITLE) as demo:
@@ -249,7 +258,7 @@ with gr.Blocks(title=TITLE) as demo:
                     )
         with gr.Row():
             examplars_plot = (
-                gr.JSON()
+                gr.Gallery()
             )  # TODO: turn this into a plot with the actual images
             demo.load(
                 show_examplars,