import gradio as gr from datasets import load_dataset import numpy as np gender_labels = ['man', 'non-binary', 'woman', 'no_gender_specified', ] ethnicity_labels = ['African-American', 'American_Indian', 'Black', 'Caucasian', 'East_Asian', 'First_Nations', 'Hispanic', 'Indigenous_American', 'Latino', 'Latinx', 'Multiracial', 'Native_American', 'Pacific_Islander', 'South_Asian', 'Southeast_Asian', 'White', 'no_ethnicity_specified'] models = ['DallE', 'SD_14', 'SD_2'] nos = [1,2,3,4,5,6,7,8,9,10] indexes = [768, 1536, 10752] ds = load_dataset("tti-bias/identities", split="train") def get_nearest_64(gender, ethnicity, model, no, index): df = ds.remove_columns(["image","image_path"]).to_pandas() index = np.load(f"indexes/knn_{index}_65.npy") ix = df.loc[(df['ethnicity'] == ethnicity) & (df['gender'] == gender) & (df['no'] == no) & (df['model'] == model)].index[0] image = ds.select([index[ix][0]])["image"][0] neighbors = ds.select(index[ix][1:25]) neighbor_images = neighbors["image"] neighbor_captions = [caption.split("/")[-1] for caption in neighbors["image_path"]] neighbor_captions = [' '.join(caption.split("_")[4:-3]) for caption in neighbor_captions] neighbor_models = neighbors["model"] neighbor_captions = [f"{a} {b}" for a,b in zip(neighbor_captions,neighbor_models)] return image, list(zip(neighbor_images, neighbor_captions)) with gr.Blocks() as demo: gr.Markdown("# BoVW Nearest Neighbors Explorer") gr.Markdown("### TF-IDF index of the _identities_ dataset of images generated by 3 models using a visual vocabulary of 10,752 words.") gr.Markdown("#### Choose one of the generated identity images to see its nearest neighbors according to a bag-of-visual-words model.") with gr.Row(): with gr.Column(): model = gr.Radio(models, label="Model") index = gr.Radio(indexes, label="Visual vocabulary size") gender = gr.Radio(gender_labels, label="Gender label") with gr.Column(): ethnicity = gr.Radio(ethnicity_labels, label="Ethnicity label") no = gr.Radio(nos, label="Image number") button = gr.Button(value="Get nearest neighbors") with gr.Row(): image = gr.Image() gallery = gr.Gallery().style(grid=4) button.click(get_nearest_64, inputs=[gender, ethnicity, model, no, index], outputs=[image, gallery]) demo.launch()