Spaces:

hugging-apps
/

codecsplat

Running on Zero

codecsplat / src /model /model_wrapper.py
CodecSplat demo
7405e85 verified 7 days ago
48.3 kB
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Protocol, runtime_checkable

import torch
import wandb
from einops import rearrange, repeat, einsum
from jaxtyping import Float
from pytorch_lightning import LightningModule
from pytorch_lightning.loggers.wandb import WandbLogger
from pytorch_lightning.utilities import rank_zero_only
from torch import Tensor, nn, optim
import numpy as np
import json
import os
import time
from tqdm import tqdm
import torch.nn.functional as F
import math

from ..dataset.data_module import get_data_shim
from ..dataset.types import BatchedExample
from ..dataset import DatasetCfg
from ..evaluation.metrics import compute_lpips, compute_psnr, compute_ssim
from ..global_cfg import get_cfg
from ..loss import Loss, LossLpips
from ..misc.benchmarker import Benchmarker
from ..misc.image_io import prep_image, save_image, save_video
from ..misc.LocalLogger import LocalLogger
from ..misc.step_tracker import StepTracker
from ..visualization.annotation import add_label
from ..visualization.camera_trajectory.interpolation import (
    interpolate_extrinsics,
    interpolate_intrinsics,
)
from ..visualization.camera_trajectory.wobble import (
    generate_wobble,
    generate_wobble_transformation,
)
from ..visualization.color_map import apply_color_map_to_image
from ..visualization.layout import add_border, hcat, vcat
from .decoder.decoder import Decoder, DepthRenderingMode, DecoderOutput
from .encoder import Encoder
from .types import Gaussians
from src.visualization.vis_depth import viz_depth_tensor
from PIL import Image
from ..misc.stabilize_camera import render_stabilization_path
from .ply_export import save_gaussian_ply


@dataclass
class OptimizerCfg:
    lr: float
    warm_up_steps: int
    lr_monodepth: float
    weight_decay: float
    div_factor: float
    final_div_factor: float


@dataclass
class TestCfg:
    compute_scores: bool
    save_image: bool
    save_video: bool
    eval_time_skip_steps: int
    save_gt_image: bool
    save_input_images: bool
    save_depth: bool
    save_depth_concat_img: bool
    save_depth_npy: bool
    save_gaussian: bool
    save_bitstream: bool
    render_chunk_size: int | None
    stabilize_camera: bool
    stab_camera_kernel: int
    gaussian_export_all_points: bool
    gaussian_export_align_to_view: bool


@dataclass
class TrainCfg:
    depth_mode: DepthRenderingMode | None
    extended_visualization: bool
    print_log_every_n_steps: int
    eval_model_every_n_val: int
    eval_data_length: int
    eval_time_skip_steps: int
    eval_save_model: bool
    l1_loss: bool
    intermediate_loss_weight: float
    no_viz_video: bool
    train_ignore_large_loss: float


@runtime_checkable
class TrajectoryFn(Protocol):
    def __call__(
        self,
        t: Float[Tensor, " t"],
    ) -> tuple[
        Float[Tensor, "batch view 4 4"],  # extrinsics
        Float[Tensor, "batch view 3 3"],  # intrinsics
    ]:
        pass


