code/evaluate.py

"""Evaluation script for RS-IMLE (CIFAR-10 and CelebA-HQ-256).

Computes FID (5000 samples) and Precision/Recall (1000 samples), and saves a
sample grid. Cycle counts (H, L, refinement_steps) can be overridden at test time.
"""

import argparse
import json
import os
import shutil
import time
from distutils.util import strtobool

import imageio
import numpy as np
import torch
import torchvision

from data import set_up_data
from hps import Hyperparams, add_imle_arguments, parse_args_and_update_hparams
from models import IMLE
from torch import autocast


# ──────────────────────────────────────────────────────────────────────────
# cleanfid monkey-patch 
# ──────────────────────────────────────────────────────────────────────────
from cleanfid import fid
import cleanfid.features as _cleanfid_feat
import cleanfid.inception_torchscript as _cleanfid_incept

_inception_cache = os.path.join(os.path.expanduser("~"), ".cache", "cleanfid")
os.makedirs(_inception_cache, exist_ok=True)
_orig_feature_extractor = _cleanfid_feat.feature_extractor


def _patched_feature_extractor(name="torchscript_inception",
                               device=torch.device("cuda"),
                               resize_inside=False, use_dataparallel=True):
    if name == "torchscript_inception":
        model = _cleanfid_incept.InceptionV3W(
            _inception_cache, download=True, resize_inside=resize_inside
        ).to(device)
        model.eval()
        if use_dataparallel:
            model = torch.nn.DataParallel(model)
        return lambda x: model(x)
    return _orig_feature_extractor(name, device, resize_inside, use_dataparallel)


_cleanfid_feat.feature_extractor = _patched_feature_extractor


# ──────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────

def generate_images_to_dir(imle_model, latent_dim, num_images, out_dir,
                           batch_size=16):
    """Generate num_images PNG files into out_dir for FID/P-R evaluation."""
    os.makedirs(out_dir, exist_ok=True)
    device = next(imle_model.parameters()).device
    idx = 0
    with torch.no_grad():
        while idx < num_images:
            bs = min(batch_size, num_images - idx)
            z = torch.randn(bs, latent_dim, device=device)
            with autocast(device_type="cuda"):
                imgs = imle_model(z, None)
            # Model output is in [-1, 1], convert to [0, 255] uint8
            imgs = (imgs + 1.0) * 127.5
            imgs = imgs.clamp(0, 255).permute(0, 2, 3, 1)
            imgs = imgs.cpu().numpy().astype(np.uint8)
            for j in range(bs):
                imageio.imwrite(os.path.join(out_dir, f"{idx}.png"), imgs[j])
                idx += 1


def generate_grid_image(imle_model, latent_dim, num_samples, nrow, out_path):
    """Generate a grid of num_samples images and save to out_path."""
    device = next(imle_model.parameters()).device
    with torch.no_grad():
        z = torch.randn(num_samples, latent_dim, device=device)
        with autocast(device_type="cuda"):
            imgs = imle_model(z, None)
        # Convert [-1, 1] -> [0, 1] for torchvision grid
        imgs = (imgs + 1.0) / 2.0
        imgs = imgs.clamp(0.0, 1.0)
        grid = torchvision.utils.make_grid(imgs, nrow=nrow, padding=2)
        grid_pil = torchvision.transforms.functional.to_pil_image(
            grid.cpu())
        grid_pil.save(out_path)
    print(
        f"Saved sample grid ({num_samples} images, {nrow} per row) to {out_path}")


