backend/examples.py

import ast
import re
from pathlib import Path

import requests


def group_files_by_index(file_paths, data_type="audio"):
    # Regular expression pattern to extract the key from each image path
    if data_type == "audio":
        pattern = r"audio_(\d+).(png|wav)"
    elif data_type == "video":
        pattern = r"video_(\d+).(png|mkv)"
    else:
        pattern = r"img_(\d+).png"
    # Dictionary to store the grouped files
    grouped_files = {}
    # Iterate over each file path
    for file_path in file_paths:
        # Extract the key using the regular expression pattern
        match = re.search(pattern, file_path)
        if match:
            key = int(match.group(1))

            # Add the file path to the corresponding group in the dictionary
            if key not in grouped_files:
                grouped_files[key] = []
            grouped_files[key].append(file_path)
    # Sort the dictionary by keys
    sorted_grouped_files = dict(sorted(grouped_files.items()))
    return sorted_grouped_files


def build_description(
    i, data_none, data_attack, quality_metrics=["psnr", "ssim", "lpips"]
):
    # TODO: handle this at data generation
    if isinstance(data_none["fake_det"], str):
        data_none["fake_det"] = ast.literal_eval(data_none["fake_det"])
    if isinstance(data_none["watermark_det"], str):
        data_none["watermark_det"] = ast.literal_eval(data_none["watermark_det"])

    if isinstance(data_attack["fake_det"], str):
        data_attack["fake_det"] = ast.literal_eval(data_attack["fake_det"])
    if isinstance(data_attack["watermark_det"], str):
        data_attack["watermark_det"] = ast.literal_eval(data_attack["watermark_det"])

    if i == 0:
        fake_det = data_none["fake_det"]

        return f"detected: {fake_det}"
    elif i == 1:
        # Fixed metrics
        det = data_none["watermark_det"]
        p_value = float(data_none["p_value"])
        bit_acc = data_none["bit_acc"]

        # Dynamic metrics
        metrics_output = []
        for metric in quality_metrics:
            value = float(data_none[metric])
            metrics_output.append(f"{metric}: {value:.2f}")

        # Fixed metrics output
        fixed_metrics_output = (
            f" detected: {det} p_value: {p_value:.2f} bit_acc: {bit_acc:.2f}"
        )

        # Combine all outputs
        return " ".join(metrics_output) + f"{fixed_metrics_output}"
    elif i == 2:
        fake_det = data_attack["fake_det"]

        return f"det: {fake_det}"
    elif i == 3:
        det = data_attack["watermark_det"]

        p_value = float(data_attack["p_value"])
        word_acc = data_attack["word_acc"]
        bit_acc = data_attack["bit_acc"]

        return f"word_acc: {word_acc:.2f} detected: {det} p_value: {p_value:.2f} bit_acc: {bit_acc:.2f}"