class ModelWrapper(LightningModule):
    logger: Optional[WandbLogger]
    encoder: nn.Module
    decoder: Decoder
    losses: nn.ModuleList
    optimizer_cfg: OptimizerCfg
    test_cfg: TestCfg
    train_cfg: TrainCfg
    step_tracker: StepTracker | None
    eval_data_cfg: Optional[DatasetCfg | None]

    def __init__(
        self,
        optimizer_cfg: OptimizerCfg,
        test_cfg: TestCfg,
        train_cfg: TrainCfg,
        encoder: Encoder,
        decoder: Decoder,
        losses: list[Loss],
        step_tracker: StepTracker | None,
        eval_data_cfg: Optional[DatasetCfg | None] = None,
    ) -> None:
        super().__init__()
        self.optimizer_cfg = optimizer_cfg
        self.test_cfg = test_cfg
        self.train_cfg = train_cfg
        self.step_tracker = step_tracker
        self.eval_data_cfg = eval_data_cfg

        # Set up the model.
        self.encoder = encoder
        self.decoder = decoder
        self.data_shim = get_data_shim(self.encoder)
        self.losses = nn.ModuleList(losses)

        # This is used for testing.
        self.benchmarker = Benchmarker()
        self.eval_cnt = 0

        if self.test_cfg.compute_scores:
            self.test_step_outputs = {}
            self.time_skip_steps_dict = {"encoder": 0, "decoder": 0}

    @torch.no_grad()
    def compute_lpips(
        self,
        ground_truth: Float[Tensor, "batch channel height width"],
        predicted: Float[Tensor, "batch channel height width"],
    ) -> Float[Tensor, " batch"]:
        for loss_fn in self.losses:
            if isinstance(loss_fn, LossLpips):
                with torch.autocast(
                    device_type=predicted.device.type,
                    enabled=False,
                ):
                    value = loss_fn.lpips_on_device(predicted.device).forward(
                        ground_truth.float(),
                        predicted.float(),
                        normalize=True,
                    )
                return value[:, 0, 0, 0]

        return compute_lpips(ground_truth, predicted)

    def _codec_loss_weight(self) -> float:
        return float(getattr(getattr(self.encoder, "cfg", None), "codec_loss_weight", 1.0))

    def _codec_depth_loss_weight(self) -> float:
        return float(
            getattr(getattr(self.encoder, "cfg", None), "codec_depth_loss_weight", 0.0)
        )

    def _codec_lmb_weight(self, codec_lmb: Tensor) -> Tensor:
        cfg = getattr(self.encoder, "cfg", None)
        codec_lmb_range = getattr(cfg, "codec_lmb_range", None)
        if codec_lmb_range is None or len(codec_lmb_range) == 0:
            return torch.ones_like(codec_lmb)

        max_lmb = float(codec_lmb_range[-1])
        if max_lmb <= 0:
            return torch.ones_like(codec_lmb)

        return (codec_lmb / max_lmb).clamp_min(0.0)

    def _unpack_encoder_output(
        self,
        encoder_output,
    ) -> tuple[
        Gaussians | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
    ]:
        codec_loss = None
        codec_lmb = None
        codec_num_bytes = None
        codec_bpp = None
        codec_depth_targets = None
        codec_depth_loss_pred = None
        codec_depth_loss_target = None
        if isinstance(encoder_output, tuple):
            if len(encoder_output) != 3:
                raise ValueError(
                    f"Unexpected encoder output tuple length: {len(encoder_output)}"
                )
            encoder_output, codec_loss, codec_lmb = encoder_output

        pred_depths = None
        gaussians = encoder_output
        if isinstance(encoder_output, dict):
            pred_depths = encoder_output.get("depths")
            gaussians = encoder_output["gaussians"]
            codec_num_bytes = encoder_output.get("codec_num_bytes")
            codec_bpp = encoder_output.get("codec_bpp")
            codec_depth_targets = encoder_output.get("codec_depth_targets")
            codec_depth_loss_pred = encoder_output.get("codec_depth_loss_pred")
            codec_depth_loss_target = encoder_output.get("codec_depth_loss_target")

        return (
            gaussians,
            pred_depths,
            codec_depth_targets,
            codec_depth_loss_pred,
            codec_depth_loss_target,
            codec_loss,
            codec_lmb,
            codec_num_bytes,
            codec_bpp,
        )

    def training_step(self, batch, batch_idx):
        batch: BatchedExample = self.data_shim(batch)
        batch_size, _, _, h, w = batch["target"]["image"].shape

        # Run the model.
        encoder_output = self.encoder(
            batch["context"], self.global_step, scene_names=batch["scene"]
        )
        (
            gaussians,
            pred_depths,
            codec_depth_targets,
            codec_depth_loss_pred,
            codec_depth_loss_target,
            codec_loss,
            codec_lmb,
            _,
            _,
        ) = self._unpack_encoder_output(encoder_output)

        supervise_intermediate_depth = False

        if gaussians.means.size(0) != batch["target"]["extrinsics"].size(0):
            supervise_intermediate_depth = True
            assert gaussians.means.size(0) % batch["target"]["extrinsics"].size(0) == 0
            num_depths = gaussians.means.size(0) // batch["target"]["extrinsics"].size(
                0
            )
            # add loss to intermediate depth predictions
            target_extrinsics = torch.cat(
                [batch["target"]["extrinsics"]] * num_depths, dim=0
            )
            target_intrinsics = torch.cat(
                [batch["target"]["intrinsics"]] * num_depths, dim=0
            )
            target_near = torch.cat([batch["target"]["near"]] * num_depths, dim=0)
            target_far = torch.cat([batch["target"]["far"]] * num_depths, dim=0)

            output_all = self.decoder.forward(
                gaussians,
                target_extrinsics,
                target_intrinsics,
                target_near,
                target_far,
                (h, w),
                depth_mode=self.train_cfg.depth_mode,
            )
            # split
            batch_size = batch["target"]["extrinsics"].size(0)
            # order: intermediate depth, final depth
            output_intermediate = DecoderOutput(
                color=output_all.color[:-batch_size],
                depth=(
                    output_all.depth[:-batch_size]
                    if output_all.depth is not None
                    else None
                ),
            )
            output = DecoderOutput(
                color=output_all.color[-batch_size:],
                depth=(
                    output_all.depth[-batch_size:]
                    if output_all.depth is not None
                    else None
                ),
            )

        else:
            output = self.decoder.forward(
                gaussians,
                batch["target"]["extrinsics"],
                batch["target"]["intrinsics"],
                batch["target"]["near"],
                batch["target"]["far"],
                (h, w),
                depth_mode=self.train_cfg.depth_mode,
            )

        target_gt = batch["target"]["image"]

        # Compute metrics.
        psnr = compute_psnr(
            rearrange(target_gt, "b v c h w -> (b v) c h w"),
            rearrange(output.color, "b v c h w -> (b v) c h w"),
        )
        self.log("train/psnr", psnr.mean())

        # Compute and log loss.
        total_loss_per_sample = torch.zeros(batch_size, device=target_gt.device, dtype=target_gt.dtype)

        valid_depth_mask = None

        for loss_fn in self.losses:
            if loss_fn.name == "mse":
                loss_per_sample = loss_fn.forward(
                    output,
                    batch,
                    gaussians,
                    self.global_step,
                    l1_loss=self.train_cfg.l1_loss,
                    clamp_large_error=self.train_cfg.train_ignore_large_loss,
                    valid_depth_mask=valid_depth_mask,
                    reduction="none",
                )
            else:
                loss_per_sample = loss_fn.forward(
                    output,
                    batch,
                    gaussians,
                    self.global_step,
                    valid_depth_mask=valid_depth_mask,
                    reduction="none",
                )
            self.log(f"loss/{loss_fn.name}", loss_per_sample.mean())
            total_loss_per_sample = total_loss_per_sample + loss_per_sample

        if (
            codec_depth_loss_pred is not None
            and codec_depth_loss_target is not None
            and codec_depth_loss_pred.shape == codec_depth_loss_target.shape
        ):
            codec_depth_loss = F.l1_loss(
                codec_depth_loss_pred,
                codec_depth_loss_target.detach(),
                reduction="none",
            ).mean(dim=(1, 2, 3))
            self.log("loss/codec_depth", codec_depth_loss.mean())
            if (
                pred_depths is not None
                and codec_depth_targets is not None
                and pred_depths.shape == codec_depth_targets.shape
            ):
                codec_depth_metric_l1 = F.l1_loss(
                    pred_depths,
                    codec_depth_targets.detach(),
                    reduction="none",
                ).mean(dim=(1, 2, 3))
                self.log("loss/codec_depth_metric_l1", codec_depth_metric_l1.mean())
            total_loss_per_sample = (
                total_loss_per_sample
                + self._codec_depth_loss_weight() * codec_depth_loss
            )

        # color loss on intermediate output
        if supervise_intermediate_depth:
            for loss_fn in self.losses:
                if output_intermediate.color.size(0) != batch_size:
                    assert output_intermediate.color.size(0) % batch_size == 0
                    num_intermediate = output_intermediate.color.size(0) // batch_size
                    intermediate_loss = torch.zeros(
                        batch_size,
                        device=target_gt.device,
                        dtype=target_gt.dtype,
                    )
                    for i in range(num_intermediate):
                        curr_output = DecoderOutput(
                            color=output_intermediate.color[
                                (batch_size * i) : (batch_size * (i + 1))
                            ],
                            depth=(
                                output_intermediate.depth[
                                    (batch_size * i) : (batch_size * (i + 1))
                                ]
                                if output_intermediate.depth is not None
                                else None
                            ),
                        )
                        curr_loss_weight = self.train_cfg.intermediate_loss_weight ** (
                            num_intermediate - i
                        )

                        if loss_fn.name == "mse":
                            loss = loss_fn.forward(
                                curr_output,
                                batch,
                                gaussians,
                                self.global_step,
                                l1_loss=self.train_cfg.l1_loss,
                                clamp_large_error=self.train_cfg.train_ignore_large_loss,
                                valid_depth_mask=valid_depth_mask,
                                reduction="none",
                            )
                        else:
                            loss = loss_fn.forward(
                                curr_output,
                                batch,
                                gaussians,
                                self.global_step,
                                valid_depth_mask=valid_depth_mask,
                                reduction="none",
                            )

                        intermediate_loss = intermediate_loss + curr_loss_weight * loss

                    self.log(f"loss/{loss_fn.name}_intermediate", intermediate_loss.mean())
                    total_loss_per_sample = total_loss_per_sample + intermediate_loss
                else:
                    if loss_fn.name == "mse":
                        loss = loss_fn.forward(
                            output_intermediate,
                            batch,
                            gaussians,
                            self.global_step,
                            l1_loss=self.train_cfg.l1_loss,
                            clamp_large_error=self.train_cfg.train_ignore_large_loss,
                            valid_depth_mask=valid_depth_mask,
                            reduction="none",
                        )
                    else:
                        loss = loss_fn.forward(
                            output_intermediate,
                            batch,
                            gaussians,
                            self.global_step,
                            valid_depth_mask=valid_depth_mask,
                            reduction="none",
                        )
                    weighted_loss = self.train_cfg.intermediate_loss_weight * loss
                    self.log(f"loss/{loss_fn.name}_intermediate", weighted_loss.mean())
                    total_loss_per_sample = (
                        total_loss_per_sample + weighted_loss
                    )

        distortion_loss_per_sample = total_loss_per_sample

        if codec_loss is not None:
            codec_loss_mean = codec_loss.mean()
            self.log("loss/codec", codec_loss_mean)
            num_context_views = batch["context"]["image"].shape[1]
            self.log("loss/codec_per_view", (codec_loss / num_context_views).mean())
            if codec_lmb is not None:
                self.log("loss/codec_lmb", codec_lmb.mean())
                codec_lmb_weight = self._codec_lmb_weight(codec_lmb)
                self.log("loss/codec_lmb_weight", codec_lmb_weight.mean())
                distortion_loss_per_sample = distortion_loss_per_sample * codec_lmb_weight
            total_loss_per_sample = distortion_loss_per_sample + self._codec_loss_weight() * codec_loss
        else:
            total_loss_per_sample = distortion_loss_per_sample

        self.log("loss/distortion", distortion_loss_per_sample.mean())
        total_loss = total_loss_per_sample.mean()

        self.log("loss/total", total_loss)

        if (
            self.global_rank == 0
            and self.global_step % self.train_cfg.print_log_every_n_steps == 0
        ):
            print(
                f"train step {self.global_step}; "
                f"scene = {[x[:20] for x in batch['scene']]}; "
                f"context = {batch['context']['index'].tolist()}; "
                f"bound = [{batch['context']['near'].detach().cpu().numpy().mean()} "
                f"{batch['context']['far'].detach().cpu().numpy().mean()}]; "
                f"loss = {total_loss:.6f}"
            )
        self.log("info/near", batch["context"]["near"].detach().cpu().numpy().mean())
        self.log("info/far", batch["context"]["far"].detach().cpu().numpy().mean())
        self.log("info/global_step", self.global_step)  # hack for ckpt monitor

        # Tell the data loader processes about the current step.
        if self.step_tracker is not None:
            self.step_tracker.set_step(self.global_step)

        return total_loss

    def test_step(self, batch, batch_idx):
        batch: BatchedExample = self.data_shim(batch)
        b, v, _, h, w = batch["target"]["image"].shape
        assert b == 1

        # save input views for visualization
        if self.test_cfg.save_input_images:
            (scene,) = batch["scene"]
            path = Path(get_cfg()["output_dir"])

            input_images = batch["context"]["image"][0]  # [V, 3, H, W]
            index = batch["context"]["index"][0]
            for idx, color in zip(index, input_images):
                save_image(color, path / "images" / scene / f"color/input_{idx:0>6}.png")

        # save depth vis
        if self.test_cfg.save_depth or self.test_cfg.save_gaussian:
            visualization_dump = {}
        else:
            visualization_dump = None

        # Render Gaussians.
        with self.benchmarker.time("encoder"):
            encoder_output = self.encoder(
                batch["context"],
                self.global_step,
                visualization_dump=visualization_dump,
                return_bitstream=self.test_cfg.save_bitstream,
            )
            (
                gaussians,
                _,
                _,
                _,
                _,
                _,
                _,
                codec_num_bytes,
                codec_bpp,
            ) = self._unpack_encoder_output(encoder_output)
            codec_extra_metrics = {}
            if isinstance(encoder_output, dict):
                for metric_name in (
                    "codec_encode_context_seconds",
                    "codec_compress_seconds",
                    "codec_decompress_seconds",
                    "codec_camera_num_bytes",
                    "codec_feature_num_bytes",
                    "codec_estimated_feature_num_bytes",
                    "codec_latent_feature_num_bytes",
                    "codec_depth_num_bytes",
                    "codec_feature_raw_num_bytes",
                    "codec_depth_raw_num_bytes",
                ):
                    metric_value = encoder_output.get(metric_name)
                    if metric_value is not None:
                        codec_extra_metrics[metric_name] = metric_value

        if self.test_cfg.save_bitstream:
            if not isinstance(encoder_output, dict) or "bitstream" not in encoder_output:
                raise RuntimeError(
                    "test.save_bitstream=true requires an enabled learned feature codec."
                )
            scene = batch["scene"][0]
            bitstream_path = (
                Path(get_cfg()["output_dir"]) / "bitstreams" / f"{scene}.bin"
            )
            bitstream_path.parent.mkdir(parents=True, exist_ok=True)
            bitstream_path.write_bytes(encoder_output["bitstream"])

        # save gaussians
        if self.test_cfg.save_gaussian:
            scene = batch["scene"][0]
            save_path = Path(get_cfg()['output_dir']) / 'gaussians' / (scene + '.ply')
            save_gaussian_ply(
                gaussians,
                visualization_dump,
                batch,
                save_path,
                save_all_gaussians=self.test_cfg.gaussian_export_all_points,
                align_to_view=self.test_cfg.gaussian_export_align_to_view,
            )

        with self.benchmarker.time("decoder", num_calls=v):
            camera_poses = batch["target"]["extrinsics"]

            if self.test_cfg.stabilize_camera:
                stable_poses = render_stabilization_path(
                    camera_poses[0].detach().cpu().numpy(),
                    k_size=self.test_cfg.stab_camera_kernel,
                )

                stable_poses = list(
                    map(
                        lambda x: np.concatenate(
                            (x, np.array([[0.0, 0.0, 0.0, 1.0]])), axis=0
                        ),
                        stable_poses,
                    )
                )
                stable_poses = torch.from_numpy(np.stack(stable_poses, axis=0)).to(
                    camera_poses
                )
                camera_poses = stable_poses.unsqueeze(0)

            if self.test_cfg.render_chunk_size is not None:
                chunk_size = self.test_cfg.render_chunk_size
                num_chunks = math.ceil(camera_poses.shape[1] / chunk_size)

                output = None
                for i in range(num_chunks):
                    start = chunk_size * i
                    end = chunk_size * (i + 1)

                    render_intrinsics = batch["target"]["intrinsics"]
                    render_near = batch["target"]["near"]
                    render_far = batch["target"]["far"]

                    curr_output = self.decoder.forward(
                        gaussians,
                        camera_poses[:, start:end],
                        render_intrinsics[:, start:end],
                        render_near[:, start:end],
                        render_far[:, start:end],
                        (h, w),
                        depth_mode=None,
                    )

                    if i == 0:
                        output = curr_output
                    else:
                        # ignore depth
                        output.color = torch.cat(
                            (output.color, curr_output.color), dim=1
                        )

            else:
                output = self.decoder.forward(
                    gaussians,
                    camera_poses,
                    batch["target"]["intrinsics"],
                    batch["target"]["near"],
                    batch["target"]["far"],
                    (h, w),
                    depth_mode=None,
                )

        (scene,) = batch["scene"]
        path = Path(get_cfg()["output_dir"])

        # save depth
        if self.test_cfg.save_depth:
            depth = (
                visualization_dump["depth"][0, :, :, :, 0, 0].cpu().detach()
            )  # [V, H, W]

            index = batch["context"]["index"][0]

            if self.test_cfg.save_depth_concat_img:
                # concat (img0, img1, depth0, depth1)
                image = batch['context']['image'][0]  # [V, 3, H, W] in [0,1]
                image = rearrange(image, "b c h w -> h (b w) c")  # [H, VW, 3]
                image_concat = (image.detach().cpu().numpy() * 255).astype(np.uint8)  # [H, VW, 3]

                depth_concat = []

            for idx, depth_i in zip(index, depth):
                depth_viz = viz_depth_tensor(
                    1.0 / depth_i, return_numpy=True
                )  # [H, W, 3]

                if self.test_cfg.save_depth_concat_img:
                    depth_concat.append(depth_viz)

                save_path = path / "images" / scene / "depth" / f"{idx:0>6}.png"
                save_dir = os.path.dirname(save_path)
                os.makedirs(save_dir, exist_ok=True)
                Image.fromarray(depth_viz).save(save_path)

                # save depth as npy
                if self.test_cfg.save_depth_npy:
                    depth_npy = depth_i.detach().cpu().numpy()
                    save_path = path / "images" / scene / "depth" / f"{idx:0>6}.npy"
                    save_dir = os.path.dirname(save_path)
                    os.makedirs(save_dir, exist_ok=True)
                    np.save(save_path, depth_npy)

            if self.test_cfg.save_depth_concat_img:
                depth_concat = np.concatenate(depth_concat, axis=1)  # [H, VW, 3]
                concat = np.concatenate((image_concat, depth_concat), axis=0)  # [2H, VW, 3]

                save_path = path / "images" / scene / "depth" /  f"img_depth_{scene}.png"
                save_dir = os.path.dirname(save_path)
                os.makedirs(save_dir, exist_ok=True)
                Image.fromarray(concat).save(save_path)

        images_prob = output.color[0]
        rgb_gt = batch["target"]["image"][0]

        # Save images.
        if self.test_cfg.save_image:
            if self.test_cfg.save_gt_image:
                for index, color, gt in zip(
                    batch["target"]["index"][0], images_prob, rgb_gt
                ):
                    save_image(color, path / "images" / scene / f"color/{index:0>6}.png")
                    save_image(gt, path / "images" / scene / f"color/{index:0>6}_gt.png")
            else:
                for index, color in zip(batch["target"]["index"][0], images_prob):
                    save_image(color, path / "images" / scene / f"color/{index:0>6}.png")

        # save video
        if self.test_cfg.save_video:
            frame_str = "_".join([str(x.item()) for x in batch["context"]["index"][0]])
            save_video(
                [a for a in images_prob],
                path / "videos" / f"{scene}_frame_{frame_str}.mp4",
            )

        # compute scores
        if self.test_cfg.compute_scores:
            if batch_idx < self.test_cfg.eval_time_skip_steps:
                self.time_skip_steps_dict["encoder"] += 1
                self.time_skip_steps_dict["decoder"] += v

            rgb = images_prob

            if f"psnr" not in self.test_step_outputs:
                self.test_step_outputs[f"psnr"] = []
            if f"ssim" not in self.test_step_outputs:
                self.test_step_outputs[f"ssim"] = []
            if f"lpips" not in self.test_step_outputs:
                self.test_step_outputs[f"lpips"] = []
            if codec_bpp is not None and f"codec_bpp" not in self.test_step_outputs:
                self.test_step_outputs[f"codec_bpp"] = []
            if (
                codec_num_bytes is not None
                and f"codec_num_bytes" not in self.test_step_outputs
            ):
                self.test_step_outputs[f"codec_num_bytes"] = []
            for metric_name in codec_extra_metrics:
                if metric_name not in self.test_step_outputs:
                    self.test_step_outputs[metric_name] = []

            self.test_step_outputs[f"psnr"].append(
                compute_psnr(rgb_gt, rgb).mean().item()
            )
            self.test_step_outputs[f"ssim"].append(
                compute_ssim(rgb_gt, rgb).mean().item()
            )
            self.test_step_outputs[f"lpips"].append(
                self.compute_lpips(rgb_gt, rgb).mean().item()
            )
            if codec_bpp is not None:
                self.test_step_outputs[f"codec_bpp"].append(codec_bpp.item())
            if codec_num_bytes is not None:
                self.test_step_outputs[f"codec_num_bytes"].append(
                    codec_num_bytes.item()
                )
            for metric_name, metric_value in codec_extra_metrics.items():
                self.test_step_outputs[metric_name].append(metric_value.item())

    def on_test_end(self) -> None:
        out_dir = Path(get_cfg()["output_dir"]) / "metrics"
        saved_scores = {}
        if self.test_cfg.compute_scores:
            self.benchmarker.dump_memory(out_dir / "peak_memory.json")
            self.benchmarker.dump(out_dir / "benchmark.json")

            for metric_name, metric_scores in self.test_step_outputs.items():
                avg_scores = sum(metric_scores) / len(metric_scores)
                saved_scores[metric_name] = avg_scores
                print(metric_name, avg_scores)
                with (out_dir / f"scores_{metric_name}_all.json").open("w") as f:
                    json.dump(metric_scores, f)
                metric_scores.clear()

            for tag, times in self.benchmarker.execution_times.items():
                times = times[int(self.time_skip_steps_dict[tag]) :]
                saved_scores[tag] = [len(times), np.mean(times)]
                print(
                    f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call"
                )
                self.time_skip_steps_dict[tag] = 0

            with (out_dir / f"scores_all_avg.json").open("w") as f:
                json.dump(saved_scores, f)
            self.benchmarker.clear_history()
        else:
            self.benchmarker.dump(out_dir / "benchmark.json")
            self.benchmarker.dump_memory(out_dir / "peak_memory.json")
            self.benchmarker.summarize()

    @rank_zero_only
    def validation_step(self, batch, batch_idx):
        batch: BatchedExample = self.data_shim(batch)

        if self.global_rank == 0:
            print(
                f"validation step {self.global_step}; "
                f"scene = {[a[:20] for a in batch['scene']]}; "
                f"context = {batch['context']['index'].tolist()}"
            )

        # Render Gaussians.
        b, _, _, h, w = batch["target"]["image"].shape
        assert b == 1
        gaussians_softmax = self.encoder(
            batch["context"],
            self.global_step,
        )
        (
            gaussians_softmax,
            pred_depths,
            _,
            _,
            _,
            _,
            _,
            codec_num_bytes,
            codec_bpp,
        ) = self._unpack_encoder_output(
            gaussians_softmax
        )

        output_softmax = self.decoder.forward(
            gaussians_softmax,
            batch["target"]["extrinsics"],
            batch["target"]["intrinsics"],
            batch["target"]["near"],
            batch["target"]["far"],
            (h, w),
        )
        rgb_softmax = output_softmax.color[0]

        # Compute validation metrics.
        rgb_gt = batch["target"]["image"][0]
        for tag, rgb in zip(("val",), (rgb_softmax,)):
            psnr = compute_psnr(rgb_gt, rgb).mean()
            self.log(f"val/psnr_{tag}", psnr)
            lpips = self.compute_lpips(rgb_gt, rgb).mean()
            self.log(f"val/lpips_{tag}", lpips)
            ssim = compute_ssim(rgb_gt, rgb).mean()
            self.log(f"val/ssim_{tag}", ssim)
        if codec_bpp is not None:
            self.log("val/codec_bpp", codec_bpp)
        if codec_num_bytes is not None:
            self.log("val/codec_num_bytes", codec_num_bytes)

        # viz depth
        if pred_depths is not None:
            # only visualize predicted depth
            pred_depths = pred_depths[0]  # [V, H, W]

            # gaussian downsample
            if pred_depths.shape[1:] != batch["context"]["image"].shape[-2:]:
                pred_depths = F.interpolate(
                    pred_depths.unsqueeze(1),
                    size=batch["context"]["image"].shape[-2:],
                    mode="bilinear",
                    align_corners=True,
                ).squeeze(1)

            inverse_depth_pred = 1.0 / pred_depths

            concat = []
            for i in range(inverse_depth_pred.size(0)):
                concat.append(inverse_depth_pred[i])

            concat = torch.cat(concat, dim=1)  # [H, W*N]

            depth_viz = viz_depth_tensor(concat.cpu().detach())  # [3, H, W*N]

            # also concat images
            input_images = batch["context"]["image"][0]  # [N, 3, H, W]
            concat_img = [img for img in input_images]
            concat_img = torch.cat(concat_img, dim=-1) * 255  # [3, H, W*N]

            concat = torch.cat(
                (concat_img.cpu().detach(), depth_viz), dim=1
            )  # [3, H*2, W*N]

            self.logger.log_image(
                "depth",
                [concat],
                step=self.global_step,
                caption=batch["scene"],
            )

        # Construct comparison image.
        comparison = hcat(
            add_label(vcat(*batch["context"]["image"][0]), "Context"),
            add_label(vcat(*rgb_gt), "Target (Ground Truth)"),
            add_label(vcat(*rgb_softmax), "Target (Prediction)"),
        )
        self.logger.log_image(
            "comparison",
            [prep_image(add_border(comparison))],
            step=self.global_step,
            caption=batch["scene"],
        )

        # Run video validation step.
        if not self.train_cfg.no_viz_video:
            self.render_video_interpolation(batch)
            self.render_video_wobble(batch)
            if self.train_cfg.extended_visualization:
                self.render_video_interpolation_exaggerated(batch)

    def on_validation_epoch_end(self) -> None:
        """hack to run the full validation"""
        if self.trainer.sanity_checking and self.global_rank == 0:
            print(self.encoder)  # log the model to wandb log files

        if (not self.trainer.sanity_checking) and (self.eval_data_cfg is not None):
            self.eval_cnt = self.eval_cnt + 1
            if self.eval_cnt % self.train_cfg.eval_model_every_n_val == 0:
                # backup current ckpt before running full test sets eval
                if self.train_cfg.eval_save_model:
                    ckpt_saved_path = (
                        self.trainer.checkpoint_callback.format_checkpoint_name(
                            dict(
                                epoch=self.trainer.current_epoch,
                                step=self.trainer.global_step,
                            )
                        )
                    )
                    backup_dir = str(
                        Path(ckpt_saved_path).parent.parent / "checkpoints_backups"
                    )
                    if self.global_rank == 0:
                        os.makedirs(backup_dir, exist_ok=True)
                    ckpt_saved_path = os.path.join(
                        backup_dir, os.path.basename(ckpt_saved_path)
                    )
                    # call save_checkpoint on ALL process as suggested by pytorch_lightning
                    self.trainer.save_checkpoint(
                        ckpt_saved_path,
                        weights_only=True,
                    )
                    if self.global_rank == 0:
                        print(f"backup model to {ckpt_saved_path}.")

                # run full test sets eval on rank=0 device
                self.run_full_test_sets_eval()

    @rank_zero_only
    def run_full_test_sets_eval(self) -> None:
        start_t = time.time()

        pred_depths = None
        depth_gt = None

        full_testsets = self.trainer.datamodule.test_dataloader(
            dataset_cfg=self.eval_data_cfg
        )
        scores_dict: dict[str, list[float]] = {}

        for score_tag in ("psnr", "ssim", "lpips"):
            scores_dict[score_tag] = []

        # evaluate depth
        for score_tag in ("abs_rel", "rmse", "a1"):
            scores_dict[score_tag] = []

        self.benchmarker.clear_history()
        time_skip_first_n_steps = min(
            self.train_cfg.eval_time_skip_steps, len(full_testsets)
        )
        time_skip_steps_dict = {"encoder": 0, "decoder": 0}
        for batch_idx, batch in tqdm(
            enumerate(full_testsets),
            total=min(len(full_testsets), self.train_cfg.eval_data_length),
        ):
            if batch_idx >= self.train_cfg.eval_data_length:
                break

            batch = self.data_shim(batch)
            batch = self.transfer_batch_to_device(batch, "cuda", dataloader_idx=0)

            # Render Gaussians.
            b, v, _, h, w = batch["target"]["image"].shape
            assert b == 1
            if batch_idx < time_skip_first_n_steps:
                time_skip_steps_dict["encoder"] += 1
                time_skip_steps_dict["decoder"] += v

            with self.benchmarker.time("encoder"):
                encoder_output = self.encoder(
                    batch["context"],
                    self.global_step,
                )
                (
                    gaussians,
                    pred_depths,
                    _,
                    _,
                    _,
                    _,
                    _,
                    codec_num_bytes,
                    codec_bpp,
                ) = self._unpack_encoder_output(encoder_output)

            with self.benchmarker.time("decoder", num_calls=v):
                output = self.decoder.forward(
                    gaussians,
                    batch["target"]["extrinsics"],
                    batch["target"]["intrinsics"],
                    batch["target"]["near"],
                    batch["target"]["far"],
                    (h, w),
                )

            # Compute validation metrics.
            rgb_gt = batch["target"]["image"][0]
            rgb = output.color[0]
            scores_dict["psnr"].append(compute_psnr(rgb_gt, rgb).mean().item())
            scores_dict["lpips"].append(
                self.compute_lpips(rgb_gt, rgb).mean().item()
            )
            scores_dict["ssim"].append(compute_ssim(rgb_gt, rgb).mean().item())

            if codec_bpp is not None:
                scores_dict.setdefault("codec_bpp", []).append(codec_bpp.item())
            if codec_num_bytes is not None:
                scores_dict.setdefault("codec_num_bytes", []).append(
                    codec_num_bytes.item()
                )

        # summarise scores and log to logger
        for score_tag, scores in scores_dict.items():
            if scores:
                self.log(f"test/{score_tag}", sum(scores) / len(scores))
        # summarise run time
        for tag, times in self.benchmarker.execution_times.items():
            times = times[int(time_skip_steps_dict[tag]) :]
            print(f"{tag}: {len(times)} calls, avg. {np.mean(times)} seconds per call")
            self.log(f"test/runtime_avg_{tag}", np.mean(times))
        self.benchmarker.clear_history()

        overall_eval_time = time.time() - start_t
        print(f"Eval total time cost: {overall_eval_time:.3f}s")
        self.log("test/runtime_all", overall_eval_time)

    @rank_zero_only
    def render_video_wobble(self, batch: BatchedExample) -> None:
        # Two views are needed to get the wobble radius.
        _, v, _, _ = batch["context"]["extrinsics"].shape
        if v != 2:
            return

        def trajectory_fn(t):
            origin_a = batch["context"]["extrinsics"][:, 0, :3, 3]
            origin_b = batch["context"]["extrinsics"][:, 1, :3, 3]
            delta = (origin_a - origin_b).norm(dim=-1)
            extrinsics = generate_wobble(
                batch["context"]["extrinsics"][:, 0],
                delta * 0.25,
                t,
            )
            intrinsics = repeat(
                batch["context"]["intrinsics"][:, 0],
                "b i j -> b v i j",
                v=t.shape[0],
            )
            return extrinsics, intrinsics

        return self.render_video_generic(batch, trajectory_fn, "wobble", num_frames=60)

    @rank_zero_only
    def render_video_interpolation(self, batch: BatchedExample) -> None:
        _, v, _, _ = batch["context"]["extrinsics"].shape

        def trajectory_fn(t):
            extrinsics = interpolate_extrinsics(
                batch["context"]["extrinsics"][0, 0],
                (
                    batch["context"]["extrinsics"][0, 1]
                    if v == 2
                    else batch["target"]["extrinsics"][0, 0]
                ),
                t,
            )
            intrinsics = interpolate_intrinsics(
                batch["context"]["intrinsics"][0, 0],
                (
                    batch["context"]["intrinsics"][0, 1]
                    if v == 2
                    else batch["target"]["intrinsics"][0, 0]
                ),
                t,
            )
            return extrinsics[None], intrinsics[None]

        return self.render_video_generic(batch, trajectory_fn, "rgb")

    @rank_zero_only
    def render_video_interpolation_exaggerated(self, batch: BatchedExample) -> None:
        # Two views are needed to get the wobble radius.
        _, v, _, _ = batch["context"]["extrinsics"].shape
        if v != 2:
            return

        def trajectory_fn(t):
            origin_a = batch["context"]["extrinsics"][:, 0, :3, 3]
            origin_b = batch["context"]["extrinsics"][:, 1, :3, 3]
            delta = (origin_a - origin_b).norm(dim=-1)
            tf = generate_wobble_transformation(
                delta * 0.5,
                t,
                5,
                scale_radius_with_t=False,
            )
            extrinsics = interpolate_extrinsics(
                batch["context"]["extrinsics"][0, 0],
                (
                    batch["context"]["extrinsics"][0, 1]
                    if v == 2
                    else batch["target"]["extrinsics"][0, 0]
                ),
                t * 5 - 2,
            )
            intrinsics = interpolate_intrinsics(
                batch["context"]["intrinsics"][0, 0],
                (
                    batch["context"]["intrinsics"][0, 1]
                    if v == 2
                    else batch["target"]["intrinsics"][0, 0]
                ),
                t * 5 - 2,
            )
            return extrinsics @ tf, intrinsics[None]

        return self.render_video_generic(
            batch,
            trajectory_fn,
            "interpolation_exaggerated",
            num_frames=300,
            smooth=False,
            loop_reverse=False,
        )

    @rank_zero_only
    def render_video_generic(
        self,
        batch: BatchedExample,
        trajectory_fn: TrajectoryFn,
        name: str,
        num_frames: int = 30,
        smooth: bool = True,
        loop_reverse: bool = True,
    ) -> None:
        encoder_output = self.encoder(batch["context"], self.global_step)
        gaussians, _, _, _, _, _, _, _, _ = self._unpack_encoder_output(encoder_output)

        t = torch.linspace(0, 1, num_frames, dtype=torch.float32, device=self.device)
        if smooth:
            t = (torch.cos(torch.pi * (t + 1)) + 1) / 2

        extrinsics, intrinsics = trajectory_fn(t)

        _, _, _, h, w = batch["context"]["image"].shape

        # Color-map the result.
        def depth_map(result):
            near = result[result > 0][:16_000_000].quantile(0.01).log()
            far = result.view(-1)[:16_000_000].quantile(0.99).log()
            result = result.log()
            result = 1 - (result - near) / (far - near)
            return apply_color_map_to_image(result, "turbo")

        near = repeat(batch["context"]["near"][:, 0], "b -> b v", v=num_frames)
        far = repeat(batch["context"]["far"][:, 0], "b -> b v", v=num_frames)
        output = self.decoder.forward(
            gaussians, extrinsics, intrinsics, near, far, (h, w), "depth"
        )
        predictions = [
            vcat(rgb, depth)
            for rgb, depth in zip(output.color[0], depth_map(output.depth[0]))
        ]

        images = [
            add_border(hcat(add_label(prediction, "Prediction")))
            for prediction in predictions
        ]

        video = torch.stack(images)
        if loop_reverse:
            video = torch.cat((video, video.flip(0)[1:-1]), dim=0)

        key = f"video/{name}"
        if isinstance(self.logger, WandbLogger):
            video_uint8 = (
                (video.clip(min=0, max=1) * 255).type(torch.uint8).cpu().numpy()
            )
            wandb.log(
                {key: wandb.Video(video_uint8[None], fps=30, format="mp4")}
            )
        else:
            assert isinstance(self.logger, LocalLogger)
            save_video(
                list(video),
                self.logger.log_path / key / f"{self.global_step:0>6}.mp4",
                fps=30,
            )

    def configure_optimizers(self):
        if self.optimizer_cfg.lr_monodepth > 0:
            pretrained_params = []
            new_params = []

            for name, param in self.named_parameters():
                if not param.requires_grad:
                    continue
                if "pretrained" in name:
                    pretrained_params.append(param)
                else:
                    new_params.append(param)

            optimizer = torch.optim.AdamW(
                [
                    {
                        "params": pretrained_params,
                        "lr": self.optimizer_cfg.lr_monodepth,
                    },
                    {"params": new_params, "lr": self.optimizer_cfg.lr},
                ],
                weight_decay=self.optimizer_cfg.weight_decay,
            )

            scheduler = torch.optim.lr_scheduler.OneCycleLR(
                optimizer,
                [self.optimizer_cfg.lr_monodepth, self.optimizer_cfg.lr],
                self.trainer.max_steps + 10,
                pct_start=0.01,
                div_factor=self.optimizer_cfg.div_factor,
                final_div_factor=self.optimizer_cfg.final_div_factor,
                cycle_momentum=False,
                anneal_strategy="cos",
            )

        else:
            optimizer = optim.AdamW(
                [param for param in self.parameters() if param.requires_grad],
                lr=self.optimizer_cfg.lr,
                weight_decay=self.optimizer_cfg.weight_decay,
            )

            scheduler = torch.optim.lr_scheduler.OneCycleLR(
                optimizer,
                self.optimizer_cfg.lr,
                self.trainer.max_steps + 10,
                pct_start=0.01,
                div_factor=self.optimizer_cfg.div_factor,
                final_div_factor=self.optimizer_cfg.final_div_factor,
                cycle_momentum=False,
                anneal_strategy="cos",
            )

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "interval": "step",
                "frequency": 1,
            },
        }