def override_model_cycles(imle_model, test_H_cycles=None, test_L_cycles=None,
                          test_refinement_steps=None):
    """Override H_cycles / L_cycles / refinement_steps inside the loaded model."""
    model = imle_model.module if hasattr(imle_model, 'module') else imle_model
    mapper = model.decoder.mapping_network

    inner = getattr(mapper, 'trm', None)
    if inner is None:
        print("WARNING: No TRM inner module found; cycle override skipped.")
        return

    old_H = inner.H_cycles
    old_L = inner.L_cycles
    old_halt = mapper.refinement_steps

    if test_H_cycles is not None:
        inner.H_cycles = test_H_cycles
    if test_L_cycles is not None:
        inner.L_cycles = test_L_cycles
    if test_refinement_steps is not None:
        mapper.refinement_steps = test_refinement_steps

    print(f"Cycle override: H_cycles {old_H}->{inner.H_cycles}, "
          f"L_cycles {old_L}->{inner.L_cycles}, "
          f"refinement_steps {old_halt}->{mapper.refinement_steps}")


def build_and_load_model(H):
    """Build IMLE model and load checkpoint weights.

    Returns (imle_model, ema_model_or_None).
    """
    imle_model = IMLE(H).cuda()
    ema_model = None

    if H.restore_path and os.path.isfile(H.restore_path):
        print(f"Loading model weights from: {H.restore_path}")
        state_dict = torch.load(H.restore_path, map_location='cuda')

        # Handle DataParallel-saved checkpoints (keys start with 'module.')
        has_module_prefix = any(k.startswith('module.')
                                for k in state_dict.keys())
        if has_module_prefix:
            # Wrap model in DataParallel first, then load
            imle_model = torch.nn.DataParallel(imle_model)
            imle_model.load_state_dict(state_dict, strict=bool(H.load_strict))
        else:
            imle_model.load_state_dict(state_dict, strict=bool(H.load_strict))
            imle_model = torch.nn.DataParallel(imle_model)
    else:
        imle_model = torch.nn.DataParallel(imle_model)
        if H.restore_path:
            print(f"WARNING: restore_path '{H.restore_path}' not found!")
        else:
            print("WARNING: No --restore_path specified, using random weights!")

    # Load EMA if requested
    if getattr(H, 'restore_ema_path', None) and os.path.isfile(H.restore_ema_path):
        print(f"Loading EMA weights from: {H.restore_ema_path}")
        ema_model = IMLE(H).cuda()
        ema_state = torch.load(H.restore_ema_path, map_location='cuda')
        ema_model.load_state_dict(ema_state, strict=bool(H.load_strict))

    return imle_model, ema_model


# ──────────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────────

