src/label_interface.py

"""Interface for labeling concepts in images.
"""

from typing import Optional

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,
    profile: gr.OAuthProfile
):
    username = profile.username
    try:
        sample_idx = random.choice(range(len(global_variables.all_metadata[split])))
        sample = global_variables.all_metadata[split][sample_idx]
        image_path = f"{ASSETS_FOLDER}/{DATASET_NAME}/data/{split}/{sample['file_name']}"
        try:
            username_votes = sample["votes"][username]
            voted_concepts = [c for c in CONCEPTS if username_votes.get(c, False)]
        except KeyError:
            voted_concepts = []

        return image_path, voted_concepts, f"{split}:{sample_idx}"
    except IndexError:
        gr.Warning("No image found for the selected filter.")
        return None, None, None

def submit_label(
    voted_concepts: list,
    current_image: Optional[str],
    profile: gr.OAuthProfile
):
    username = profile.username
    if current_image is None:
        gr.Warning("No image selected.")
        return None, None, None
    split, idx = current_image.split(":")
    idx = int(idx)
    global_variables.get_metadata(split)
    if "votes" not in global_variables.all_metadata[split][idx]:
        global_variables.all_metadata[split][idx]["votes"] = {}
    global_variables.all_metadata[split][idx]["votes"][username] = {c: c in voted_concepts for c in CONCEPTS}
    vote_sum = {c: 0 for c in CONCEPTS}
    concepts = {}
    for c in CONCEPTS:
        for vote in global_variables.all_metadata[split][idx]["votes"].values():
            if c not in vote:
                continue
            vote_sum[c] += 2 * vote[c] - 1
        concepts[c] = vote_sum[c] > 0 if vote_sum[c] != 0 else None
    global_variables.all_metadata[split][idx]["concepts"] = concepts
    global_variables.save_metadata(split)
    return get_next_image(split, profile)
    

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",],
                    value="train",
                )
            with gr.Group():
                voted_concepts = gr.CheckboxGroup(
                    label="Concepts",
                    choices=CONCEPTS,
                )
            with gr.Row():
                next_button = gr.Button(
                    value="Next",
                )
                gr.LoginButton()
                submit_button = gr.Button(
                    value="Submit",
                )

        with gr.Column():
            image = gr.Image(
                label="Image",
            )
    current_image = gr.State(None)
    next_button.click(
        get_next_image, 
        inputs=[split],
        outputs=[image, voted_concepts, current_image],
    )
    submit_button.click(
        submit_label,
        inputs=[voted_concepts, current_image],
        outputs=[image, voted_concepts, current_image],
    )