ui/layout.py

import gradio as gr
import time

from logic.data_utils import CustomHFDatasetSaver
from data.lang2eng_map import lang2eng_mapping
from gradio_modal import Modal
from logic.handlers import *
from config.settings import *
from functools import partial

from .selection_page import build_selection_page
from .main_page import build_main_page
from .main_page import sort_with_pyuca

def get_key_by_value(dictionary, value):
    for key, val in dictionary.items():
        if val == value:
            return key
    return None

def load_concepts(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _, _ = local_storage
    if category:
        if lang in words_mapping:
            category = get_key_by_value(words_mapping[lang], category) # translate back the category name
        concept_val = None
        # I'm doing this to avoid the case where the user load a previous example and the concept 'reamins' selected even if category changed
        if loading_example:
            concept_val = concept_btn if concept_btn else None
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)][category]), value=concept_val), False
    else:
        concepts_list = []
        # cats = []  # FIXME: Assumes concepts are unique across all categories
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            for cat in concepts[country][lang]:
                for concept in concepts[country][lang][cat]:
                    # cats.append(cat)
                    concepts_list.append(concept)
            return gr.update(choices=sort_with_pyuca(concepts_list), value=None), False
        # return gr.update(choices=[], value=None), False

def update_category(category, concept_btn, local_storage, loading_example, concepts):
    country, lang, _, _ = local_storage
    if not category and concept_btn:
        # FIXME: Assumes unique concepts across categories
        for cat in concepts[country][lang2eng_mapping.get(lang, lang)]:
            for concept in concepts[country][lang2eng_mapping.get(lang, lang)][cat]:
                if concept == concept_btn:
                    category = cat
        return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)]), value=category), True
    else:
        # import pdb; pdb.set_trace()
        if country is None or lang is None:
            return gr.update(choices=[], value=None), False
        else:
            return gr.update(choices=sort_with_pyuca(concepts[country][lang2eng_mapping.get(lang, lang)]), value=category), False


def switch_ui(country, language, username, password, flag=False, metadata_dict=None):
    # wait for 2 seconds
    print(f"Language: {language}, Country: {country}")

    if not country or not language:
        gr.Warning(f"⚠️ Please select Country and Language first.",  duration=3)
        return

    # Hide the selection page and show the main UI.
    if flag:
        gr.Info(f"Loading the Language Selection UI",  duration=3)
        time.sleep(2)
        local_storage = [None, None, "", ""]
        countries = sort_with_pyuca(list(metadata_dict.keys()))
        return (local_storage, gr.update(visible=flag), gr.update(visible=not flag), 
                gr.update(visible=not flag), 
                gr.update(choices=countries, label="Country", value=None), 
                gr.update(value=None, choices=[], label="Language", allow_custom_value=False, interactive=False))
    
    else:
        gr.Info(f"Loading the Main UI",  duration=3)
        time.sleep(2)
        local_storage = [country, language, username, password]
        return local_storage, gr.update(visible=flag), gr.update(visible=not flag), gr.update(visible=not flag)

