src/explore_interface.py

"""Interface to explore images and labels.
"""

import random

import gradio as gr

from src import global_variables
from src.constants import CONCEPTS, ASSETS_FOLDER, DATASET_NAME

def get_next_image(
    split: str,
    concepts: list,
    current_metadata: dict,
    selected_concepts: list,
):
    if concepts != selected_concepts:
        for split, values in global_variables.all_metadata.items():
            current_metadata[split] = [x for x in values if all([c in x["concepts"] for c in concepts])]
        selected_concepts = concepts
    try:
        if split == "all":
            sample = random.choice(
                current_metadata["train"] + current_metadata["validation"] + current_metadata["test"]
            )
        else:
            sample = random.choice(current_metadata[split])
        image_path = f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/{sample['file_name']}"
        return image_path, sample["class"], sample["concepts"], current_metadata, selected_concepts
    except IndexError:
        gr.Warning("No image found for the selected filter.")
        return None, None, None, current_metadata, selected_concepts

with gr.Blocks() as interface:        
    with gr.Row():
        with gr.Column():
            with gr.Group():
                gr.Markdown(
                    "## # Image Selection",
                )
                split = gr.Radio(
                    label="Split",
                    choices=["train", "validation", "test", "all"],
                    value="train",
                )
                concepts = gr.Dropdown(
                    label="Concepts",
                    multiselect=True,
                    choices=CONCEPTS,
                )
            button = gr.Button(
                value="Next",
            )
            with gr.Group():
                gr.Markdown(
                    "##  # Image Info",
                )
                im_class = gr.Textbox(
                    label="Class",
                )
                im_concepts = gr.JSON(
                    label="Concepts",
                )
        with gr.Column():
            image = gr.Image(
                label="Image",
            )
    current_metadata = gr.State(global_variables.all_metadata)
    selected_concepts = gr.State([])
    button.click(
        get_next_image, 
        outputs=[image, im_class, im_concepts, current_metadata, selected_concepts], 
        inputs=[split, concepts, current_metadata, selected_concepts]
    )