app.py

import json
import logging
import os
import random
from functools import partial

import gradio as gr
from datasets import Dataset, load_dataset
from dotenv import load_dotenv

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

load_dotenv()

# dataset = load_dataset("detection-datasets/coco")
it_dataset = (
    load_dataset("imagenet-1k", split="train", streaming=True, trust_remote_code=True)
    .shuffle(42)
    .skip(0)
    .take(1000)
)


def gen_from_iterable_dataset(iterable_ds):
    """
    Convert an iterable dataset to a generator
    """
    yield from iterable_ds


dataset = Dataset.from_generator(
    partial(gen_from_iterable_dataset, it_dataset), features=it_dataset.features
)

# imagenet_categories_data.json is a JSON file containing a hierarchy of ImageNet categories.
# We want to take all categories under "artifact, artefact".
# Each node has this structure:
# {
#     "id": 1,
#     "name": "entity",
#     "children": ...
# }
with open("imagenet_categories_data.json") as f:
    data = json.load(f)

    # Recursively find all categories under "artifact, artefact".
    # We want to get all the "index" values of the leaf nodes. Nodes that are not leaf nodes have a "children" key.
    def find_categories(node):
        if "children" in node:
            for child in node["children"]:
                yield from find_categories(child)
        elif "index" in node:
            yield node["index"]

    broad_categories = data["children"]
    artifact_category = next(
        filter(lambda x: x["name"] == "artifact, artefact", broad_categories)
    )
    artifact_categories = list(find_categories(artifact_category))
    # logger.info(f"Artifact categories: {artifact_categories}")


def filter_imgs_by_label(x):
    """
    Filter out the images that have label -1
    """
    return x["label"] in artifact_categories


dataset = dataset.filter(filter_imgs_by_label)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

load_dotenv()


def get_user_prompt():
    # Pick the first 3 images and labels
    images = []
    machine_labels = []
    human_labels = []
    for i in range(3):
        data = dataset[random.randint(0, len(dataset) - 1)]
        images.append(data["image"])
        # Get the label as a human readable string
        machine_labels.append(data["label"])
        human_label = dataset.features["label"].int2str(data["label"])
        human_labels.append(human_label)
    return {
        "images": images,
        "machine_labels": machine_labels,
        "human_labels": human_labels,
    }


hf_writer = gr.HuggingFaceDatasetSaver(
    hf_token=os.environ["HF_TOKEN"],
    dataset_name="acmc/maker-faire-bot",
    private=True,
)
csv_writer = gr.CSVLogger()

theme = gr.themes.Default(primary_hue="cyan", secondary_hue="fuchsia")

translation_table = {
    "Maker Faire Bot": "Maker Faire Bot",
    "**Think about these objects...**": "**Tænk på disse objekter...**",
    "We want to build a Maker Faire Bot that can generate creative ideas. Help us by providing ideas on what you'd build with the following three objects!": "Vi vil bygge en Maker Faire Bot, der kan generere kreative ideer. Hjælp os ved at give ideer til, hvad du ville bygge med de følgende tre objekter!",
    "Change": "Skift",
    "What would you build with these 3 things?": "Hvad ville du bygge med disse 3 ting?",
    "For example, if you have a roll of string, a camera, and a loudspeaker, you could build an electronic guitar. If you can write in Danish, that's great!": "For eksempel, hvis du har en rulle snor, et kamera og en højttaler, kunne du bygge en elektronisk guitar. Hvis du kan skrive på dansk, er det fantastisk!",
    "It doesn't need to be a very long explanation, just a few sentences to help the bot understand your idea.": "Det behøver ikke være en meget lang forklaring, bare et par sætninger for at hjælpe robotten med at forstå din idé.",
    "Submit": "Indsend",
    "New Prompt": "Ny opgave",
    "How would you build it?": "Hvordan ville du bygge det?",
    "This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our policy.": "Dette er et eksperimentelt projekt. Dine data er anonyme og vil blive brugt til at træne en AI-model. Ved at bruge dette værktøj accepterer du vores politik.",
    "(example): An digital electronic guitar": "(eksempel): En digital elektronisk guitar",
    """I would use the roll of string to create the strings of the guitar, and the camera to analyze the hand movements. Then, I would use an AI model to predict the chords and play the sound through the loudspeaker.""": """Jeg ville bruge snoren til at skabe guitarens strenge, og kameraet til at analysere håndbevægelserne. Derefter ville jeg bruge en AI-model til at forudsige akkorderne og afspille lyden gennem højttaleren.""",
}


def get_bilingual_string(key):
    return f"{translation_table[key]}  //  {key}"


