scripts/eval_taic.py

#!/usr/bin/env python3
"""Test / eval for TAIC: codec stats + optional full task-network metrics.

Codec-only (default):
  bpp, feature distortion D, loss; optional --actual-bpp

With task metrics (--with-metrics):
  also load official task-network config/checkpoint, run truncated task net from h,
  report mAP-bbox / mAP-mask / mIoU / PQ / mAP-OKS (task-dependent).

Examples:
  python scripts/eval_taic.py -c configs/eval/taic_detection.yaml
  python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics
  python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics --max-batches 50
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from datetime import datetime

import torch
from torch.utils.data import DataLoader

REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

from flexicm.data import (
    COCOImageDataset,
    COCOWholeBodyImageDataset,
    ImageFolderDataset,
    build_test_transform,
)
from flexicm.data.coco_eval import COCOEvalDataset, TASK_ANN_FILES, coco_eval_collate
from flexicm.models import TAIC
from flexicm.tasks import TASK_META, build_teacher
from flexicm.tasks.losses import TAICCriterion
from flexicm.tasks.metric_eval import run_task_metric_eval
from flexicm.tasks.metric_runners import (
    DEFAULT_TASK_NET_CKPTS,
    DEFAULT_TASK_NET_CONFIGS,
    build_metric_runner,
)
from flexicm.utils.codec_test import resolve_ckpt, test_taic_loader
from flexicm.utils.train_utils import load_checkpoint_dict, load_yaml_config, set_seed


def parse_args(argv):
    parser = argparse.ArgumentParser("Test FlexICM TAIC (codec + optional task metrics)")
    parser.add_argument("-c", "--config", required=True, help="configs/eval/taic_*.yaml")
    given, remaining = parser.parse_known_args(argv)
    cfg_path = given.config if os.path.isabs(given.config) else os.path.join(REPO_ROOT, given.config)
    cfg = load_yaml_config(cfg_path)
    # -c already consumed by the first parse; keep it as a default for the second pass
    parser.set_defaults(config=cfg_path, **cfg)
    for action in parser._actions:
        if "--config" in action.option_strings:
            action.required = False
            break
    parser.add_argument("--actual-bpp", action="store_true")
    parser.add_argument("--with-metrics", action="store_true", help="Run full task-network metrics")
    parser.add_argument("--max-batches", type=int, default=None)
    parser.add_argument("--split", type=str, default=None)
    parser.add_argument(
        "--eval-size",
        type=int,
        default=cfg.get("eval_size", None),
        help="If set (e.g. 256), resize eval images to eval_size×eval_size",
    )
    args = parser.parse_args(remaining)
    args.config = cfg_path
    if "--actual-bpp" in argv:
        args.actual_bpp = True
    if "--with-metrics" in argv:
        args.with_metrics = True
    return args


def build_codec_loader(args, device):
    split = args.split or getattr(args, "split", None) or "val2017"
    eval_size = getattr(args, "eval_size", None)
    tf = build_test_transform(eval_size=eval_size)
    root = args.dataset_path
    split_dir = os.path.join(root, split)
    if os.path.isdir(split_dir):
        if args.task == "pose":
            dataset = COCOWholeBodyImageDataset(root, split, tf)
        else:
            dataset = COCOImageDataset(root, split, tf)
    else:
        dataset = ImageFolderDataset(root, tf)
    return DataLoader(
        dataset,
        batch_size=getattr(args, "test_batch_size", 1),
        shuffle=False,
        num_workers=getattr(args, "num_workers", 4),
        pin_memory=(device == "cuda"),
    )


def build_metric_loader(args, device):
    split = args.split or getattr(args, "split", None) or "val2017"
    ann_rel = getattr(args, "ann_file", None) or TASK_ANN_FILES.get(args.task)
    if ann_rel is None:
        raise ValueError(f"No ann_file for task={args.task}")
    ann_file = ann_rel if os.path.isabs(ann_rel) else os.path.join(args.dataset_path, ann_rel)
    eval_size = getattr(args, "eval_size", None)
    dataset = COCOEvalDataset(
        args.dataset_path,
        ann_file=ann_file,
        image_prefix=split,
        transform=build_test_transform(eval_size=eval_size),
    )
    loader = DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        num_workers=getattr(args, "num_workers", 4),
        pin_memory=(device == "cuda"),
        collate_fn=coco_eval_collate,
    )
    return loader, ann_file


def main(argv):
    args = parse_args(argv)
    set_seed(getattr(args, "seed", 42))

    os.environ["CUDA_VISIBLE_DEVICES"] = str(getattr(args, "gpu_id", 0))
    device = "cuda" if getattr(args, "cuda", True) and torch.cuda.is_available() else "cpu"

    task = args.task
    meta = TASK_META[task]
    out_channels = getattr(args, "out_channels", meta["out_channels"])
    align_mode = getattr(args, "align_mode", meta["align_mode"])
    lmbda = getattr(args, "lmbda", 0.0035)

    ckpt = resolve_ckpt(args.checkpoint, REPO_ROOT, label="TAIC checkpoint")
    print(f"Loading TAIC checkpoint: {ckpt}")

    net = TAIC(N=128, M=192, out_channels=out_channels).to(device)
    state, _ = load_checkpoint_dict(ckpt, map_location=device)
    missing = net.load_state_dict(state, strict=False)
    print(f"load_state_dict: missing={len(missing.missing_keys)} unexpected={len(missing.unexpected_keys)}")
    net.eval()

    # ---- codec test ----
    teacher = build_teacher(
        task,
        pretrained_backbone=getattr(args, "pretrained_backbone", True),
        use_official_teacher=getattr(args, "use_official_teacher", True),
        task_config=getattr(args, "task_config", None),
        task_checkpoint=getattr(args, "task_checkpoint", None),
        device=device,
    )
    teacher = teacher.to(device).eval()
    criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode)
    codec_loader = build_codec_loader(args, device)
    eval_size = getattr(args, "eval_size", None)
    size_msg = f"{eval_size}x{eval_size}" if eval_size else "original"
    print(f"[codec] test set size: {len(codec_loader.dataset)}  device={device}  task={task}  input={size_msg}")

    codec_result = test_taic_loader(
        net,
        teacher,
        codec_loader,
        criterion,
        device,
        align_divisor=256,
        run_actual_bpp=bool(getattr(args, "actual_bpp", False)),
        max_batches=args.max_batches,
    )
    print("==== TAIC codec test summary ====")
    for k, v in codec_result.items():
        print(f"  {k}: {v:.6f}" if isinstance(v, float) else f"  {k}: {v}")

    payload = {
        "task": task,
        "checkpoint": ckpt,
        "config": args.config,
        "codec_result": codec_result,
    }

    # ---- optional task metrics ----
    if getattr(args, "with_metrics", False):
        task_cfg = getattr(args, "task_config", None) or DEFAULT_TASK_NET_CONFIGS[task]
        task_ckpt = getattr(args, "task_checkpoint", None) or DEFAULT_TASK_NET_CKPTS[task]
        if not os.path.isabs(task_cfg):
            task_cfg = os.path.join(REPO_ROOT, task_cfg)
        task_ckpt = resolve_ckpt(task_ckpt, REPO_ROOT, label=f"{task} task-network checkpoint")

        print(f"[metric] loading task network:\n  config={task_cfg}\n  ckpt={task_ckpt}")
        runner = build_metric_runner(task, device=device)
        runner.load(task_cfg, task_ckpt)

        metric_loader, ann_file = build_metric_loader(args, device)
        print(f"[metric] COCO eval images: {len(metric_loader.dataset)}  ann={ann_file}")
        metrics = run_task_metric_eval(
            net,
            runner,
            metric_loader,
            device,
            ann_file=ann_file,
            use_condition=False,
            base_codec=None,
            max_batches=args.max_batches,
            finalize_kwargs={
                "gt_folder": getattr(args, "panoptic_gt_folder", None),
                "pred_folder": getattr(args, "panoptic_pred_folder", None),
                "num_classes": getattr(args, "num_classes", 133),
            },
        )
        print("==== TAIC task metric summary ====")
        for k, v in metrics.items():
            print(f"  {k}: {v}")
        payload["task_config"] = task_cfg
        payload["task_checkpoint"] = task_ckpt
        payload["task_metrics"] = metrics

    out_dir = getattr(args, "result_dir", None) or os.path.join(
        REPO_ROOT, "logs", "eval_taic", task, str(getattr(args, "quality_level", 1))
    )
    os.makedirs(out_dir, exist_ok=True)
    out_json = os.path.join(out_dir, f"eval_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
    with open(out_json, "w") as f:
        json.dump(payload, f, indent=2, default=str)
    print(f"Wrote {out_json}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv[1:]))