|
import gradio as gr |
|
from datasets import load_from_disk |
|
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] |
|
|
|
ds = load_from_disk("color-sorted") |
|
|
|
def get_nearest_64(gender, ethnicity, model, no, index): |
|
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"][0] |
|
neighbors = ds.select(max(ix-10, 0), min(ix+10, len(ds)-1)) |
|
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("# Colorfulness Nearest Neighbors Explorer") |
|
gr.Markdown("### TF-IDF index of the _identities_ dataset of images generated by 3 models using colorfulness") |
|
gr.Markdown("#### Choose one of the generated identity images to see its nearest neighbors according to colorfulness") |
|
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() |