def build_infos(abs_path: Path, datatype: str, dataset_name: str, db_key: str):
    def generate_file_patterns(prefixes, extensions):
        indices = [0, 1, 3, 4, 5]
        return [
            f"{prefix}_{index:05d}.{ext}"
            for prefix in prefixes
            for index in indices
            for ext in extensions
        ]

    if datatype == "audio":
        quality_metrics = ["snr", "sisnr", "stoi", "pesq"]
        extensions = ["png", "wav"]
        datatype_abbr = "audio"
        eval_results_path = abs_path + f"{dataset_name}_1k/examples_eval_results.json"
    elif datatype == "image":
        quality_metrics = ["psnr", "ssim", "lpips"]
        extensions = ["png"]
        datatype_abbr = "img"
        eval_results_path = abs_path + f"{dataset_name}_1k/examples_eval_results.json"
    elif datatype == "video":
        quality_metrics = ["psnr", "ssim", "lpips", "msssim", "vmaf"]
        extensions = ["mkv"]
        datatype_abbr = "video"
        eval_results_path = abs_path + f"{dataset_name}/examples_eval_results.json"

    response = requests.get(eval_results_path)
    if response.status_code == 200:
        results_data = response.json()
    else:
        return {}

    dataset = results_data["eval"][db_key]

    prefixes = [
        f"attacked_{datatype_abbr}",
        f"attacked_wmd_{datatype_abbr}",
        f"{datatype_abbr}",
        f"wmd_{datatype_abbr}",
    ]

    file_patterns = generate_file_patterns(prefixes, extensions)

    infos = {}
    for model_name in dataset.keys():
        model_infos = {}

        default_attack_name = "none"
        if datatype == "audio":
            default_attack_name = "identity"
        elif datatype == "video":
            default_attack_name = "Identity"

        identity_attack_rows = dataset[model_name][default_attack_name]["default"]

        for attack_name, attack_variants_data in dataset[model_name].items():
            for attack_variant, attack_rows in attack_variants_data.items():
                if attack_variant == "default":
                    attack = attack_name
                else:
                    attack = f"{attack_name}_{attack_variant}"

                if len(attack_rows) == 0:
                    model_infos[attack] = []
                    continue

                if datatype == "video":
                    file_paths = [
                        f"{abs_path}{dataset_name}/examples/{datatype}/{model_name}/{attack}/{pattern}"
                        for pattern in file_patterns
                    ]
                else:
                    file_paths = [
                        f"{abs_path}{dataset_name}_1k/examples/{datatype}/{model_name}/{attack}/{pattern}"
                        for pattern in file_patterns
                    ]

                all_files = []

                for i, files in group_files_by_index(
                    file_paths,
                    data_type=datatype,
                ).items():
                    data_none = [e for e in identity_attack_rows if e["idx"] == i][0]
                    data_attack = [e for e in attack_rows if e["idx"] == i][0]

                    files = sorted(
                        [(f, Path(f).stem) for f in files], key=lambda x: x[1]
                    )
                    files = files[2:] + files[:2]

                    files = [
                        {
                            "image_url": f,
                            "description": f"{n}\n{build_description(i, data_none, data_attack, quality_metrics)}",
                            **(
                                {"audio_url": f.replace(".png", ".wav")}
                                if datatype == "audio" and f.endswith(".png")
                                else {}
                            ),
                        }
                        for i, (f, n) in enumerate(files)
                    ]

                    all_files.extend(files)

                model_infos[attack] = all_files

        infos[model_name] = model_infos

    return infos


def image_examples_tab(abs_path: Path):
    dataset_name = "coco_val2014"
    datatype = "image"
    db_key = "coco_val2014"

    image_infos = build_infos(
        abs_path,
        datatype=datatype,
        dataset_name=dataset_name,
        db_key=db_key,
    )

    # First combo box (category selection)
    # model_choice = gr.Dropdown(
    #     choices=list(image_infos.keys()),
    #     label="Select a Model",
    #     value=None,
    # )
    # Second combo box (subcategory selection)
    # Initialize with options from the first category by default
    # attack_choice = gr.Dropdown(
    #     choices=list(image_infos["wam"].keys()),
    #     label="Select an Attack",
    #     value=None,
    # )

    # # Gallery component to display images
    # gallery = gr.Gallery(
    #     label="Image Gallery",
    #     columns=4,
    #     rows=1,
    # )

    # Update options for the second combo box when the first one changes
    # def update_subcategories(selected_category):
    #     values = list(image_infos[selected_category].keys())
    #     values = [(v, v) for v in values]
    #     attack_choice.choices = values
    #     # return gr.Dropdown.update(choices=list(image_infos[selected_category].keys()))

    # # Function to load images based on selections from both combo boxes
    # def load_images(category, subcategory):
    #     return image_infos.get(category, {}).get(subcategory, [])

    # # Update gallery based on both combo box selections
    # model_choice.change(
    #     fn=update_subcategories, inputs=model_choice, outputs=attack_choice
    # )
    # attack_choice.change(
    #     fn=load_images, inputs=[model_choice, attack_choice], outputs=gallery
    # )
    return image_infos