with gr.Blocks(theme=theme) as demo:
    with gr.Row() as header:
        gr.Image(
            "maker-faire-logo.webp",
            show_download_button=False,
            show_label=False,
            show_share_button=False,
            container=False,
            # height=100,
            scale=0.2,
        )
        gr.Markdown(
            get_bilingual_string("Maker Faire Bot"),
            visible=False,
        )

    user_prompt = gr.State(get_user_prompt())

    gr.Markdown(get_bilingual_string("**Think about these objects...**"))
    gr.Markdown(
        get_bilingual_string(
            "We want to build a Maker Faire Bot that can generate creative ideas. Help us by providing ideas on what you'd build with the following three objects!"
        )
    )
    image_components = []
    with gr.Row(variant="panel") as row:
        def change_image(this_i, user_prompt):
            logger.info(
                f"Current user prompt: {user_prompt}, current image index: {this_i}"
            )
            data = dataset[random.randint(0, len(dataset) - 1)]
            new_user_prompt = user_prompt.copy()
            new_user_prompt["images"][this_i] = data["image"]
            new_user_prompt["machine_labels"][this_i] = data["label"]
            new_user_prompt["human_labels"][this_i] = dataset.features[
                "label"
            ].int2str(data["label"])
            logger.info(f"New user prompt: {new_user_prompt}")
            return (
                new_user_prompt,
                new_user_prompt["images"][this_i],
                gr.update(
                    label=new_user_prompt["human_labels"][this_i],
                ),
            )

        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][0],
                label=user_prompt.value["human_labels"][0],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change"), variant="secondary")

            btn.click(
                lambda *args: change_image(0, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )
        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][1],
                label=user_prompt.value["human_labels"][1],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change"), variant="secondary")

            btn.click(
                lambda *args: change_image(1, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )
        with gr.Column(variant="default") as col:
            img = gr.Image(
                user_prompt.value["images"][2],
                label=user_prompt.value["human_labels"][2],
                interactive=False,
                show_download_button=False,
                show_share_button=False,
            )
            image_components.append(img)
            btn = gr.Button(get_bilingual_string("Change"), variant="secondary")

            btn.click(
                lambda *args: change_image(2, *args),
                inputs=[user_prompt],
                outputs=[user_prompt, img, img],
                preprocess=True,
                postprocess=True,
            )

    user_answer_object = gr.Textbox(
        autofocus=True,
        placeholder=get_bilingual_string("(example): An digital electronic guitar"),
        label=get_bilingual_string("What would you build with these 3 things?"),
        info=get_bilingual_string("For example, if you have a roll of string, a camera, and a loudspeaker, you could build an electronic guitar. If you can write in Danish, that's great!")
    )
    user_answer_explanation = gr.TextArea(
        autofocus=True,
        label=get_bilingual_string("How would you build it?"),
        # The example uses a roll of string, a camera, and a loudspeaker to build an electronic guitar.
        placeholder=get_bilingual_string(
            """I would use the roll of string to create the strings of the guitar, and the camera to analyze the hand movements. Then, I would use an AI model to predict the chords and play the sound through the loudspeaker."""
        ),
        info=get_bilingual_string("It doesn't need to be a very long explanation, just a few sentences to help the bot understand your idea.")
    )

    csv_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_csv",
    )
    hf_writer.setup(
        components=[user_prompt, user_answer_object, user_answer_explanation],
        flagging_dir="user_data_hf",
    )

    submit_btn = gr.Button(get_bilingual_string("Submit"), variant="primary")

    def log_results(prompt, object, explanation):
        logger.info(f"logging - Prompt: {prompt}")
        # csv_writer.flag(
        #    [
        #        {
        #            "machine_labels": prompt["machine_labels"],
        #            "human_labels": prompt["human_labels"],
        #        },
        #        object,
        #        explanation,
        #    ]
        # )
        hf_writer.flag(
            [
                {
                    "machine_labels": prompt["machine_labels"],
                    "human_labels": prompt["human_labels"],
                },
                object,
                explanation,
            ]
        )
        return ["", ""]  # Clear the textboxes

    submit_btn.click(
        log_results,
        inputs=[user_prompt, user_answer_object, user_answer_explanation],
        outputs=[user_answer_object, user_answer_explanation],
        preprocess=True,
    )

    # def renew_prompt(image_components):
    #    new_prompt = get_user_prompt()
    #    for i in range(len(new_prompt["images"])):
    #        image_components[i].update(
    #            url=new_prompt["images"][i],
    #            label=new_prompt["human_labels"][i],
    #        )
    #    return new_prompt

    # new_prompt_btn = gr.Button(get_bilingual_string("New Prompt"), variant="secondary")
    # new_prompt_btn.click(
    #    renew_prompt,
    #    inputs=image_components,
    #    outputs=[user_prompt],
    #    # preprocess=True,
    # )

    gr.Markdown(
        get_bilingual_string(
            "This is an experimental project. Your data is anonymous and will be used to train an AI model. By using this tool, you agree to our policy."
        )
    )
if __name__ == "__main__":
    demo.launch()