def main():
    # ── Parse arguments ──
    H = Hyperparams()
    parser = argparse.ArgumentParser(
        description="RS-IMLE Evaluation: FID, Precision/Recall, Sample Grid",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser = add_imle_arguments(parser)

    # Test-specific arguments
    parser.add_argument('--test_H_cycles', type=int, default=None,
                        help='Override H_cycles at test time (default: use trained value)')
    parser.add_argument('--test_L_cycles', type=int, default=None,
                        help='Override L_cycles at test time (default: use trained value)')
    # Note: --test_refinement_steps and --num_fid_samples are registered by
    # add_imle_arguments() above; only test-time-only knobs are added here.
    parser.add_argument('--num_pr_samples', type=int, default=1000,
                        help='Number of generated samples for Precision/Recall')
    parser.add_argument('--num_grid_samples', type=int, default=40,
                        help='Number of samples in the output grid image')
    parser.add_argument('--grid_nrow', type=int, default=8,
                        help='Number of images per row in grid (default 8, so 8x5=40)')
    parser.add_argument('--output_dir', type=str, default=None,
                        help='Output directory (default: {save_dir}/test_results)')
    parser.add_argument('--use_ema', default=False,
                        type=lambda x: bool(strtobool(x)),
                        help='Evaluate EMA model instead of main model')
    parser.add_argument('--skip_fid', default=False,
                        type=lambda x: bool(strtobool(x)),
                        help='Skip FID computation')
    parser.add_argument('--skip_pr', default=False,
                        type=lambda x: bool(strtobool(x)),
                        help='Skip Precision/Recall computation')
    parser.add_argument('--test_batch_size', type=int, default=128,
                        help='Batch size for generating images during test')
    parser.add_argument('--dump_samples_dir', type=str, default=None,
                        help='If set, dump --dump_samples_n PNGs here. '
                             'Combine with --skip_fid True --skip_pr True to skip metrics.')
    parser.add_argument('--dump_samples_n', type=int, default=256,
                        help='Number of PNGs to save when --dump_samples_dir is set.')

    parse_args_and_update_hparams(H, parser)

    # ── Setup output directory ──
    if H.output_dir:
        output_dir = H.output_dir
    else:
        output_dir = os.path.join(H.save_dir, 'test_results')
    os.makedirs(output_dir, exist_ok=True)

    # ── Setup data (needed for FID reference path) ──
    H, data_train, data_valid, preprocess_fn = set_up_data(H)

    print("=" * 60)
    print("RS-IMLE Evaluation")
    print("=" * 60)
    print(f"  Dataset:        {H.dataset}")
    print(f"  Data root:      {H.data_root}")
    print(f"  Checkpoint:     {H.restore_path}")
    print(f"  use_rtm:        {getattr(H, 'use_rtm', False)}")
    print(f"  H_cycles:       {H.H_cycles}")
    print(f"  L_cycles:       {H.L_cycles}")
    print(f"  refinement_steps: {H.refinement_steps}")
    if H.test_H_cycles is not None:
        print(f"  test_H_cycles:  {H.test_H_cycles}  (OVERRIDE)")
    if H.test_L_cycles is not None:
        print(f"  test_L_cycles:  {H.test_L_cycles}  (OVERRIDE)")
    if H.test_refinement_steps is not None:
        print(f"  test_refinement_steps: {H.test_refinement_steps}  (OVERRIDE)")
    print(f"  Output dir:     {output_dir}")
    print(f"  FID samples:    {H.num_fid_samples}")
    print(f"  P/R samples:    {H.num_pr_samples}")
    print(f"  Grid samples:   {H.num_grid_samples}")
    print("=" * 60)

    # ── Build and load model ──
    imle_model, ema_model = build_and_load_model(H)

    # Decide which model to evaluate
    if H.use_ema and ema_model is not None:
        eval_model = ema_model
        print("Evaluating EMA model")
    else:
        eval_model = imle_model
        if H.use_ema and ema_model is None:
            print("WARNING: --use_ema True but no EMA weights loaded, using main model")
        print("Evaluating main model")

    # ── Override cycles at test time ──
    override_model_cycles(
        eval_model,
        test_H_cycles=H.test_H_cycles,
        test_L_cycles=H.test_L_cycles,
        test_refinement_steps=H.test_refinement_steps,
    )

    eval_model.eval()

    batch_size = min(getattr(H, 'test_batch_size', 128), H.num_fid_samples)
    results = {}

    if H.dump_samples_dir:
        print(f"\nDumping {H.dump_samples_n} PNG samples to {H.dump_samples_dir}")
        generate_images_to_dir(eval_model, H.latent_dim, H.dump_samples_n,
                               H.dump_samples_dir, batch_size=batch_size)

    # ── 1. Sample Grid ──
    print("\n[1/3] Generating sample grid...")
    grid_path = os.path.join(output_dir, "samples_grid.png")
    generate_grid_image(eval_model, H.latent_dim, H.num_grid_samples,
                        H.grid_nrow, grid_path)

    # ── 2. FID ──
    if not H.skip_fid:
        print(f"\n[2/3] Computing FID ({H.num_fid_samples} samples)...")
        fid_dir = os.path.join(output_dir, "fid")
        os.makedirs(fid_dir, exist_ok=True)
        t0 = time.time()
        generate_images_to_dir(eval_model, H.latent_dim, H.num_fid_samples,
                               fid_dir, batch_size=batch_size)
        ref_dir = f'{H.data_root}/img'
        print(f"  Reference dir: {ref_dir}")
        print(f"  Generated dir: {fid_dir}")
        cur_fid = fid.compute_fid(
            ref_dir, fid_dir, verbose=False, num_workers=0)
        results['fid'] = cur_fid
        print(f"  FID = {cur_fid:.4f}  ({time.time() - t0:.1f}s)")
        shutil.rmtree(fid_dir, ignore_errors=True)
    else:
        print("\n[2/3] FID skipped (--skip_fid True)")

    # ── 3. Precision / Recall ──
    if not H.skip_pr:
        print(
            f"\n[3/3] Computing Precision/Recall ({H.num_pr_samples} samples)...")
        pr_dir = os.path.join(output_dir, "prec_rec")
        os.makedirs(pr_dir, exist_ok=True)
        t0 = time.time()
        generate_images_to_dir(eval_model, H.latent_dim, H.num_pr_samples,
                               pr_dir, batch_size=batch_size)
        try:
            from helpers.improved_precision_recall import compute_prec_recall
            ref_dir = f'{H.data_root}/img'
            precision, recall = compute_prec_recall(ref_dir, pr_dir)
            results['precision'] = precision
            results['recall'] = recall
            print(f"  Precision = {precision:.4f}")
            print(f"  Recall    = {recall:.4f}")
            print(f"  ({time.time() - t0:.1f}s)")
        except ImportError:
            print("  WARNING: helpers.improved_precision_recall not found, skipping P/R")
        shutil.rmtree(pr_dir, ignore_errors=True)
    else:
        print("\n[3/3] Precision/Recall skipped (--skip_pr True)")

    # ── Save results ──
    # Get final cycle values (after override)
    model_unwrapped = eval_model.module if hasattr(
        eval_model, 'module') else eval_model
    mapper = model_unwrapped.decoder.mapping_network
    inner = getattr(mapper, 'trm', None)

    results['config'] = {
        'restore_path': H.restore_path,
        'dataset': H.dataset,
        'data_root': H.data_root,
        'use_rtm': bool(getattr(H, 'use_rtm', False)),
        'latent_dim': H.latent_dim,
        'H_cycles_trained': H.H_cycles,
        'L_cycles_trained': H.L_cycles,
        'refinement_steps_trained': H.refinement_steps,
        'H_cycles_eval': inner.H_cycles if inner else H.H_cycles,
        'L_cycles_eval': inner.L_cycles if inner else H.L_cycles,
        'refinement_steps_eval': mapper.refinement_steps if hasattr(mapper, 'refinement_steps') else H.refinement_steps,
        'num_fid_samples': H.num_fid_samples,
        'num_pr_samples': H.num_pr_samples,
    }

    results_path = os.path.join(output_dir, "test_metrics.json")
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nResults saved to {results_path}")

    print("\n" + "=" * 60)
    print("RESULTS SUMMARY")
    print("=" * 60)
    if 'fid' in results:
        print(f"  FID:       {results['fid']:.4f}")
    if 'precision' in results:
        print(f"  Precision: {results['precision']:.4f}")
    if 'recall' in results:
        print(f"  Recall:    {results['recall']:.4f}")
    if 'ada_fid' in results:
        print(f"  Ada FID:   {results['ada_fid']:.5f}")
    if 'ada_sfid' in results:
        print(f"  Ada sFID:  {results['ada_sfid']:.5f}")
    if 'ada_inception_score' in results:
        print(f"  Ada IS:    {results['ada_inception_score']:.5f}")
    if 'ada_precision' in results:
        print(f"  Ada Prec:  {results['ada_precision']:.4f}")
    if 'ada_recall' in results:
        print(f"  Ada Recall:{results['ada_recall']:.4f}")
    print(f"  Grid:      {grid_path}")
    print("=" * 60)


if __name__ == "__main__":
    main()