diffusionsfm/eval/eval_jobs.py

"""
python -m diffusionsfm.eval.eval_jobs --eval_path output/multi_diffusionsfm_dense --use_submitit
"""

import os
import json
import submitit
import argparse
import itertools
from glob import glob

import numpy as np
from tqdm.auto import tqdm

from diffusionsfm.dataset.co3d_v2 import TEST_CATEGORIES, TRAINING_CATEGORIES
from diffusionsfm.eval.eval_category import save_results
from diffusionsfm.utils.slurm import submitit_job_watcher


def evaluate_diffusionsfm(eval_path, use_submitit, mode):
    JOB_PARAMS = {
        "output_dir": [eval_path],
        "checkpoint": [800_000],
        "num_images": [2, 3, 4, 5, 6, 7, 8],
        "sample_num": [0, 1, 2, 3, 4],
        "category": TEST_CATEGORIES, # TRAINING_CATEGORIES + TEST_CATEGORIES,
        "calculate_additional_timesteps": [True],
    }
    if mode == "test":
        JOB_PARAMS["category"] = TEST_CATEGORIES
    elif mode == "train1":
        JOB_PARAMS["category"] = TRAINING_CATEGORIES[:len(TRAINING_CATEGORIES) // 2]
    elif mode == "train2":
        JOB_PARAMS["category"] = TRAINING_CATEGORIES[len(TRAINING_CATEGORIES) // 2:]
    keys, values = zip(*JOB_PARAMS.items())
    job_configs = [dict(zip(keys, p)) for p in itertools.product(*values)]

    if use_submitit:
        log_output = "./slurm_logs"
        executor = submitit.AutoExecutor(
            cluster=None, folder=log_output, slurm_max_num_timeout=10
        )
        # Use your own parameters
        executor.update_parameters(
            slurm_additional_parameters={
                "nodes": 1,
                "cpus-per-task": 5,
                "gpus": 1,
                "time": "6:00:00",
                "partition": "all",
                "exclude": "grogu-1-9, grogu-1-14,"
            }
        )
        jobs = []
        with executor.batch():
            # This context manager submits all jobs at once at the end.
            for params in job_configs:
                job = executor.submit(save_results, **params)
                job_param = f"{params['category']}_N{params['num_images']}_{params['sample_num']}"
                jobs.append((job_param, job))
        jobs = {f"{job_param}_{job.job_id}": job for job_param, job in jobs}
        submitit_job_watcher(jobs)
    else:
        for job_config in tqdm(job_configs):
            # This is much slower.
            save_results(**job_config)


def process_predictions(eval_path, pred_index, checkpoint=800_000, threshold_R=15, threshold_CC=0.1):
    """
    pred_index should be 1 (corresponding to T=90)
    """
    def aggregate_per_category(categories, metric_key, num_images, sample_num, threshold=None):
        """
        Aggregates one metric over all data points in a prediction file and then across categories.
        - For R_error and CC_error: use mean to threshold-based accuracy
        - For CD and CD_Object: use median to reduce the effect of outliers
        """
        per_category_values = []

        for category in tqdm(categories, desc=f"Sample {sample_num}, N={num_images}, {metric_key}"):
            per_pred_values = []

            data_path = glob(
                os.path.join(eval_path, "eval", f"{category}_{num_images}_{sample_num}_ckpt{checkpoint}*.json")
            )[0]

            with open(data_path) as f:
                eval_data = json.load(f)

            for preds in eval_data.values():
                if metric_key in ["R_error", "CC_error"]:
                    vals = np.array(preds[pred_index][metric_key])
                    per_pred_values.append(np.mean(vals < threshold))
                else: 
                    per_pred_values.append(preds[pred_index][metric_key])

            # Aggregate over all predictions within this category
            per_category_values.append(
                np.mean(per_pred_values) if metric_key in ["R_error", "CC_error"]
                else np.median(per_pred_values)  # CD or CD_Object — use median to filter outliers
            )

        if metric_key in ["R_error", "CC_error"]:
            return np.mean(per_category_values)
        else:
            return np.median(per_category_values)

    def aggregate_metric(categories, metric_key, num_images, threshold=None):
        """Aggregates one metric over 5 random samples per category and returns the final mean"""
        return np.mean([
            aggregate_per_category(categories, metric_key, num_images, sample_num, threshold=threshold)
            for sample_num in range(5)
        ])

    # Output containers
    all_seen_acc_R, all_seen_acc_CC = [], []
    all_seen_CD, all_seen_CD_Object = [], []
    all_unseen_acc_R, all_unseen_acc_CC = [], []
    all_unseen_CD, all_unseen_CD_Object = [], []

    for num_images in range(2, 9):
        # Seen categories
        all_seen_acc_R.append(
            aggregate_metric(TRAINING_CATEGORIES, "R_error", num_images, threshold=threshold_R)
        )
        all_seen_acc_CC.append(
            aggregate_metric(TRAINING_CATEGORIES, "CC_error", num_images, threshold=threshold_CC)
        )
        all_seen_CD.append(
            aggregate_metric(TRAINING_CATEGORIES, "CD", num_images)
        )
        all_seen_CD_Object.append(
            aggregate_metric(TRAINING_CATEGORIES, "CD_Object", num_images)
        )

        # Unseen categories
        all_unseen_acc_R.append(
            aggregate_metric(TEST_CATEGORIES, "R_error", num_images, threshold=threshold_R)
        )
        all_unseen_acc_CC.append(
            aggregate_metric(TEST_CATEGORIES, "CC_error", num_images, threshold=threshold_CC)
        )
        all_unseen_CD.append(
            aggregate_metric(TEST_CATEGORIES, "CD", num_images)
        )
        all_unseen_CD_Object.append(
            aggregate_metric(TEST_CATEGORIES, "CD_Object", num_images)
        )

    # Print the results in formatted rows
    print("N=           ", " ".join(f"{i: 5}" for i in range(2, 9)))
    print("Seen R       ", " ".join([f"{x:0.3f}" for x in all_seen_acc_R]))
    print("Seen CC      ", " ".join([f"{x:0.3f}" for x in all_seen_acc_CC]))
    print("Seen CD      ", " ".join([f"{x:0.3f}" for x in all_seen_CD]))
    print("Seen CD_Obj  ", " ".join([f"{x:0.3f}" for x in all_seen_CD_Object]))
    print("Unseen R     ", " ".join([f"{x:0.3f}" for x in all_unseen_acc_R]))
    print("Unseen CC    ", " ".join([f"{x:0.3f}" for x in all_unseen_acc_CC]))
    print("Unseen CD    ", " ".join([f"{x:0.3f}" for x in all_unseen_CD]))
    print("Unseen CD_Obj", " ".join([f"{x:0.3f}" for x in all_unseen_CD_Object]))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--eval_path", type=str, default=None)
    parser.add_argument("--use_submitit", action="store_true")
    parser.add_argument("--mode", type=str, default="test")
    args = parser.parse_args()

    eval_path = "output/multi_diffusionsfm_dense" if args.eval_path is None else args.eval_path
    use_submitit = args.use_submitit
    mode = args.mode

    evaluate_diffusionsfm(eval_path, use_submitit, mode)
    process_predictions(eval_path, 1)