def video_examples_tab(abs_path: Path):
    dataset_name = "sav_val_full"
    datatype = "video"
    db_key = "sa-v_sav_val_videos"

    image_infos = build_infos(
        abs_path,
        datatype=datatype,
        dataset_name=dataset_name,
        db_key=db_key,
    )

    return image_infos

    # First combo box (category selection)
    # model_choice = gr.Dropdown(
    #     choices=list(image_infos.keys()),
    #     label="Select a Model",
    #     value=None,
    # )
    # Second combo box (subcategory selection)
    # Initialize with options from the first category by default
    # attack_choice = gr.Dropdown(
    #     choices=list(image_infos["videoseal_0.0"].keys()),
    #     label="Select an Attack",
    #     value=None,
    # )

    # Gallery component to display images
    # gallery = gr.Gallery(
    #     label="Video Gallery",
    #     columns=4,
    #     rows=1,
    # )

    # Update options for the second combo box when the first one changes
    # def update_subcategories(selected_category):
    #     values = list(image_infos[selected_category].keys())
    #     values = [(v, v) for v in values]
    #     attack_choice.choices = values
    #     # return gr.Dropdown.update(choices=list(image_infos[selected_category].keys()))

    # Function to load images based on selections from both combo boxes
    # def load_images(category, subcategory):
    #     return image_infos.get(category, {}).get(subcategory, [])

    # # Update gallery based on both combo box selections
    # model_choice.change(
    #     fn=update_subcategories, inputs=model_choice, outputs=attack_choice
    # )
    # attack_choice.change(
    #     fn=load_images, inputs=[model_choice, attack_choice], outputs=gallery
    # )


def audio_examples_tab(abs_path: Path):
    dataset_name = "voxpopuli"
    datatype = "audio"
    db_key = "voxpopuli"

    audio_infos = build_infos(
        abs_path,
        datatype=datatype,
        dataset_name=dataset_name,
        db_key=db_key,
    )

    return audio_infos

    print(audio_infos)

    # First combo box (category selection)
    # model_choice = gr.Dropdown(
    #     choices=list(audio_infos.keys()),
    #     label="Select a Model",
    #     value=None,
    # )
    # Second combo box (subcategory selection)
    # Initialize with options from the first category by default
    attack_choice = gr.Dropdown(
        choices=list(audio_infos["audioseal"].keys()),
        label="Select an Attack",
        value=None,
    )

    # Gallery component to display images
    gallery = gr.Gallery(
        label="Image Gallery", columns=4, rows=1, object_fit="scale-down"
    )

    audio_player = gr.Audio(visible=False)
    audio_map_state = gr.State({})

    # Update options for the second combo box when the first one changes
    def update_subcategories(selected_category):
        values = list(audio_infos[selected_category].keys())
        values = [(v, v) for v in values]
        attack_choice.choices = values
        # return gr.Dropdown.update(choices=list(image_infos[selected_category].keys()))

    # Function to load images based on selections from both combo boxes
    def load_audios(category, subcategory):
        files = audio_infos.get(category, {}).get(subcategory, [])
        images = [f for f in files if f[0].endswith(".png")]
        audios = {f[0]: f[0].replace(".png", ".wav") for f in images}
        return images, audios

    def play_audio(selected_image, audios):
        image_path = selected_image["image"]["path"]
        audio_file = audios.get(image_path)
        return gr.update(value=audio_file, visible=audio_file is not None)

    def hide_audio_player():
        # Hide the audio player when the preview is closed
        return gr.update(visible=False)

    def get_selected_image(select_data: gr.SelectData, audios):
        if select_data is None:
            return gr.update(visible=False)
        selected_image = select_data.value
        return play_audio(selected_image, audios)

    # Update gallery based on both combo box selections
    model_choice.change(
        fn=update_subcategories, inputs=model_choice, outputs=attack_choice
    )
    attack_choice.change(
        fn=load_audios,
        inputs=[model_choice, attack_choice],
        outputs=[gallery, audio_map_state],
    )
    gallery.select(
        fn=get_selected_image,
        inputs=[audio_map_state],
        outputs=audio_player,
    )
    gallery.preview_close(
        fn=hide_audio_player,
        outputs=audio_player,
    )
    return gr.Column([model_choice, attack_choice, gallery, audio_player])