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, }, }