train_gs.py

"""
Stage 3.2: Train a 2DGS or 3DGS model on the aligned canonical point cloud.

Initialises SH0 from canonical RGB and renders via gsplat.
The ``--config.renderer`` flag selects between 2DGS and 3DGS backends.
"""

from __future__ import annotations

import json
import math
import os
import subprocess
import sys
from typing import TYPE_CHECKING

import cv2
import numpy as np
import open3d as o3d
import torch
import tyro
from tqdm.auto import tqdm

from configs.stage3_gs import GSConfig
from data.checkpoint_loading import (
    AlignmentDataParams,
    load_aligned_point_cloud,
    load_alignment_data_params,
    load_deformation_checkpoints,
    load_inverse_local_from_checkpoint,
)
from data.data_loading import load_da3_original_images_from_folder, load_data
from losses import init_lpips
from models.canonical_gs_model import CanonicalGSModel, sh0_to_rgb
from utils.confidence import build_pixel_confidence_weights, resolve_confidence_weighting, weighted_mean
from utils.density_control import (
    collect_gaussian_parameters,
    create_density_strategy,
    enforce_default_strategy_budget,
    resize_model_gaussians_from_state,
    robust_scene_scale,
    sync_gaussian_parameters,
)
from utils.downsample import downsample_to_target
from utils.dynamic_mask import (
    build_visibility_aware_residual_weights,
    dynamic_mask_ramp,
    update_residual_ema_,
    validate_dynamic_mask_params,
)
from utils.eval_metrics import flow_guided_temporal_l1, image_psnr, image_ssim, interleaved_frame_split
from utils.export_checkpoint_to_ply import (
    ExportGSCheckpointToPlyConfig,
)
from utils.export_checkpoint_to_ply import (
    main as export_gs_checkpoint_to_ply,
)
from utils.knn import query_knn_with_backend
from utils.logging import get_logger, tb_log_hparams, try_create_tensorboard_writer
from utils.normals import estimate_normals
from utils.reconstruction_losses import dssim_loss, edge_l1_loss, multiscale_l1_loss
from utils.topology_reliability import (
    build_motion_consistency_maps,
    combine_reliability_maps,
    gate_density_gradients,
    load_feature_consistency_maps,
)

logger = get_logger(__name__)

if TYPE_CHECKING:
    from torch.utils.tensorboard import SummaryWriter


# ==========================================================================
# Main
# ==========================================================================


