eval_all.py

import argparse
import glob
import json
import os
import re

this_dir = os.path.abspath(os.path.dirname(__file__))


modality2dataset = {
    "image": [
        "ImageNet-1K",
        "N24News",
        "HatefulMemes",
        "VOC2007",
        "SUN397",
        "Place365",
        "ImageNet-A",
        "ImageNet-R",
        "ObjectNet",
        "Country211",
        "OK-VQA",
        "A-OKVQA",
        "DocVQA",
        "InfographicsVQA",
        "ChartQA",
        "Visual7W",
        "ScienceQA",
        "VizWiz",
        "GQA",
        "TextVQA",
        "VisDial",
        "CIRR",
        "VisualNews_t2i",
        "VisualNews_i2t",
        "MSCOCO_t2i",
        "MSCOCO_i2t",
        "NIGHTS",
        "WebQA",
        "FashionIQ",
        "Wiki-SS-NQ",
        "OVEN",
        "EDIS",
        "MSCOCO",
        "RefCOCO",
        "RefCOCO-Matching",
        "Visual7W-Pointing",
    ],
    "video": [
        "K700",
        "SmthSmthV2",
        "HMDB51",
        "UCF101",
        "Breakfast",
        "MVBench",
        "Video-MME",
        "NExTQA",
        "EgoSchema",
        "ActivityNetQA",
        "DiDeMo",
        "MSR-VTT",
        "MSVD",
        "VATEX",
        "YouCook2",
        "QVHighlight",
        "Charades-STA",
        "MomentSeeker",
    ],
    "visdoc": [
        "ViDoRe_arxivqa",
        "ViDoRe_docvqa",
        "ViDoRe_infovqa",
        "ViDoRe_tabfquad",
        "ViDoRe_tatdqa",
        "ViDoRe_shiftproject",
        "ViDoRe_syntheticDocQA_artificial_intelligence",
        "ViDoRe_syntheticDocQA_energy",
        "ViDoRe_syntheticDocQA_government_reports",
        "ViDoRe_syntheticDocQA_healthcare_industry",
        "ViDoRe_esg_reports_human_labeled_v2",
        "ViDoRe_biomedical_lectures_v2_multilingual",
        "ViDoRe_economics_reports_v2_multilingual",
        "ViDoRe_esg_reports_v2_multilingual",
        "VisRAG_ArxivQA",
        "VisRAG_ChartQA",
        "VisRAG_MP-DocVQA",
        "VisRAG_SlideVQA",
        "VisRAG_InfoVQA",
        "VisRAG_PlotQA",
        "ViDoSeek-page",
        "ViDoSeek-doc",
        "MMLongBench-page",
        "MMLongBench-doc",
    ],
}
modality2metric = {
    "image": "hit@1",
    "video": "hit@1",
    "visdoc": "ndcg_linear@5",
}
modalities = ["image", "video", "visdoc"]  # Process in this order

OUTPUT_BASEDIR = os.path.join(this_dir, "./outputs/evals")

# ===== Dataset-to-category mapping =====
CLS_DATASETS = {
    "ImageNet-1K",
    "N24News",
    "HatefulMemes",
    "VOC2007",
    "SUN397",
    "Place365",
    "ImageNet-A",
    "ImageNet-R",
    "ObjectNet",
    "Country211",
}
QA_DATASETS = {
    "OK-VQA",
    "A-OKVQA",
    "DocVQA",
    "InfographicsVQA",
    "ChartQA",
    "Visual7W",
    "ScienceQA",
    "VizWiz",
    "GQA",
    "TextVQA",
}
RET_DATASETS = {
    # i -> t
    "MSCOCO_i2t",
    "VisualNews_i2t",
    # t -> i
    "VisDial",
    "MSCOCO_t2i",
    "VisualNews_t2i",
    "WebQA",
    "EDIS",
    "Wiki-SS-NQ",
    # i -> i, grouped under RET
    "CIRR",
    "NIGHTS",
    "OVEN",
    "FashionIQ",
}
GROUNDING_DATASETS = {"MSCOCO", "RefCOCO", "RefCOCO-Matching", "Visual7W-Pointing"}


