merge_data.py

from pathlib import Path
import torch
from collections import defaultdict

from modules.constants import OPTIONS, SAL_METHODS
from modules.parser import DATA_ROOT
from plot_saliency import DISTINGUISHERS
from saliency import extract_alpha_from_filename


def read_and_merge_files(data_root: Path, groups: list, option: str, saliency_method: str):
    # Initialize nested dict: distinguisher -> metric -> alpha -> value
    merged_data = {d: defaultdict(dict) for d in DISTINGUISHERS}
    all_metrics = set()

    for group in groups:
        input_directory = data_root / group / option / saliency_method
        if not input_directory.exists():
            print(f"WARNING: {input_directory} does not exist. Skipping.")
            continue

        for file in input_directory.iterdir():
            for distinguisher in DISTINGUISHERS:
                if distinguisher in file.name:
                    content = torch.load(file)

                    for metric, values in content.items():
                        all_metrics.add(metric)
                        alpha = extract_alpha_from_filename(str(file))
                        merged_data[distinguisher][metric][alpha] = values

    return merged_data, sorted(all_metrics)


def save_merged_group(output_root: Path, merged_group_name: str, option: str, saliency_method: str, merged_data: dict):
    save_dir = output_root / merged_group_name / option / saliency_method
    save_dir.mkdir(parents=True, exist_ok=True)

    for distinguisher, metric_dict in merged_data.items():
        # metric_dict: metric -> alpha -> value
        all_alphas = set()
        for metric in metric_dict.values():
            all_alphas.update(metric.keys())

        for alpha in all_alphas:
            single_alpha_dict = {}
            for metric, alpha_dict in metric_dict.items():
                if alpha in alpha_dict:
                    single_alpha_dict[metric] = alpha_dict[alpha]
            if single_alpha_dict:  # Only save if there is data
                save_path = save_dir / f"{distinguisher}_{alpha}.pt"
                torch.save(single_alpha_dict, save_path)




GROUPS = ["faithfulness_saliency_evaluation", "new_lle_complexity_saliency_evaluation", "max_sensitivity_saliency_evaluation"]
MERGED_GROUP = "MERGED_saliency_evaluation"

for option in OPTIONS:
    for saliency_method in SAL_METHODS:
        data, metrics = read_and_merge_files(DATA_ROOT, GROUPS, option, saliency_method)
        save_merged_group(DATA_ROOT, MERGED_GROUP, option, saliency_method, data)