def main(config: GSConfig):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    logger.info("Device: %s", device)

    torch.manual_seed(config.seed)
    np.random.seed(config.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config.seed)
    logger.info("Random seed: %d", config.seed)

    if config.rasterize_eps2d <= 0:
        raise ValueError("rasterize_eps2d must be positive")
    if config.mip_filter_3d:
        if config.renderer != "3dgs":
            raise ValueError("Mip 3D filtering is only supported by renderer='3dgs'")
        if config.mip_filter_variance < 0 or config.mip_filter_update_every < 0:
            raise ValueError("Mip filter variance/update interval must be non-negative")
    if config.appearance_mode not in {"off", "affine", "embedding"}:
        raise ValueError(f"Unknown appearance_mode: {config.appearance_mode!r}")
    if config.appearance_mode != "off":
        if config.appearance_embedding_dim <= 0 or config.appearance_hidden_dim <= 0:
            raise ValueError("appearance embedding dimensions must be positive")
        if config.lr_appearance < 0 or config.appearance_reg_weight < 0:
            raise ValueError("appearance learning rate and regularisation must be non-negative")
    if config.transient_mode not in {"off", "low_rank"}:
        raise ValueError(f"Unknown transient_mode: {config.transient_mode!r}")
    if config.transient_mode != "off":
        if config.renderer != "3dgs":
            raise ValueError("static/transient decomposition currently requires renderer='3dgs'")
        if config.density_control != "off":
            raise ValueError("freeze adaptive density before enabling the transient branch")
        if config.transient_rank <= 0 or config.lr_transient < 0:
            raise ValueError("transient rank must be positive and its learning rate non-negative")
        if min(
            config.transient_sparsity_weight,
            config.transient_temporal_weight,
            config.transient_feature_weight,
        ) < 0:
            raise ValueError("transient regularisation weights must be non-negative")
        if config.transient_feature_weight > 0 and not config.transient_feature_consistency_path:
            raise ValueError("transient feature supervision requires a consistency-map path")
    if min(
        config.l1_weight,
        config.lpips_weight,
        config.dssim_weight,
        config.multiscale_l1_weight,
        config.edge_weight,
    ) < 0:
        raise ValueError("reconstruction loss weights must be non-negative")
    if config.multiscale_levels <= 0:
        raise ValueError("multiscale_levels must be positive")

    if config.dynamic_mask_mode not in {"off", "residual_ema"}:
        raise ValueError(f"Unknown dynamic_mask_mode: {config.dynamic_mask_mode!r}")
    if config.dynamic_mask_mode != "off":
        validate_dynamic_mask_params(
            keep_ratio=config.dynamic_mask_keep_ratio,
            floor=config.dynamic_mask_floor,
            smooth_kernel=config.dynamic_mask_smooth_kernel,
            alpha_threshold=config.dynamic_mask_alpha_threshold,
            ema_decay=config.dynamic_mask_ema_decay,
        )
        if config.dynamic_mask_warmup_iters < 0 or config.dynamic_mask_ramp_iters < 0:
            raise ValueError("dynamic mask warmup/ramp iterations must be non-negative")
    if not 0 <= config.start_iter < config.num_iters:
        raise ValueError(
            f"start_iter must satisfy 0 <= start_iter < num_iters, got {config.start_iter} and {config.num_iters}"
        )
    if config.frames_per_iter <= 0:
        raise ValueError("frames_per_iter must be positive")
    if config.eval_holdout_every != 0:
        if config.eval_holdout_every < 2:
            raise ValueError("eval_holdout_every must be 0 or at least 2")
        if not 0 <= config.eval_holdout_offset < config.eval_holdout_every:
            raise ValueError("eval_holdout_offset must be in [0, eval_holdout_every)")
    if config.density_control != "off":
        if config.renderer != "3dgs":
            raise ValueError("adaptive density control is only supported by renderer='3dgs'")
        if not config.optimize_positions:
            raise ValueError("adaptive density control requires optimize_positions=True")
        if config.frames_per_iter != 1:
            raise ValueError("adaptive density control currently requires frames_per_iter=1")
        if not 0 <= config.density_start_iter < config.density_stop_iter <= config.num_iters:
            raise ValueError("density iterations must satisfy 0 <= start < stop <= num_iters")
        if config.density_refine_every <= 0 or config.density_reset_every <= 0:
            raise ValueError("density refine/reset intervals must be positive")
        if config.density_max_gaussians <= 0:
            raise ValueError("density_max_gaussians must be positive")
    if config.topology_reliability_mode != "off":
        if config.density_control == "off":
            raise ValueError("topology reliability requires adaptive density control")
        if config.density_control == "mcmc":
            raise ValueError("topology reliability gating currently supports default/absgrad ADC, not MCMC")
        if not 0.0 <= config.topology_reliability_threshold <= 1.0:
            raise ValueError("topology_reliability_threshold must be in [0, 1]")
        if config.topology_reliability_max_weight <= 0:
            raise ValueError("topology_reliability_max_weight must be positive")
        if config.topology_motion_max_side <= 0 or config.topology_motion_fb_sigma <= 0:
            raise ValueError("topology motion parameters must be positive")
    if config.covariance_transport not in {"local_se3", "polar", "partial_jacobian", "full_jacobian"}:
        raise ValueError(f"Unknown covariance_transport: {config.covariance_transport!r}")
    if config.covariance_transport != "local_se3":
        if config.renderer != "3dgs":
            raise ValueError("Jacobian covariance transport is only supported by renderer='3dgs'")
        if not config.deform_inverse_rotations:
            raise ValueError("Jacobian covariance transport requires deform_inverse_rotations=True")
        if not 0.0 <= config.jacobian_stretch_beta <= 1.0:
            raise ValueError("jacobian_stretch_beta must be in [0, 1]")
        if not 0.0 < config.jacobian_fd_min_step <= config.jacobian_fd_max_step:
            raise ValueError("Jacobian finite-difference steps must be positive and ordered")
        if config.jacobian_fd_scale <= 0 or config.jacobian_chunk_size <= 0:
            raise ValueError("Jacobian finite-difference scale and chunk size must be positive")
        if not 0.0 < config.jacobian_singular_min <= config.jacobian_singular_max:
            raise ValueError("Jacobian singular-value bounds must be positive and ordered")
        if config.jacobian_condition_max <= 1.0 or config.jacobian_det_min <= 0:
            raise ValueError("Jacobian condition bound must exceed 1 and determinant bound must be positive")

    # ------------------------------------------------------------------
    # 1. Load data & checkpoint
    # ------------------------------------------------------------------
    nrba_dir = os.path.join(config.root_path, config.run, config.global_opt_subdir)
    if not os.path.exists(nrba_dir):
        raise FileNotFoundError(f"Global optimization directory not found: {nrba_dir}")

    convention_path = os.path.join(nrba_dir, "convention.json")
    if os.path.exists(convention_path):
        with open(convention_path) as f:
            conv = json.load(f)
        if conv.get("global_deform_is") != "c2w":
            raise ValueError(f"Expected c2w convention, got: {conv}")
        logger.info("Convention verified: c2w")
    else:
        logger.warning("No convention.json — assuming c2w convention")

    logger.info("Loading aligned point cloud...")
    canonical_pts, canonical_cols = load_aligned_point_cloud(nrba_dir, device)
    logger.info("Loaded %d canonical points", canonical_pts.shape[0])

    effective_target_points = config.target_num_points
    if config.density_control != "off":
        effective_target_points = min(effective_target_points, config.density_max_gaussians)
    if canonical_pts.shape[0] > effective_target_points:
        logger.info(
            "Downsampling from %d to ~%d points...",
            canonical_pts.shape[0],
            effective_target_points,
        )
        canonical_pts, canonical_cols = downsample_to_target(
            canonical_pts,
            canonical_cols,
            target_count=effective_target_points,
        )
        logger.info("After downsampling: %d points", canonical_pts.shape[0])

    per_frame_global_deform, per_frame_local_deform, bbox_min, bbox_max = load_deformation_checkpoints(
        nrba_dir,
        device,
        first_local="none",
        allow_rigid_fallback=True,
    )
    num_deform_frames = len(per_frame_global_deform)

    # Reuse the exact data loading / confidence-filtering configuration from the
    # original frame_to_model_icp Stage 1 run so that Stage 3.2 is consistent
    # with the alignment stage.
    align_params: AlignmentDataParams = load_alignment_data_params(
        root_path=config.root_path,
        run=config.run,
    )

    load_data_kwargs: dict = dict(
        conf_thresh_percentile=align_params.conf_thresh_percentile,
        conf_mode=align_params.conf_mode,
        conf_local_percentile=align_params.conf_local_percentile,
        conf_global_percentile=align_params.conf_global_percentile,
        voxel_size=align_params.conf_voxel_size,
        voxel_min_count_percentile=align_params.conf_voxel_min_count_percentile,
        offset=align_params.offset,
    )
    logger.info(
        "Using alignment data params for GS training: "
        "num_frames=%d, stride=%d, offset=%d, conf_thresh_percentile=%.1f, "
        "conf_mode=%s, conf_local_percentile=%s, conf_global_percentile=%s, "
        "conf_voxel_size=%.4f, conf_voxel_min_count_percentile=%s",
        align_params.num_frames,
        align_params.stride,
        align_params.offset,
        align_params.conf_thresh_percentile,
        align_params.conf_mode,
        str(align_params.conf_local_percentile),
        str(align_params.conf_global_percentile),
        align_params.conf_voxel_size,
        str(align_params.conf_voxel_min_count_percentile),
    )

    (
        pcls,
        extrinsics_np,
        intrinsics_np,
        images,
        valid_pixel_indices,
        depth_conf,
        _depth_maps,
        _orig_images,
        _orig_intrinsics,
    ) = load_data(
        config.root_path,
        align_params.num_frames,
        align_params.stride,
        device,
        **load_data_kwargs,
        load_original_images_and_intrinsics=False,
    )
    if config.original_images_dir:
        folder_images, folder_intrinsics = load_da3_original_images_from_folder(
            config.root_path,
            config.original_images_dir,
            num_frames=align_params.num_frames,
            stride=align_params.stride,
            device=device,
        )
        images = folder_images
        intrinsics_np = folder_intrinsics
    num_frames = min(len(pcls), num_deform_frames)
    H, W = images.shape[2], images.shape[3]
    logger.info("Loaded %d frames (%d x %d)", num_frames, H, W)
    train_frame_indices, heldout_frame_indices = interleaved_frame_split(
        num_frames,
        config.eval_holdout_every,
        config.eval_holdout_offset,
    )
    logger.info(
        "Stage-3 frame protocol: %d train, %d held out%s",
        len(train_frame_indices),
        len(heldout_frame_indices),
        " (alignment/deformation may still use held-out frames)" if heldout_frame_indices else "",
    )

    intrinsics_list: list[torch.Tensor] = []
    for i in range(num_frames):
        intrinsics_list.append(torch.from_numpy(intrinsics_np[i]).to(device).float())

    gt_images = images[:num_frames]  # (N, 3, H, W)

    alignment_confidence_weighting = align_params.use_confidence_weighting or align_params.use_confidence_weighted_gs
    use_confidence_weighted_gs = resolve_confidence_weighting(
        alignment_confidence_weighting,
        config.confidence_weighting_mode,
    )
    confidence_weight_maps: torch.Tensor | None = None
    if use_confidence_weighted_gs:
        if depth_conf is None:
            raise ValueError("Confidence weighting is enabled, but no depth confidence maps were loaded.")
        confidence_weight_maps = build_pixel_confidence_weights(
            depth_conf[:num_frames],
            (H, W),
            device=device,
            floor=align_params.confidence_weight_floor,
            gamma=align_params.confidence_weight_gamma,
            low_quantile=align_params.confidence_weight_low_quantile,
            high_quantile=align_params.confidence_weight_high_quantile,
        )
        logger.info(
            "Enabled confidence-weighted GS L1: floor=%.3f, gamma=%.3f, quantiles=(%.3f, %.3f), "
            "weight range=(%.3f, %.3f)",
            align_params.confidence_weight_floor,
            align_params.confidence_weight_gamma,
            align_params.confidence_weight_low_quantile,
            align_params.confidence_weight_high_quantile,
            float(confidence_weight_maps.min()),
            float(confidence_weight_maps.max()),
        )

    topology_reliability_maps: torch.Tensor | None = None
    if config.topology_reliability_mode != "off":
        if depth_conf is None:
            raise ValueError("Topology reliability is enabled, but no depth confidence maps were loaded.")
        reliability_components = [
            build_pixel_confidence_weights(
                depth_conf[:num_frames],
                (H, W),
                device=device,
                floor=0.0,
                gamma=align_params.confidence_weight_gamma,
                low_quantile=align_params.confidence_weight_low_quantile,
                high_quantile=align_params.confidence_weight_high_quantile,
            )
        ]
        if config.topology_reliability_mode in {
            "depth_motion",
            "depth_motion_residual",
            "depth_motion_feature",
        }:
            logger.info("Computing forward/backward motion consistency maps...")
            reliability_components.append(
                build_motion_consistency_maps(
                    gt_images,
                    max_side=config.topology_motion_max_side,
                    fb_sigma=config.topology_motion_fb_sigma,
                )
            )
        if config.topology_reliability_mode == "depth_motion_feature":
            reliability_components.append(
                load_feature_consistency_maps(
                    config.topology_feature_consistency_path,
                    num_frames=num_frames,
                    output_size=(H, W),
                    device=device,
                    dtype=gt_images.dtype,
                )
            )
        topology_reliability_maps = combine_reliability_maps(reliability_components)
        logger.info(
            "Topology reliability: mode=%s, gate=%s, threshold=%.3f, map_mean=%.3f",
            config.topology_reliability_mode,
            config.topology_reliability_gate,
            config.topology_reliability_threshold,
            float(topology_reliability_maps.mean()),
        )

    transient_feature_maps: torch.Tensor | None = None
    if config.transient_mode != "off" and config.transient_feature_weight > 0:
        transient_feature_maps = load_feature_consistency_maps(
            config.transient_feature_consistency_path,
            num_frames=num_frames,
            output_size=(H, W),
            device=device,
            dtype=gt_images.dtype,
        )
        logger.info(
            "Loaded transient feature consistency maps: mean=%.3f",
            float(transient_feature_maps.mean()),
        )

    if not config.inverse_deform_dir:
        raise ValueError("--config.inverse-deform-dir is required")
    inverse_deform_net, _inv_cfg = load_inverse_local_from_checkpoint(config.inverse_deform_dir, device=device)
    inverse_deform_net.eval()
    for p in inverse_deform_net.parameters():
        p.requires_grad = False

    # ------------------------------------------------------------------
    # 2. Build model
    # ------------------------------------------------------------------
    init_normals = None
    knn_dists = None
    kdtree = None

    if config.normal_k > 0:
        logger.info(
            "Estimating normals (k=%d) for %d canonical points...",
            config.normal_k,
            canonical_pts.shape[0],
        )
        init_normals, kdtree = estimate_normals(
            canonical_pts,
            k=config.normal_k,
            backend="cpu_kdtree",
        )
        init_normals = init_normals.cpu()
        logger.info("Normal estimation done.")

    if config.scale_init == "knn":
        logger.info(
            "Computing KNN distances (K=%d) for %d points...",
            config.knn_neighbors,
            canonical_pts.shape[0],
        )
        _, d2 = query_knn_with_backend(
            canonical_pts,
            canonical_pts,
            K=config.knn_neighbors,
            backend="cpu_kdtree",
            cpu_tree=kdtree,
        )
        knn_dists = d2[:, 1:].mean(dim=-1).sqrt().cpu()
        logger.info(
            "KNN dists: median=%.5f, min=%.5f, max=%.5f",
            knn_dists.median().item(),
            knn_dists.min().item(),
            knn_dists.max().item(),
        )

    model = CanonicalGSModel(
        canonical_points=canonical_pts,
        canonical_colors=canonical_cols,
        per_frame_global_deform=per_frame_global_deform[:num_frames],
        per_frame_local_deform=per_frame_local_deform[:num_frames],
        bbox_min=bbox_min,
        bbox_max=bbox_max,
        height=H,
        width=W,
        renderer=config.renderer,
        rasterize_mode=config.rasterize_mode,
        rasterize_eps2d=config.rasterize_eps2d,
        mip_filter_3d=config.mip_filter_3d,
        mip_filter_variance=config.mip_filter_variance,
        mip_filter_opacity_compensation=config.mip_filter_opacity_compensation,
        appearance_mode=config.appearance_mode,
        appearance_embedding_dim=config.appearance_embedding_dim,
        appearance_hidden_dim=config.appearance_hidden_dim,
        transient_mode=config.transient_mode,
        transient_rank=config.transient_rank,
        transient_color_scale=config.transient_color_scale,
        transient_opacity_scale=config.transient_opacity_scale,
        optimize_cams=config.optimize_cams,
        optimize_positions=config.optimize_positions,
        deform_rotations=config.deform_inverse_rotations,
        covariance_transport=config.covariance_transport,
        jacobian_stretch_beta=config.jacobian_stretch_beta,
        jacobian_fd_scale=config.jacobian_fd_scale,
        jacobian_fd_min_step=config.jacobian_fd_min_step,
        jacobian_fd_max_step=config.jacobian_fd_max_step,
        jacobian_singular_min=config.jacobian_singular_min,
        jacobian_singular_max=config.jacobian_singular_max,
        jacobian_condition_max=config.jacobian_condition_max,
        jacobian_det_min=config.jacobian_det_min,
        jacobian_chunk_size=config.jacobian_chunk_size,
        initial_opacity=config.initial_opacity,
        initial_scale=config.initial_scale,
        initial_flat_ratio=config.initial_flat_ratio,
        near_plane=0.01,
        far_plane=1e10,
        inverse_deform_net=inverse_deform_net,
        knn_dists=knn_dists,
        init_normals=init_normals,
        sh_degree=config.sh_degree,
    ).to(device)

    if config.init_checkpoint:
        if not os.path.isfile(config.init_checkpoint):
            raise FileNotFoundError(f"Initial GS checkpoint not found: {config.init_checkpoint}")
        logger.info(
            "Warm-starting model from %s at absolute iteration %d (optimizer state is reset)",
            config.init_checkpoint,
            config.start_iter,
        )
        state = torch.load(config.init_checkpoint, map_location="cpu")
        resize_model_gaussians_from_state(model, state)
        if "mip_filter_3d" not in state:
            state["mip_filter_3d"] = torch.zeros_like(model.mip_filter_3d)
        incompatible = model.load_state_dict(state, strict=False)
        allowed_missing_prefixes = (
            "appearance_log_gain",
            "appearance_bias",
            "appearance_embeddings",
            "appearance_mlp.",
            "transient_logits",
            "transient_color_basis",
            "transient_opacity_basis",
            "transient_codes",
            "transient_supervised_frames",
        )
        unexpected_missing = [
            key for key in incompatible.missing_keys if not key.startswith(allowed_missing_prefixes)
        ]
        if unexpected_missing or incompatible.unexpected_keys:
            raise RuntimeError(
                "Incompatible warm-start checkpoint: "
                f"missing={unexpected_missing}, unexpected={incompatible.unexpected_keys}"
            )
        if incompatible.missing_keys:
            logger.info("Initialising new checkpoint fields: %s", incompatible.missing_keys)
        del state
        # Dynamic-topology resizing follows the checkpoint's CPU tensors and
        # replaces registered Parameters. Move the resized model back to the
        # selected device before optimizers are constructed.
        model.to(device)

    if model.mip_filter_enabled:
        model.update_mip_filter_3d(intrinsics_list)
    if model.transient_mode != "off":
        model.set_transient_supervised_frames(train_frame_indices)

    # ------------------------------------------------------------------
    # 3. Optimiser
    # ------------------------------------------------------------------
    gaussian_params = collect_gaussian_parameters(model)
    gaussian_lrs = {
        "means": config.lr_positions,
        "scales": config.lr_scales,
        "quats": config.lr_quats,
        "opacities": config.lr_opacities,
        "sh_dc": config.lr_sh0,
        "sh_rest": config.lr_shN,
    }
    gaussian_optimizers = {
        name: torch.optim.Adam([{"params": [param], "lr": gaussian_lrs[name], "name": name}])
        for name, param in gaussian_params.items()
    }
    all_optimizers: dict[str, torch.optim.Optimizer] = dict(gaussian_optimizers)
    if config.optimize_cams:
        cam_params = [p for p in model.per_frame_c2w.parameters() if p.requires_grad]
        if cam_params:
            all_optimizers["cams"] = torch.optim.Adam([{"params": cam_params, "lr": config.lr_cams, "name": "cams"}])
    appearance_params = model.appearance_parameters()
    if appearance_params:
        all_optimizers["appearance"] = torch.optim.Adam(
            [{"params": appearance_params, "lr": config.lr_appearance, "name": "appearance"}]
        )
    transient_params = model.transient_parameters()
    if transient_params:
        all_optimizers["transient"] = torch.optim.Adam(
            [{"params": transient_params, "lr": config.lr_transient, "name": "transient"}]
        )

    def _lr_multiplier(local_step: int) -> float:
        # LambdaLR counts from zero after construction. Offset it for a
        # warm-start so continued runs do not accidentally restart at the
        # maximum learning rate.
        absolute_step = min(config.start_iter + local_step, config.num_iters)
        return config.lr_decay + (1.0 - config.lr_decay) * 0.5 * (
            1.0 + math.cos(math.pi * absolute_step / config.num_iters)
        )

    def _new_module_lr_multiplier(local_step: int) -> float:
        # Newly introduced modules (currently appearance) have no checkpoint
        # history, so give them a full local continuation schedule.
        continuation_steps = max(config.num_iters - config.start_iter, 1)
        progress = min(local_step, continuation_steps) / continuation_steps
        return config.lr_decay + (1.0 - config.lr_decay) * 0.5 * (1.0 + math.cos(math.pi * progress))

    schedulers = {
        name: torch.optim.lr_scheduler.LambdaLR(
            optimizer,
            lr_lambda=_new_module_lr_multiplier if name in {"appearance", "transient"} else _lr_multiplier,
        )
        for name, optimizer in all_optimizers.items()
    }

    scene_scale = config.density_scene_scale if config.density_scene_scale > 0 else robust_scene_scale(canonical_pts)
    density_strategy, density_state = create_density_strategy(config, scene_scale)
    if density_strategy is not None:
        density_strategy.check_sanity(gaussian_params, gaussian_optimizers)
        model.rasterization_absgrad = config.density_control == "absgrad"
        logger.info(
            "Density control: mode=%s, start=%d, stop=%d, every=%d, cap=%d, scene_scale=%.5f",
            config.density_control,
            config.density_start_iter,
            config.density_stop_iter,
            config.density_refine_every,
            config.density_max_gaussians,
            scene_scale,
        )

    lpips_fn = init_lpips(device) if config.lpips_weight > 0 or config.eval_lpips else None

    # ------------------------------------------------------------------
    # 4. Output & logging
    # ------------------------------------------------------------------
    if config.out_dir is None:
        config.out_dir = os.path.join(
            config.root_path,
            config.run,
            f"gs_{config.renderer}",
        )
    os.makedirs(config.out_dir, exist_ok=True)

    writer = None
    if config.tensorboard:
        tb_dir = os.path.join(config.out_dir, "tensorboard")
        writer = try_create_tensorboard_writer(tb_dir)

    with open(os.path.join(config.out_dir, "config.json"), "w") as f:
        json.dump(
            {
                k: str(v) if not isinstance(v, (int, float, bool, str, type(None))) else v
                for k, v in vars(config).items()
            },
            f,
            indent=2,
        )

    logger.info("Output directory: %s", config.out_dir)
    logger.info("Renderer: %s, sh_degree=%d", config.renderer, config.sh_degree)
    logger.info("Trainable parameters:")
    for name, optimizer in all_optimizers.items():
        n_params = sum(p.numel() for group in optimizer.param_groups for p in group["params"])
        logger.info("  %s: %d params, lr=%.2e", name, n_params, optimizer.param_groups[0]["lr"])

    if writer is not None:
        tb_log_hparams(
            writer,
            {
                "root_path": config.root_path,
                "run": config.run,
                "renderer": config.renderer,
                "rasterize_mode": config.rasterize_mode,
                "rasterize_eps2d": config.rasterize_eps2d,
                "mip_filter_3d": config.mip_filter_3d,
                "mip_filter_variance": config.mip_filter_variance,
                "mip_filter_opacity_compensation": config.mip_filter_opacity_compensation,
                "mip_filter_update_every": config.mip_filter_update_every,
                "appearance_mode": config.appearance_mode,
                "appearance_embedding_dim": config.appearance_embedding_dim,
                "appearance_hidden_dim": config.appearance_hidden_dim,
                "lr_appearance": config.lr_appearance,
                "appearance_reg_weight": config.appearance_reg_weight,
                "transient_mode": config.transient_mode,
                "transient_rank": config.transient_rank,
                "transient_color_scale": config.transient_color_scale,
                "transient_opacity_scale": config.transient_opacity_scale,
                "lr_transient": config.lr_transient,
                "transient_sparsity_weight": config.transient_sparsity_weight,
                "transient_temporal_weight": config.transient_temporal_weight,
                "transient_feature_weight": config.transient_feature_weight,
                "transient_feature_consistency_path": config.transient_feature_consistency_path,
                "sh_degree": config.sh_degree,
                "target_num_points": config.target_num_points,
                "optimize_cams": config.optimize_cams,
                "optimize_positions": config.optimize_positions,
                "lr_colors": config.lr_colors,
                "lr_opacities": config.lr_opacities,
                "lr_scales": config.lr_scales,
                "lr_quats": config.lr_quats,
                "lr_sh0": config.lr_sh0,
                "lr_shN": config.lr_shN,
                "l1_weight": config.l1_weight,
                "lpips_weight": config.lpips_weight,
                "dssim_weight": config.dssim_weight,
                "multiscale_l1_weight": config.multiscale_l1_weight,
                "multiscale_levels": config.multiscale_levels,
                "edge_weight": config.edge_weight,
                "num_iters": config.num_iters,
                "frames_per_iter": config.frames_per_iter,
                "num_frames": num_frames,
                "initial_num_gaussians": len(model.canonical_points),
                "stride": align_params.stride,
                "conf_thresh_percentile": align_params.conf_thresh_percentile,
                "use_confidence_weighting": align_params.use_confidence_weighting,
                "use_confidence_weighted_gs": use_confidence_weighted_gs,
                "confidence_weighting_mode": config.confidence_weighting_mode,
                "confidence_weight_floor": align_params.confidence_weight_floor,
                "confidence_weight_gamma": align_params.confidence_weight_gamma,
                "confidence_weight_low_quantile": align_params.confidence_weight_low_quantile,
                "confidence_weight_high_quantile": align_params.confidence_weight_high_quantile,
                "dynamic_mask_mode": config.dynamic_mask_mode,
                "dynamic_mask_warmup_iters": config.dynamic_mask_warmup_iters,
                "dynamic_mask_ramp_iters": config.dynamic_mask_ramp_iters,
                "dynamic_mask_keep_ratio": config.dynamic_mask_keep_ratio,
                "dynamic_mask_floor": config.dynamic_mask_floor,
                "dynamic_mask_smooth_kernel": config.dynamic_mask_smooth_kernel,
                "dynamic_mask_alpha_threshold": config.dynamic_mask_alpha_threshold,
                "dynamic_mask_ema_decay": config.dynamic_mask_ema_decay,
                "seed": config.seed,
                "init_checkpoint": config.init_checkpoint,
                "start_iter": config.start_iter,
                "eval_holdout_every": config.eval_holdout_every,
                "eval_holdout_offset": config.eval_holdout_offset,
                "eval_lpips": config.eval_lpips,
                "eval_temporal_consistency": config.eval_temporal_consistency,
                "density_control": config.density_control,
                "density_start_iter": config.density_start_iter,
                "density_stop_iter": config.density_stop_iter,
                "density_refine_every": config.density_refine_every,
                "density_grad_threshold": config.density_grad_threshold,
                "density_revised_opacity": config.density_revised_opacity,
                "density_max_gaussians": config.density_max_gaussians,
                "density_scene_scale": scene_scale,
                "topology_reliability_mode": config.topology_reliability_mode,
                "topology_reliability_gate": config.topology_reliability_gate,
                "topology_reliability_threshold": config.topology_reliability_threshold,
                "topology_reliability_max_weight": config.topology_reliability_max_weight,
                "topology_feature_consistency_path": config.topology_feature_consistency_path,
                "covariance_transport": config.covariance_transport,
                "jacobian_stretch_beta": config.jacobian_stretch_beta,
                "jacobian_fd_scale": config.jacobian_fd_scale,
                "jacobian_singular_min": config.jacobian_singular_min,
                "jacobian_singular_max": config.jacobian_singular_max,
                "jacobian_condition_max": config.jacobian_condition_max,
            },
            step=0,
        )

    # ------------------------------------------------------------------
    # 5. Training loop
    # ------------------------------------------------------------------
    is_2dgs = config.renderer == "2dgs"
    use_dynamic_mask = config.dynamic_mask_mode == "residual_ema"
    needs_residual_ema = use_dynamic_mask or config.topology_reliability_mode == "depth_motion_residual"
    residual_ema_maps = (
        torch.zeros((num_frames, 1, H, W), device=device, dtype=gt_images.dtype) if needs_residual_ema else None
    )
    residual_ema_initialized = [False] * num_frames if needs_residual_ema else None
    if use_dynamic_mask:
        logger.info(
            "Enabled visibility-aware residual EMA mask: warmup=%d, ramp=%d, keep_ratio=%.3f, "
            "floor=%.3f, smooth_kernel=%d, alpha_threshold=%.3f, ema_decay=%.3f",
            config.dynamic_mask_warmup_iters,
            config.dynamic_mask_ramp_iters,
            config.dynamic_mask_keep_ratio,
            config.dynamic_mask_floor,
            config.dynamic_mask_smooth_kernel,
            config.dynamic_mask_alpha_threshold,
            config.dynamic_mask_ema_decay,
        )
    model.train()
    pbar = tqdm(range(config.start_iter, config.num_iters), desc=f"GS training ({config.renderer})")
    defer_means_freeze = bool(
        density_strategy is not None
        and config.sh_freeze_means_when_full_sh
        and 0 < config.sh_full_from_iter < config.density_stop_iter
    )

    for it in pbar:
        if config.sh_full_from_iter > 0 and it == config.sh_full_from_iter:
            for p in model.per_frame_c2w.parameters():
                p.requires_grad = False
            if config.sh_freeze_means_when_full_sh and not defer_means_freeze:
                model.canonical_points.requires_grad = False
            logger.info(
                "At iter %d: froze cams (means frozen=%s), full SH enabled.",
                it,
                config.sh_freeze_means_when_full_sh and not defer_means_freeze,
            )
        if defer_means_freeze and it == config.density_stop_iter:
            model.canonical_points.requires_grad = False
            logger.info("At iter %d: density control stopped and canonical means were frozen.", it)

        for optimizer in all_optimizers.values():
            optimizer.zero_grad(set_to_none=True)

        sampled_train_indices = torch.randint(0, len(train_frame_indices), (config.frames_per_iter,))
        total_loss = torch.tensor(0.0, device=device)
        render_result = None
        sh_reg_loss = torch.tensor(0.0, device=device)
        dynamic_mask_mean = torch.tensor(1.0, device=device)
        density_reliability_map = None
        topology_stats = None

        for sampled_idx in sampled_train_indices:
            fi = train_frame_indices[sampled_idx.item()]
            K = intrinsics_list[fi]

            render_result = model.render_frame(fi, K, use_inverse_deform=True)
            density_active = bool(
                density_strategy is not None and config.density_start_iter <= it < config.density_stop_iter
            )
            if density_active:
                assert density_state is not None
                density_strategy.step_pre_backward(
                    gaussian_params,
                    gaussian_optimizers,
                    density_state,
                    it,
                    render_result["info"],
                )
            rendered_rgb = render_result["rgb"]
            gt = gt_images[fi : fi + 1]

            pixel_l1 = (rendered_rgb - gt).abs().mean(dim=1, keepdim=True)
            loss_weights = confidence_weight_maps[fi : fi + 1] if confidence_weight_maps is not None else None
            dynamic_weights = None
            if residual_ema_maps is not None and residual_ema_initialized is not None:
                ema = residual_ema_maps[fi : fi + 1]
                update_residual_ema_(
                    ema,
                    pixel_l1,
                    decay=config.dynamic_mask_ema_decay,
                    initialized=residual_ema_initialized[fi],
                )
                residual_ema_initialized[fi] = True
                mask_strength = dynamic_mask_ramp(
                    it,
                    warmup_iters=config.dynamic_mask_warmup_iters,
                    ramp_iters=config.dynamic_mask_ramp_iters,
                )
                dynamic_weights = build_visibility_aware_residual_weights(
                    ema,
                    render_result["alpha"],
                    keep_ratio=config.dynamic_mask_keep_ratio,
                    floor=config.dynamic_mask_floor,
                    smooth_kernel=config.dynamic_mask_smooth_kernel,
                    alpha_threshold=config.dynamic_mask_alpha_threshold,
                    strength=mask_strength,
                )
                dynamic_mask_mean = dynamic_weights.mean()
                if use_dynamic_mask:
                    loss_weights = dynamic_weights if loss_weights is None else loss_weights * dynamic_weights
            if topology_reliability_maps is not None:
                density_reliability_map = topology_reliability_maps[fi : fi + 1]
                if config.topology_reliability_mode == "depth_motion_residual":
                    assert dynamic_weights is not None
                    density_reliability_map = combine_reliability_maps(
                        [density_reliability_map, dynamic_weights.detach()]
                    )
            loss_l1 = weighted_mean(pixel_l1, loss_weights) if loss_weights is not None else pixel_l1.mean()
            loss_lpips = torch.tensor(0.0, device=device)
            if lpips_fn is not None and config.lpips_weight > 0:
                loss_lpips = lpips_fn(
                    rendered_rgb.clamp(0, 1) * 2 - 1,
                    gt.clamp(0, 1) * 2 - 1,
                ).mean()

            loss_dssim = (
                dssim_loss(rendered_rgb, gt) if config.dssim_weight > 0 else torch.tensor(0.0, device=device)
            )
            loss_multiscale = (
                multiscale_l1_loss(rendered_rgb, gt, levels=config.multiscale_levels)
                if config.multiscale_l1_weight > 0
                else torch.tensor(0.0, device=device)
            )
            loss_edge = (
                edge_l1_loss(rendered_rgb, gt) if config.edge_weight > 0 else torch.tensor(0.0, device=device)
            )

            frame_loss = (
                config.l1_weight * loss_l1
                + config.lpips_weight * loss_lpips
                + config.dssim_weight * loss_dssim
                + config.multiscale_l1_weight * loss_multiscale
                + config.edge_weight * loss_edge
            )

            # ---------- Regularisation ----------
            reg_loss = torch.tensor(0.0, device=device)

            if config.scale_reg_weight > 0:
                scales = model.log_scales.exp()
                reg_loss = reg_loss + config.scale_reg_weight * scales.mean()

            if is_2dgs:
                if config.opacity_reg_weight > 0:
                    opacities = torch.sigmoid(model.logit_opacities)
                    reg_loss = reg_loss + config.opacity_reg_weight * (opacities * (1 - opacities)).mean()

                if config.normal_consistency_weight > 0 and render_result is not None and "normals" in render_result:
                    from losses.gaussian import normal_consistency_loss

                    normals = render_result["normals"]
                    surf_normals = render_result["surf_normals"]
                    alphas = render_result.get("alpha")
                    if normals is not None and surf_normals is not None and alphas is not None:
                        nc_loss = normal_consistency_loss(
                            normals,
                            surf_normals.unsqueeze(0),
                            alphas.permute(0, 2, 3, 1),
                        )
                        reg_loss = reg_loss + config.normal_consistency_weight * nc_loss

                if config.distortion_weight > 0 and render_result is not None and "distort" in render_result:
                    from losses import distortion_loss

                    distort = render_result["distort"]
                    if distort is not None:
                        dl = distortion_loss(distort)
                        reg_loss = reg_loss + config.distortion_weight * dl

                if config.alpha_reg_weight > 0 and render_result is not None and "alpha" in render_result:
                    alpha = render_result["alpha"]
                    if alpha is not None:
                        reg_loss = reg_loss + config.alpha_reg_weight * (1.0 - alpha).mean()

            if config.sh_reg_weight > 0.0 and hasattr(model, "sh_rest") and model.sh_rest is not None:
                sh_reg = (model.sh_rest**2).mean()
                sh_reg_loss = sh_reg_loss + sh_reg / config.frames_per_iter
                reg_loss = reg_loss + config.sh_reg_weight * sh_reg

            if config.appearance_reg_weight > 0 and model.appearance_mode != "off":
                reg_loss = reg_loss + config.appearance_reg_weight * model.appearance_regularization(fi)

            if model.transient_mode != "off":
                transient_sparsity, transient_temporal = model.transient_regularization()
                reg_loss = reg_loss + config.transient_sparsity_weight * transient_sparsity
                reg_loss = reg_loss + config.transient_temporal_weight * transient_temporal
                if transient_feature_maps is not None and config.transient_feature_weight > 0:
                    from utils.topology_reliability import sample_projected_reliability

                    gaussian_ids = render_result["info"].get("gaussian_ids")
                    if gaussian_ids is None:
                        raise RuntimeError("gsplat packed rasterization did not return gaussian_ids")
                    consistency = sample_projected_reliability(
                        render_result["info"], transient_feature_maps[fi : fi + 1]
                    )
                    target_transient = 1.0 - consistency
                    assert model.transient_logits is not None
                    predicted_transient = torch.sigmoid(model.transient_logits[gaussian_ids].squeeze(-1))
                    transient_feature_loss = torch.nn.functional.binary_cross_entropy(
                        predicted_transient,
                        target_transient,
                    )
                    reg_loss = reg_loss + config.transient_feature_weight * transient_feature_loss

            frame_loss = frame_loss + reg_loss
            total_loss = total_loss + frame_loss / config.frames_per_iter
        total_loss.backward()

        if density_active and density_reliability_map is not None:
            assert render_result is not None
            topology_stats = gate_density_gradients(
                render_result["info"],
                density_reliability_map,
                threshold=config.topology_reliability_threshold,
                gate=config.topology_reliability_gate,
                absgrad=config.density_control == "absgrad",
                max_weight=config.topology_reliability_max_weight,
            )

        # SH gradient freezing schedule
        if (
            hasattr(model, "sh_rest")
            and model.sh_rest is not None
            and model.sh_rest.grad is not None
            and config.sh_degree > 0
        ):
            if config.sh_full_from_iter > 0 and it < config.sh_full_from_iter:
                model.sh_rest.grad.zero_()
            elif config.sh_increase_every > 0 and config.sh_full_from_iter == 0:
                L = int(config.sh_degree)
                total_rest = model.sh_rest.shape[1]
                if total_rest == L * (L + 2):
                    bands_unlocked = min(max(it // config.sh_increase_every, 0), L)
                    active_non_dc = bands_unlocked * (bands_unlocked + 2)
                    if active_non_dc < total_rest:
                        model.sh_rest.grad[:, active_non_dc:, :].zero_()

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        for optimizer in all_optimizers.values():
            optimizer.step()
        for scheduler in schedulers.values():
            scheduler.step()

        if density_active:
            assert density_strategy is not None and density_state is not None and render_result is not None
            if config.density_control == "mcmc":
                density_strategy.step_post_backward(
                    gaussian_params,
                    gaussian_optimizers,
                    density_state,
                    it,
                    render_result["info"],
                    lr=schedulers["means"].get_last_lr()[0],
                )
            else:
                density_strategy.step_post_backward(
                    gaussian_params,
                    gaussian_optimizers,
                    density_state,
                    it,
                    render_result["info"],
                    packed=True,
                )
                removed_for_budget = enforce_default_strategy_budget(
                    gaussian_params,
                    gaussian_optimizers,
                    density_state,
                    config.density_max_gaussians,
                )
                if removed_for_budget:
                    logger.info("Density budget pruned %d lowest-opacity Gaussians", removed_for_budget)
            sync_gaussian_parameters(model, gaussian_params)

        mip_filter_due = bool(
            model.mip_filter_enabled
            and (
                len(model.mip_filter_3d) != len(model.canonical_points)
                or (
                    config.mip_filter_update_every > 0
                    and it > config.start_iter
                    and it % config.mip_filter_update_every == 0
                )
            )
        )
        if mip_filter_due:
            model.update_mip_filter_3d(intrinsics_list)

        # ---------- Logging ----------
        if it % config.log_every == 0 or it == config.num_iters - 1:
            postfix = {"loss": f"{total_loss.item():.5f}"}
            postfix["n_gs"] = str(len(model.canonical_points))
            if use_dynamic_mask:
                postfix["mask_w"] = f"{dynamic_mask_mean.item():.3f}"
            if topology_stats is not None:
                postfix["topo_w"] = f"{topology_stats['projected_mean']:.3f}"
                if config.topology_reliability_gate == "binary":
                    postfix["topo_keep"] = f"{topology_stats['projected_keep_ratio']:.3f}"
            if model.last_jacobian_stats is not None:
                postfix["jac_ok"] = f"{model.last_jacobian_stats['valid_ratio'].item():.3f}"
            pbar.set_postfix(postfix)

            if writer is not None:
                writer.add_scalar("train/loss", total_loss.item(), it)
                if config.sh_reg_weight > 0.0:
                    writer.add_scalar("train/sh_reg_loss", sh_reg_loss.item(), it)
                for name, scheduler in schedulers.items():
                    writer.add_scalar(f"train/lr_{name}", scheduler.get_last_lr()[0], it)
                writer.add_scalar("train/num_gaussians", len(model.canonical_points), it)
                if use_dynamic_mask:
                    writer.add_scalar("train/dynamic_mask_mean", dynamic_mask_mean.item(), it)
                if topology_stats is not None:
                    for name, value in topology_stats.items():
                        writer.add_scalar(f"train/topology_{name}", value, it)
                if model.last_jacobian_stats is not None:
                    for name, value in model.last_jacobian_stats.items():
                        writer.add_scalar(f"train/jacobian_{name}", value.item(), it)

                if it % (config.log_every * 5) == 0 and render_result is not None:
                    with torch.no_grad():
                        vis_fi = 0
                        vis_K = intrinsics_list[vis_fi]
                        vis_result = model.render_frame(vis_fi, vis_K, use_inverse_deform=True)
                        writer.add_image("vis/rendered", vis_result["rgb"][0].clamp(0, 1), it)
                        writer.add_image("vis/gt", gt_images[vis_fi], it)
                        if residual_ema_maps is not None and residual_ema_initialized is not None:
                            vis_residual = residual_ema_maps[vis_fi : vis_fi + 1]
                            vis_weights = build_visibility_aware_residual_weights(
                                vis_residual,
                                vis_result["alpha"],
                                keep_ratio=config.dynamic_mask_keep_ratio,
                                floor=config.dynamic_mask_floor,
                                smooth_kernel=config.dynamic_mask_smooth_kernel,
                                alpha_threshold=config.dynamic_mask_alpha_threshold,
                                strength=dynamic_mask_ramp(
                                    it,
                                    warmup_iters=config.dynamic_mask_warmup_iters,
                                    ramp_iters=config.dynamic_mask_ramp_iters,
                                ),
                            )
                            writer.add_image("vis/dynamic_mask", vis_weights[0], it)

        # ---------- Save ----------
        if (it > 0 and it % config.save_every == 0) or it == config.num_iters - 1:
            ckpt_path = os.path.join(config.out_dir, f"checkpoint_{it:06d}.pt")
            torch.save(model.state_dict(), ckpt_path)
            logger.info("Saved checkpoint: %s", ckpt_path)

        # ---------- Eval ----------
        if (it > 0 and it % config.eval_every == 0) or it == config.num_iters - 1:
            _evaluate(
                model,
                gt_images,
                intrinsics_list,
                num_frames,
                config,
                it,
                writer,
                lpips_fn,
                train_frame_indices,
                heldout_frame_indices,
                num_frames_to_save=5 if it < config.num_iters - 1 else 1000,
            )

    # ------------------------------------------------------------------
    # 6. Final save
    # ------------------------------------------------------------------
    final_path = os.path.join(config.out_dir, "model_final.pt")
    torch.save(model.state_dict(), final_path)
    logger.info("Training complete.  Final model: %s", final_path)

    final_pts = model.canonical_points.detach().cpu().numpy()
    if hasattr(model, "sh_coeffs") and model.sh_coeffs is not None:
        sh0 = model.sh_coeffs[:, 0, :]
        final_cols = sh0_to_rgb(sh0).detach().cpu().numpy().clip(0, 1)
    else:
        final_cols = sh0_to_rgb(model.sh_dc.detach()[:, 0, :]).cpu().numpy()
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(final_pts)
    pcd.colors = o3d.utility.Vector3dVector(final_cols)
    o3d.io.write_point_cloud(os.path.join(config.out_dir, "canonical_points_final.ply"), pcd)

    if writer is not None:
        writer.close()

    # If we trained a 3DGS model, automatically export a 3DGS-compatible PLY
    # checkpoint alongside the final model for convenient downstream viewing.
    if config.renderer == "3dgs":
        try:
            export_cfg = ExportGSCheckpointToPlyConfig(
                root_path=config.root_path,
                run=config.run,
                checkpoint_dir=config.out_dir,
                global_opt_subdir=config.global_opt_subdir,
            )
            export_gs_checkpoint_to_ply(export_cfg)
        except Exception as exc:
            logger.error("Automatic 3DGS PLY export failed: %s", exc)

    logger.info("Done.")

    # ------------------------------------------------------------------
    # 9. Optional automatic evaluation
    # ------------------------------------------------------------------
    if config.auto_eval:
        cmd = [
            sys.executable,
            "-m",
            "eval_gs",
            "--config.root-path",
            config.root_path,
            "--config.run",
            config.run,
            "--config.checkpoint-dir",
            config.out_dir,
        ]
        if config.original_images_dir:
            cmd.extend(["--config.original-images-dir", config.original_images_dir])
        try:
            logger.info("Running automatic eval via: %s", " ".join(cmd))
            project_root = os.path.dirname(os.path.abspath(__file__))
            subprocess.run(cmd, check=True, cwd=project_root)
        except Exception as e:
            logger.error("Automatic eval failed: %s", e)


@torch.no_grad()
def _evaluate(
    model: CanonicalGSModel,
    gt_images: torch.Tensor,
    intrinsics_list: list[torch.Tensor],
    num_frames: int,
    config: GSConfig,
    global_step: int,
    writer: SummaryWriter | None,
    lpips_fn,
    train_frame_indices: list[int],
    heldout_frame_indices: list[int],
    num_frames_to_save: int = 5,
):
    """Evaluate input fitting and an optional Stage-3 held-out appearance split."""
    model.eval()
    psnrs: list[float] = []
    ssims: list[float] = []
    lpips_values: list[float] = []
    rendered_frames: list[np.ndarray] = []
    target_frames: list[np.ndarray] = []
    alpha_frames: list[np.ndarray] = []

    eval_dir = os.path.join(config.out_dir, f"eval_{global_step:06d}")
    os.makedirs(eval_dir, exist_ok=True)

    for fi in range(num_frames):
        K = intrinsics_list[fi]
        result = model.render_frame(fi, K, use_inverse_deform=True)
        rendered = result["rgb"]  # (1, 3, H, W)
        gt = gt_images[fi : fi + 1]

        psnr = image_psnr(rendered, gt).item()
        ssim = image_ssim(rendered, gt).item()
        psnrs.append(psnr)
        ssims.append(ssim)
        if config.eval_lpips and lpips_fn is not None:
            lpips_values.append(lpips_fn(rendered.clamp(0, 1) * 2 - 1, gt.clamp(0, 1) * 2 - 1).mean().item())
        else:
            lpips_values.append(float("nan"))

        rendered_frames.append(rendered[0].clamp(0, 1).permute(1, 2, 0).cpu().numpy())
        target_frames.append(gt[0].clamp(0, 1).permute(1, 2, 0).cpu().numpy())
        alpha_frames.append(result["alpha"][0, 0].clamp(0, 1).cpu().numpy())

        if fi < num_frames_to_save:
            img_np = (rendered[0].clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            cv2.imwrite(
                os.path.join(eval_dir, f"rendered_{fi:05d}.png"),
                cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR),
            )
            gt_np = (gt[0].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            cv2.imwrite(
                os.path.join(eval_dir, f"gt_{fi:05d}.png"),
                cv2.cvtColor(gt_np, cv2.COLOR_RGB2BGR),
            )

            rendered_canon = model.render_frame(fi, K, use_inverse_deform=False)["rgb"]
            rendered_canon_np = (rendered_canon[0].clamp(0, 1).permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
            cv2.imwrite(
                os.path.join(eval_dir, f"rendered_canon_{fi:05d}.png"),
                cv2.cvtColor(rendered_canon_np, cv2.COLOR_RGB2BGR),
            )

    all_frame_indices = list(range(num_frames))

    def _mean_for(values: list[float], indices: list[int]) -> float:
        if not indices:
            return float("nan")
        selected = np.asarray([values[i] for i in indices], dtype=np.float64)
        return float(selected[np.isfinite(selected)].mean()) if np.isfinite(selected).any() else float("nan")

    avg_psnr = _mean_for(psnrs, all_frame_indices)
    avg_ssim = _mean_for(ssims, all_frame_indices)
    avg_lpips = _mean_for(lpips_values, all_frame_indices)
    train_psnr = _mean_for(psnrs, train_frame_indices)
    train_ssim = _mean_for(ssims, train_frame_indices)
    train_lpips = _mean_for(lpips_values, train_frame_indices)
    heldout_psnr = _mean_for(psnrs, heldout_frame_indices)
    heldout_ssim = _mean_for(ssims, heldout_frame_indices)
    heldout_lpips = _mean_for(lpips_values, heldout_frame_indices)
    logger.info(
        "Eval [step %d]: input-fitting PSNR=%.2f dB, SSIM=%.4f, LPIPS=%.4f",
        global_step,
        avg_psnr,
        avg_ssim,
        avg_lpips,
    )
    if heldout_frame_indices:
        logger.info(
            "Eval split: train PSNR/SSIM/LPIPS=%.2f/%.4f/%.4f, held-out appearance=%.2f/%.4f/%.4f",
            train_psnr,
            train_ssim,
            train_lpips,
            heldout_psnr,
            heldout_ssim,
            heldout_lpips,
        )

    temporal = None
    if config.eval_temporal_consistency:
        temporal = flow_guided_temporal_l1(rendered_frames, target_frames, alpha_frames)
        logger.info(
            "Temporal diagnostic: render_warp_l1=%.5f, target_warp_l1=%.5f, valid_ratio=%.3f",
            temporal["render_warp_l1"],
            temporal["target_warp_l1"],
            temporal["valid_ratio"],
        )

    if writer is not None:
        writer.add_scalar("eval/psnr_avg", avg_psnr, global_step)
        writer.add_scalar("eval/ssim_avg", avg_ssim, global_step)
        writer.add_scalar("eval/lpips_avg", avg_lpips, global_step)
        writer.add_scalar("eval/train_psnr", train_psnr, global_step)
        writer.add_scalar("eval/train_ssim", train_ssim, global_step)
        writer.add_scalar("eval/train_lpips", train_lpips, global_step)
        if heldout_frame_indices:
            writer.add_scalar("eval/heldout_appearance_psnr", heldout_psnr, global_step)
            writer.add_scalar("eval/heldout_appearance_ssim", heldout_ssim, global_step)
            writer.add_scalar("eval/heldout_appearance_lpips", heldout_lpips, global_step)
        if temporal is not None:
            for name, value in temporal.items():
                writer.add_scalar(f"eval/temporal_{name}", value, global_step)
        for fi, p in enumerate(psnrs):
            writer.add_scalar(f"eval/psnr_frame_{fi:03d}", p, global_step)

    def _finite_or_none(value: float) -> float | None:
        return value if np.isfinite(value) else None

    metrics = {
        "step": global_step,
        "protocol": "input_fitting" if not heldout_frame_indices else "stage3_interleaved_holdout",
        "alignment_deformation_may_use_heldout": bool(heldout_frame_indices),
        "all": {
            "psnr": _finite_or_none(avg_psnr),
            "ssim": _finite_or_none(avg_ssim),
            "lpips": _finite_or_none(avg_lpips),
        },
        "train": {
            "indices": train_frame_indices,
            "psnr": _finite_or_none(train_psnr),
            "ssim": _finite_or_none(train_ssim),
            "lpips": _finite_or_none(train_lpips),
        },
        "heldout_appearance": {
            "indices": heldout_frame_indices,
            "psnr": _finite_or_none(heldout_psnr),
            "ssim": _finite_or_none(heldout_ssim),
            "lpips": _finite_or_none(heldout_lpips),
        },
        "temporal_flow_diagnostic": temporal,
    }
    with open(os.path.join(eval_dir, "metrics.json"), "w") as f:
        json.dump(metrics, f, indent=2)

    model.train()


if __name__ == "__main__":
    tyro.cli(main)