def build_ui(concepts_dict, metadata_dict, HF_API_TOKEN, HF_DATASET_NAME):
    hf_writer = CustomHFDatasetSaver(HF_API_TOKEN, HF_DATASET_NAME, private=True)
    custom_css = """
        .compact-container {
            max-width: 600px;
            margin: auto;
            padding: 20px;
        }
        .compact-btn {
        }
        #image_inp img {
            object-fit: contain;        /* make sure the full image shows */
            height: 460px;           /* set a fixed height */
        }
    """
    ############################################################################
    with gr.Blocks(css=custom_css) as ui:
        local_storage = gr.State([None, None, "", ""])
        loading_example = gr.State(False) # to check if the values are loaded from a user click on an example in 
        # First page: selection

        selection_page, country_choice, language_choice, proceed_btn, username, password, intro_markdown = build_selection_page(metadata_dict)
        
        # Second page
        cmp_main_ui = build_main_page(concepts_dict, metadata_dict, local_storage)

        main_ui_placeholder = cmp_main_ui["main_ui_placeholder"]
        country_inp = cmp_main_ui["country_inp"]
        language_inp = cmp_main_ui["language_inp"]
        image_inp = cmp_main_ui["image_inp"]
        image_url_inp = cmp_main_ui["image_url_inp"]
        long_caption_inp = cmp_main_ui["long_caption_inp"]
        num_words_inp = cmp_main_ui["num_words_inp"]
        category_btn = cmp_main_ui["category_btn"]
        concept_btn = cmp_main_ui["concept_btn"]
        category_concept_dropdowns = cmp_main_ui["category_concept_dropdowns"]
        instruct_btn = cmp_main_ui["instruct_btn"]
        clear_btn = cmp_main_ui["clear_btn"]
        hide_faces_btn = cmp_main_ui["hide_faces_btn"]
        hide_all_faces_btn = cmp_main_ui["hide_all_faces_btn"]
        unhide_faces_btn = cmp_main_ui["unhide_faces_btn"]
        submit_btn = cmp_main_ui["submit_btn"]
        timestamp_btn = cmp_main_ui["timestamp_btn"]
        exampleid_btn = cmp_main_ui["exampleid_btn"]
        username_inp = cmp_main_ui["username_inp"]
        password_inp = cmp_main_ui["password_inp"]
        modal_saving = cmp_main_ui["modal_saving"]
        modal_data_saved = cmp_main_ui["modal_data_saved"]
        modal = cmp_main_ui["modal"]
        exit_btn = cmp_main_ui["exit_btn"]
        intro_text_inp = cmp_main_ui["intro_text_inp"] # Note intro_text is actually the "Task" text from metadata
        intro_text_inst_inp = cmp_main_ui["intro_text_inst_inp"] # This is the "Instructions" text from metadata
        modal_saving_text = cmp_main_ui["modal_saving_text"]
        modal_data_saved_text = cmp_main_ui["modal_data_saved_text"]
        exclude_btn = cmp_main_ui["exclude_btn"]
        modal_exclude_confirm = cmp_main_ui["modal_exclude_confirm"]
        cancel_exclude_btn = cmp_main_ui["cancel_exclude_btn"]
        confirm_exclude_btn = cmp_main_ui["confirm_exclude_btn"]


        ### Category button
        category_btn.change(
            fn=partial(load_concepts, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[concept_btn, loading_example]
        )
        concept_btn.change(
            fn=partial(update_category, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, loading_example],
            outputs=[category_btn, loading_example]
        ).then(
            fn=partial(load_concepts, concepts=concepts_dict),
            inputs=[category_btn, concept_btn, local_storage, gr.State(True)],
            outputs=[concept_btn, loading_example]
        )
        ####

        # is_blurred = gr.State(False)  # Initialize as False

        # Notification image is false
        with Modal(visible=False) as modal_img_url:
            gr.Markdown("The image URL is not valid. Please provide a valid image URL.")
            # gr.Warning(f"⚠️ The image URL is not valid. Please provide a valid image URL.")

        # Event listeners
        gr.on(
            triggers=[country_choice.change, language_choice.change],
            fn=validate_metadata,
            inputs=[country_choice, language_choice],
            outputs=[proceed_btn],
        )
        gr.on(
            triggers=[image_url_inp.change],
            fn=update_image,
            inputs=[image_url_inp],
            outputs=[image_inp, modal_img_url]
        )
        gr.on(
            triggers=[language_choice.change],
            fn=partial(update_intro_language, metadata=metadata_dict),
            inputs=[country_choice, language_choice, intro_markdown],
            outputs=[intro_markdown]
        )

        gr.on(
            triggers=[image_inp.change, long_caption_inp.change],
            fn=update_timestamp,
            outputs=[timestamp_btn]
        )
        ori_img = gr.State(None)
        gr.on(
            triggers=[image_inp.change],
            fn=validate_inputs,
            inputs=[image_inp, ori_img, concept_btn], # is_blurred
            outputs=[submit_btn, image_inp, ori_img], # is_blurred
        )

        gr.on(
            triggers=[long_caption_inp.change],
            fn=count_words,
            inputs=[long_caption_inp, language_choice],
            outputs=[num_words_inp],
        )

        # Clear Button
        clear_btn.click(
            fn=clear_data,
            outputs=[
                image_inp, image_url_inp, long_caption_inp, exampleid_btn, 
                category_btn, concept_btn, 
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ],
        )

        #============= Face Blurring ============= #
        with Modal(visible=False) as modal_faces:
            with gr.Column():
                face_img = gr.Image(label="Image Faces", elem_id="image_faces", format="png", height=512, width=768)
                with gr.Row():
                    faces_count = gr.Textbox(label="Face Counts", elem_id="face_counts", interactive=False)
                    blur_faces_ids = gr.Dropdown(
                        [], value=[], multiselect=True, label="Please select the faces IDs you want to blur.", elem_id="blur_faces_ids")
                    # blur_faces_ids = gr.Textbox(label="Specify faces ids to blur by comma", elem_id="blur_faces_ids", interactive=True)
                    submit_btn_face = gr.Button("Submit", variant="primary", interactive=True, elem_id="submit_btn_face")
        
        with Modal(visible=False) as modal_faces_2:
            gr.Markdown("The current image does not have any faces to hide.")
    
        faces_info = gr.State(None)
        hide_faces_btn.click(
            fn=select_faces_to_hide,
            inputs=[image_inp, blur_faces_ids],
            outputs=[image_inp, modal_faces, modal_faces_2, face_img, faces_count, faces_info, blur_faces_ids]
        )

        submit_btn_face.click(
            fn=blur_selected_faces,
            inputs=[image_inp, blur_faces_ids, faces_info, face_img, faces_count], # is_blurred 
            outputs=[image_inp, modal_faces, face_img, faces_count, blur_faces_ids] # is_blurred
        )

        hide_all_faces_btn.click(
            fn=blur_all_faces,
            inputs=[image_inp],
            outputs=[image_inp, modal_faces_2] # is_blurred
        )

        unhide_faces_btn.click(
            fn=unhide_faces,
            inputs=[image_inp, ori_img], # is_blurred
            outputs=[image_inp] # is_blurred
        )

        # ===============================
        with gr.Column(visible=False, elem_id="browse_data") as browse_data_placeholder:
            # Browse Data
            browse_data_text = gr.Markdown("## Browse Data")

            loading_msg = gr.Markdown("**Loading your data, please wait ...**")
            for contrib_type in ["Your data", "Other data"]:
                with gr.Tab(contrib_type):
                    if contrib_type == "Your data":
                        # Show user's past data points
                        user_examples = gr.Dataset(
                            samples=[],
                            components=['image','textbox','textbox','textbox','textbox', 
                                       'textbox','textbox','textbox', 'textbox'],
                            headers=['Image', 'Image URL (Optional, if not uploading an image)', 'Description', 'Country', 'Language', 
                                    'Category', 'Concept', 'Additional Concepts', 'ID'],
                        )
                        # Handle clicking on an example
                        user_examples.click(
                            fn=partial(handle_click_example, concepts_dict=concepts_dict),
                            inputs=[user_examples],
                            outputs=[
                                image_inp, image_url_inp, long_caption_inp, exampleid_btn, 
                                category_btn, concept_btn,
                                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                                category_concept_dropdowns[3], category_concept_dropdowns[4], loading_example
                            ],
                        )

                    elif contrib_type == "Other data":
                        # Show others' images
                        pass

        # ============================================ #
        # Submit Button Click events

        proceed_btn.click(
            fn=partial(switch_ui, flag=False),
            inputs=[country_choice, language_choice, username, password],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_data_placeholder],
        ).then(
            fn=partial(update_language, metadata_dict=metadata_dict, concepts_dict=concepts_dict),
            inputs=[local_storage],
            outputs=[
                country_inp, language_inp, username_inp, password_inp, category_btn, concept_btn, image_inp, 
                image_url_inp, long_caption_inp, intro_text_inp, intro_text_inst_inp, instruct_btn, clear_btn,
                submit_btn, modal_saving_text, modal_data_saved_text, timestamp_btn, exit_btn, browse_data_text,
                loading_msg, hide_all_faces_btn, hide_faces_btn, unhide_faces_btn, exclude_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ]
        ).then(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path = LOCAL_DS_DIRECTORY_PATH),
            inputs=[username_inp, password_inp, country_choice, language_choice],
            outputs=[user_examples, loading_msg],
        )


        # Exit Button
        exit_btn.click(
            fn=exit,
            outputs=[
                image_inp, image_url_inp, long_caption_inp, user_examples, loading_msg, 
                username, password, local_storage, exampleid_btn, category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ],
        ).success(
            fn=partial(switch_ui, flag=True, metadata_dict=metadata_dict),
            inputs=[country_choice, language_choice, username, password],
            outputs=[local_storage, selection_page, main_ui_placeholder, browse_data_placeholder, country_choice, language_choice],
        )

        # Disable button while saving
        # Note: I think this is not longer needed as we clear all the inputs and disable the submit button if the data is saved correctly
        # def disable_submit():
        #     return gr.update(interactive=False)

        # def enable_submit():
        #     return gr.update(interactive=True)

        # STEP 1: show modal
        # submit_btn.click(lambda: Modal(visible=True), None, modal_saving)
        
        # STEP 2: disable button
        # submit_btn.click(disable_submit, None, [submit_btn], queue=False)

        #STEP 3: perform save_data
        data_outputs = {
            "image": image_inp,
            "image_url": image_url_inp,
            "caption": long_caption_inp,
            "country": country_inp,
            "language": language_inp,
            "category": category_btn,
            "concept": concept_btn,
            "category_1_concepts": category_concept_dropdowns[0],
            "category_2_concepts": category_concept_dropdowns[1],
            "category_3_concepts": category_concept_dropdowns[2],
            "category_4_concepts": category_concept_dropdowns[3],
            "category_5_concepts": category_concept_dropdowns[4],
            "timestamp": timestamp_btn,
            "username": username_inp,
            "password": password_inp,
            "id": exampleid_btn,
            "excluded": gr.State(value=False),
            "concepts_dict": gr.State(value=concepts_dict),
            "country_lang_map": gr.State(value=lang2eng_mapping),
            # "is_blurred": is_blurred
        }
        # data_outputs = [image_inp, image_url_inp, long_caption_inp,
        #         country_inp, language_inp, category_btn, concept_btn,
        #         timestamp_btn, username_inp, password_inp, exampleid_btn]
        hf_writer.setup(list(data_outputs.keys()), local_ds_folder = LOCAL_DS_DIRECTORY_PATH)

        # STEP 4: Chain save_data, then update_user_data, then re-enable button, hide modal, and clear
        submit_btn.click(
            hf_writer.save,
            list(data_outputs.values()),
            None,
        ).success(
            fn=partial(clear_data, "submit"),
            outputs=[
                image_inp, image_url_inp, long_caption_inp, exampleid_btn, 
                category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ],
        # ).success(enable_submit, 
        #    None, [submit_btn]
        # ).success(lambda: Modal(visible=False), 
        #    None, modal_saving
        # ).success(lambda: Modal(visible=True), 
        #        None, modal_data_saved
        ).success(
            # set loading msg
            lambda: gr.update(value="**Loading your data, please wait ...**"),
            None, loading_msg
        ).success(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path = LOCAL_DS_DIRECTORY_PATH),
            inputs=[username_inp, password_inp, country_choice, language_choice],
            outputs=[user_examples, loading_msg]
        )

        # ============================================ #
        # instructions button 
        instruct_btn.click(lambda: Modal(visible=True), None, modal)
        # ============================================ #

        # # Load saved values from local storage (browser storage)
        # @ui.load(inputs=[local_storage], outputs=[country_choice, language_choice, username, password])
        # def load_from_local_storage(saved_values):
        #     print("loading from local storage", saved_values)
        #     return saved_values[0], saved_values[1], saved_values[2], saved_values[3]
        
        # ============================================= #
        # Exclude button 
        # ============================================= #

        exclude_status = gr.Markdown(visible=False)

        # Show confirmation modal when exclude button is clicked
        exclude_btn.click(
            fn=check_exclude_fn,
            inputs=[image_inp],
            outputs=[modal_exclude_confirm]
        )

        # Close modal when cancel button is clicked
        cancel_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        )

        # Handle confirm exclusion
        confirm_exclude_btn.click(
            fn=lambda: gr.update(visible=False),
            outputs=[modal_exclude_confirm]
        ).success(
            fn=hf_writer.save,
            inputs=[
                image_inp, image_url_inp, long_caption_inp, country_inp, language_inp, 
                category_btn, concept_btn,
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4],
                timestamp_btn, username_inp, password_inp, exampleid_btn, gr.State(value=True),
                gr.State(value=concepts_dict), gr.State(value=lang2eng_mapping)
            ],
            outputs=None
        ).success(
            fn=partial(clear_data, "remove"),
            outputs=[
                image_inp, image_url_inp, long_caption_inp, exampleid_btn, 
                category_btn, concept_btn, 
                category_concept_dropdowns[0], category_concept_dropdowns[1], category_concept_dropdowns[2], 
                category_concept_dropdowns[3], category_concept_dropdowns[4]
            ]
        ).success(
            fn=lambda: gr.update(value="Example excluded successfully", visible=True),
            outputs=exclude_status
        ).success(
            fn=lambda: gr.update(value="**Refreshing your data, please wait...**"),
            outputs=loading_msg
        ).success(
            fn=partial(update_user_data, HF_DATASET_NAME=HF_DATASET_NAME, local_ds_directory_path=LOCAL_DS_DIRECTORY_PATH),
            inputs=[username_inp, password_inp, country_choice, language_choice],
            outputs=[user_examples, loading_msg]
        )

    return ui