def get_category(dataset_name: str) -> str:
    if dataset_name in CLS_DATASETS:
        return "CLS"
    if dataset_name in QA_DATASETS:
        return "QA"
    if dataset_name in RET_DATASETS:
        return "RET"
    if dataset_name in GROUNDING_DATASETS:
        return "Grounding"
    return "Other"


def compute_category_means(
    modality: str, modality_data: dict, modality2dataset: dict, main_metric_key: str
) -> dict:

    category_scores = {}

    for dataset_name in modality2dataset.get(modality, []):
        score_info = modality_data.get(dataset_name)
        if isinstance(score_info, dict):
            metric_value = score_info.get(main_metric_key)
            if isinstance(metric_value, (int, float)):
                category = get_category(dataset_name)
                if category not in category_scores:
                    category_scores[category] = []
                category_scores[category].append(metric_value)

    category_means = {}
    for category, values in category_scores.items():
        if values:
            category_means[category] = sum(values) / len(values)
        else:
            category_means[category] = "N/A"

    return category_means


def summarize_experiment_scores(model_name, checkpoint_name):
    base_path = os.path.join(OUTPUT_BASEDIR, model_name, checkpoint_name)
    experiment_name_for_log = model_name + "_" + checkpoint_name

    current_experiment_scores = {}
    for modality in modalities:
        current_experiment_scores[modality] = {}
        modality_specific_result_dir = os.path.join(base_path, modality)

        for dataset_name in modality2dataset.get(modality, []):
            current_experiment_scores[modality][dataset_name] = "FILE_N/A"  # Initialize

        if not os.path.isdir(modality_specific_result_dir):
            print(f"    Directory not found: {modality_specific_result_dir}")
            for dataset_name in modality2dataset.get(modality, []):
                current_experiment_scores[modality][dataset_name] = "DIR_N/A"
            continue

        for filename in os.listdir(modality_specific_result_dir):
            if filename.endswith("_score.json"):
                score_file_path = os.path.join(modality_specific_result_dir, filename)
                dataset_name_from_file = None
                for known_dataset in modality2dataset.get(modality, []):
                    if filename == f"{known_dataset}_score.json":
                        dataset_name_from_file = known_dataset
                        break

                if dataset_name_from_file:
                    try:
                        with open(score_file_path, "r") as f:
                            score_data = json.load(f)
                            current_experiment_scores[modality][
                                dataset_name_from_file
                            ] = score_data
                    except json.JSONDecodeError:
                        print(f"      Error decoding JSON from {score_file_path}")
                        current_experiment_scores[modality][
                            dataset_name_from_file
                        ] = "JSON_ERROR"
                    except Exception as e:
                        print(f"      Error reading file {score_file_path}: {e}")
                        current_experiment_scores[modality][
                            dataset_name_from_file
                        ] = "READ_ERROR"

    # --- Construct and Save the Final JSON Report ---
    final_output = {
        "model_name": model_name,
        "checkpoint_name": checkpoint_name,
        "metrics": current_experiment_scores,
    }

    # --- Print average scores and missing datasets per modality ---
    # print(f"\n  --- Summary for Experiment: {experiment_name_for_log} ---")
    for modality in modalities:
        modality_info = {}

        if modality not in current_experiment_scores:
            print(f"    Modality '{modality.upper()}' not processed.")
            continue
        main_metric_key = modality2metric[modality]
        modality_data = current_experiment_scores[modality]
        collected_metric_values = []
        datasets_missing_score_file = []
        datasets_file_found_metric_missing = []

        for dataset_name in modality2dataset.get(modality, []):
            score_info = modality_data.get(dataset_name)
            if isinstance(score_info, dict):
                metric_value = score_info.get(main_metric_key)
                if isinstance(metric_value, (int, float)):
                    collected_metric_values.append(metric_value)
                else:
                    datasets_file_found_metric_missing.append(
                        f"{dataset_name} (metric '{main_metric_key}' missing/invalid)"
                    )
            else:
                datasets_missing_score_file.append(
                    f"{dataset_name} (status: {score_info if score_info else 'Not Processed'})"
                )

        modality_info["n_datasets"] = len(collected_metric_values)
        if collected_metric_values:
            average_score = sum(collected_metric_values) / len(collected_metric_values)
            modality_info[main_metric_key] = average_score
        else:
            modality_info[main_metric_key] = "N/A"

        if datasets_missing_score_file:
            modality_info["missing_files"] = datasets_missing_score_file
        if datasets_file_found_metric_missing:
            modality_info["missing_metrics"] = datasets_file_found_metric_missing

        if modality.upper() == "IMAGE" and not datasets_missing_score_file:
            modality_info["category_means"] = compute_category_means(
                modality, modality_data, modality2dataset, main_metric_key
            )

        final_output[modality.upper()] = modality_info

    output_json_path = os.path.join(base_path, f"final_results.json")
    try:
        with open(output_json_path, "w") as f:
            json.dump(final_output, f, indent=4)
        print(f"  Report for '{experiment_name_for_log}' saved to: {output_json_path}")
    except Exception as e:
        print(
            f"  Error saving JSON report for '{experiment_name_for_log}' to {output_json_path}: {e}"
        )


def eval_cmd(model_name, checkpoint_name):
    if "Qwen2.5vl" in model_name:
        model_backbone = "qwen2_5_vl"
    elif "Qwen2vl" in model_name:
        model_backbone = "qwen2_vl"
    else:
        raise NotImplementedError

    cmd_parameters = [
        f'MODEL_NAME="{model_name}"',
        f'CHECKPOINT_NAME="{checkpoint_name}"',
        f'MODEL_BACKBONE="{model_backbone}"',
    ]

    if "TgtInstruction" in model_name:
        cmd_parameters.append(f"TGT_PREFIX_INSTRUCTION=true")

    if "PreMLP" in model_name:
        cmd_parameters.append(f"POOLING=pre_mlp")

    if "Residual" in model_name:
        cmd_parameters.append(f"RESIDUAL_EMBEDDING=true")
        if "ResidualAvg" in model_name:
            cmd_parameters.append(f"RESIDUAL_EMBEDDING_METHOD=avg")
        if "ResidualMax" in model_name:
            cmd_parameters.append(f"RESIDUAL_EMBEDDING_METHOD=max")
        if "ResidualResidual" in model_name:
            cmd_parameters.append(f"RESIDUAL_EMBEDDING_METHOD=residual")
        if "ResidualGem" in model_name:
            cmd_parameters.append(f"RESIDUAL_EMBEDDING_METHOD=gem")

    if "ChatTemplate" in model_name:
        cmd_parameters.append(f"CHAT_TEMPLATE=true")

    match = re.search(r"Queries(\d+)", model_name)
    if match:
        num_queries = match.group(1)
        cmd_parameters.append("LEARNABLE_QUERIES=true")
        cmd_parameters.append(f"NUM_QUERIES={num_queries}")

    cmd_str = " ".join(cmd_parameters)

    project_root = this_dir
    os.system(
        f"""cd {project_root} && {cmd_str} bash experiments/public/eval/eval_rematch.sh"""
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_name", type=str, default=None)
    parser.add_argument("--checkpoint_name", type=str, default=None)
    args = parser.parse_args()

    if args.model_name and args.checkpoint_name:
        try:
            for checkpoint_name in args.checkpoint_name.split(","):
                eval_cmd(args.model_name, checkpoint_name)
                summarize_experiment_scores(args.model_name, checkpoint_name)
        except Exception as e:
            print(e)
        return

    outputs_dir = os.path.join(this_dir, "./outputs")
    all_model_checkpoints = glob.glob(os.path.join(outputs_dir, "*/*"))
    for intermediate_checkpoint in all_model_checkpoints:
        model_name = intermediate_checkpoint.split(os.sep)[-2]
        checkpoint_name = intermediate_checkpoint.split(os.sep)[-1]

        if checkpoint_name.startswith("checkpoint"):
            try:
                eval_cmd(model_name, checkpoint_name)
                summarize_experiment_scores(model_name, checkpoint_name)
            except Exception as e:
                print(e)


if __name__ == "__main__":
    main()