diff --git a/.gitignore b/.gitignore index 5b7992b..e244cb1 100644 --- a/.gitignore +++ b/.gitignore @@ -184,4 +184,16 @@ app.log # data /data/ -*.png \ No newline at end of file +*.png + +# machine-specific run / inspection scripts (private paths, accounts, node names) +scripts/*.sbatch +scripts/export_pointcloud.py +scripts/preview_pointcloud.py +scripts/smoke_erayzer_load.py +scripts/viz_erayzer.py + +# fine-tune config holds a machine-specific init_checkpoint path (kept local) +configs/multiview/erayzer_cheese3d_ft.yaml +# local model checkpoints (machine-specific, large) +checkpoints/ diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 711d493..b54b23e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -15,7 +15,8 @@ See our [code of conduct](CODE_OF_CONDUCT.md) for more information. install in editable mode with dev dependencies: ```bash -pip install -e ".[dev]" +pip install lightning poetry-core +pip install -e ".[dev]" --no-build-isolation pre-commit install ``` diff --git a/README.md b/README.md index 6291637..e66c2c6 100644 --- a/README.md +++ b/README.md @@ -45,15 +45,23 @@ For Github cloning: ```commandline git clone https://github.com/paninski-lab/beast cd beast -pip install -e . +pip install lightning poetry-core +pip install -e . --no-build-isolation ``` For installation through PyPI: ```commandline -pip install beast-backbones +pip install lightning poetry-core +pip install beast-backbones --no-build-isolation ``` +> **Note:** `beast` depends on a custom fork of +> [gsplat](https://github.com/QitaoZhao/gsplat) that must be compiled from source. The +> `gsplat` build requires `torch` (provided by `lightning`) and the build backend requires +> `poetry-core`. Installing these first and using `--no-build-isolation` lets the build find +> them in your environment. + ## Usage > The commands below are for the **single-view BEAST model**. diff --git a/beast/api/model.py b/beast/api/model.py index 84044da..67efaad 100644 --- a/beast/api/model.py +++ b/beast/api/model.py @@ -10,14 +10,12 @@ import torch -from beast.config import BeastConfig +from beast.config import get_beast_config_class from beast.inference import predict_images, predict_video from beast.io import load_config from beast.logging import log_step from beast.models.base import BaseLightningModel -from beast.models.resnets import ResnetAutoencoder -from beast.models.vits import VisionTransformer -from beast.train import train +from beast.models.registry import MODEL_REGISTRY, TRAIN_REGISTRY _logger = logging.getLogger(__name__) @@ -46,12 +44,6 @@ class Model: This class manages both the model and the training/inference processes. """ - MODEL_REGISTRY = { - 'vit': VisionTransformer, - 'resnet': ResnetAutoencoder, - # Add more models as needed - } - def __init__( self, model: BaseLightningModel, @@ -83,11 +75,11 @@ def from_dir(cls, model_dir: str | Path) -> 'Model': config = load_config(config_path) model_type = config['model'].get('model_class', '').lower() - if model_type not in cls.MODEL_REGISTRY: + if model_type not in MODEL_REGISTRY: raise ValueError(f'Unknown model type: {model_type}') # Initialize the LightningModule - model_class = cls.MODEL_REGISTRY[model_type] + model_class = MODEL_REGISTRY[model_type] model = model_class(config) _logger.info(f'Loaded a {model_class} model') @@ -116,14 +108,15 @@ def from_config(cls, config_path: str | Path | dict) -> 'Model': if not isinstance(config_path, dict): config = load_config(config_path) else: - config = BeastConfig.model_validate(config_path).model_dump() + model_class = (config_path.get('model') or {}).get('model_class', '') + config = get_beast_config_class(model_class).model_validate(config_path).model_dump() model_type = config['model'].get('model_class', '').lower() - if model_type not in cls.MODEL_REGISTRY: + if model_type not in MODEL_REGISTRY: raise ValueError(f'Unknown model type: {model_type}') # Initialize the LightningModule - model_class = cls.MODEL_REGISTRY[model_type] + model_class = MODEL_REGISTRY[model_type] log_step(f"Creating {model_type} model instance", level='debug') log_step( f"About to call {model_class.__name__}.__init__() - this may take several" @@ -148,8 +141,10 @@ def train(self, output_dir: str | Path = 'runs/default') -> None: """ self.model_dir = Path(output_dir) + model_type = self.config['model'].get('model_class', '').lower() + train_fn = TRAIN_REGISTRY[model_type] with chdir(self.model_dir): - self.model = train(self.config, self.model, output_dir=self.model_dir) + self.model = train_fn(self.config, self.model, output_dir=self.model_dir) def predict_images( self, diff --git a/beast/config.py b/beast/config.py index 27b0e07..9164221 100644 --- a/beast/config.py +++ b/beast/config.py @@ -1,4 +1,32 @@ -"""Pydantic models for BEAST configuration validation.""" +"""Pydantic schemas for BEAST configuration validation and per-model dispatch. + +Top-level config classes +------------------------ +BeastConfig + Used for models that share the standard training loop (resnet, vit). + Contains TrainingConfig, OptimizerConfig, and DataConfig. + +ERayZerBeastConfig + Used for ERayZer and its subclasses (e.g. BEAST3D if they share ERayZer's + training schema). Contains ERayZerTrainingConfig and ERayZerOptimizerConfig. + +Shared schemas (reusable for new models) +----------------------------------------- +TrainingConfig — standard epoch-based training fields (batch sizes, epochs, imgaug, …) +OptimizerConfig — AdamW/Adam + cosine/step scheduler fields +DataConfig — data_dir only; extend or replace for richer data configs + +Per-model dispatch +------------------ +get_beast_config_class(model_class) returns the correct top-level config class for a +given model_class string. Both beast.io.load_config and beast.api.model.Model.from_config +use this function so that YAML files and raw dicts are validated against the right schema. + +To add a new model with a divergent training schema: + 1. Define TrainingConfig and OptimizerConfig in the model's *_config.py + 2. Define BeastConfig here (avoids circular imports since DataConfig lives here) + 3. Add 'model_class': BeastConfig to _MODEL_CONFIG_CLASSES +""" from __future__ import annotations @@ -7,68 +35,21 @@ from pydantic import BaseModel, Field +from beast.models.beast_resnet.beast_resnet_config import ResnetModelConfig +from beast.models.beast_vit.beast_vit_config import VitModelConfig +from beast.models.erayzer.erayzer_config import ( + ERayZerModelConfig, + ERayZerOptimizerConfig, + ERayZerTrainingConfig, +) + class BeastConfig(BaseModel): model: ModelConfig training: TrainingConfig optimizer: OptimizerConfig data: DataConfig - - -class ResnetModelConfig(BaseModel): - model_class: Literal['resnet'] - model_params: ResnetModelParams - seed: int = 0 - checkpoint: str | None = None - - -class ResnetModelParams(BaseModel): - backbone: Literal[ - 'resnet18', - 'resnet34', - 'resnet50', - 'resnet101', - 'resnet152', - ] = 'resnet18' - num_latents: int | None = None - image_size: int = 224 - num_channels: int = 3 - - -class VitModelConfig(BaseModel): - model_class: Literal['vit'] - model_params: VitModelParams - seed: int = 0 - checkpoint: str | None = None - - -class VitModelParams(BaseModel): - hidden_size: int = 768 - num_hidden_layers: int = 12 - num_attention_heads: int = 12 - intermediate_size: int = 3072 - hidden_act: str = 'gelu' - hidden_dropout_prob: float = 0.0 - attention_probs_dropout_prob: float = 0.0 - initializer_range: float = 0.02 - layer_norm_eps: float = 1e-12 - image_size: int = 224 - patch_size: int = 16 - num_channels: int = 3 - qkv_bias: bool = True - decoder_num_attention_heads: int = 16 - decoder_hidden_size: int = 512 - decoder_num_hidden_layers: int = 8 - decoder_intermediate_size: int = 2048 - mask_ratio: float = 0.75 - norm_pix_loss: bool = False - embed_size: int = 768 - temp_scale: bool = False - random_init: bool = False - use_infoNCE: bool = False - infoNCE_weight: float = 0.03 - use_perceptual_loss: bool = False - lambda_perceptual: float = 10.0 + inference: bool = False ModelConfig = Annotated[ @@ -113,4 +94,34 @@ class DataConfig(BaseModel): data_dir: str | Path +class ERayZerBeastConfig(BaseModel): + """Complete top-level config for ERayZer training runs.""" + + model: ERayZerModelConfig + training: ERayZerTrainingConfig + optimizer: ERayZerOptimizerConfig + data: DataConfig + inference: bool = False + + +_MODEL_CONFIG_CLASSES: dict[str, type[BaseModel]] = { + 'erayzer': ERayZerBeastConfig, +} + + +def get_beast_config_class(model_class: str) -> type[BaseModel]: + """Return the top-level Pydantic config class for the given model_class identifier. + + Parameters + ---------- + model_class: model type identifier string (e.g., 'resnet', 'vit', 'erayzer') + + Returns + ------- + Pydantic model class for the full top-level config + + """ + return _MODEL_CONFIG_CLASSES.get(model_class, BeastConfig) + + BeastConfig.model_rebuild() diff --git a/beast/data/datamodules.py b/beast/data/datamodules.py index d90ee6b..cc42d3d 100644 --- a/beast/data/datamodules.py +++ b/beast/data/datamodules.py @@ -4,6 +4,7 @@ import logging import multiprocessing import os +from pathlib import Path from typing import cast import lightning.pytorch as pl @@ -12,7 +13,7 @@ from lightning.pytorch.utilities import rank_zero_only from torch.utils.data import DataLoader, Subset, random_split -from beast.data.datasets import BaseDataset +from beast.data.datasets import BaseDataset, MultiViewDataset from beast.data.samplers import ContrastBatchSampler, contrastive_collate_fn _logger = logging.getLogger(__name__) @@ -332,3 +333,140 @@ def split_sizes_from_probabilities( ) return [train_number, val_number, test_number] + + +class MultiViewDataModule(pl.LightningDataModule): + """Lightning data module for BEAST3D multi-view self-supervised training. + + Splits unique frame IDs sequentially into train and val sets (no shuffle before + splitting, so later sessions end up in val — avoids temporal leakage). Each split + gets its own ``MultiViewDataset`` instance with the appropriate mode so that + view-order shuffling is only applied during training. + """ + + def __init__( + self, + data_dir: str | Path, + image_size: int, + train_batch_size: int = 4, + val_batch_size: int = 4, + train_fraction: float = 0.9, + use_mask: bool = False, + normalize_cameras: bool = True, + num_workers: int | None = None, + seed: int = 42, + ) -> None: + """Initialize the multi-view data module. + + Parameters + ---------- + data_dir: path to the ``dataset/`` directory produced by ``beast extract_3d``. + image_size: square size to resize images to. + train_batch_size: batch size for the training dataloader. + val_batch_size: batch size for the validation dataloader. + train_fraction: fraction of unique frames used for training; the rest go to val. + use_mask: if True, load binary segmentation masks. + normalize_cameras: if True, normalize camera poses per sample. + num_workers: dataloader worker count; defaults to SLURM_CPUS_PER_TASK or cpu_count. + seed: random seed for the training dataloader shuffle. + + """ + if not 0 < train_fraction <= 1: + raise ValueError(f'train_fraction must be in (0, 1], got {train_fraction}') + + super().__init__() + self.data_dir = Path(data_dir) + self.image_size = image_size + self.train_batch_size = train_batch_size + self.val_batch_size = val_batch_size + self.train_fraction = train_fraction + self.use_mask = use_mask + self.normalize_cameras = normalize_cameras + self.seed = seed + + if num_workers is not None: + self.num_workers = num_workers + else: + slurm_cpus = os.getenv('SLURM_CPUS_PER_TASK') + self.num_workers = int(slurm_cpus) if slurm_cpus else (os.cpu_count() or 0) + + self.train_dataset: MultiViewDataset | None = None + self.val_dataset: MultiViewDataset | None = None + + def setup(self, stage: str | None = None) -> None: + """Build train and val datasets by splitting unique frame IDs. + + Parameters + ---------- + stage: Lightning stage string (unused; both splits are always built). + + """ + # build a temporary dataset just to enumerate all frame IDs + full = MultiViewDataset( + data_dir=self.data_dir, + image_size=self.image_size, + mode='test', + use_mask=False, + normalize_cameras=self.normalize_cameras, + ) + all_ids = full.unique_frame_ids + n_train = max(1, int(len(all_ids) * self.train_fraction)) + train_ids = all_ids[:n_train] + val_ids = all_ids[n_train:] or all_ids[-1:] # at least one val sample + + self.train_dataset = MultiViewDataset( + data_dir=self.data_dir, + image_size=self.image_size, + mode='train', + use_mask=self.use_mask, + normalize_cameras=self.normalize_cameras, + frame_ids=train_ids, + ) + self.val_dataset = MultiViewDataset( + data_dir=self.data_dir, + image_size=self.image_size, + mode='test', + use_mask=self.use_mask, + normalize_cameras=self.normalize_cameras, + frame_ids=val_ids, + ) + + if rank_zero_only.rank == 0: + _logger.info( + f'MultiViewDataModule: {len(train_ids)} train frames, ' + f'{len(val_ids)} val frames' + ) + + def train_dataloader(self) -> DataLoader: + """Return the training dataloader.""" + if self.train_dataset is None: + raise RuntimeError('call setup() before train_dataloader()') + return DataLoader( + self.train_dataset, + batch_size=self.train_batch_size, + shuffle=True, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + pin_memory=True, + drop_last=True, + generator=torch.Generator().manual_seed(self.seed), + multiprocessing_context=( + multiprocessing.get_context('spawn') if self.num_workers > 0 else None + ), + ) + + def val_dataloader(self) -> DataLoader: + """Return the validation dataloader.""" + if self.val_dataset is None: + raise RuntimeError('call setup() before val_dataloader()') + return DataLoader( + self.val_dataset, + batch_size=self.val_batch_size, + shuffle=False, + num_workers=self.num_workers, + persistent_workers=self.num_workers > 0, + pin_memory=True, + multiprocessing_context=( + multiprocessing.get_context('spawn') if self.num_workers > 0 else None + ), + ) diff --git a/beast/data/datasets.py b/beast/data/datasets.py index 5d46f98..fe07fd0 100644 --- a/beast/data/datasets.py +++ b/beast/data/datasets.py @@ -1,5 +1,7 @@ """Dataset objects store images and augmentation pipeline.""" +import json +import logging import time from collections.abc import Callable from pathlib import Path @@ -11,9 +13,17 @@ from PIL import Image from torchvision import transforms -from beast.data.types import ExampleDict +from beast.data.types import ExampleDict, MultiViewExampleDict +from beast.geometry.camera import ( + intrinsics_to_fxfycxcy, + normalize_camera_sequence, + scale_intrinsics, + w2c_to_c2w, +) from beast.logging import log_step +_logger = logging.getLogger(__name__) + _IMAGENET_MEAN = [0.485, 0.456, 0.406] _IMAGENET_STD = [0.229, 0.224, 0.225] @@ -144,3 +154,197 @@ def _get_single_item(self, idx: int) -> ExampleDict: idx=idx, image_path=str(img_path), ) + + +class MultiViewDataset(torch.utils.data.Dataset): + """Multi-view dataset for BEAST3D training. + + Each item is a single time point from one session, containing synchronized + images from all available cameras along with their camera parameters. + + The dataset reads from the output directory produced by ``beast extract_3d``. + Each item returns: + + - ``image``: float32 tensor of shape ``(V, 3, H, W)`` in ``[0, 1]``. + - ``c2w``: float32 tensor of shape ``(V, 4, 4)`` camera-to-world matrices. + - ``fxfycxcy``: float32 tensor of shape ``(V, 4)`` intrinsics at ``image_size`` + resolution in absolute pixel units. + - ``view_names``: list of V camera name strings. + - ``video_id``, ``frame_id``: metadata strings. + - ``input_mask`` (optional): float32 tensor of shape ``(V, 1, H, W)`` binary + foreground masks, only present when ``use_mask=True``. + """ + + def __init__( + self, + data_dir: str | Path, + image_size: int, + mode: str = 'train', + use_mask: bool = False, + normalize_cameras: bool = True, + frame_ids: list[str] | None = None, + ) -> None: + """Initialize the multi-view dataset. + + Parameters + ---------- + data_dir: path to the ``dataset/`` directory produced by ``beast extract_3d``. + image_size: square size to resize images to (both height and width). + mode: ``'train'`` shuffles view order randomly; ``'test'`` uses sorted order. + use_mask: if True, load binary segmentation masks alongside images. + normalize_cameras: if True, re-center cameras so camera 0 is at the origin + and scale translations so the mean camera distance is 1. + frame_ids: optional pre-selected list of ``'{video_id}/{frame_filename}'`` + strings; if None, all frames from all sessions are used. + + Raises + ------ + ValueError + if data_dir does not exist, contains no sessions, or use_mask is True + but a mask file is missing for any frame. + + """ + self.data_dir = Path(data_dir) + if not self.data_dir.is_dir(): + raise ValueError(f'{self.data_dir} is not a directory') + + self.image_size = image_size + self.mode = mode + self.use_mask = use_mask + self.normalize_cameras = normalize_cameras + + info_path = self.data_dir / 'info.json' + if not info_path.exists(): + raise ValueError(f'info.json not found in {self.data_dir}') + with open(info_path) as f: + info = json.load(f) + self.available_views: list[str] = sorted(info['available_views']) + + if frame_ids is not None: + self.unique_frame_ids = frame_ids + else: + csv_files = sorted(self.data_dir.rglob('selected_frames.csv')) + if not csv_files: + raise ValueError(f'No selected_frames.csv files found under {self.data_dir}') + ids: list[str] = [] + for csv_path in csv_files: + video_id = csv_path.parent.name + frames = csv_path.read_text().splitlines() + ids.extend(f'{video_id}/{f}' for f in frames if f) + self.unique_frame_ids = sorted(set(ids)) + + if not self.unique_frame_ids: + raise ValueError(f'{self.data_dir} contains no multi-view frame data') + + # ToTensor makes both pipelines return a Tensor; Compose's stub keeps the + # input type, so annotate the callables to reflect the real output + self.img_transform: Callable[[Image.Image], torch.Tensor] = transforms.Compose([ + transforms.Resize( + (image_size, image_size), + interpolation=transforms.InterpolationMode.BICUBIC, + antialias=True, + ), + transforms.ToTensor(), + ]) + self.mask_transform: Callable[[Image.Image], torch.Tensor] = transforms.Compose([ + transforms.Resize( + (image_size, image_size), + interpolation=transforms.InterpolationMode.NEAREST, + ), + transforms.ToTensor(), + ]) + + _logger.info( + f'MultiViewDataset: {len(self.unique_frame_ids)} frames × ' + f'{len(self.available_views)} views (mode={mode}, use_mask={use_mask})' + ) + + def __len__(self) -> int: + """Return number of frames in the dataset.""" + return len(self.unique_frame_ids) + + def __getitem__(self, idx: int) -> MultiViewExampleDict: + """Return all camera views for a single time point. + + Parameters + ---------- + idx: index into unique_frame_ids. + + Returns + ------- + MultiViewExampleDict with stacked tensors for all V cameras. + + """ + return self._get_single_item(idx) + + def _get_single_item(self, idx: int) -> MultiViewExampleDict: + """Load images and camera params for one (session, frame) pair.""" + unique_frame_id = self.unique_frame_ids[idx] + video_id, frame_id = unique_frame_id.split('/', 1) + + images: list[torch.Tensor] = [] + extrinsics: list[torch.Tensor] = [] + fxfycxcy_list: list[torch.Tensor] = [] + masks: list[torch.Tensor] = [] + + for view in self.available_views: + img_path = self.data_dir / video_id / view / frame_id + cam_path = img_path.with_suffix('.npy') + + image = Image.open(img_path).convert('RGB') + images.append(self.img_transform(image)) + + if not cam_path.exists(): + raise FileNotFoundError( + f'camera file missing for {img_path.name} ' + f'(expected {cam_path})' + ) + cam_info = np.load(cam_path, allow_pickle=True).item() + K = torch.from_numpy(cam_info['intrinsics']).float() + scale_w = self.image_size / cam_info['width'] + scale_h = self.image_size / cam_info['height'] + K = scale_intrinsics(K, scale_w, scale_h) + fxfycxcy_list.append(intrinsics_to_fxfycxcy(K)) + extrinsics.append(torch.from_numpy(cam_info['extrinsics']).float()) + + if self.use_mask: + mask_filename = frame_id.replace('img', 'mask', 1) + mask_path = self.data_dir / video_id / view / mask_filename + mask = Image.open(mask_path).convert('L') + mask_tensor = self.mask_transform(mask) + masks.append((mask_tensor > 0.5).float()) + + images_t = torch.stack(images) # (V, 3, H, W) + w2c_t = torch.stack(extrinsics) # (V, 4, 4) + fxfycxcy_t = torch.stack(fxfycxcy_list) # (V, 4) + + if self.normalize_cameras: + c2w_t = normalize_camera_sequence(w2c_t) + else: + c2w_t = w2c_to_c2w(w2c_t) + + view_names = list(self.available_views) + + if self.mode == 'train': + perm = torch.randperm(len(self.available_views)) + images_t = images_t[perm] + c2w_t = c2w_t[perm] + fxfycxcy_t = fxfycxcy_t[perm] + view_names = [view_names[i] for i in perm.tolist()] + if self.use_mask: + masks_t = torch.stack(masks)[perm] + else: + if self.use_mask: + masks_t = torch.stack(masks) + + result: MultiViewExampleDict = { + 'image': images_t, + 'c2w': c2w_t, + 'fxfycxcy': fxfycxcy_t, + 'view_names': view_names, + 'video_id': video_id, + 'frame_id': frame_id, + } + if self.use_mask: + result['input_mask'] = masks_t + return result diff --git a/beast/data/types.py b/beast/data/types.py index 889b147..960d7d4 100644 --- a/beast/data/types.py +++ b/beast/data/types.py @@ -4,6 +4,7 @@ from jaxtyping import Float from torch import Tensor +from typing_extensions import NotRequired class ExampleDict(TypedDict): @@ -12,3 +13,21 @@ class ExampleDict(TypedDict): video: str | list[str] idx: int | list[int] image_path: str | list[str] + + +class MultiViewExampleDict(TypedDict): + """Return type when calling MultiViewDataset.__getitem__(). + + Required fields: image, view_names, video_id, frame_id. + Optional fields: c2w, fxfycxcy (absent when calibration files are missing), + input_mask (present only when use_mask=True). + """ + image: Float[Tensor, 'views channels height width'] + view_names: list[str] + video_id: str + frame_id: str + # the three fields below are only used by BEAST3D (GT cameras + mask loss); + # may become a separate subclass dict later + c2w: NotRequired[Float[Tensor, 'views 4 4']] + fxfycxcy: NotRequired[Float[Tensor, 'views 4']] + input_mask: NotRequired[Float[Tensor, 'views 1 height width']] diff --git a/beast/geometry/__init__.py b/beast/geometry/__init__.py new file mode 100644 index 0000000..9702511 --- /dev/null +++ b/beast/geometry/__init__.py @@ -0,0 +1 @@ +"""Geometry utilities: camera math, positional encodings, and point cloud alignment.""" diff --git a/beast/geometry/camera.py b/beast/geometry/camera.py new file mode 100644 index 0000000..19ff326 --- /dev/null +++ b/beast/geometry/camera.py @@ -0,0 +1,464 @@ +"""Camera transformation helpers, Plucker ray encoding, and pose utilities. + +Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. +Licensed under the Apache License, Version 2.0. + +""" + +import math + +import numpy as np +import numpy.typing as npt +import torch +from einops import rearrange +from jaxtyping import Float +from torch import Tensor + +_EPS = np.finfo(float).eps * 4.0 + + +def unit_vector(data: np.ndarray, axis: int | None = None) -> np.ndarray: + """Return ndarray normalized by Euclidean norm along axis. + + Parameters + ---------- + data: input array. + axis: the axis along which to normalize into unit vector. + + """ + data = np.array(data, dtype=np.float64, copy=True) + if data.ndim == 1: + data /= math.sqrt(np.dot(data, data)) + return data + length = np.atleast_1d(np.sum(data * data, axis)) + np.sqrt(length, length) + if axis is not None: + length = np.expand_dims(length, axis) + data /= length + return data + + +def quaternion_from_matrix(matrix: np.ndarray, isprecise: bool = False) -> np.ndarray: + """Return quaternion from rotation matrix. + + Parameters + ---------- + matrix: rotation matrix to obtain quaternion. + isprecise: if True, input matrix is assumed to be a precise rotation matrix + and a faster algorithm is used. + + """ + M = np.asarray(matrix, dtype=np.float64)[:4, :4] + if isprecise: + q = np.empty((4,)) + t = np.trace(M) + if t > M[3, 3]: + q[0] = t + q[3] = M[1, 0] - M[0, 1] + q[2] = M[0, 2] - M[2, 0] + q[1] = M[2, 1] - M[1, 2] + else: + i, j, k = 1, 2, 3 + if M[1, 1] > M[0, 0]: + i, j, k = 2, 3, 1 + if M[2, 2] > M[i, i]: + i, j, k = 3, 1, 2 + t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3] + q[i] = t + q[j] = M[i, j] + M[j, i] + q[k] = M[k, i] + M[i, k] + q[3] = M[k, j] - M[j, k] + q *= 0.5 / math.sqrt(t * M[3, 3]) + else: + m00 = M[0, 0] + m01 = M[0, 1] + m02 = M[0, 2] + m10 = M[1, 0] + m11 = M[1, 1] + m12 = M[1, 2] + m20 = M[2, 0] + m21 = M[2, 1] + m22 = M[2, 2] + K = [ + [m00 - m11 - m22, 0.0, 0.0, 0.0], + [m01 + m10, m11 - m00 - m22, 0.0, 0.0], + [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0], + [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22], + ] + K = np.array(K) + K /= 3.0 + w, V = np.linalg.eigh(K) + q = V[np.array([3, 0, 1, 2]), np.argmax(w)] + if q[0] < 0.0: + np.negative(q, q) + return q + + +def quaternion_slerp( + quat0: np.ndarray, + quat1: np.ndarray, + fraction: float, + spin: int = 0, + shortestpath: bool = True, +) -> np.ndarray: + """Return spherical linear interpolation between two quaternions. + + Parameters + ---------- + quat0: first quaternion. + quat1: second quaternion. + fraction: interpolation parameter (0 → quat0, 1 → quat1). + spin: additional spin to place on the interpolation. + shortestpath: whether to return the short or long path to rotation. + + """ + q0 = unit_vector(quat0[:4]) + q1 = unit_vector(quat1[:4]) + if q0 is None or q1 is None: + raise ValueError('Input quaternions invalid.') + if fraction == 0.0: + return q0 + if fraction == 1.0: + return q1 + d = np.dot(q0, q1) + if abs(abs(d) - 1.0) < _EPS: + return q0 + if shortestpath and d < 0.0: + d = -d + np.negative(q1, q1) + angle = math.acos(d) + spin * math.pi + if abs(angle) < _EPS: + return q0 + isin = 1.0 / math.sin(angle) + q0 *= math.sin((1.0 - fraction) * angle) * isin + q1 *= math.sin(fraction * angle) * isin + q0 += q1 + return q0 + + +def quaternion_matrix(quaternion: npt.ArrayLike) -> np.ndarray: + """Return homogeneous rotation matrix from quaternion. + + Parameters + ---------- + quaternion: value to convert to matrix. + + """ + q = np.array(quaternion, dtype=np.float64, copy=True) + n = np.dot(q, q) + if n < _EPS: + return np.identity(4) + q *= math.sqrt(2.0 / n) + q = np.outer(q, q) + return np.array( + [ + [1.0 - q[2, 2] - q[3, 3], q[1, 2] - q[3, 0], q[1, 3] + q[2, 0], 0.0], + [q[1, 2] + q[3, 0], 1.0 - q[1, 1] - q[3, 3], q[2, 3] - q[1, 0], 0.0], + [q[1, 3] - q[2, 0], q[2, 3] + q[1, 0], 1.0 - q[1, 1] - q[2, 2], 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + + +def get_interpolated_poses( + pose_a: np.ndarray, + pose_b: np.ndarray, + steps: int = 10, +) -> list[np.ndarray]: + """Return interpolation of poses with the specified number of steps. + + Parameters + ---------- + pose_a: first pose. + pose_b: second pose. + steps: number of steps the interpolated pose path should contain. + + """ + quat_a = quaternion_from_matrix(pose_a[:3, :3]) + quat_b = quaternion_from_matrix(pose_b[:3, :3]) + + ts = np.linspace(0, 1, steps) + quats = [quaternion_slerp(quat_a, quat_b, t) for t in ts] + trans = [(1 - t) * pose_a[:3, 3] + t * pose_b[:3, 3] for t in ts] + + poses_ab = [] + for quat, tran in zip(quats, trans, strict=True): + pose = np.identity(4) + pose[:3, :3] = quaternion_matrix(quat)[:3, :3] + pose[:3, 3] = tran + poses_ab.append(pose[:3]) + return poses_ab + + +def get_interpolated_k( + k_a: Float[Tensor, '3 3'], + k_b: Float[Tensor, '3 3'], + steps: int = 10, +) -> list[Float[Tensor, '3 4']]: + """Return interpolated path between two camera intrinsic matrices. + + Parameters + ---------- + k_a: camera matrix 1. + k_b: camera matrix 2. + steps: number of steps the interpolated path should contain. + + Returns + ------- + list of interpolated camera matrices. + + """ + Ks: list[Float[Tensor, '3 3']] = [] + ts = np.linspace(0, 1, steps) + for t in ts: + new_k = k_a * (1.0 - t) + k_b * t + Ks.append(new_k) + return Ks + + +def get_ordered_poses_and_k( + poses: Float[Tensor, 'num_poses 3 4'], + Ks: Float[Tensor, 'num_poses 3 3'], +) -> tuple[Float[Tensor, 'num_poses 3 4'], Float[Tensor, 'num_poses 3 3']]: + """Return poses and intrinsics ordered by Euclidean distance between poses. + + Parameters + ---------- + poses: list of camera poses. + Ks: list of camera intrinsics. + + Returns + ------- + tuple of ordered poses and intrinsics. + + """ + poses_num = len(poses) + + ordered_poses = torch.unsqueeze(poses[0], 0) + ordered_ks = torch.unsqueeze(Ks[0], 0) + + poses = poses[1:] + Ks = Ks[1:] + + for _ in range(poses_num - 1): + distances = torch.norm(ordered_poses[-1][:, 3] - poses[:, :, 3], dim=1) + idx = torch.argmin(distances) + ordered_poses = torch.cat((ordered_poses, torch.unsqueeze(poses[idx], 0)), dim=0) + ordered_ks = torch.cat((ordered_ks, torch.unsqueeze(Ks[idx], 0)), dim=0) + poses = torch.cat((poses[0:idx], poses[idx + 1:]), dim=0) + Ks = torch.cat((Ks[0:idx], Ks[idx + 1:]), dim=0) + + return ordered_poses, ordered_ks + + +def get_interpolated_poses_many( + poses: Float[Tensor, 'num_poses 3 4'], + Ks: Float[Tensor, 'num_poses 3 3'], + steps_per_transition: int = 10, + order_poses: bool = False, +) -> tuple[Float[Tensor, 'num_poses 3 4'], Float[Tensor, 'num_poses 3 3']]: + """Return interpolated poses for many camera poses. + + Parameters + ---------- + poses: list of camera poses. + Ks: list of camera intrinsics. + steps_per_transition: number of steps per transition. + order_poses: whether to order poses by Euclidean distance. + + Returns + ------- + tuple of (new poses, intrinsics). + + """ + traj = [] + k_interp = [] + + if poses.shape[0] == 0: + raise ValueError('get_interpolated_poses_many requires at least one pose') + if poses.shape[0] == 1: + traj = [poses[0].cpu().numpy()] + k_interp = [Ks[0]] + else: + if order_poses: + poses, Ks = get_ordered_poses_and_k(poses, Ks) + for idx in range(poses.shape[0] - 1): + pose_a = poses[idx].cpu().numpy() + pose_b = poses[idx + 1].cpu().numpy() + poses_ab = get_interpolated_poses(pose_a, pose_b, steps=steps_per_transition) + traj += poses_ab + k_interp += get_interpolated_k(Ks[idx], Ks[idx + 1], steps=steps_per_transition) + + traj = np.stack(traj, axis=0) + k_interp = torch.stack(k_interp, dim=0) + + return torch.tensor(traj, dtype=torch.float32), torch.tensor(k_interp, dtype=torch.float32) + + +def w2c_to_c2w(w2c: torch.Tensor) -> torch.Tensor: + """Convert world-to-camera matrices to camera-to-world using the analytical SE3 inverse. + + For a rigid body transform [R | t; 0 | 1], the inverse is [R^T | -R^T t; 0 | 1]. + Handles any leading batch dimensions. + + Parameters + ---------- + w2c: world-to-camera matrices of shape (..., 4, 4). + + Returns + ------- + camera-to-world matrices of shape (..., 4, 4). + + """ + R = w2c[..., :3, :3] + t = w2c[..., :3, 3:] + R_T = R.transpose(-2, -1) + t_new = -torch.matmul(R_T, t) + c2w = torch.zeros_like(w2c) + c2w[..., :3, :3] = R_T + c2w[..., :3, 3:] = t_new + c2w[..., 3, 3] = 1.0 + return c2w + + +def intrinsics_to_fxfycxcy(K: torch.Tensor) -> torch.Tensor: + """Extract [fx, fy, cx, cy] from a 3x3 camera intrinsics matrix. + + Parameters + ---------- + K: intrinsics matrix of shape (..., 3, 3). + + Returns + ------- + tensor of shape (..., 4) containing [fx, fy, cx, cy]. + + """ + return torch.stack( + [K[..., 0, 0], K[..., 1, 1], K[..., 0, 2], K[..., 1, 2]], + dim=-1, + ) + + +def scale_intrinsics(K: torch.Tensor, scale_w: float, scale_h: float) -> torch.Tensor: + """Scale camera intrinsics to account for an image resize. + + Parameters + ---------- + K: intrinsics matrix of shape (..., 3, 3). + scale_w: horizontal scale factor (new_width / orig_width). + scale_h: vertical scale factor (new_height / orig_height). + + Returns + ------- + scaled intrinsics matrix of shape (..., 3, 3). + + """ + K_scaled = K.clone() + K_scaled[..., 0, 0] = K_scaled[..., 0, 0] * scale_w # fx + K_scaled[..., 0, 2] = K_scaled[..., 0, 2] * scale_w # cx + K_scaled[..., 1, 1] = K_scaled[..., 1, 1] * scale_h # fy + K_scaled[..., 1, 2] = K_scaled[..., 1, 2] * scale_h # cy + return K_scaled + + +def normalize_camera_sequence(extrinsics: torch.Tensor) -> torch.Tensor: + """Normalize a sequence of w2c camera matrices and return c2w. + + Applies two normalizations in sequence: + 1. Re-center: transforms coordinates so camera 0 is at the world origin. + 2. Scale: divides all translations so the mean camera distance from origin is 1. + + Parameters + ---------- + extrinsics: world-to-camera matrices of shape (V, 4, 4). + + Returns + ------- + camera-to-world matrices of shape (V, 4, 4) in the normalized coordinate frame. + + """ + first_c2w = w2c_to_c2w(extrinsics[0:1]).squeeze(0) # (4, 4) + ex_norm = torch.matmul(extrinsics, first_c2w) # (V, 4, 4) + + c2w = w2c_to_c2w(ex_norm) # (V, 4, 4) + scale = c2w[:, :3, 3].norm(dim=-1).mean() + + if scale > 1e-8: + ex_norm = ex_norm.clone() + ex_norm[:, :3, 3] = ex_norm[:, :3, 3] / scale + c2w = w2c_to_c2w(ex_norm) + + return c2w + + +def cam_info_to_plucker( + c2w: torch.Tensor, + fxfycxcy: torch.Tensor, + target_imgs_info: dict, + normalized: bool = True, + return_moment: bool = True, +) -> torch.Tensor: + """Compute per-pixel Plucker ray embeddings from camera parameters. + + Parameters + ---------- + c2w: camera-to-world matrices of shape [b, 4, 4] or [b, v, 4, 4]. + fxfycxcy: intrinsics of shape [b, 4] or [b, v, 4]. + target_imgs_info: dict with keys 'height' and 'width'. + normalized: if True, scale fxfycxcy by image resolution before use. + return_moment: if True, return moment+direction encoding; otherwise + origin+direction. + + Returns + ------- + Plucker ray tensor of shape [b, 6, h, w]. + + """ + if len(c2w.shape) == 3: + b = c2w.shape[0] + elif len(c2w.shape) == 4: + c2w = rearrange(c2w.clone(), 'b v n d -> (b v) n d') + fxfycxcy = rearrange(fxfycxcy.clone(), 'b v d -> (b v) d') + b = c2w.shape[0] + + h, w = target_imgs_info['height'], target_imgs_info['width'] + + fxfycxcy = fxfycxcy.clone() + if normalized: + fxfycxcy[:, 0] *= w + fxfycxcy[:, 1] *= h + fxfycxcy[:, 2] *= w + fxfycxcy[:, 3] *= h + + y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing='ij') + y, x = y.to(c2w), x.to(c2w) + x = x[None, :, :].expand(b, -1, -1).reshape(b, -1) + y = y[None, :, :].expand(b, -1, -1).reshape(b, -1) + x = (x + 0.5 - fxfycxcy[:, 2:3]) / fxfycxcy[:, 0:1] + y = (y + 0.5 - fxfycxcy[:, 3:4]) / fxfycxcy[:, 1:2] + z = torch.ones_like(x) + ray_d = torch.stack([x, y, z], dim=2) # [b, h*w, 3] + ray_d = torch.bmm(ray_d, c2w[:, :3, :3].transpose(1, 2)) # [b, h*w, 3] + ray_d = ray_d / torch.norm(ray_d, dim=2, keepdim=True) # [b, h*w, 3] + ray_o = c2w[:, :3, 3][:, None, :].expand_as(ray_d) # [b, h*w, 3] + + ray_o = ray_o.reshape(b, h, w, 3).permute(0, 3, 1, 2) # [b, 3, h, w] + ray_d = ray_d.reshape(b, h, w, 3).permute(0, 3, 1, 2) + + if return_moment: + plucker = torch.cat( + [ + torch.cross(ray_o, ray_d, dim=1), + ray_d, + ], + dim=1, + ) + else: + plucker = torch.cat( + [ + ray_o, + ray_d, + ], + dim=1, + ) + return plucker # [b, 6, h, w] diff --git a/beast/geometry/positional_encoding.py b/beast/geometry/positional_encoding.py new file mode 100644 index 0000000..b3600a7 --- /dev/null +++ b/beast/geometry/positional_encoding.py @@ -0,0 +1,98 @@ +"""Sinusoidal positional embedding utilities for 1D and 2D grids.""" + +import torch + + +def get_1d_sincos_pos_emb_from_grid( + embed_dim: int, + pos: torch.Tensor, + device: str | torch.device = 'cpu', +) -> torch.Tensor: + """Generate 1D sinusoidal positional embeddings from grid positions. + + Parameters + ---------- + embed_dim: the embedding dimension (must be even). + pos: the grid positions, shape [b * gh * gw] or [batch_size, sequence_length]. + device: device for the output tensor. + + Returns + ------- + sinusoidal positional embeddings of shape [len(pos), embed_dim]. + + """ + assert embed_dim % 2 == 0, 'Embedding dimension must be even for sine and cosine.' + + pos = pos.float() + + dim = torch.arange(embed_dim // 2, dtype=torch.float32, device=device) + freq = 1.0 / (10000 ** (dim / (embed_dim // 2))) + + pos_emb_sin = torch.sin(pos[:, None] * freq) + pos_emb_cos = torch.cos(pos[:, None] * freq) + + pos_emb = torch.cat([pos_emb_sin, pos_emb_cos], dim=-1) + + return pos_emb + + +def get_2d_sincos_pos_embed( + embed_dim: int, + grid_size: tuple[int, int], + device: str | torch.device = 'cpu', +) -> torch.Tensor: + """Generate 2D sine-cosine positional embeddings with separate grid height and width. + + Parameters + ---------- + embed_dim: the embedding dimension. + grid_size: tuple specifying the grid height and width (grid_h, grid_w). + device: the device to place the embeddings on. + + Returns + ------- + positional embeddings of shape [grid_h*grid_w, embed_dim]. + + """ + grid_h_size, grid_w_size = grid_size + + grid_h = torch.arange(grid_h_size, dtype=torch.float32, device=device) + grid_w = torch.arange(grid_w_size, dtype=torch.float32, device=device) + grid = torch.meshgrid(grid_w, grid_h, indexing='ij') + grid = torch.stack(grid, dim=0) + + grid = grid.view(2, 1, grid.size(1), grid.size(2)) + + pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, device=device) + + return pos_embed + + +def get_2d_sincos_pos_embed_from_grid( + embed_dim: int, + grid: torch.Tensor, + device: str | torch.device = 'cpu', +) -> torch.Tensor: + """Generate 2D sine-cosine positional embeddings from a grid. + + Parameters + ---------- + embed_dim: the embedding dimension. + grid: the grid of shape [2, 1, grid_h, grid_w]. + device: the device to place the embeddings on. + + Returns + ------- + positional embeddings of shape [grid_h*grid_w, embed_dim]. + + """ + assert embed_dim % 2 == 0, 'Embedding dimension must be even.' + + grid_h = grid[0].view(-1) + grid_w = grid[1].view(-1) + + emb_h = get_1d_sincos_pos_emb_from_grid(embed_dim // 2, grid_h, device=device) + emb_w = get_1d_sincos_pos_emb_from_grid(embed_dim // 2, grid_w, device=device) + + pos_embed = torch.cat([emb_h, emb_w], dim=-1) + return pos_embed diff --git a/beast/geometry/rotations.py b/beast/geometry/rotations.py new file mode 100644 index 0000000..c13ea6a --- /dev/null +++ b/beast/geometry/rotations.py @@ -0,0 +1,51 @@ +"""Rotation representation conversion utilities.""" + +import torch +import torch.nn.functional as F + + +def rot6d2mat(rot_6d: torch.Tensor) -> torch.Tensor: + """Convert a 6D rotation representation to a 3×3 rotation matrix. + + Uses the Gram-Schmidt orthonormalization procedure from Zhou et al., + "On the Continuity of Rotation Representations in Neural Networks" (CVPR 2019). + The first three values encode the first column and the second three encode + the second column of the rotation matrix; the third column is their cross product. + + Parameters + ---------- + rot_6d: rotation tensor of shape (..., 6). + + Returns + ------- + rotation matrices of shape (..., 3, 3) with columns [a1, a2, a3]. + + """ + a1 = F.normalize(rot_6d[..., :3], dim=-1) + a2 = rot_6d[..., 3:] + a2 = F.normalize(a2 - (a1 * a2).sum(dim=-1, keepdim=True) * a1, dim=-1) + a3 = torch.linalg.cross(a1, a2) + return torch.stack([a1, a2, a3], dim=-1) + + +def quat2mat(quat: torch.Tensor) -> torch.Tensor: + """Convert unit quaternions to 3×3 rotation matrices. + + Parameters + ---------- + quat: quaternion tensor of shape (..., 4) in (w, x, y, z) order. + Need not be pre-normalized. + + Returns + ------- + rotation matrices of shape (..., 3, 3). + + """ + quat = F.normalize(quat, dim=-1) + w, x, y, z = quat.unbind(dim=-1) + mat = torch.stack([ + 1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y), + 2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x), + 2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y), + ], dim=-1) + return mat.reshape(*quat.shape[:-1], 3, 3) diff --git a/beast/io.py b/beast/io.py index 2a4a102..12c0fe2 100644 --- a/beast/io.py +++ b/beast/io.py @@ -7,7 +7,7 @@ import yaml from aniposelib.cameras import CameraGroup as CameraGroupAnipose -from beast.config import BeastConfig +from beast.config import get_beast_config_class from beast.preprocess.config_3d import Beast3DConfig _logger = logging.getLogger(__name__) @@ -34,9 +34,11 @@ def load_config(path: str | Path) -> dict: with open(path) as file: raw = yaml.safe_load(file) - # validate against the schema; raises ValidationError on missing required - # fields, wrong types, or invalid Literal values - validated = BeastConfig.model_validate(raw) + # dispatch to the appropriate config class based on model_class, then validate; + # raises ValidationError on missing required fields, wrong types, or invalid Literals + model_class = (raw.get('model') or {}).get('model_class', '') + config_cls = get_beast_config_class(model_class) + validated = config_cls.model_validate(raw) # convert back to a plain nested dict so callers don't depend on pydantic types; # this also fills in any fields that have defaults but were absent from the yaml diff --git a/beast/losses.py b/beast/losses.py new file mode 100644 index 0000000..c5b5b44 --- /dev/null +++ b/beast/losses.py @@ -0,0 +1,31 @@ +"""Loss functions shared across BEAST models.""" + +import torch +import torch.nn.functional as F + + +def masked_mse_loss( + rendering: torch.Tensor, + target: torch.Tensor, + pixel_mask: torch.Tensor, +) -> torch.Tensor: + """Compute weighted MSE over RGB pixels using a foreground mask. + + Parameters + ---------- + rendering: predicted image tensor + target: target image tensor + pixel_mask: binary mask tensor; loss is sum(w*(p-t)^2) / sum(w), w repeated across channels + + Returns + ------- + scalar masked MSE loss + + """ + m = pixel_mask.to(dtype=rendering.dtype, device=rendering.device) + if m.ndim == 3: + m = m.unsqueeze(1) + valid_mask = m.expand_as(rendering) + loss = F.mse_loss(rendering, target, reduction='none') * valid_mask + normalizer = valid_mask.sum() + return loss.sum() / normalizer diff --git a/beast/models/__init__.py b/beast/models/__init__.py index 9aeaedc..28d430b 100644 --- a/beast/models/__init__.py +++ b/beast/models/__init__.py @@ -1 +1,16 @@ -"""Model architectures for BEAST.""" +"""Model architectures for BEAST. + +Each model lives in its own sub-package under beast/models// containing: + + _config.py — Pydantic schemas for model (and optionally training/optimizer) + _model.py — LightningModule subclass + _train.py — training entry point (delegates to beast.train.train or custom loop) + __init__.py — re-exports model class, config classes, and train function + +The central registry (beast.models.registry) maps model_class strings to their +implementation classes and training functions. All models inherit from +BaseLightningModel (beast.models.base) and must implement get_model_outputs(), +compute_loss(), and predict_step(). + +See docs/developer_guide.md for a full walkthrough of adding a new model. +""" diff --git a/beast/models/beast_resnet/__init__.py b/beast/models/beast_resnet/__init__.py new file mode 100644 index 0000000..4bb0fca --- /dev/null +++ b/beast/models/beast_resnet/__init__.py @@ -0,0 +1,7 @@ +"""ResNet autoencoder model package.""" + +from beast.models.beast_resnet.beast_resnet_config import ResnetModelConfig, ResnetModelParams +from beast.models.beast_resnet.beast_resnet_model import ResnetAutoencoder +from beast.models.beast_resnet.beast_resnet_train import train + +__all__ = ['ResnetAutoencoder', 'ResnetModelConfig', 'ResnetModelParams', 'train'] diff --git a/beast/models/beast_resnet/beast_resnet_config.py b/beast/models/beast_resnet/beast_resnet_config.py new file mode 100644 index 0000000..b2bae44 --- /dev/null +++ b/beast/models/beast_resnet/beast_resnet_config.py @@ -0,0 +1,29 @@ +"""Pydantic config schemas for the ResNet autoencoder model.""" + +from typing import Literal + +from pydantic import BaseModel + + +class ResnetModelParams(BaseModel): + """Parameters for the ResNet backbone and bottleneck.""" + + backbone: Literal[ + 'resnet18', + 'resnet34', + 'resnet50', + 'resnet101', + 'resnet152', + ] = 'resnet18' + num_latents: int | None = None + image_size: int = 224 + num_channels: int = 3 + + +class ResnetModelConfig(BaseModel): + """Top-level model-section config for the ResNet autoencoder.""" + + model_class: Literal['resnet'] + model_params: ResnetModelParams + seed: int = 0 + checkpoint: str | None = None diff --git a/beast/models/resnets.py b/beast/models/beast_resnet/beast_resnet_model.py similarity index 99% rename from beast/models/resnets.py rename to beast/models/beast_resnet/beast_resnet_model.py index caf08c9..58ce9ae 100644 --- a/beast/models/resnets.py +++ b/beast/models/beast_resnet/beast_resnet_model.py @@ -1,4 +1,4 @@ -"""Resnet-based autoencoder implementation. +"""ResNet autoencoder implementation. Adapted from https://github.com/Horizon2333/imagenet-autoencoder diff --git a/beast/models/beast_resnet/beast_resnet_train.py b/beast/models/beast_resnet/beast_resnet_train.py new file mode 100644 index 0000000..501ced7 --- /dev/null +++ b/beast/models/beast_resnet/beast_resnet_train.py @@ -0,0 +1,9 @@ +"""Training entry point for the ResNet autoencoder. + +Delegates to the shared beast.train.train function, which handles +BaseDataset / BaseDataModule and epoch-based training. +""" + +from beast.train import train + +__all__ = ['train'] diff --git a/beast/models/beast_vit/__init__.py b/beast/models/beast_vit/__init__.py new file mode 100644 index 0000000..d539f44 --- /dev/null +++ b/beast/models/beast_vit/__init__.py @@ -0,0 +1,7 @@ +"""Vision Transformer autoencoder model package.""" + +from beast.models.beast_vit.beast_vit_config import VitModelConfig, VitModelParams +from beast.models.beast_vit.beast_vit_model import VisionTransformer +from beast.models.beast_vit.beast_vit_train import train + +__all__ = ['VisionTransformer', 'VitModelConfig', 'VitModelParams', 'train'] diff --git a/beast/models/beast_vit/beast_vit_config.py b/beast/models/beast_vit/beast_vit_config.py new file mode 100644 index 0000000..c0c958c --- /dev/null +++ b/beast/models/beast_vit/beast_vit_config.py @@ -0,0 +1,45 @@ +"""Pydantic config schemas for the Vision Transformer autoencoder model.""" + +from typing import Literal + +from pydantic import BaseModel + + +class VitModelParams(BaseModel): + """Parameters for the ViT-MAE backbone and decoder.""" + + hidden_size: int = 768 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + intermediate_size: int = 3072 + hidden_act: str = 'gelu' + hidden_dropout_prob: float = 0.0 + attention_probs_dropout_prob: float = 0.0 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-12 + image_size: int = 224 + patch_size: int = 16 + num_channels: int = 3 + qkv_bias: bool = True + decoder_num_attention_heads: int = 16 + decoder_hidden_size: int = 512 + decoder_num_hidden_layers: int = 8 + decoder_intermediate_size: int = 2048 + mask_ratio: float = 0.75 + norm_pix_loss: bool = False + embed_size: int = 768 + temp_scale: bool = False + random_init: bool = False + use_infoNCE: bool = False + infoNCE_weight: float = 0.03 + use_perceptual_loss: bool = False + lambda_perceptual: float = 10.0 + + +class VitModelConfig(BaseModel): + """Top-level model-section config for the Vision Transformer autoencoder.""" + + model_class: Literal['vit'] + model_params: VitModelParams + seed: int = 0 + checkpoint: str | None = None diff --git a/beast/models/vits.py b/beast/models/beast_vit/beast_vit_model.py similarity index 99% rename from beast/models/vits.py rename to beast/models/beast_vit/beast_vit_model.py index f5c9004..11aaa39 100644 --- a/beast/models/vits.py +++ b/beast/models/beast_vit/beast_vit_model.py @@ -13,7 +13,7 @@ from beast.logging import log_step from beast.models.base import BaseLightningModel -from beast.models.perceptual import AlexPerceptual +from beast.nn.perceptual import AlexPerceptual class BatchNormProjector(nn.Module): diff --git a/beast/models/beast_vit/beast_vit_train.py b/beast/models/beast_vit/beast_vit_train.py new file mode 100644 index 0000000..784d304 --- /dev/null +++ b/beast/models/beast_vit/beast_vit_train.py @@ -0,0 +1,9 @@ +"""Training entry point for the Vision Transformer autoencoder. + +Delegates to the shared beast.train.train function, which handles +BaseDataset / BaseDataModule and epoch-based training. +""" + +from beast.train import train + +__all__ = ['train'] diff --git a/beast/models/erayzer/__init__.py b/beast/models/erayzer/__init__.py new file mode 100644 index 0000000..8ba573f --- /dev/null +++ b/beast/models/erayzer/__init__.py @@ -0,0 +1,45 @@ +"""ERayZer multi-view 3DGS model package.""" + +from beast.models.erayzer.erayzer_config import ( + ERayZerGaussiansConfig, + ERayZerImageTokenizerConfig, + ERayZerModelConfig, + ERayZerOptimizerConfig, + ERayZerPoseLatentConfig, + ERayZerRangeSettingConfig, + ERayZerTargetImageConfig, + ERayZerTrainingConfig, + ERayZerTransformerConfig, +) +from beast.models.erayzer.erayzer_model import ( + ERayZer, + GaussiansUpsampler, + LossComputer, + PoseEstimator, + build_transformer_blocks, + get_cam_se3, + get_point_range_func, + sanitize, +) +from beast.models.erayzer.erayzer_train import train + +__all__ = [ + 'ERayZer', + 'ERayZerGaussiansConfig', + 'ERayZerImageTokenizerConfig', + 'ERayZerModelConfig', + 'ERayZerOptimizerConfig', + 'ERayZerPoseLatentConfig', + 'ERayZerRangeSettingConfig', + 'ERayZerTargetImageConfig', + 'ERayZerTrainingConfig', + 'ERayZerTransformerConfig', + 'GaussiansUpsampler', + 'LossComputer', + 'PoseEstimator', + 'build_transformer_blocks', + 'get_cam_se3', + 'get_point_range_func', + 'sanitize', + 'train', +] diff --git a/beast/models/erayzer/erayzer_config.py b/beast/models/erayzer/erayzer_config.py new file mode 100644 index 0000000..c4a93a7 --- /dev/null +++ b/beast/models/erayzer/erayzer_config.py @@ -0,0 +1,153 @@ +"""Pydantic config schemas for the ERayZer multi-view 3DGS model.""" + +from typing import Literal + +from pydantic import BaseModel + + +class ERayZerImageTokenizerConfig(BaseModel): + """Patch tokenizer configuration.""" + + image_size: int = 256 + patch_size: int = 16 + in_channels: int = 3 + + +class ERayZerTargetImageConfig(BaseModel): + """Novel-view render resolution.""" + + height: int = 256 + width: int = 256 + + +class ERayZerTransformerConfig(BaseModel): + """Transformer architecture settings.""" + + d: int + d_head: int + encoder_n_layer: int = 0 + encoder_geom_n_layer: int + use_qk_norm: bool = False + special_init: bool = False + depth_init: bool = False + fix_decoder: bool = False + + +class ERayZerPoseLatentConfig(BaseModel): + """Camera-prediction pose parameterization settings.""" + + canonical: Literal['first', 'middle', 'unordered'] = 'first' + mode: Literal['pairwise', 'global'] = 'pairwise' + representation: Literal['6d', 'quat'] = '6d' + per_view_focal: bool = False + + +class ERayZerRangeSettingConfig(BaseModel): + """Depth range mapping settings for the Gaussian z-coordinate. + + ``near`` and ``far`` are only used for linear_depth, log_depth, and + disparity types. object_centric_depth ignores them. + """ + + type: Literal[ + 'object_centric_depth', + 'linear_depth', + 'log_depth', + 'disparity', + ] = 'object_centric_depth' + near: float = 0.0 + far: float = 500.0 + + +class ERayZerGaussiansConfig(BaseModel): + """Gaussian splatting settings.""" + + sh_degree: int = 3 + range_setting: ERayZerRangeSettingConfig = ERayZerRangeSettingConfig() + + +class ERayZerModelConfig(BaseModel): + """Complete model-section config for ERayZer. + + ``checkpoint`` resumes a full Lightning training state. ``init_checkpoint`` + instead loads model weights only (``strict=False``, dropping non-model keys + such as the perceptual loss network) to warm-start fine-tuning from a + pretrained ERayZer checkpoint. + """ + + model_class: Literal['erayzer'] + seed: int = 0 + checkpoint: str | None = None + init_checkpoint: str | None = None + + image_tokenizer: ERayZerImageTokenizerConfig = ERayZerImageTokenizerConfig() + target_image: ERayZerTargetImageConfig = ERayZerTargetImageConfig() + transformer: ERayZerTransformerConfig + pose_latent: ERayZerPoseLatentConfig = ERayZerPoseLatentConfig() + gaussians: ERayZerGaussiansConfig = ERayZerGaussiansConfig() + + hard_pixelalign: bool = False + input_with_pe: bool = True + mask_ratio: float = 0.0 + scaling_bias: float = -2.3 + scaling_max: float = -1.2 + opacity_bias: float = -2.0 + near_plane: float = 0.2 + use_deferred_rendering: bool = False + + +class ERayZerTrainingConfig(BaseModel): + """Training configuration for ERayZer.""" + + train_batch_size: int + val_batch_size: int + test_batch_size: int = 128 + num_epochs: int = 200 + seed: int = 0 + num_workers: int = 8 + num_gpus: int = 1 + num_nodes: int = 1 + log_every_n_steps: int = 10 + check_val_every_n_epoch: int = 1 + ckpt_every_n_epochs: int | None = None + viz_every_n_epochs: int = 1 + precision: str = '32-true' + num_views: int + num_input_views: int + num_target_views: int + random_num_input_views: bool = False + min_input_views: int = 2 + max_input_views: int = 5 + target_view_range: list[int] | None = None + freeze_focal_steps: int = 0 + max_fwdbwd_passes: int + grad_checkpoint_every: int = 1 + train_fraction: float = 0.9 + random_split: bool = False + random_inputs: bool = False + render_interpolate: bool = False + l2_loss_weight: float = 1.0 + gs_reg_loss_weight: float = 0.0 + perceptual_loss_weight: float = 0.0 + pose_consistency_reg_weight: float = 0.0 + + +class ERayZerOptimizerConfig(BaseModel): + """Optimizer configuration for ERayZer (AdamW + LR schedule). + + ``schedule`` selects the LR schedule: ``onecycle`` (cosine warmup+decay) or + ``constant`` (fixed LR, matching the multi-view recipe). ``scale_lr_by_batch`` + applies the linear rule ``lr * global_batch_size / 256`` so the effective LR + tracks the global batch across GPUs/nodes. + """ + + lr: float + beta1: float = 0.9 + beta2: float = 0.95 + wd: float = 0.05 + warmup: int = 3000 + div_factor: float = 1.0 + final_div_factor: float = 1.0 + accumulate_grad_batches: int = 1 + schedule: Literal['onecycle', 'constant'] = 'onecycle' + scale_lr_by_batch: bool = False diff --git a/beast/models/erayzer/erayzer_model.py b/beast/models/erayzer/erayzer_model.py new file mode 100644 index 0000000..205d28f --- /dev/null +++ b/beast/models/erayzer/erayzer_model.py @@ -0,0 +1,1949 @@ +"""ERayZer model: 3DGS renderer, loss computer, and Lightning model. + +ERayZer is the base class that includes a camera-prediction branch +(transformer_encoder + PoseEstimator). Subclasses can override +``_resolve_cameras`` to substitute a different camera source: + +- ``BEAST3D``: drops the prediction modules and reads GT cameras from the + batch dict. +- ``SABLE`` (future): inherits prediction and adds depth/pointcloud + initialization. +""" + +import copy +import logging +from pathlib import Path +from types import SimpleNamespace +from typing import cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from einops import rearrange, repeat +from einops.layers.torch import Rearrange +from torch.nn.modules.module import _IncompatibleKeys + +from beast.geometry.camera import cam_info_to_plucker, get_interpolated_poses_many +from beast.geometry.positional_encoding import get_2d_sincos_pos_embed +from beast.geometry.rotations import quat2mat, rot6d2mat +from beast.losses import masked_mse_loss +from beast.models.base import BaseLightningModel +from beast.models.erayzer.visualize import ( + camera_intrinsic_stats, + export_gaussian_glb, + make_camera_pose_image, + make_render_grid, + viz_is_due, +) +from beast.nn.perceptual import VGGPerceptual +from beast.nn.transformer import ( + QK_Norm_TransformerBlock, + _init_weights, + _init_weights_layerwise, +) +from beast.rendering.gaussians_renderer import ( + GaussianModel, + deferred_gaussian_render, + render_opencv_cam_gsplat, +) + +_logger = logging.getLogger(__name__) + +_INIT_CKPT_CONTAINER_KEYS = ('state_dict', 'model', 'model_state_dict') +_INIT_CKPT_DROP_PREFIXES = ('loss_computer.',) + + +def _filter_init_state_dict( + raw: dict, + drop_prefixes: tuple[str, ...] = _INIT_CKPT_DROP_PREFIXES, +) -> dict: + """Extract a model weight dict from a raw checkpoint and drop ignored keys. + + Handles both bare state dicts and checkpoints that nest the weights under a + container key (``state_dict``, ``model``, or ``model_state_dict``). Keys + beginning with any of ``drop_prefixes`` are removed; the default drops the + perceptual loss network (``loss_computer.*``), which is not part of the + renderable model and is rebuilt from torchvision on demand. + + Parameters + ---------- + raw: object returned by ``torch.load`` on a checkpoint file. + drop_prefixes: key prefixes to exclude from the returned dict. + + Returns + ------- + Flat ``{name: tensor}`` state dict ready for ``load_state_dict``. + + """ + state_dict = raw + for key in _INIT_CKPT_CONTAINER_KEYS: + if isinstance(raw, dict) and isinstance(raw.get(key), dict): + state_dict = raw[key] + break + return { + name: tensor + for name, tensor in state_dict.items() + if not any(name.startswith(prefix) for prefix in drop_prefixes) + } + + +def _parse_hf_checkpoint_url(url: str) -> tuple[str, str, str]: + """Parse a huggingface.co URL into ``(repo_id, revision, filename)``. + + Accepts both the ``...//blob|resolve|raw//`` form and the + bare ``...//`` form (defaulting the revision to ``main``). + """ + from urllib.parse import urlparse + + segments = [s for s in urlparse(url).path.split('/') if s] + if len(segments) < 3: + raise ValueError(f'malformed Hugging Face URL: {url}') + repo_id = '/'.join(segments[:2]) + if segments[2] in ('blob', 'resolve', 'raw'): + if len(segments) < 5: + raise ValueError(f'missing file path in Hugging Face URL: {url}') + return repo_id, segments[3], '/'.join(segments[4:]) + return repo_id, 'main', '/'.join(segments[2:]) + + +def _resolve_init_checkpoint(spec: str) -> Path: + """Resolve an ``init_checkpoint`` spec to a local file path. + + A ``huggingface.co`` URL is downloaded (and cached) via ``huggingface_hub``; + any other value is treated as a local filesystem path. This mirrors the + multi-view loader's ``pretrained_url`` so a pretrained checkpoint can be + fetched from the Hub without a manual download. + + Parameters + ---------- + spec: a local path or a huggingface.co checkpoint URL. + + Returns + ------- + local path to the checkpoint file. + + """ + if spec.startswith(('http://', 'https://')) and 'huggingface.co' in spec: + from huggingface_hub import hf_hub_download # lazy: only needed for URLs + + repo_id, revision, filename = _parse_hf_checkpoint_url(spec) + return Path(hf_hub_download(repo_id=repo_id, filename=filename, revision=revision)) + return Path(spec) + + +def sanitize(t: torch.Tensor) -> torch.Tensor: + """Replace non-finite entries so downstream losses stay valid. + + Parameters + ---------- + t: input tensor. + + Returns + ------- + tensor with NaN → 0, ±Inf → ±1e6, clamped to [-1e6, 1e6]. + + """ + _safe = 1e6 + return torch.nan_to_num(t, nan=0.0, posinf=_safe, neginf=-_safe).clamp(-_safe, _safe) + + +def build_transformer_blocks( + num_layers: int, + d: int, + d_head: int, + use_qk_norm: bool, + special_init: bool = False, + depth_init: bool = False, +) -> nn.ModuleList: + """Build a ModuleList of QK_Norm_TransformerBlock layers. + + Parameters + ---------- + num_layers: number of transformer layers. + d: embedding dimension. + d_head: per-head dimension. + use_qk_norm: whether to apply QK normalization. + special_init: if True, apply layerwise weight scaling. + depth_init: if True, scale std by layer depth rather than total layers. + + Returns + ------- + initialized ModuleList of transformer blocks. + + """ + layers = [ + QK_Norm_TransformerBlock(d, d_head, use_qk_norm=use_qk_norm) + for _ in range(num_layers) + ] + if special_init: + for idx, layer in enumerate(layers): + std = 0.02 / (2 * (idx + 1)) ** 0.5 if depth_init else 0.02 / (2 * num_layers) ** 0.5 + layer.apply(lambda module, s=std: _init_weights_layerwise(module, s)) + return nn.ModuleList(layers) + + +# --------------------------------------------------------------------------- +# Camera-prediction helpers +# --------------------------------------------------------------------------- + +def get_cam_se3( + cam_info: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Decode camera tokens into c2w matrices and fxfycxcy intrinsics. + + Parameters + ---------- + cam_info: predicted camera parameters of shape [B, D] where D is either + 13 (6D rotation + 3D translation + 4 intrinsics) or + 11 (quaternion rotation + 3D translation + 4 intrinsics). + + Returns + ------- + tuple of (c2w [B, 4, 4], fxfycxcy [B, 4]). + + Raises + ------ + NotImplementedError + if D is neither 11 nor 13. + + """ + b, n = cam_info.shape + if n == 13: + R = rot6d2mat(cam_info[:, :6]) + t = cam_info[:, 6:9].unsqueeze(-1) + fxfycxcy = cam_info[:, 9:] + elif n == 11: + R = quat2mat(cam_info[:, :4]) + t = cam_info[:, 4:7].unsqueeze(-1) + fxfycxcy = cam_info[:, 7:] + else: + raise NotImplementedError(f'cam_info width {n} is neither 11 nor 13') + + bottom = torch.tensor([0, 0, 0, 1], dtype=R.dtype, device=R.device) + bottom = bottom.view(1, 1, 4).expand(b, -1, -1) + c2w = torch.cat([torch.cat([R, t], dim=2), bottom], dim=1) + return c2w, fxfycxcy + + +class PoseEstimator(nn.Module): + """Predict per-view SE3 poses and focal lengths from camera tokens. + + Supports three canonical modes (``first``, ``middle``, ``unordered``) and + two pose parameterizations (``6d`` continuous rotation, ``quat``ernion). + """ + + _FOCAL_BIAS = 1.25 # default predicted focal length offset (normalized) + + def __init__(self, config: dict) -> None: + """Initialize PoseEstimator. + + Parameters + ---------- + config: full model config dict; reads config['model']['pose_latent'] + for canonical, mode, and representation settings. + + """ + super().__init__() + self.config = config + pose_cfg = config['model']['pose_latent'] + self.canonical = pose_cfg.get('canonical', 'first') + if self.canonical not in ('first', 'middle', 'unordered'): + raise ValueError(f'Unknown canonical mode: {self.canonical!r}') + self.is_pairwise = pose_cfg.get('mode', 'pairwise') == 'pairwise' + # per_view_focal=False (default) predicts one focal from the canonical + # view and broadcasts it (matches the pretrained ERayZer checkpoint); + # True applies the focal head independently to every view + self.per_view_focal = pose_cfg.get('per_view_focal', False) + self.pose_consistency_reg_weight = config['training'].get( + 'pose_consistency_reg_weight', 0.0, + ) + self.pose_rep = pose_cfg.get('representation', '6d') + if self.pose_rep == '6d': + self.num_pose_element = 6 + elif self.pose_rep == 'quat': + self.num_pose_element = 4 + else: + raise NotImplementedError(f'Unknown pose representation: {self.pose_rep!r}') + + d = config['model']['transformer']['d'] + rel_head_in = d * 2 if self.is_pairwise else d + self.rel_head = nn.Sequential( + nn.Linear(rel_head_in, d, bias=True), + nn.SiLU(), + nn.Linear(d, self.num_pose_element + 3, bias=True), + ) + self.rel_head.apply(_init_weights) + + self.canonical_k_head = nn.Sequential( + nn.Linear(d, d, bias=True), + nn.SiLU(), + nn.Linear(d, 1, bias=False), + ) + self.canonical_k_head.apply(_init_weights) + + def forward( + self, + x: torch.Tensor, + v: int, + ) -> torch.Tensor: + """Predict camera parameters for all views. + + Parameters + ---------- + x: per-view camera tokens of shape [B*V, D] or [B, V, D]. + v: number of views V. + + Returns + ------- + camera parameter tensor of shape [B*V, pose_dim] where pose_dim is + either 13 (6d) or 11 (quat). + + """ + return_flat = x.ndim == 2 + if return_flat: + x = rearrange(x, '(b v) d -> b v d', v=v) + + if self.is_pairwise: + if v == 1: + out = self._forward_single_view(x, v) + else: + out = self._forward_canonical(x, v, self.canonical) + else: + out = self._forward_global(x, v, self.canonical) + + return rearrange(out, 'b v d -> (b v) d') if return_flat else out + + # ------------------------------------------------------------------ + # Private forward variants + # ------------------------------------------------------------------ + + def _identity_rt( + self, + b: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Return an identity rotation+translation canonical tensor [B, 1, num_pose_element+3].""" + if self.pose_rep == '6d': + # first two columns of I₃ then zero translation + vals = [1, 0, 0, 0, 1, 0, 0, 0, 0] + else: + # w=1, x=y=z=0 quaternion, then zero translation + vals = [1, 0, 0, 0, 0, 0, 0] + return torch.tensor(vals, device=device, dtype=dtype).reshape(1, 1, -1).expand(b, 1, -1) + + def _predict_fxfy( + self, + x: torch.Tensor, + v: int, + cano_idx: int = 0, + ) -> torch.Tensor: + """Predict focal lengths from camera tokens. + + With ``per_view_focal`` the focal head is applied independently to every + view, so each camera gets its own focal length. Otherwise the head is + applied to the canonical view only and broadcast to all views, matching + how the pretrained ERayZer checkpoint was trained. + + Parameters + ---------- + x: per-view camera tokens of shape [B, V, D]. + v: number of views to broadcast to. + cano_idx: index of the canonical view (used when not per-view). + + Returns + ------- + fxfy tensor of shape [B, V, 2] (fx == fy per view). + + """ + if self.per_view_focal: + fxfy = self.canonical_k_head(x) + self._FOCAL_BIAS # [B, V', 1] + if fxfy.shape[1] == 1 and v > 1: + fxfy = fxfy.expand(-1, v, -1) # broadcast single view + else: + x_cano = x[:, cano_idx:cano_idx + 1] # [B, 1, D] + fxfy = self.canonical_k_head(x_cano) + self._FOCAL_BIAS # [B, 1, 1] + fxfy = fxfy.expand(-1, v, -1) # [B, V, 1] + return fxfy.expand(-1, -1, 2) # [B, V, 2] + + def _append_cxcy(self, info: torch.Tensor) -> torch.Tensor: + """Append fixed cx=0.5, cy=0.5 principal point to all views.""" + b, v = info.shape[:2] + cxcy = info.new_tensor([0.5, 0.5]).reshape(1, 1, 2).expand(b, v, -1) + return torch.cat([info, cxcy], dim=-1) + + def _forward_canonical( + self, + x: torch.Tensor, + v: int, + canonical: str, + ) -> torch.Tensor: + """Canonical-based pairwise prediction (first or middle view as anchor).""" + b = x.shape[0] + device, dtype = x.device, x.dtype + + if canonical == 'first': + cano_idx = 0 + rel_indices = torch.arange(1, v, device=device) + elif canonical == 'middle': + cano_idx = v // 2 + rel_indices = torch.cat([ + torch.arange(cano_idx, device=device), + torch.arange(cano_idx + 1, v, device=device), + ]) + else: + raise NotImplementedError + + x_canonical = x[:, cano_idx:cano_idx + 1] # [B, 1, D] + x_rel = x[:, rel_indices] # [B, V-1, D] + + # identity canonical pose + fxfy (per-view or canonical-broadcast) + rt_cano = self._identity_rt(b, device, dtype) # [B, 1, num_pose+3] + fxfy = self._predict_fxfy(x, v, cano_idx) # [B, V, 2] + info = torch.cat([rt_cano.expand(b, v, -1), fxfy], dim=-1) # [B, V, num_pose+3+2] + + # relative pose residuals + feat_rel = torch.cat([x_canonical.expand(-1, v - 1, -1), x_rel], dim=-1) + residuals = self.rel_head(feat_rel) # [B, V-1, num_pose+3] + info[:, rel_indices, :self.num_pose_element + 3] = ( + info[:, rel_indices, :self.num_pose_element + 3] + residuals + ) + + return self._append_cxcy(info) + + def _forward_single_view(self, x: torch.Tensor, v: int) -> torch.Tensor: + """Degenerate single-view case: identity pose broadcast to all views.""" + b = x.shape[0] + rt_cano = self._identity_rt(b, x.device, x.dtype) + fxfy = self._predict_fxfy(x, v, cano_idx=0) + info = torch.cat([rt_cano.expand(b, v, -1), fxfy], dim=-1) + return self._append_cxcy(info) + + def _forward_global( + self, + x: torch.Tensor, + v: int, + canonical: str, + ) -> torch.Tensor: + """Global (non-pairwise) prediction: each view gets its own pose offset from identity.""" + b = x.shape[0] + device, dtype = x.device, x.dtype + + rt_canonical = self._identity_rt(b, device, dtype).expand(b, v, -1) + residuals = self.rel_head(x) # [B, V, num_pose+3] + + if canonical == 'unordered': + cano_idx = 0 + rt_final = rt_canonical + residuals + else: + if canonical == 'first': + cano_idx = 0 + elif canonical == 'middle': + cano_idx = v // 2 + else: + raise ValueError(f'Unknown canonical mode: {canonical!r}') + mask = torch.ones(1, v, 1, device=device) + mask[:, cano_idx, :] = 0.0 + rt_final = rt_canonical + residuals * mask + + fxfy = self._predict_fxfy(x, v, cano_idx) # [B, V, 2] + info = torch.cat([rt_final, fxfy], dim=-1) + return self._append_cxcy(info) + + +# --------------------------------------------------------------------------- +# Gaussian splatting helpers +# --------------------------------------------------------------------------- + +class GaussiansUpsampler(nn.Module): + """Project token features to per-Gaussian attributes.""" + + def __init__(self, config: dict) -> None: + """Initialize. + + Parameters + ---------- + config: full model config dict. + + """ + super().__init__() + self.config = config + self.scaling_bias = config['model'].get('scaling_bias', -2.3) + self.scaling_max = config['model'].get('scaling_max', -1.2) + self.opacity_bias = config['model'].get('opacity_bias', -2.0) + + def to_gs( + self, + gaussians: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Split and activate raw Gaussian attribute predictions. + + Parameters + ---------- + gaussians: raw feature tensor of shape [B, N, d_out]. + + Returns + ------- + tuple of (xyz, features, scaling, rotation, opacity) tensors. + + """ + sh_degree = self.config['model']['gaussians']['sh_degree'] + n_sh = (sh_degree + 1) ** 2 * 3 + xyz, features, scaling, rotation, opacity = gaussians.split( + [3, n_sh, 3, 4, 1], dim=2, + ) + + if not self.config['model']['hard_pixelalign']: + xyz = xyz.clamp(-500.0, 500.0) + + features = features.reshape( + features.size(0), + features.size(1), + (sh_degree + 1) ** 2, + 3, + ) + + scaling = (scaling + self.scaling_bias).clamp(max=self.scaling_max).clamp(min=-10.0) + opacity = (opacity + self.opacity_bias).clamp(min=-10.0) + return xyz, features, scaling, rotation, opacity + + +class Renderer(nn.Module): + """Batched Gaussian splat renderer using config-driven dispatch.""" + + def __init__(self, config: dict) -> None: + """Initialize renderer from model config. + + Parameters + ---------- + config: full model config dict; reads config['model']['gaussians']. + + """ + super().__init__() + self.config = config + self.sh_degree = config['model']['gaussians']['sh_degree'] + self.gaussians_model = GaussianModel(config['model']['gaussians']['sh_degree']) + + @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32) + def forward( + self, + xyz, + features, + scaling, + rotation, + opacity, + height, + width, + C2W, + fxfycxcy, + ): + """Render Gaussians for all batch items and views. + + Parameters + ---------- + xyz: [b, n_gaussians, 3]. + features: [b, n_gaussians, (sh_degree+1)^2, 3]. + scaling: [b, n_gaussians, 3]. + rotation: [b, n_gaussians, 4]. + opacity: [b, n_gaussians, 1]. + height: render height in pixels. + width: render width in pixels. + C2W: camera-to-world matrices of shape [b, v, 4, 4]. + fxfycxcy: intrinsics of shape [b, v, 4]. + + Returns + ------- + SimpleNamespace with fields render, depth, alpha. + + """ + if self.config['model'].get('use_deferred_rendering', False): + renderings = deferred_gaussian_render( + xyz, features, scaling, rotation, opacity, height, width, C2W, fxfycxcy, + ) + b, v = C2W.size(0), C2W.size(1) + depth = torch.zeros(b, v, 1, height, width, dtype=torch.float32, device=xyz.device) + alpha = torch.zeros(b, v, 1, height, width, dtype=torch.float32, device=xyz.device) + else: + b, v = C2W.size(0), C2W.size(1) + kw = {'dtype': torch.float32, 'device': xyz.device} + renderings = torch.zeros(b, v, 3, height, width, **kw) + depth = torch.zeros(b, v, 1, height, width, **kw) + alpha = torch.zeros(b, v, 1, height, width, **kw) + + for i in range(b): + pc = self.gaussians_model.set_data( + xyz[i], features[i], scaling[i], rotation[i], opacity[i], + ) + near_plane = self.config['model'].get('near_plane', 0.2) + buffers = render_opencv_cam_gsplat( + pc, height, width, C2W[i], fxfycxcy[i], self.sh_degree, + near_plane=near_plane, + ) + render = buffers['render'] + assert render is not None + renderings[i] = render + if 'depth' in buffers and buffers['depth'] is not None: + depth[i] = buffers['depth'] + if 'alpha' in buffers and buffers['alpha'] is not None: + alpha[i] = buffers['alpha'] + + return SimpleNamespace(render=renderings, depth=depth, alpha=alpha) + + +def get_point_range_func(gaussians_config: dict): + """Return a depth range mapping function based on the gaussians config. + + Parameters + ---------- + gaussians_config: config dict for the gaussians sub-section. + + Returns + ------- + callable mapping raw network output to depth values. + + Raises + ------ + NotImplementedError + if range_setting type is unrecognised. + + """ + range_setting = gaussians_config.get('range_setting', {'type': 'object_centric_depth'}) + if range_setting['type'] == 'object_centric_depth': + def range_object_centric(t): + return (2.0 * torch.sigmoid(t) - 1.0) * 1.5 + 2.7 + return range_object_centric + elif range_setting['type'] == 'linear_depth': + near = range_setting.get('near', 0.0) + far = range_setting.get('far', 500.0) + def range_linear(t): + return torch.sigmoid(t) * (far - near) + near + return range_linear + elif range_setting['type'] == 'log_depth': + near = range_setting.get('near', -6.2) + far = range_setting.get('far', 6.2) + def range_log(t): + return torch.exp(torch.sigmoid(t) * (far - near) + near) + return range_log + elif range_setting['type'] == 'disparity': + near = range_setting.get('near', 0.1) + far = range_setting.get('far', 500.0) + def range_disparity(t): + return 1.0 / (torch.sigmoid(t) * (1.0 / near - 1.0 / far) + 1.0 / far) + return range_disparity + else: + raise NotImplementedError(f'Unknown range_setting type: {range_setting["type"]}') + + +class LossComputer(nn.Module): + """Composite render loss: L2 + gs_reg + optional LPIPS + optional perceptual.""" + + def __init__(self, config: dict, device: str = 'cpu') -> None: + """Initialize. + + Parameters + ---------- + config: full training config dict; reads config['training'] for loss weights. + device: device to place loss modules on. + + """ + super().__init__() + self.config = config + self.device = device + + self.perceptual_loss_module = None + if self.config['training'].get('perceptual_loss_weight', 0.0) > 0.0: + self.perceptual_loss_module = VGGPerceptual(device).eval() + + def forward( + self, + rendering: torch.Tensor, + target: torch.Tensor, + xyz_norm: torch.Tensor | None, + xyz_init_norm: torch.Tensor | None, + pixel_mask: torch.Tensor | None = None, + ) -> SimpleNamespace: + """Compute composite render loss. + + Parameters + ---------- + rendering: predicted images of shape [B, V, 3, H, W]. + target: ground-truth images of shape [B, V, 3 or 4, H, W]. + xyz_norm: normalized xyz coordinates of shape [B, N, 3] or None. + xyz_init_norm: normalized initial xyz coordinates of shape [B, N, 3] or None. + pixel_mask: optional foreground mask of shape [B, V, H, W] or [B*V, 1, H, W]. + + Returns + ------- + SimpleNamespace with fields: loss, l2_loss, psnr, gs_reg_loss, perceptual_loss. + + """ + batch_size, num_views, _, height, width = rendering.shape + rendering = rendering.reshape(batch_size * num_views, 3, height, width) + target = target.reshape(batch_size * num_views, target.shape[2], height, width) + + if target.shape[1] == 4: + target = target[:, :3] + + if pixel_mask is not None: + pixel_mask = pixel_mask.reshape(batch_size * num_views, 1, height, width) + l2_loss = masked_mse_loss(rendering, target, pixel_mask) + else: + l2_loss = F.mse_loss(rendering, target) + + gs_reg_loss = ( + F.mse_loss(xyz_norm, xyz_init_norm) + if xyz_norm is not None and xyz_init_norm is not None + else rendering.new_zeros(()) + ) + + psnr = -10.0 * torch.log10(l2_loss.clamp_min(1e-8)) + + perceptual_loss = rendering.new_zeros(()) + if self.perceptual_loss_module is not None: + perceptual_loss = self.perceptual_loss_module(rendering, target) + + total_loss = ( + self.config['training'].get('l2_loss_weight', 1.0) * l2_loss + + self.config['training'].get('gs_reg_loss_weight', 0.0) * gs_reg_loss + + self.config['training'].get('perceptual_loss_weight', 0.0) * perceptual_loss + ) + + return SimpleNamespace( + loss=total_loss, + l2_loss=l2_loss, + gs_reg_loss=gs_reg_loss, + psnr=psnr, + perceptual_loss=perceptual_loss, + ) + + +# --------------------------------------------------------------------------- +# ERayZer base model +# --------------------------------------------------------------------------- + +class ERayZer(BaseLightningModel): + """ERayZer: encode multi-view images into 3D Gaussians and render. + + Camera parameters are obtained via ``_resolve_cameras``, which by default + runs a learned pose-prediction branch (transformer_encoder + PoseEstimator). + Subclasses override ``_resolve_cameras`` to use a different source (e.g. + ground-truth cameras in ``BEAST3D``). + + The camera prediction branch is only constructed when + ``config['model']['transformer']['encoder_n_layer'] > 0``. + """ + + _NUM_REGISTER_TOKENS = 4 + + def __init__(self, config: dict) -> None: + """Initialize ERayZer model architecture. + + Parameters + ---------- + config: full training/model config dict. + + """ + super().__init__(config) + + self.d = config['model']['transformer']['d'] + self.d_head = config['model']['transformer']['d_head'] + img_size = config['model']['image_tokenizer']['image_size'] + patch_size = config['model']['image_tokenizer']['patch_size'] + self.hh = self.ww = img_size // patch_size + self.ph = self.pw = patch_size + + # patch tokenizer: images → per-patch embedding vectors + self.image_tokenizer = nn.Sequential( + Rearrange( + 'b v c (hh ph) (ww pw) -> (b v) (hh ww) (ph pw c)', + ph=self.ph, + pw=self.pw, + ), + nn.Linear( + config['model']['image_tokenizer']['in_channels'] * self.ph * self.pw, + self.d, + bias=False, + ), + ) + self.image_tokenizer.apply(_init_weights) + + # spatial positional embedding + self.use_pe_embedding_layer = config['model'].get('input_with_pe', True) + if self.use_pe_embedding_layer: + self.pe_embedder = nn.Sequential( + nn.Linear(self.d, self.d), + nn.SiLU(), + nn.Linear(self.d, self.d), + ) + self.pe_embedder.apply(_init_weights) + + self.pe_embedder_plucker = nn.Sequential( + nn.Linear(self.d, self.d), + nn.SiLU(), + nn.Linear(self.d, self.d), + ) + self.pe_embedder_plucker.apply(_init_weights) + + use_qk_norm = config['model']['transformer'].get('use_qk_norm', False) + special_init = config['model']['transformer'].get('special_init', False) + depth_init = config['model']['transformer'].get('depth_init', False) + + # camera prediction branch (only when encoder_n_layer is set) + encoder_n_layer = config['model']['transformer'].get('encoder_n_layer', 0) + if encoder_n_layer > 0: + if encoder_n_layer % 2 != 0: + raise ValueError( + f'encoder_n_layer must be even (alternating frame/global attention), ' + f'got {encoder_n_layer}' + ) + self.num_register_tokens = self._NUM_REGISTER_TOKENS + self.camera_token = nn.Parameter(torch.empty(1, 1, self.d)) + self.register_token = nn.Parameter(torch.empty(1, self.num_register_tokens, self.d)) + nn.init.normal_(self.camera_token, std=1e-6) + nn.init.normal_(self.register_token, std=1e-6) + self.transformer_encoder = build_transformer_blocks( + num_layers=encoder_n_layer, + d=self.d, + d_head=self.d_head, + use_qk_norm=use_qk_norm, + special_init=special_init, + depth_init=depth_init, + ) + self.pose_predictor = PoseEstimator(config) + + # geometry encoder (alternating frame / global attention) + encoder_geom_n_layer = config['model']['transformer']['encoder_geom_n_layer'] + if encoder_geom_n_layer % 2 != 0: + raise ValueError( + f'encoder_geom_n_layer must be even (alternating frame/global attention), ' + f'got {encoder_geom_n_layer}' + ) + self.transformer_encoder_geom = build_transformer_blocks( + num_layers=encoder_geom_n_layer, + d=self.d, + d_head=self.d_head, + use_qk_norm=use_qk_norm, + special_init=special_init, + depth_init=depth_init, + ) + + # Plucker ray tokenizer: 6-channel ray map → per-patch d-dim embedding + self.input_pose_tokenizer = nn.Sequential( + Rearrange( + 'b v (hh ph) (ww pw) c -> (b v) (hh ww) (ph pw c)', + ph=self.ph, + pw=self.pw, + ), + nn.Linear(6 * self.ph * self.pw, self.d, bias=False), + ) + self.input_pose_tokenizer.apply(_init_weights) + + # fusion MLP: concatenated (image + plucker) tokens → d-dim tokens + self.mlp_fuse = nn.Sequential( + nn.LayerNorm(self.d * 2, bias=False), + nn.Linear(self.d * 2, self.d, bias=True), + nn.SiLU(), + nn.Linear(self.d, self.d, bias=True), + ) + self.mlp_fuse.apply(_init_weights) + + self.num_views = config['training']['num_views'] + self.mask_ratio = float(config['model'].get('mask_ratio', 0.0)) + + # Gaussian attribute decoder + sh_degree = config['model']['gaussians']['sh_degree'] + n_gs_channels = 3 + (sh_degree + 1) ** 2 * 3 + 3 + 4 + 1 + self.image_token_decoder = nn.Sequential( + nn.LayerNorm(self.d, bias=False), + nn.Linear(self.d, self.ph * self.pw * n_gs_channels, bias=False), + ) + self.image_token_decoder.apply(_init_weights) + + self.upsampler = GaussiansUpsampler(config) + self.range_func = get_point_range_func(config['model']['gaussians']) + self.renderer = Renderer(config) + self.loss_computer = LossComputer(config) + + self.config_bk = copy.deepcopy(config) + self.render_interpolate = config['training'].get('render_interpolate', False) + + if config.get('inference', False): + self.random_index = config['training'].get('random_inputs', False) + else: + self.random_index = config['training'].get('random_split', False) + + # warm-start fine-tuning from a pretrained ERayZer checkpoint, if given + init_checkpoint = config['model'].get('init_checkpoint') + if init_checkpoint: + self._load_init_checkpoint(init_checkpoint) + + def _load_init_checkpoint(self, ckpt_path: str) -> _IncompatibleKeys: + """Load pretrained model weights (only) from a checkpoint file. + + Reads the weights, drops non-model keys (e.g. the perceptual loss + network), and loads with ``strict=False`` so a partially-overlapping + checkpoint (different layer count, extra heads) still warm-starts the + shared parameters. ``ckpt_path`` may be a local path or a huggingface.co + URL (auto-downloaded), matching the multi-view ``pretrained_url``. + + Parameters + ---------- + ckpt_path: local path or huggingface.co URL of a weight file. + + Returns + ------- + the ``load_state_dict`` result with ``missing_keys``/``unexpected_keys``. + + Raises + ------ + FileNotFoundError: if ``ckpt_path`` does not resolve to a file. + + """ + path = _resolve_init_checkpoint(ckpt_path) + if not path.is_file(): + raise FileNotFoundError(f'init_checkpoint not found: {path}') + raw = torch.load(str(path), map_location='cpu', weights_only=False) + state_dict = _filter_init_state_dict(raw) + result = self.load_state_dict(state_dict, strict=False) + _logger.info( + f'loaded init_checkpoint {path.name}: {len(state_dict)} tensors ' + f'({len(result.missing_keys)} missing, ' + f'{len(result.unexpected_keys)} unexpected)' + ) + return result + + # ------------------------------------------------------------------ + # Camera resolution hook + # ------------------------------------------------------------------ + + def _resolve_cameras( + self, + img_tokens: torch.Tensor, + b: int, + v_all: int, + n: int, + data: dict, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor, bool]: + """Predict camera parameters from image tokens. + + Override this in subclasses to substitute a different camera source. + The base implementation runs the learned camera-prediction branch. + + Parameters + ---------- + img_tokens: spatial-PE-added image tokens of shape [B*V, n, d]. + b: batch size. + v_all: total number of views (may include padding). + n: number of patch tokens per view. + data: batch dict (unused by the prediction branch but available to overrides). + device: compute device. + + Returns + ------- + tuple of: + c2w_all [B, V, 4, 4] camera-to-world matrices, + fxfycxcy_all [B, V, 4] intrinsics, + normalized (bool) — True means fxfycxcy values are divided by image dims. + + Raises + ------ + RuntimeError + if the camera prediction modules were not initialized (encoder_n_layer == 0). + + """ + if not hasattr(self, 'transformer_encoder'): + raise RuntimeError( + 'ERayZer camera prediction modules are not initialized ' + '(encoder_n_layer was 0 or not set). ' + 'Either set encoder_n_layer > 0 in config, or use BEAST3D for GT cameras.' + ) + + # build per-view token sequences: [cam_token, register_tokens, img_tokens] + cam_tokens = repeat(self.camera_token, '1 one d -> bv one d', bv=b * v_all) + reg_tokens = repeat(self.register_token, '1 r d -> bv r d', bv=b * v_all) + all_tokens = torch.cat([cam_tokens, reg_tokens, img_tokens], dim=1) + all_tokens = rearrange(all_tokens, '(b v) n d -> b (v n) d', b=b) + + # camera encoder (alternating frame / global attention) + all_tokens = self.run_vggt_encoder(all_tokens, b, v_all) + all_tokens = rearrange(all_tokens, 'b (v n) d -> (b v) n d', v=v_all) + + # extract camera token (first token of each view) + cam_out = all_tokens[:, 0] # [B*V, d] + + # predict poses + cam_info = self.pose_predictor(cam_out, v_all) # [B*V, pose_dim] + pred_c2w, pred_fxfycxcy = get_cam_se3(cam_info) + pred_c2w = rearrange(pred_c2w, '(b v) h w -> b v h w', b=b) + pred_fxfycxcy = rearrange(pred_fxfycxcy, '(b v) d -> b v d', b=b) + + # learn intrinsics end-to-end through the render loss (requires the + # gsplat fork's v_Ks); freeze them for an initial warmup so geometry can + # stabilize first (freeze_focal_steps=0 disables the freeze) + freeze_steps = self.config['training'].get('freeze_focal_steps', 0) + if self.global_step < freeze_steps: + pred_fxfycxcy = pred_fxfycxcy.detach() + + return pred_c2w, pred_fxfycxcy, True + + # ------------------------------------------------------------------ + # BaseLightningModel interface + # ------------------------------------------------------------------ + + def get_model_outputs(self, batch_dict: dict) -> dict: + """Run ERayZer forward pass and return a plain dict. + + Parameters + ---------- + batch_dict: batch from the dataloader. + + Returns + ------- + dict of model outputs converted from SimpleNamespace. + + """ + return vars(self(batch_dict)) + + def compute_loss( + self, + stage: str, + **kwargs, + ) -> tuple[torch.Tensor, list[dict]]: + """Compute composite render + regularization loss. + + Parameters + ---------- + stage: one of 'train', 'val', 'test'. + **kwargs: model output fields from get_model_outputs. + + Returns + ------- + tuple of (total_loss, log_list). + + """ + rendering = kwargs['render'] + target = kwargs['target_image'] + xyz_norm = kwargs.get('xyz_norm') + xyz_init_norm = kwargs.get('xyz_init_norm') + pixel_mask = kwargs.get('pixel_mask') + result = self.loss_computer(rendering, target, xyz_norm, xyz_init_norm, pixel_mask) + log_list = [ + {'name': f'{stage}_l2', 'value': result.l2_loss}, + {'name': f'{stage}_psnr', 'value': result.psnr, 'prog_bar': True}, + {'name': f'{stage}_gs_reg', 'value': result.gs_reg_loss}, + ] + if result.perceptual_loss.item() != 0.0: + log_list.append({'name': f'{stage}_perceptual', 'value': result.perceptual_loss}) + return result.loss, log_list + + def predict_step(self, batch_dict: dict, batch_idx: int) -> dict: + """Run inference and return model outputs. + + Parameters + ---------- + batch_dict: batch from the dataloader. + batch_idx: batch index (unused). + + Returns + ------- + dict of model outputs. + + """ + return self.get_model_outputs(batch_dict) + + def validation_step(self, batch_dict: dict, batch_idx: int) -> None: + """Run validation and, on cadence, log render/camera/intrinsic visuals. + + Extends the base validation (loss logging) by logging an input/GT/pred + render grid, a 3D predicted-camera plot, and intrinsic/Gaussian scalars + to TensorBoard for the first validation batch every ``viz_every_n_epochs``. + + Parameters + ---------- + batch_dict: batch from the validation dataloader. + batch_idx: index of the current validation batch. + + """ + self.evaluate_batch(batch_dict, 'val') + every = self.config['training'].get('viz_every_n_epochs', 1) + if viz_is_due(batch_idx, self.current_epoch, every): + self._log_validation_visuals(batch_dict) + + @torch.no_grad() + def _log_validation_visuals(self, batch_dict: dict) -> None: + """Log render grid, camera-pose plot, and scalar stats to TensorBoard. + + Best-effort: any failure is logged as a warning and never interrupts + training. No-ops when no image-capable logger is attached. + + Parameters + ---------- + batch_dict: batch from the validation dataloader. + + """ + experiment = getattr(self.logger, 'experiment', None) + if experiment is None or not hasattr(experiment, 'add_image'): + return + try: + out = self.get_model_outputs(batch_dict) + step = self.global_step + # one merged grid: input GT / input render / GT target / pred target + experiment.add_image( + 'val/render_grid', + make_render_grid( + out['input_image'], + out['target_image'], + out['render'], + render_input=out.get('render_input'), + ), + step, + ) + experiment.add_image( + 'val/pred_cameras', + make_camera_pose_image(out['c2w_input'], out['c2w_target']), + step, + ) + image_size = self.config['model']['image_tokenizer']['image_size'] + stats = camera_intrinsic_stats(out['fxfycxcy_target'], image_size) + for name, value in stats.items(): + experiment.add_scalar(f'val/{name}', value, step) + + gaussians = out.get('gaussians') + if gaussians: + opacity = torch.cat([g.get_opacity.flatten() for g in gaussians]).mean() + scaling = torch.cat([g.get_scaling.flatten() for g in gaussians]).mean() + experiment.add_scalar('val/opacity_mean', opacity.item(), step) + experiment.add_scalar('val/scale_mean', scaling.item(), step) + self._export_validation_pointcloud(gaussians[0], step) + except Exception as exc: # noqa: BLE001 — viz must never break training + _logger.warning(f'validation visualization failed: {exc}') + + def _export_validation_pointcloud(self, gaussian_model, step: int) -> None: + """Export the predicted Gaussian point cloud as a GLB under the log dir. + + Writes ``/pointclouds/val_step.glb``. Best-effort: a + missing log dir or export failure is logged and otherwise ignored. + + Parameters + ---------- + gaussian_model: GaussianModel for the sample to export. + step: current global step (used in the filename). + + """ + log_dir = getattr(self.logger, 'log_dir', None) + if log_dir is None: + return + path = Path(log_dir) / 'pointclouds' / f'val_step{step:07d}.glb' + # drop near-transparent Gaussians (opacity <= 0.05) from the export + n = export_gaussian_glb(gaussian_model, path, opacity_threshold=0.05) + _logger.info(f'saved validation point cloud ({n} pts): {path}') + + def configure_optimizers(self) -> dict: + """Configure the AdamW optimizer and LR schedule. + + Optionally scales the LR by the global batch size (linear rule, + ``lr * global_batch_size / 256``) and selects either a constant schedule + (matching the multi-view recipe) or OneCycleLR (cosine warmup+decay). + + Returns + ------- + Lightning optimizer/scheduler config dict. + + """ + cfg = self.config['optimizer'] + lr = cfg['lr'] + if cfg.get('scale_lr_by_batch', False): + tcfg = self.config['training'] + global_batch_size = ( + tcfg['train_batch_size'] * tcfg['num_gpus'] * tcfg['num_nodes'] + ) + lr = lr * global_batch_size / 256 + + optimizer = torch.optim.AdamW( + filter(lambda p: p.requires_grad, self.parameters()), + lr=lr, + betas=(cfg['beta1'], cfg['beta2']), + weight_decay=cfg['wd'], + ) + + schedule = cfg.get('schedule', 'onecycle') + if schedule == 'constant': + # fixed LR for the whole run (multi-view recipe) + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda _step: 1.0) + elif schedule == 'onecycle': + total_steps = self.config['training']['max_fwdbwd_passes'] + scheduler = torch.optim.lr_scheduler.OneCycleLR( + optimizer, + max_lr=lr, + total_steps=total_steps, + pct_start=cfg['warmup'] / total_steps, + anneal_strategy='cos', + div_factor=cfg.get('div_factor', 1.0), + final_div_factor=cfg.get('final_div_factor', 1.0), + ) + else: + raise ValueError(f'unknown optimizer.schedule: {schedule!r}') + + return { + 'optimizer': optimizer, + 'lr_scheduler': {'scheduler': scheduler, 'interval': 'step'}, + 'monitor': 'val_loss', + } + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def forward( + self, + data: dict, + render_video: bool = False, + ) -> SimpleNamespace: + """Encode posed multi-view images into 3D Gaussians and render. + + Parameters + ---------- + data: batch dict with keys: + - image: [B, V, 3, H, W] in [0, 1] + - c2w / fxfycxcy: only required when using a subclass that reads + GT cameras (e.g. BEAST3D). + - input_indices (optional): [B, n_in] + - target_indices (optional): [B, n_tgt] + render_video: whether to render an interpolated video trajectory. + + Returns + ------- + SimpleNamespace with rendered images and auxiliary fields. + + """ + image_all = data['image'] * 2.0 - 1.0 # [B, V, 3, H, W], rescale [0,1] → [-1,1] + + b, v_real, c, h, w = image_all.shape + device = image_all.device + batch_idx = torch.arange(b, device=device).unsqueeze(1) + + input_idx, target_idx = self.resolve_view_indices(data, v_real, device) + if self.training: + input_idx, target_idx = self.maybe_randomize_view_indices( + input_idx, target_idx, device, + ) + + v_input = input_idx.shape[1] + + # pad to 10 views (repeat last view) if fewer are available + pad_views = max(0, 10 - v_real) + if pad_views: + last_view = image_all[:, -1:, ...].expand(-1, pad_views, -1, -1, -1) + image_all = torch.cat([image_all, last_view], dim=1) + v_all = 10 + else: + v_all = v_real + + # tokenize all views and add spatial PE + img_tokens = self.image_tokenizer(image_all) # [B*V, n, d] + _, n, d = img_tokens.shape + if self.use_pe_embedding_layer: + img_tokens = self.add_spatial_pe( + img_tokens, b, v_all, self.hh, self.ww, embedder=self.pe_embedder, + ) + + # resolve cameras (predicted or GT, depending on subclass) + c2w_all, fxfycxcy_all, normalized = self._resolve_cameras( + img_tokens, b, v_all, n, data, device, + ) + + # select input and target cameras + c2w_input = c2w_all[batch_idx, input_idx, ...] # [B, v_input, 4, 4] + fxfycxcy_input = fxfycxcy_all[batch_idx, input_idx, ...] # [B, v_input, 4] + c2w_target = c2w_all[batch_idx, target_idx, ...] # [B, v_target, 4, 4] + fxfycxcy_target = fxfycxcy_all[batch_idx, target_idx, ...].clone() + + # Plucker ray encoding for input views + plucker_rays_input = cam_info_to_plucker( + c2w_input, + fxfycxcy_input, + self.config['model']['target_image'], + normalized=normalized, + return_moment=True, + ) + plucker_rays_input = rearrange( + plucker_rays_input, '(b v) c h w -> b v h w c', b=b, v=v_input, + ) + plucker_emb_input = self.input_pose_tokenizer(plucker_rays_input) + if self.use_pe_embedding_layer: + plucker_emb_input = self.add_spatial_pe( + plucker_emb_input, b, v_input, self.hh, self.ww, + embedder=self.pe_embedder_plucker, + ) + plucker_emb_input = rearrange(plucker_emb_input, '(b v) n d -> b (v n) d', v=v_input) + + # select and optionally mask input image tokens + img_tokens_all = rearrange(img_tokens, '(b v) n d -> b v n d', b=b, v=v_all) + img_tokens_input = img_tokens_all[batch_idx, input_idx, ...] + if self.training and self.mask_ratio > 0: + keep = torch.rand(img_tokens_input.shape[:-1], device=device) >= self.mask_ratio + img_tokens_input = img_tokens_input * keep.unsqueeze(-1).to(img_tokens_input.dtype) + + # fuse image tokens with Plucker embeddings + img_tokens_input = rearrange(img_tokens_input, 'b v n d -> b (v n) d') + img_tokens_input = torch.cat([img_tokens_input, plucker_emb_input], dim=-1) + all_tokens = self.mlp_fuse(img_tokens_input) + + all_tokens = self.run_vggt_encoder_geom(all_tokens, b, v_input) + + # decode tokens to per-pixel Gaussian attributes + img_aligned_gaussians = self.image_token_decoder(all_tokens) + img_aligned_gaussians = rearrange( + img_aligned_gaussians, 'b (v n) d -> b v n d', v=v_input, + )[:, :v_real] + img_aligned_gaussians = rearrange( + img_aligned_gaussians, + 'b v n (ph pw c) -> b (v n ph pw) c', + ph=self.ph, + pw=self.pw, + ) + + xyz, features, scaling, rotation, opacity = self.upsampler.to_gs(img_aligned_gaussians) + + img_aligned_xyz = rearrange( + xyz, + 'b (v hh ww ph pw) c -> b v c (hh ph) (ww pw)', + v=v_input, + hh=self.hh, + ww=self.ww, + ph=self.ph, + pw=self.pw, + ) + + ray_o = None + if self.config['model']['hard_pixelalign']: + img_aligned_xyz = img_aligned_xyz.mean(dim=2, keepdim=True) + img_aligned_xyz = self.range_func(img_aligned_xyz) + plucker_rays = cam_info_to_plucker( + c2w_input, + fxfycxcy_input, + self.config['model']['target_image'], + normalized=normalized, + return_moment=False, + ) + plucker_rays = rearrange(plucker_rays, '(b v) c h w -> b v c h w', b=b) + ray_o, ray_d = plucker_rays.split([3, 3], dim=2) + img_aligned_xyz = ray_o + img_aligned_xyz * ray_d + + xyz = rearrange( + img_aligned_xyz, + 'b v c (hh ph) (ww pw) -> b (v hh ww ph pw) c', + ph=self.ph, + pw=self.pw, + ) + + xyz, features, scaling, rotation, opacity = map( + sanitize, (xyz, features, scaling, rotation, opacity), + ) + + gaussian_attrs = SimpleNamespace( + xyz=xyz, + features=features, + scaling=scaling, + rotation=rotation, + opacity=opacity, + ) + + height, width = int(h), int(w) + if height <= 0 or width <= 0: + raise ValueError( + f'Invalid image dimensions from batch: h={h}, w={w}. ' + 'Check dataset preprocessing and target_image config.' + ) + if normalized: + scale = fxfycxcy_target.new_tensor([width, height, width, height]) + fxfycxcy_target = fxfycxcy_target * scale + + render = self.renderer( + xyz, features, scaling, rotation, opacity, + height, width, C2W=c2w_target, fxfycxcy=fxfycxcy_target, + ) + + # input-view reconstruction: render the Gaussians back at the reference + # cameras as a self-consistency check; computed at val/inference only + render_input = None + if not self.training or self.config.get('inference', False): + fxfycxcy_input_px = fxfycxcy_input + if normalized: + fxfycxcy_input_px = fxfycxcy_input * fxfycxcy_input.new_tensor( + [width, height, width, height], + ) + with torch.no_grad(): + render_input = self.renderer( + xyz, features, scaling, rotation, opacity, + height, width, C2W=c2w_input, fxfycxcy=fxfycxcy_input_px, + ).render + + vis_only_results = None + should_render_video = ( + render_video + or not self.training + or self.config.get('inference', False) + or data.get('return_render_video', False) + ) + if should_render_video: + with torch.no_grad(): + vis_only_results = self.render_images_video( + gaussian_attrs, c2w_target, fxfycxcy_target, normalized=False, + ) + + gaussians = [] + pixelalign_xyz = [] + gaussians_usage = [] + gaussians_scale = [] + gaussians_opacity = [] + for b_i in range(xyz.size(0)): + self.renderer.gaussians_model.empty() + gaussians_model = copy.deepcopy(self.renderer.gaussians_model) + gaussians.append( + gaussians_model.set_data( + xyz[b_i].detach().float(), + features[b_i].detach().float(), + scaling[b_i].detach().float(), + rotation[b_i].detach().float(), + opacity[b_i].detach().float(), + ) + ) + + usage_mask = gaussians[-1].get_opacity > 0.05 + usage = usage_mask.sum() / usage_mask.numel() + gaussians_usage.append(usage.item() if torch.is_tensor(usage) else usage) + + mean_scale = gaussians[-1].get_scaling.mean() + gaussians_scale.append( + mean_scale.item() if torch.is_tensor(mean_scale) else mean_scale, + ) + + mean_opacity = gaussians[-1].get_opacity.mean() + gaussians_opacity.append( + mean_opacity.item() if torch.is_tensor(mean_opacity) else mean_opacity, + ) + + pa_xyz = gaussians[-1].get_xyz + pa_xyz = rearrange( + pa_xyz, + '(v hh ww ph pw) c -> v c (hh ph) (ww pw)', + v=v_input, + hh=self.hh, + ww=self.ww, + ph=self.ph, + pw=self.pw, + ) + pixelalign_xyz.append(pa_xyz) + pixelalign_xyz = torch.stack(pixelalign_xyz, dim=0) + + return SimpleNamespace( + ray_o=ray_o, + gaussians=gaussians, + pixelalign_xyz=pixelalign_xyz, + image=data['image'], + input_image=data['image'][batch_idx, input_idx.clamp(max=v_real - 1), ...], + target_image=data['image'][batch_idx, target_idx, ...], + render=render.render, + render_input=render_input, + c2w_input=c2w_input, + fxfycxcy_input=fxfycxcy_input, + c2w_target=c2w_target, + fxfycxcy_target=fxfycxcy_target, + input_indices=input_idx, + target_indices=target_idx, + xyz_norm=None, + xyz_init_norm=None, + render_video=( + vis_only_results.rendered_images_video.detach().clamp(0, 1) + if vis_only_results is not None + else None + ), + ) + + def predict_frame_from_all_tokens( + self, + all_tokens: torch.Tensor, + c2w_input: torch.Tensor, + fxfycxcy_input: torch.Tensor, + c2w_target: torch.Tensor, + fxfycxcy_target: torch.Tensor, + data: dict, + ) -> SimpleNamespace: + """Decode Gaussians and render from pre-computed tokens. + + Parameters + ---------- + all_tokens: per-view patch tokens of shape [B, v_input, n, d]. + c2w_input: input camera-to-world matrices [B, v_input, 4, 4]. + fxfycxcy_input: input intrinsics [B, v_input, 4]. + c2w_target: target camera-to-world matrices [B, v_target, 4, 4]. + fxfycxcy_target: target intrinsics [B, v_target, 4]. + data: batch dict (used for image shape metadata). + + Returns + ------- + SimpleNamespace with render, gaussians, pixelalign_xyz, and image fields. + + """ + b, v_input, n, d = all_tokens.shape + h = w = int(self.config['model']['image_tokenizer']['image_size']) + + all_tokens = rearrange(all_tokens, 'b v n d -> b (v n) d', v=v_input) + + img_aligned_gaussians = self.image_token_decoder(all_tokens) + img_aligned_gaussians = rearrange( + img_aligned_gaussians, 'b (v n) d -> b v n d', v=v_input, + ) + img_aligned_gaussians = rearrange( + img_aligned_gaussians, + 'b v n (ph pw c) -> b (v n ph pw) c', + ph=self.ph, + pw=self.pw, + ) + + if img_aligned_gaussians.shape[0] != b: + img_aligned_gaussians = rearrange( + img_aligned_gaussians, '(b v) n c -> b (v n) c', b=b, v=v_input, + ) + + xyz, features, scaling, rotation, opacity = self.upsampler.to_gs(img_aligned_gaussians) + + img_aligned_xyz = rearrange( + xyz, + 'b (v hh ww ph pw) c -> b v c (hh ph) (ww pw)', + v=v_input, + hh=self.hh, + ww=self.ww, + ph=self.ph, + pw=self.pw, + ) + + # this method receives pre-computed cameras so normalized=False + normalized = False + ray_o = None + if self.config['model']['hard_pixelalign']: + img_aligned_xyz = img_aligned_xyz.mean(dim=2, keepdim=True) + img_aligned_xyz = self.range_func(img_aligned_xyz) + plucker_rays = cam_info_to_plucker( + c2w_input, + fxfycxcy_input, + self.config['model']['target_image'], + normalized=normalized, + return_moment=False, + ) + plucker_rays = rearrange(plucker_rays, '(b v) c h w -> b v c h w', b=b) + ray_o, ray_d = plucker_rays.split([3, 3], dim=2) + img_aligned_xyz = ray_o + img_aligned_xyz * ray_d + + xyz = rearrange( + img_aligned_xyz, + 'b v c (hh ph) (ww pw) -> b (v hh ww ph pw) c', + ph=self.ph, + pw=self.pw, + ) + + xyz, features, scaling, rotation, opacity = map( + sanitize, (xyz, features, scaling, rotation, opacity), + ) + + height, width = int(h), int(w) + if height <= 0 or width <= 0: + raise ValueError( + f'Invalid image dimensions: h={h}, w={w}. ' + 'Check target_image config.' + ) + + render = self.renderer( + xyz, features, scaling, rotation, opacity, + height, width, C2W=c2w_target, fxfycxcy=fxfycxcy_target, + ) + + gaussians = [] + pixelalign_xyz = [] + for b_i in range(xyz.size(0)): + self.renderer.gaussians_model.empty() + gaussians_model = copy.deepcopy(self.renderer.gaussians_model) + gaussians.append( + gaussians_model.set_data( + xyz[b_i].detach().float(), + features[b_i].detach().float(), + scaling[b_i].detach().float(), + rotation[b_i].detach().float(), + opacity[b_i].detach().float(), + ) + ) + pa_xyz = gaussians[-1].get_xyz + pa_xyz = rearrange( + pa_xyz, + '(v hh ww ph pw) c -> v c (hh ph) (ww pw)', + v=v_input, + hh=self.hh, + ww=self.ww, + ph=self.ph, + pw=self.pw, + ) + pixelalign_xyz.append(pa_xyz) + pixelalign_xyz = torch.stack(pixelalign_xyz, dim=0) + + return SimpleNamespace( + ray_o=ray_o, + gaussians=gaussians, + image=data['image'], + pixelalign_xyz=pixelalign_xyz, + render=render.render, + ) + + # ------------------------------------------------------------------ + # Helper methods + # ------------------------------------------------------------------ + + def resolve_view_indices( + self, + data: dict[str, torch.Tensor], + num_real_views: int, + device: str | torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Resolve input and target view indices from the batch or config. + + Parameters + ---------- + data: batch dict potentially containing explicit view index tensors. + num_real_views: total number of views in the current batch. + device: device for output tensors. + + Returns + ------- + tuple of (input_indices [B, n_in], target_indices [B, n_tgt]). + + """ + batch_size = data['image'].shape[0] + + for input_key in ('input_indices', 'context_indices'): + if input_key in data and 'target_indices' in data: + input_idx = data[input_key].to(device=device, dtype=torch.long) + target_idx = data['target_indices'].to(device=device, dtype=torch.long) + return input_idx, target_idx + + # multi-view-style sampling (matches the multi-view sampler): each batch + # draws a random TOTAL view count in [2, num_views], holds out a small + # number of target views (target_view_range), and feeds the rest as + # input. The per-sample camera shuffle in the dataset makes which + # physical cameras land in each slot differ across samples. + target_view_range = self.config['training'].get('target_view_range') + if self.training and target_view_range: + total = int(torch.randint(2, num_real_views + 1, (1,), device=device).item()) + lo = int(target_view_range[0]) + hi = min(int(target_view_range[1]), total - 1) + hi = max(hi, lo) + num_target = int(torch.randint(lo, hi + 1, (1,), device=device).item()) + num_input = total - num_target + perm = torch.randperm(num_real_views, device=device) + input_idx = perm[:num_input].unsqueeze(0).repeat(batch_size, 1) + target_idx = perm[num_input:total].unsqueeze(0).repeat(batch_size, 1) + return input_idx, target_idx + + # variable reference-view-count sampling: pick a random number of input + # views in [min, max] (per batch, so all samples share a shape) and use + # the remaining views as targets. The dataset shuffles which physical + # cameras occupy each slot per sample, so this is random-N-random-which. + if self.training and self.config['training'].get('random_num_input_views', False): + lo = self.config['training'].get('min_input_views', 2) + hi = min(self.config['training'].get('max_input_views', 5), num_real_views - 1) + hi = max(hi, lo) + num_input = int(torch.randint(lo, hi + 1, (1,), device=device).item()) + perm = torch.randperm(num_real_views, device=device) + input_idx = perm[:num_input].unsqueeze(0).repeat(batch_size, 1) + target_idx = perm[num_input:].unsqueeze(0).repeat(batch_size, 1) + return input_idx, target_idx + + num_input = self.config['training'].get('num_input_views', num_real_views) + num_target = self.config['training'].get('num_target_views', num_real_views) + + can_split = ( + not self.config.get('inference', False) + and num_real_views == self.config['training']['num_views'] + and num_input + num_target <= num_real_views + ) + + if can_split and self.random_index: + perm = torch.randperm(num_real_views, device=device) + input_idx = perm[:num_input].unsqueeze(0).repeat(batch_size, 1) + target_idx = perm[num_input:num_input + num_target].unsqueeze(0).repeat(batch_size, 1) + return input_idx, target_idx + + if can_split: + input_idx = torch.arange(num_input, device=device).unsqueeze(0).repeat(batch_size, 1) + target_idx = torch.arange( + num_input, num_input + num_target, device=device, + ).unsqueeze(0).repeat(batch_size, 1) + return input_idx, target_idx + + input_count = min(num_input, num_real_views) + input_idx = torch.arange( + input_count, device=device, dtype=torch.long, + ).unsqueeze(0).repeat(batch_size, 1) + target_idx = torch.arange( + num_real_views, device=device, dtype=torch.long, + ).unsqueeze(0).repeat(batch_size, 1) + return input_idx, target_idx + + def maybe_randomize_view_indices( + self, + input_idx: torch.Tensor, + target_idx: torch.Tensor, + device: str | torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Randomly swap input and target view indices for two-view training. + + Parameters + ---------- + input_idx: input view indices [B, n_in]. + target_idx: target view indices [B, n_tgt]. + device: device for the swap mask. + + Returns + ------- + tuple of (input_idx, target_idx) possibly swapped per batch item. + + """ + b = input_idx.shape[0] + swap = (torch.rand(b, device=device) < 0.5).unsqueeze(1) + if input_idx.shape[1] == 1: + return ( + torch.where(swap, target_idx, input_idx), + torch.where(swap, input_idx, target_idx), + ) + return input_idx, target_idx + + def add_spatial_pe( + self, + tokens: torch.Tensor, + b: int, + v: int, + h_tokens: int, + w_tokens: int, + embedder: nn.Module, + ) -> torch.Tensor: + """Add 2D sinusoidal positional encoding to patch tokens. + + Parameters + ---------- + tokens: patch token tensor of shape [B*V, n, d]. + b: batch size. + v: number of views. + h_tokens: number of token rows (H / patch_size). + w_tokens: number of token columns (W / patch_size). + embedder: learnable projection applied to the positional encoding. + + Returns + ------- + tokens with spatial positional encoding added, shape [B*V, n, d]. + + """ + bv, n, d = tokens.shape + if h_tokens * w_tokens != n: + raise ValueError(f'Token count {n} != h_tokens*w_tokens {h_tokens}*{w_tokens}') + + spatial_pe = get_2d_sincos_pos_embed( + embed_dim=d, + grid_size=(h_tokens, w_tokens), + device=tokens.device, + ).to(tokens.dtype) # [n, d] + + spatial_pe = spatial_pe.reshape(1, 1, n, d).expand(b, v, -1, -1).reshape(bv, n, d) + return tokens + embedder(spatial_pe) + + def run_vggt_encoder( + self, + all_tokens: torch.Tensor, + b: int, + v: int, + ) -> torch.Tensor: + """Run the camera-prediction encoder with alternating frame/global attention. + + Parameters + ---------- + all_tokens: token tensor of shape [B, V*n, d]. + b: batch size. + v: number of views. + + Returns + ------- + updated token tensor of shape [B, V*n, d]. + + """ + checkpoint_every = self.config['training']['grad_checkpoint_every'] + for i in range(0, len(self.transformer_encoder), checkpoint_every): + if i % 2 == 0: + all_tokens = rearrange(all_tokens, 'b (v n) d -> (b v) n d', v=v) + else: + all_tokens = rearrange(all_tokens, '(b v) n d -> b (v n) d', b=b) + all_tokens = cast( + torch.Tensor, + torch.utils.checkpoint.checkpoint( + self._run_layers_encoder(i, i + 1), + all_tokens, + use_reentrant=False, + ), + ) + if checkpoint_every > 1: + all_tokens = self._run_layers_encoder(i + 1, i + checkpoint_every)(all_tokens) + return all_tokens + + def _run_layers_encoder(self, start: int, end: int): + """Return a closure applying camera-encoder layers [start, end). + + Parameters + ---------- + start: index of the first layer to apply. + end: one past the index of the last layer to apply. + + Returns + ------- + callable that applies the specified layer range. + + """ + def custom_forward(tokens): + for i in range(start, min(end, len(self.transformer_encoder))): + tokens = self.transformer_encoder[i](tokens) + return tokens + return custom_forward + + def run_vggt_encoder_geom( + self, + all_tokens: torch.Tensor, + b: int, + v: int, + ) -> torch.Tensor: + """Run the geometry encoder with alternating frame/global attention. + + Parameters + ---------- + all_tokens: token tensor of shape [B, V*n, d]. + b: batch size. + v: number of views. + + Returns + ------- + updated token tensor of shape [B, V*n, d]. + + """ + checkpoint_every = self.config['training']['grad_checkpoint_every'] + for i in range(0, len(self.transformer_encoder_geom), checkpoint_every): + if i % 2 == 0: + all_tokens = rearrange(all_tokens, 'b (v n) d -> (b v) n d', v=v) + else: + all_tokens = rearrange(all_tokens, '(b v) n d -> b (v n) d', b=b) + all_tokens = cast( + torch.Tensor, + torch.utils.checkpoint.checkpoint( + self._run_layers_geom(i, i + 1), + all_tokens, + use_reentrant=False, + ), + ) + if checkpoint_every > 1: + all_tokens = self._run_layers_geom(i + 1, i + checkpoint_every)(all_tokens) + return all_tokens + + def _run_layers_geom(self, start: int, end: int): + """Return a closure applying geometry-encoder layers [start, end). + + Parameters + ---------- + start: index of the first layer to apply. + end: one past the index of the last layer to apply. + + Returns + ------- + callable that applies the specified layer range. + + """ + def custom_forward(tokens): + for i in range(start, min(end, len(self.transformer_encoder_geom))): + tokens = self.transformer_encoder_geom[i](tokens) + return tokens + return custom_forward + + def render_images_video( + self, + gaussian_attrs: SimpleNamespace, + c2w_all: torch.Tensor, + fxfycxcy_all: torch.Tensor, + normalized: bool = False, + ) -> SimpleNamespace: + """Render an interpolated video trajectory from a set of Gaussians. + + Parameters + ---------- + gaussian_attrs: SimpleNamespace with xyz, features, scaling, rotation, opacity. + c2w_all: camera-to-world matrices [B, V, 4, 4]. + fxfycxcy_all: intrinsics [B, V, 4]. + normalized: whether intrinsics are normalized (will be scaled by image size). + + Returns + ------- + SimpleNamespace with field rendered_images_video of shape [B, V', 3, H, W]. + + """ + with torch.no_grad(): + xyz = gaussian_attrs.xyz.detach() + features = gaussian_attrs.features.detach() + scaling = gaussian_attrs.scaling.detach() + rotation = gaussian_attrs.rotation.detach() + opacity = gaussian_attrs.opacity.detach() + c2w_all = c2w_all.detach() + fxfycxcy_all = fxfycxcy_all.detach() + + b, v, _, _ = c2w_all.shape + device = xyz.device + num_frames = 30 + + all_renderings = [] + for i in range(b): + c2ws = c2w_all[i] + fxfycxcy = fxfycxcy_all[i] + Ks = torch.zeros((c2ws.shape[0], 3, 3), device=device) + Ks[:, 0, 0] = fxfycxcy[:, 0] + Ks[:, 1, 1] = fxfycxcy[:, 1] + Ks[:, 0, 2] = fxfycxcy[:, 2] + Ks[:, 1, 2] = fxfycxcy[:, 3] + c2ws_interp, Ks_interp = get_interpolated_poses_many( + c2ws[:, :3, :4], Ks, num_frames, + ) + frame_c2ws = torch.cat( + [ + c2ws_interp.to(device), + torch.tensor([[[0, 0, 0, 1]]], device=device).expand( + c2ws_interp.shape[0], -1, -1, + ), + ], + dim=1, + ) + frame_fxfycxcy = torch.zeros((c2ws_interp.shape[0], 4), device=device) + frame_fxfycxcy[:, 0] = Ks_interp[:, 0, 0] + frame_fxfycxcy[:, 1] = Ks_interp[:, 1, 1] + frame_fxfycxcy[:, 2] = Ks_interp[:, 0, 2] + frame_fxfycxcy[:, 3] = Ks_interp[:, 1, 2] + + img_size = self.config['model']['image_tokenizer']['image_size'] + num_views_traj = frame_c2ws.shape[0] + renderings = [] + for start in range(0, num_views_traj, 5): + end = min(start + 5, num_views_traj) + rendered_batch = self.renderer( + xyz, features, scaling, rotation, opacity, + img_size, img_size, + C2W=frame_c2ws[start:end].unsqueeze(0), + fxfycxcy=frame_fxfycxcy[start:end].unsqueeze(0), + ).render.squeeze(0) + renderings.append(rendered_batch) + + all_renderings.append(torch.cat(renderings, dim=0)) + + all_renderings = torch.stack(all_renderings) # [B, V', 3, H, W] + + return SimpleNamespace(rendered_images_video=all_renderings) diff --git a/beast/models/erayzer/erayzer_train.py b/beast/models/erayzer/erayzer_train.py new file mode 100644 index 0000000..b0fc80f --- /dev/null +++ b/beast/models/erayzer/erayzer_train.py @@ -0,0 +1,9 @@ +"""Training entry point for ERayZer. + +Delegates to the shared beast.train.train function, which handles +MultiViewDataModule and step-based training for the erayzer model class. +""" + +from beast.train import train + +__all__ = ['train'] diff --git a/beast/models/erayzer/visualize.py b/beast/models/erayzer/visualize.py new file mode 100644 index 0000000..286900c --- /dev/null +++ b/beast/models/erayzer/visualize.py @@ -0,0 +1,390 @@ +"""Validation-time visualization for ERayZer. + +Builds TensorBoard-ready images so training quality can be inspected live: + +- ``make_render_grid``: a 2-row GT-vs-render comparison over all views, with + colored borders marking input (reference) vs target (novel) views. +- ``make_camera_pose_image``: a 3D plot of all predicted camera poses as + frustums, colored by input vs target (frustum style ported from the + multi-view fig_teaser_cam_erayzer figure). +- ``camera_intrinsic_stats``: scalar summaries of the predicted intrinsics. +- ``export_gaussian_glb``: write the predicted Gaussians as a GLB point cloud. +- ``viz_is_due``: cadence helper deciding when to log during validation. + +Image builders return ``[3, H, W]`` float tensors in ``[0, 1]`` for +``SummaryWriter.add_image`` with the default ``CHW`` data format. +""" + +import numpy as np +import torch +from matplotlib.backends.backend_agg import FigureCanvasAgg +from matplotlib.figure import Figure +from mpl_toolkits.mplot3d import Axes3D # noqa: F401 (registers the 3d projection) + +# shared color code: input (reference) views vs target (novel) views +_INPUT_RGB = (0.20, 0.80, 0.36) # green +_TARGET_RGB = (0.95, 0.35, 0.20) # orange-red +_SH_C0 = 0.28209479177387814 # 0th-order SH coefficient (SH DC -> RGB) + + +def viz_is_due(batch_idx: int, current_epoch: int, every_n_epochs: int) -> bool: + """Decide whether validation visuals should be logged for this batch. + + Visuals are logged only for the first validation batch of an epoch, and only + every ``every_n_epochs`` epochs. + + Parameters + ---------- + batch_idx: index of the current validation batch. + current_epoch: current training epoch. + every_n_epochs: log cadence in epochs; ``<= 0`` disables logging. + + Returns + ------- + True if visuals should be logged now. + + """ + if every_n_epochs <= 0: + return False + if batch_idx != 0: + return False + return current_epoch % every_n_epochs == 0 + + +# --------------------------------------------------------------------------- +# Render grid +# --------------------------------------------------------------------------- + + +def _select_sample(t: torch.Tensor, sample_idx: int) -> torch.Tensor: + """Return a single sample [V, ...] from a [B, V, ...] or [V, ...] tensor.""" + if t.ndim >= 5: + return t[sample_idx] + return t + + +def _add_border( + img: torch.Tensor, + color: tuple[float, float, float], + width: int = 4, +) -> torch.Tensor: + """Paint a solid colored border of ``width`` pixels around a [3, H, W] image.""" + out = img.clone() + c = torch.tensor(color, dtype=out.dtype).view(3, 1, 1) + out[:, :width, :] = c + out[:, -width:, :] = c + out[:, :, :width] = c + out[:, :, -width:] = c + return out + + +def _row_from_views( + views: torch.Tensor, + sample_idx: int, + border_color: tuple[float, float, float], + pad: int = 2, +) -> torch.Tensor: + """Tile a sample's views horizontally, each with a colored border.""" + v = _select_sample(views, sample_idx).detach().cpu().float().clamp(0.0, 1.0) + h = v.shape[2] + sep = torch.ones(3, h, pad) + cols = [] + for i in range(v.shape[0]): + cols.append(_add_border(v[i], border_color)) + if i < v.shape[0] - 1: + cols.append(sep) + return torch.cat(cols, dim=2) + + +def _build_row( + groups: list[tuple[torch.Tensor, tuple[float, float, float]]], + sample_idx: int, + group_gap: int = 10, +) -> torch.Tensor: + """Concatenate per-group bordered view strips into one [3, H, W] row.""" + parts = [] + for i, (views, color) in enumerate(groups): + row = _row_from_views(views, sample_idx, color) + parts.append(row) + if i < len(groups) - 1: + parts.append(torch.ones(3, row.shape[1], group_gap)) # gap between groups + return torch.cat(parts, dim=2) + + +def _pad_width(row: torch.Tensor, width: int, fill: float = 1.0) -> torch.Tensor: + """Right-pad a [3, H, W] row to ``width`` with a constant fill.""" + cur = row.shape[2] + if cur >= width: + return row + pad = torch.full((3, row.shape[1], width - cur), fill) + return torch.cat([row, pad], dim=2) + + +def _stack_rows(rows: list[torch.Tensor]) -> torch.Tensor: + """Right-pad rows to a common width and stack them with grey separators.""" + width = max(row.shape[2] for row in rows) + rows = [_pad_width(row, width) for row in rows] + sep = torch.full((3, 3, width), 0.5) # grey separator between rows + stacked = [] + for i, row in enumerate(rows): + stacked.append(row) + if i < len(rows) - 1: + stacked.append(sep) + return torch.cat(stacked, dim=1).clamp(0.0, 1.0).float() + + +def make_render_grid( + input_image: torch.Tensor, + target_image: torch.Tensor, + render: torch.Tensor, + render_input: torch.Tensor | None = None, + sample_idx: int = 0, +) -> torch.Tensor: + """Build a 2-row GT-vs-render comparison grid over all views. + + Row 0 is ground truth, row 1 is the model's render, with one column per view + in the order input views then target views. Each image gets a colored + border: green for input (reference) views, red for target (novel) views, so + the two view types are distinguishable at a glance. When ``render_input`` is + omitted only the target views are shown (the input views cannot be rendered). + + Inputs are stacks of per-view images shaped ``[B, V, 3, H, W]`` (or + ``[V, 3, H, W]``). + + Parameters + ---------- + input_image: reference views fed to the model. + target_image: ground-truth novel (target) views. + render: predicted renders of the target views. + render_input: optional renders of the input views (input-view NVS). + sample_idx: which batch element to visualize. + + Returns + ------- + a ``[3, H, W]`` float image in ``[0, 1]``: row 0 GT, row 1 render. + + """ + if render_input is not None: + gt_row = _build_row( + [(input_image, _INPUT_RGB), (target_image, _TARGET_RGB)], sample_idx, + ) + render_row = _build_row( + [(render_input, _INPUT_RGB), (render, _TARGET_RGB)], sample_idx, + ) + else: + gt_row = _build_row([(target_image, _TARGET_RGB)], sample_idx) + render_row = _build_row([(render, _TARGET_RGB)], sample_idx) + return _stack_rows([gt_row, render_row]) + + +# --------------------------------------------------------------------------- +# Camera-pose plot +# --------------------------------------------------------------------------- + + +def _select_cameras(c2w: torch.Tensor, sample_idx: int) -> np.ndarray: + """Return a [V, 4, 4] numpy array of camera matrices for one sample.""" + if c2w.ndim == 4: + c2w = c2w[sample_idx] + return c2w.detach().cpu().to(torch.float32).numpy() + + +def _frustum_segments(c2w: np.ndarray, scale: float) -> list[tuple[np.ndarray, np.ndarray]]: + """Line segments of one camera frustum in world coords (OpenCV convention). + + Apex at the camera center, base ahead at +z, plus a short stub on the top + edge encoding the 'up' direction so orientation flips are visible. + """ + s = scale + pts_cam = np.array([ + [0.0, 0.0, 0.0], # apex + [-s, -s, 1.5 * s], # base top-left + [s, -s, 1.5 * s], # base top-right + [s, s, 1.5 * s], # base bottom-right + [-s, s, 1.5 * s], # base bottom-left + ]) + rot = c2w[:3, :3] + trans = c2w[:3, 3] + pts_w = (rot @ pts_cam.T).T + trans + apex, base = pts_w[0], pts_w[1:] + segs = [(apex, base[i]) for i in range(4)] + segs += [(base[0], base[1]), (base[1], base[2]), (base[2], base[3]), (base[3], base[0])] + top_mid = 0.5 * (base[0] + base[1]) + segs.append((top_mid, top_mid + 0.4 * s * (-rot[:, 1]))) # up stub + return segs + + +def _frustum_scale_for(*c2w_arrays: np.ndarray) -> float: + """Pick a frustum size from the spread of camera centers.""" + centers = np.concatenate([c[:, :3, 3] for c in c2w_arrays], axis=0) + radii = np.linalg.norm(centers - np.median(centers, axis=0), axis=1) + median_radius = float(np.median(radii)) if radii.size else 1.0 + return 0.15 * (median_radius if median_radius > 1e-6 else 1.0) + + +def _draw_camera_frustums( + ax: Axes3D, + c2w_np: np.ndarray, + color: tuple[float, float, float], + label: str, + scale: float, +) -> None: + """Plot one set of camera frustums plus center markers on ``ax``.""" + centers = c2w_np[:, :3, 3] + # Axes3D.scatter accepts an array for zs at runtime; its stub types it as int + ax.scatter( + centers[:, 0], + centers[:, 1], + centers[:, 2], # pyright: ignore[reportArgumentType] + color=color, s=24, depthshade=False, label=label, + ) + for v in range(c2w_np.shape[0]): + for p0, p1 in _frustum_segments(c2w_np[v], scale): + ax.plot([p0[0], p1[0]], [p0[1], p1[1]], [p0[2], p1[2]], + color=color, linewidth=1.3) + + +def _set_axes_equal(ax: Axes3D, centers: np.ndarray, pad: float = 0.3) -> None: + """Tight equal-aspect bounding cube around all camera centers.""" + lo, hi = centers.min(axis=0), centers.max(axis=0) + mid = 0.5 * (lo + hi) + extent = float(np.max(hi - lo)) + half = (1.0 + pad) * 0.5 * (extent if extent > 0 else 1.0) + ax.set_xlim(mid[0] - half, mid[0] + half) + ax.set_ylim(mid[1] - half, mid[1] + half) + ax.set_zlim(mid[2] - half, mid[2] + half) + ax.set_box_aspect((1, 1, 1)) + + +def make_camera_pose_image( + c2w_input: torch.Tensor, + c2w_target: torch.Tensor | None = None, + sample_idx: int = 0, +) -> torch.Tensor: + """Render a 3D plot of all predicted camera poses to an image tensor. + + Input (reference) cameras are drawn as green frustums and target (novel) + cameras as red frustums, so every camera in the batch is shown and the two + types are color-coded. Uses the Agg backend so it works headless. + + Parameters + ---------- + c2w_input: input camera-to-world matrices ``[B, V, 4, 4]`` or ``[V, 4, 4]``. + c2w_target: optional target camera-to-world matrices, same layout. + sample_idx: which batch element to visualize. + + Returns + ------- + a ``[3, H, W]`` float image in ``[0, 1]``. + + """ + cams_input = _select_cameras(c2w_input, sample_idx) + arrays = [cams_input] + cams_target = None + if c2w_target is not None: + cams_target = _select_cameras(c2w_target, sample_idx) + arrays.append(cams_target) + scale = _frustum_scale_for(*arrays) + + fig = Figure(figsize=(6, 5)) + canvas = FigureCanvasAgg(fig) + ax = fig.add_subplot(111, projection='3d') + ax.view_init(elev=15, azim=-70) + _draw_camera_frustums(ax, cams_input, _INPUT_RGB, 'input', scale) + if cams_target is not None: + _draw_camera_frustums(ax, cams_target, _TARGET_RGB, 'target', scale) + + centers = np.concatenate([c[:, :3, 3] for c in arrays], axis=0) + _set_axes_equal(ax, centers) + ax.set_xlabel('X') + ax.set_ylabel('Y') + ax.set_zlabel('Z') + ax.set_xticklabels([]) + ax.set_yticklabels([]) + # set_zticklabels is wrapped by a descriptor the stub types as non-callable + ax.set_zticklabels([]) # pyright: ignore[reportCallIssue] + ax.legend(loc='upper right', fontsize=9) + ax.set_title('predicted cameras (green=input, red=target)') + + canvas.draw() + rgb = np.asarray(canvas.buffer_rgba())[..., :3].copy() # [H, W, 3] uint8 + img = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0 + return img.clamp(0.0, 1.0) + + +# --------------------------------------------------------------------------- +# Scalars + point cloud +# --------------------------------------------------------------------------- + + +def camera_intrinsic_stats(fxfycxcy: torch.Tensor, image_size: int) -> dict[str, float]: + """Summarize predicted intrinsics as normalized scalar statistics. + + Parameters + ---------- + fxfycxcy: intrinsics tensor of shape ``[..., 4]`` in pixels (fx, fy, cx, cy). + image_size: image side length used to normalize pixel intrinsics. + + Returns + ------- + dict of plain floats: fx/fy means, fx spread (std), and cx/cy means, all + normalized to ``[0, 1]`` image fractions. + + """ + norm = fxfycxcy.detach().float() / image_size + fx = norm[..., 0].flatten() + fy = norm[..., 1].flatten() + cx = norm[..., 2].flatten() + cy = norm[..., 3].flatten() + return { + 'focal_fx_mean': fx.mean().item(), + 'focal_fx_std': fx.std(unbiased=False).item(), + 'focal_fy_mean': fy.mean().item(), + 'cx_mean': cx.mean().item(), + 'cy_mean': cy.mean().item(), + } + + +def export_gaussian_glb(gaussian_model, path, opacity_threshold: float = 0.0) -> int: + """Export a predicted Gaussian field as a colored point-cloud GLB. + + Reads positions and SH-DC color from a ``GaussianModel`` (anything exposing + ``get_xyz``, ``get_features``, ``get_opacity``), optionally drops near-empty + Gaussians, and writes a GLB point cloud. ``trimesh`` is imported lazily so the + core model import does not depend on it. + + Parameters + ---------- + gaussian_model: a GaussianModel-like object with get_xyz/get_features/get_opacity. + path: output ``.glb`` path. + opacity_threshold: drop Gaussians with activated opacity at or below this. + + Returns + ------- + number of points written (0 if nothing passed the filter; no file is written). + + """ + from pathlib import Path + + # lazy import keeps trimesh optional for the core model import + import trimesh # pyright: ignore[reportMissingImports] + + xyz = gaussian_model.get_xyz.detach().cpu().float() + feat = gaussian_model.get_features.detach().cpu().float() # [N, K, 3] + opacity = gaussian_model.get_opacity.detach().cpu().float().reshape(-1) + rgb = (feat[:, 0, :] * _SH_C0 + 0.5).clamp(0.0, 1.0) # SH DC -> RGB + + if opacity_threshold > 0.0: + mask = opacity > opacity_threshold + xyz, rgb = xyz[mask], rgb[mask] + + n = int(xyz.shape[0]) + if n == 0: + return 0 + + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + colors = (rgb.numpy() * 255).astype(np.uint8) + rgba = np.concatenate([colors, np.full((n, 1), 255, np.uint8)], axis=1) + trimesh.PointCloud(vertices=xyz.numpy(), colors=rgba).export(str(path)) + return n diff --git a/beast/models/perceptual.py b/beast/models/perceptual.py deleted file mode 100644 index 3307e85..0000000 --- a/beast/models/perceptual.py +++ /dev/null @@ -1,70 +0,0 @@ -"""Perceptual loss modules using pretrained feature extractors.""" - -from typing import Any - -import torch -import torchvision -from torch import nn - -# https://github.com/MLReproHub/SMAE/blob/main/src/loss/perceptual.py - - -class Perceptual(nn.Module): - """Base perceptual loss module that compares feature representations.""" - - def __init__(self, *, network: nn.Module, criterion: nn.Module) -> None: - """Initialize perceptual loss module. - - Parameters - ---------- - network: feature extractor that maps input images to feature tensors - criterion: loss function applied to extracted features (e.g. MSELoss) - """ - super().__init__() - self.net = network - self.criterion = criterion - self.sigmoid = nn.Sigmoid() - - def forward(self, x_hat: torch.Tensor, x: torch.Tensor) -> torch.Tensor: - """Compute perceptual loss between reconstructed and target images. - - Parameters - ---------- - x_hat: reconstructed image batch - x: target image batch - - Returns - ------- - scalar loss tensor - - """ - x_hat_features = self.sigmoid(self.net(x_hat)) - x_features = self.sigmoid(self.net(x)) - loss = self.criterion(x_hat_features, x_features) - return loss - - -class AlexPerceptual(Perceptual): - """Perceptual loss using the first five layers of a pretrained AlexNet.""" - - def __init__(self, *, device: str | torch.device, **kwargs: Any) -> None: - """Perceptual loss using pretrained AlexNet features [Pihlgren et al. 2020]. - - Extracts features from the first five layers of AlexNet (pretrained on ImageNet) - and computes loss between reconstructed and target feature maps. - - Parameters - ---------- - device: device to run the feature extractor on (e.g. 'cuda', 'cpu') - **kwargs: passed to parent; must include criterion (e.g. nn.MSELoss()) - """ - # Load alex net pretrained on IN1k - alex_net = torchvision.models.alexnet(weights='IMAGENET1K_V1') - # Extract features after second relu activation - # Append sigmoid layer to normalize features - perceptual_net = alex_net.features[:5].to(device) - # don't record gradients for the perceptual net; gradients will still propagate through - for parameter in perceptual_net.parameters(): - parameter.requires_grad = False - - super().__init__(network=perceptual_net, **kwargs) diff --git a/beast/models/registry.py b/beast/models/registry.py new file mode 100644 index 0000000..c1c7f9c --- /dev/null +++ b/beast/models/registry.py @@ -0,0 +1,100 @@ +"""Central registries for BEAST model classes, training functions, and config schemas. + +Three dicts are populated by _register_all() at import time: + + MODEL_REGISTRY model_class str → LightningModule subclass + TRAIN_REGISTRY model_class str → train(config, model, output_dir) callable + CONFIG_REGISTRY model_class str → model-section Pydantic config class + +To register a new model, add three lines inside _register_all(): + MODEL_REGISTRY[''] = + TRAIN_REGISTRY[''] = + CONFIG_REGISTRY[''] = + +See docs/developer_guide.md for the full step-by-step checklist. +""" + +from collections.abc import Callable +from typing import Any + +from pydantic import BaseModel + +MODEL_REGISTRY: dict[str, type] = {} +TRAIN_REGISTRY: dict[str, Callable] = {} +CONFIG_REGISTRY: dict[str, type[BaseModel]] = {} + + +def _register_all() -> None: + """Import all model packages so they can populate the registries.""" + from beast.models.beast_resnet.beast_resnet_config import ResnetModelConfig + from beast.models.beast_resnet.beast_resnet_model import ResnetAutoencoder + from beast.models.beast_resnet.beast_resnet_train import train as resnet_train + from beast.models.beast_vit.beast_vit_config import VitModelConfig + from beast.models.beast_vit.beast_vit_model import VisionTransformer + from beast.models.beast_vit.beast_vit_train import train as vit_train + from beast.models.erayzer.erayzer_config import ERayZerModelConfig + from beast.models.erayzer.erayzer_model import ERayZer + from beast.models.erayzer.erayzer_train import train as erayzer_train + + MODEL_REGISTRY['resnet'] = ResnetAutoencoder + TRAIN_REGISTRY['resnet'] = resnet_train + CONFIG_REGISTRY['resnet'] = ResnetModelConfig + + MODEL_REGISTRY['vit'] = VisionTransformer + TRAIN_REGISTRY['vit'] = vit_train + CONFIG_REGISTRY['vit'] = VitModelConfig + + MODEL_REGISTRY['erayzer'] = ERayZer + TRAIN_REGISTRY['erayzer'] = erayzer_train + CONFIG_REGISTRY['erayzer'] = ERayZerModelConfig + + +def get_model_class(model_class: str) -> type: + """Return the model class for the given model_class string. + + Parameters + ---------- + model_class: model type identifier string (e.g., 'resnet', 'vit') + + Returns + ------- + model class + + Raises + ------ + KeyError: if model_class is not registered + + """ + if model_class not in MODEL_REGISTRY: + raise KeyError( + f'Unknown model class {model_class!r}. ' + f'Registered: {sorted(MODEL_REGISTRY)}' + ) + return MODEL_REGISTRY[model_class] + + +def get_train_fn(model_class: str) -> Callable[..., Any]: + """Return the training function for the given model_class string. + + Parameters + ---------- + model_class: model type identifier string + + Returns + ------- + training callable + + Raises + ------ + KeyError: if model_class is not registered + + """ + if model_class not in TRAIN_REGISTRY: + raise KeyError( + f'No training function registered for {model_class!r}. ' + f'Registered: {sorted(TRAIN_REGISTRY)}' + ) + return TRAIN_REGISTRY[model_class] + + +_register_all() diff --git a/beast/nn/__init__.py b/beast/nn/__init__.py new file mode 100644 index 0000000..793e0d0 --- /dev/null +++ b/beast/nn/__init__.py @@ -0,0 +1,8 @@ +"""Reusable neural network building blocks shared across BEAST models. + +Modules +------- +perceptual — AlexNet-based perceptual loss network +transformer — attention layers, MLP, and transformer blocks (used by ERayZer) +dino — DINOv3 feature extractor wrapper +""" diff --git a/beast/nn/dino.py b/beast/nn/dino.py new file mode 100644 index 0000000..3d40004 --- /dev/null +++ b/beast/nn/dino.py @@ -0,0 +1,61 @@ +"""DINOv3 feature extractor for patch and CLS token extraction.""" + +import torch.nn as nn +from transformers import AutoModel + + +class DinoV3(nn.Module): + """DINOv3 feature extractor returning patch and CLS tokens.""" + + def __init__( + self, + model_name: str = 'facebook/dinov3-vitb16-pretrain-lvd1689m', + freeze: bool = True, + ) -> None: + """Initialize DINOv3. + + Parameters + ---------- + model_name: HuggingFace model identifier. + freeze: whether to freeze all parameters (default True). + + """ + super().__init__() + + self.model = AutoModel.from_pretrained(model_name) + + self.embed_dim = self.model.config.hidden_size + + if freeze: + for p in self.model.parameters(): + p.requires_grad = False + + def forward(self, images): + """Extract patch and CLS tokens from multi-view images. + + Parameters + ---------- + images: float tensor of shape (B, V, 3, H, W) in [0, 1]. + + Returns + ------- + tuple of (patch_tokens [B, V, N, embed_dim], cls_tokens [B, V, embed_dim]). + + """ + B, V = images.shape[:2] + + x = images.view(B * V, *images.shape[2:]) + + outputs = self.model(pixel_values=x) + + hidden = outputs.last_hidden_state + + cls_tokens = hidden[:, 0] + patch_tokens = hidden[:, 5:] + + N = patch_tokens.shape[1] + + patch_tokens = patch_tokens.view(B, V, N, self.embed_dim) + cls_tokens = cls_tokens.view(B, V, self.embed_dim) + + return patch_tokens, cls_tokens diff --git a/beast/nn/perceptual.py b/beast/nn/perceptual.py new file mode 100644 index 0000000..e317138 --- /dev/null +++ b/beast/nn/perceptual.py @@ -0,0 +1,167 @@ +"""Perceptual loss modules using pretrained feature extractors.""" + +from typing import cast + +import torch +import torch.nn as nn +import torchvision +from torchvision.models import VGG19_Weights, vgg19 + +# https://github.com/MLReproHub/SMAE/blob/main/src/loss/perceptual.py + + +class Perceptual(nn.Module): + """Base class for perceptual loss modules. + + Subclasses must implement ``forward(x_hat, x) -> Tensor`` where ``x_hat`` + is the reconstruction and ``x`` is the target, both in [0, 1]. + """ + + +class AlexPerceptual(Perceptual): + """Perceptual loss using the first five layers of a pretrained AlexNet.""" + + def __init__(self, *, device: str | torch.device, criterion: nn.Module) -> None: + """Perceptual loss using pretrained AlexNet features [Pihlgren et al. 2020]. + + Extracts features from the first five layers of AlexNet (pretrained on ImageNet) + and computes loss between reconstructed and target feature maps. + + Parameters + ---------- + device: device to run the feature extractor on (e.g. 'cuda', 'cpu') + criterion: loss function applied to sigmoid-normalised features (e.g. nn.MSELoss()) + + """ + super().__init__() + alex_net = torchvision.models.alexnet(weights='IMAGENET1K_V1') + perceptual_net = alex_net.features[:5].to(device) + for parameter in perceptual_net.parameters(): + parameter.requires_grad = False + self.net = perceptual_net + self.criterion = criterion + self.sigmoid = nn.Sigmoid() + + def forward(self, x_hat: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Compute perceptual loss between reconstructed and target images. + + Parameters + ---------- + x_hat: reconstructed image batch + x: target image batch + + Returns + ------- + scalar loss tensor + + """ + x_hat_features = self.sigmoid(self.net(x_hat)) + x_features = self.sigmoid(self.net(x)) + return self.criterion(x_hat_features, x_features) + + +class VGGPerceptual(Perceptual): + """RayZer-style multi-scale perceptual loss using VGG19 features. + + Extracts features at five spatial scales and combines weighted L1 losses, + matching the RayZer perceptual loss formulation. Attempts to load + matconvnet weights; falls back to torchvision ImageNet weights. + """ + + def __init__(self, device: str | torch.device = 'cpu') -> None: + """Initialize VGG19 feature extractor. + + Parameters + ---------- + device: device to place the VGG model on + + """ + super().__init__() + self.device = device + self.vgg = self._build_vgg() + self._setup_feature_blocks() + + def _build_vgg(self) -> nn.Module: + """Build VGG19 with ImageNet weights, replacing MaxPool layers with AvgPool.""" + model = vgg19(weights=VGG19_Weights.IMAGENET1K_V1) + features = cast(nn.Sequential, model.features) + for idx, layer in enumerate(features): + if isinstance(layer, nn.MaxPool2d): + features[idx] = nn.AvgPool2d(kernel_size=2, stride=2) + model = model.to(self.device).eval() + for param in model.parameters(): + param.requires_grad = False + return model + + def _setup_feature_blocks(self) -> None: + """Build sequential feature extraction blocks from VGG19 layers.""" + output_indices = [0, 4, 9, 14, 23, 32] + self.blocks = nn.ModuleList() + features = cast(nn.Sequential, self.vgg.features) + for i in range(len(output_indices) - 1): + layers = cast(nn.Sequential, features[output_indices[i]:output_indices[i + 1]]) + block = nn.Sequential(*layers) + block = block.to(self.device).eval() + for param in block.parameters(): + param.requires_grad = False + self.blocks.append(block) + + def _extract_features(self, x: torch.Tensor) -> list[torch.Tensor]: + """Run input through all feature blocks and collect intermediate outputs. + + Parameters + ---------- + x: preprocessed image tensor + + Returns + ------- + list of feature tensors, one per block + + """ + features = [] + for block in self.blocks: + x = block(x) + features.append(x) + return features + + def _preprocess(self, x: torch.Tensor) -> torch.Tensor: + """Subtract ImageNet channel means from images scaled to [0, 255]. + + Parameters + ---------- + x: image tensor in [0, 1] + + Returns + ------- + preprocessed image tensor + + """ + mean = torch.tensor([123.68, 116.779, 103.939], device=x.device).view(1, 3, 1, 1) + return x * 255.0 - mean + + def forward(self, x_hat: torch.Tensor, x: torch.Tensor) -> torch.Tensor: + """Compute multi-scale perceptual loss between reconstructed and target images. + + Parameters + ---------- + x_hat: reconstructed image tensor in [0, 1] + x: target image tensor in [0, 1] + + Returns + ------- + scalar perceptual loss + + """ + x_hat = self._preprocess(x_hat) + x = self._preprocess(x) + + x_hat_f = self._extract_features(x_hat) + x_f = self._extract_features(x) + + e0 = torch.mean(torch.abs(x - x_hat)) + e1 = torch.mean(torch.abs(x_f[0] - x_hat_f[0])) / 2.6 + e2 = torch.mean(torch.abs(x_f[1] - x_hat_f[1])) / 4.8 + e3 = torch.mean(torch.abs(x_f[2] - x_hat_f[2])) / 3.7 + e4 = torch.mean(torch.abs(x_f[3] - x_hat_f[3])) / 5.6 + e5 = torch.mean(torch.abs(x_f[4] - x_hat_f[4])) * 10.0 / 1.5 + return (e0 + e1 + e2 + e3 + e4 + e5) / 255.0 diff --git a/beast/nn/transformer.py b/beast/nn/transformer.py new file mode 100644 index 0000000..dbcaf58 --- /dev/null +++ b/beast/nn/transformer.py @@ -0,0 +1,236 @@ +"""Transformer building blocks: attention layers, MLP, and transformer blocks.""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +def _init_weights(module: nn.Module) -> None: + """Apply standard weight initialisation to Linear, Embedding, and Conv2d layers. + + Parameters + ---------- + module: module to initialise. + + """ + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + elif isinstance(module, nn.Conv2d): + torch.nn.init.kaiming_normal_(module.weight, mode='fan_out', nonlinearity='relu') + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + + +def _init_weights_layerwise(module: nn.Module, weight_init_std: float) -> None: + """Apply layerwise weight initialisation with a custom standard deviation. + + Parameters + ---------- + module: module to initialise. + weight_init_std: standard deviation for normal initialisation. + + """ + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=weight_init_std) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=weight_init_std) + + +class MLP(nn.Module): + """MLP layer. + + Reference: https://github.com/facebookresearch/dino/blob/7c446df5b9f45747937fb0d72314eb9f7b66930a/vision_transformer.py#L49-L65 + + """ + + def __init__( + self, + d: int, + mlp_ratio: int = 4, + mlp_bias: bool = False, + mlp_dropout: float = 0.0, + mlp_dim: int | None = None, + ) -> None: + """Initialize. + + Parameters + ---------- + d: token dimension. + mlp_ratio: hidden dimension multiplier when mlp_dim is None. + mlp_bias: whether to include bias in linear layers. + mlp_dropout: dropout probability after the output projection. + mlp_dim: optional explicit hidden dimension; overrides mlp_ratio. + + """ + super().__init__() + if mlp_dim is None: + mlp_dim = d * mlp_ratio + self.mlp = nn.Sequential( + nn.Linear(d, mlp_dim, bias=mlp_bias), + nn.GELU(), + nn.Linear(mlp_dim, d, bias=mlp_bias), + nn.Dropout(mlp_dropout), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + x = self.mlp(x) + return x + + +class RMSNorm(nn.Module): + """Root mean square layer normalisation.""" + + def __init__(self, dim: int, eps: float = 1e-5) -> None: + """Initialize. + + Parameters + ---------- + dim: feature dimension to normalise over. + eps: epsilon added to variance for numerical stability. + + """ + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def _norm(self, x: torch.Tensor) -> torch.Tensor: + """Apply RMS normalisation without the learnable scale. + + Parameters + ---------- + x: input tensor. + + Returns + ------- + normalised tensor with the same shape as x. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + output = self._norm(x.float()).type_as(x) + return output * self.weight.type_as(x) + + +class QK_Norm_SelfAttention(nn.Module): + """Self-attention with optional QK normalisation via RMSNorm.""" + + def __init__( + self, + d: int, + d_head: int, + attn_qkv_bias: bool = False, + attn_fc_bias: bool = True, + attn_dropout: float = 0.0, + attn_fc_dropout: float = 0.0, + use_qk_norm: bool = True, + ) -> None: + """Initialize. + + Parameters + ---------- + d: token dimension. + d_head: per-head dimension; must evenly divide d. + attn_qkv_bias: whether to include bias in the QKV projection. + attn_fc_bias: whether to include bias in the output projection. + attn_dropout: dropout probability on attention weights. + attn_fc_dropout: dropout probability after the output projection. + use_qk_norm: whether to apply RMSNorm to Q and K before attention. + + """ + super().__init__() + assert ( + d % d_head == 0 + ), f'Token dimension {d} should be divisible by head dimension {d_head}' + self.d = d + self.d_head = d_head + self.attn_dropout = attn_dropout + + self.to_qkv = nn.Linear(d, 3 * d, bias=attn_qkv_bias) + self.fc = nn.Linear(d, d, bias=attn_fc_bias) + self.attn_fc_dropout = nn.Dropout(attn_fc_dropout) + self.use_qk_norm = use_qk_norm + + if self.use_qk_norm: + self.q_norm = RMSNorm(d_head) + self.k_norm = RMSNorm(d_head) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass. + + Parameters + ---------- + x: input tensor of shape (b, l, d). + + """ + q, k, v = self.to_qkv(x).split(self.d, dim=2) + + q, k, v = (rearrange(t, 'b l (nh dh) -> b nh l dh', dh=self.d_head) for t in (q, k, v)) + + if self.use_qk_norm: + q = self.q_norm(q) + k = self.k_norm(k) + + dropout_p = self.attn_dropout if self.training else 0.0 + x = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p) + x = rearrange(x, 'b nh l dh -> b l (nh dh)') + x = self.attn_fc_dropout(self.fc(x)) + return x + + +class QK_Norm_TransformerBlock(nn.Module): + """Pre-norm transformer block with QK-normalised self-attention.""" + + def __init__( + self, + d: int, + d_head: int, + ln_bias: bool = False, + attn_qkv_bias: bool = False, + attn_dropout: float = 0.0, + attn_fc_bias: bool = False, + attn_fc_dropout: float = 0.0, + mlp_ratio: int = 4, + mlp_bias: bool = False, + mlp_dropout: float = 0.0, + use_qk_norm: bool = True, + ) -> None: + """Initialize. + + Parameters + ---------- + d: token dimension. + d_head: per-head dimension. + ln_bias: whether to include bias in LayerNorm. + attn_qkv_bias: whether to include bias in the QKV projection. + attn_dropout: dropout probability on attention weights. + attn_fc_bias: whether to include bias in the attention output projection. + attn_fc_dropout: dropout probability after the attention output projection. + mlp_ratio: MLP hidden dimension multiplier. + mlp_bias: whether to include bias in MLP linear layers. + mlp_dropout: dropout probability in the MLP. + use_qk_norm: whether to apply RMSNorm to Q and K before attention. + + """ + super().__init__() + self.norm1 = nn.LayerNorm(d, bias=ln_bias) + self.attn = QK_Norm_SelfAttention( + d, d_head, attn_qkv_bias, attn_fc_bias, attn_dropout, attn_fc_dropout, use_qk_norm, + ) + self.norm2 = nn.LayerNorm(d, bias=ln_bias) + self.mlp = MLP(d, mlp_ratio, mlp_bias, mlp_dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x diff --git a/beast/rendering/__init__.py b/beast/rendering/__init__.py new file mode 100644 index 0000000..8b76ff4 --- /dev/null +++ b/beast/rendering/__init__.py @@ -0,0 +1 @@ +"""Rendering utilities: Gaussian splatting renderer.""" diff --git a/beast/rendering/gaussians_renderer.py b/beast/rendering/gaussians_renderer.py new file mode 100644 index 0000000..3f556bf --- /dev/null +++ b/beast/rendering/gaussians_renderer.py @@ -0,0 +1,1039 @@ +"""Gaussian splatting renderer and 3D Gaussian model utilities.""" + +import logging +import math +from collections import OrderedDict +from pathlib import Path +from typing import TypeVar + +import matplotlib +import numpy as np +import torch +from gsplat import rasterization +from plyfile import PlyData, PlyElement + +_logger = logging.getLogger(__name__) + + +def strip_lowerdiag(L: torch.Tensor) -> torch.Tensor: + """Extract lower-diagonal elements of a batch of matrices. + + Parameters + ---------- + L: input tensor of shape [N, 3, 3]. + + Returns + ------- + uncertainty tensor of shape [N, 6]. + + """ + uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device=L.device) + + uncertainty[:, 0] = L[:, 0, 0] + uncertainty[:, 1] = L[:, 0, 1] + uncertainty[:, 2] = L[:, 0, 2] + uncertainty[:, 3] = L[:, 1, 1] + uncertainty[:, 4] = L[:, 1, 2] + uncertainty[:, 5] = L[:, 2, 2] + return uncertainty + + +def strip_symmetric(sym: torch.Tensor) -> torch.Tensor: + """Extract symmetric lower-diagonal elements. + + Parameters + ---------- + sym: symmetric matrix batch of shape [N, 3, 3]. + + Returns + ------- + lower-diagonal elements of shape [N, 6]. + + """ + return strip_lowerdiag(sym) + + +def build_rotation(r: torch.Tensor) -> torch.Tensor: + """Build rotation matrices from quaternions. + + Parameters + ---------- + r: quaternion tensor of shape [N, 4]. + + Returns + ------- + rotation matrices of shape [N, 3, 3]. + + """ + norm = torch.sqrt( + r[:, 0] * r[:, 0] + r[:, 1] * r[:, 1] + r[:, 2] * r[:, 2] + r[:, 3] * r[:, 3] + ) + + q = r / norm[:, None] + + R = torch.zeros((q.size(0), 3, 3), device=r.device) + + r = q[:, 0] + x = q[:, 1] + y = q[:, 2] + z = q[:, 3] + + R[:, 0, 0] = 1 - 2 * (y * y + z * z) + R[:, 0, 1] = 2 * (x * y - r * z) + R[:, 0, 2] = 2 * (x * z + r * y) + R[:, 1, 0] = 2 * (x * y + r * z) + R[:, 1, 1] = 1 - 2 * (x * x + z * z) + R[:, 1, 2] = 2 * (y * z - r * x) + R[:, 2, 0] = 2 * (x * z - r * y) + R[:, 2, 1] = 2 * (y * z + r * x) + R[:, 2, 2] = 1 - 2 * (x * x + y * y) + return R + + +def build_scaling_rotation(s: torch.Tensor, r: torch.Tensor) -> torch.Tensor: + """Build scaling+rotation covariance matrix. + + Parameters + ---------- + s: scale tensor of shape [N, 3]. + r: rotation quaternion tensor of shape [N, 4]. + + Returns + ------- + covariance factor tensor of shape [N, 3, 3]. + + """ + L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device=s.device) + R = build_rotation(r) + + L[:, 0, 0] = s[:, 0] + L[:, 1, 1] = s[:, 1] + L[:, 2, 2] = s[:, 2] + + L = R @ L + return L + + +C0 = 0.28209479177387814 + +# SH<->RGB conversions are pure arithmetic and work on either tensors or arrays; +# the TypeVar preserves the input type in the return so callers keep ndarray methods +TensorOrArray = TypeVar('TensorOrArray', torch.Tensor, np.ndarray) + + +def RGB2SH(rgb: torch.Tensor) -> torch.Tensor: + """Convert RGB to SH coefficients. + + Parameters + ---------- + rgb: RGB tensor in [0, 1]. + + Returns + ------- + SH coefficient tensor. + + """ + return (rgb - 0.5) / C0 + + +def SH2RGB(sh: TensorOrArray) -> TensorOrArray: + """Convert SH coefficients to RGB. + + Parameters + ---------- + sh: SH coefficient tensor or array. + + Returns + ------- + RGB tensor or array (matching the input type). + + """ + return sh * C0 + 0.5 + + + +class GaussianModel: + """3D Gaussian splat model with per-Gaussian attributes.""" + + def setup_functions(self) -> None: + """Set up activation functions for Gaussian parameters.""" + def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): + """Build covariance matrix from scaling and rotation.""" + L = build_scaling_rotation(scaling_modifier * scaling, rotation) + actual_covariance = L @ L.transpose(1, 2) + symm = strip_symmetric(actual_covariance) + return symm + + self.scaling_activation = torch.exp + self.inv_scaling_activation = torch.log + self.rotation_activation = torch.nn.functional.normalize + self.opacity_activation = torch.sigmoid + self.covariance_activation = build_covariance_from_scaling_rotation + + def __init__(self, sh_degree: int, scaling_modifier: float | None = None) -> None: + """Initialize with SH degree and optional scaling modifier. + + Parameters + ---------- + sh_degree: spherical harmonics degree. + scaling_modifier: optional global scale factor applied at render time. + + """ + self.sh_degree = sh_degree + self._xyz = torch.empty(0) + self._features_dc = torch.empty(0) + if self.sh_degree > 0: + self._features_rest = torch.empty(0) + else: + self._features_rest = None + self._scaling = torch.empty(0) + self._rotation = torch.empty(0) + self._opacity = torch.empty(0) + self.setup_functions() + + self.scaling_modifier = scaling_modifier + + def empty(self) -> None: + """Reset to empty state.""" + self.__init__(self.sh_degree, self.scaling_modifier) + + def set_data( + self, + xyz: torch.Tensor, + features: torch.Tensor, + scaling: torch.Tensor, + rotation: torch.Tensor, + opacity: torch.Tensor, + ) -> 'GaussianModel': + """Set Gaussian data from tensors. + + Parameters + ---------- + xyz: positions of shape (N, 3). + features: SH features of shape (N, (sh_degree + 1) ** 2, 3). + scaling: log-scale values of shape (N, 3). + rotation: quaternions of shape (N, 4). + opacity: logit-opacity of shape (N, 1). + + Returns + ------- + self (for chaining). + + """ + self._xyz = xyz + self._features_dc = features[:, :1, :].contiguous() + if self.sh_degree > 0: + self._features_rest = features[:, 1:, :].contiguous() + else: + self._features_rest = None + self._scaling = scaling + self._rotation = rotation + self._opacity = opacity + return self + + def to(self, device: str | torch.device) -> 'GaussianModel': + """Move all tensors to device. + + Parameters + ---------- + device: target device. + + Returns + ------- + self (for chaining). + + """ + self._xyz = self._xyz.to(device) + self._features_dc = self._features_dc.to(device) + if self.sh_degree > 0: + assert self._features_rest is not None + self._features_rest = self._features_rest.to(device) + self._scaling = self._scaling.to(device) + self._rotation = self._rotation.to(device) + self._opacity = self._opacity.to(device) + return self + + def filter(self, valid_mask: torch.Tensor) -> 'GaussianModel': + """Keep only the Gaussians indicated by valid_mask. + + Parameters + ---------- + valid_mask: boolean tensor of shape (N,). + + Returns + ------- + self (for chaining). + + """ + self._xyz = self._xyz[valid_mask] + self._features_dc = self._features_dc[valid_mask] + if self.sh_degree > 0: + assert self._features_rest is not None + self._features_rest = self._features_rest[valid_mask] + self._scaling = self._scaling[valid_mask] + self._rotation = self._rotation[valid_mask] + self._opacity = self._opacity[valid_mask] + return self + + def crop(self, crop_bbx: list[float] | None = None) -> 'GaussianModel': + """Remove Gaussians outside the given bounding box. + + Parameters + ---------- + crop_bbx: [x_min, x_max, y_min, y_max, z_min, z_max]; defaults to unit cube. + + Returns + ------- + self (for chaining). + + """ + if crop_bbx is None: + crop_bbx = [-1, 1, -1, 1, -1, 1] + x_min, x_max, y_min, y_max, z_min, z_max = crop_bbx + xyz = self._xyz + invalid_mask = ( + (xyz[:, 0] < x_min) + | (xyz[:, 0] > x_max) + | (xyz[:, 1] < y_min) + | (xyz[:, 1] > y_max) + | (xyz[:, 2] < z_min) + | (xyz[:, 2] > z_max) + ) + valid_mask = ~invalid_mask + + return self.filter(valid_mask) + + def prune(self, opacity_thres: float = 0.05) -> 'GaussianModel': + """Remove low-opacity Gaussians. + + Parameters + ---------- + opacity_thres: opacity threshold below which Gaussians are removed. + + Returns + ------- + self (for chaining). + + """ + opacity = self.get_opacity.squeeze(1) + valid_mask = opacity > opacity_thres + + return self.filter(valid_mask) + + def prune_by_nearfar( + self, + cam_origins: torch.Tensor, + nearfar_percent: tuple[float, float] = (0.01, 0.99), + ) -> 'GaussianModel': + """Remove Gaussians that are too close or too far from any camera. + + Parameters + ---------- + cam_origins: camera origin positions of shape [V, 3]. + nearfar_percent: (near_pct, far_pct) quantile thresholds in [0, 1]. + + Returns + ------- + self (for chaining). + + """ + assert len(nearfar_percent) == 2 + assert nearfar_percent[0] < nearfar_percent[1] + assert nearfar_percent[0] >= 0 and nearfar_percent[1] <= 1 + + device = self._xyz.device + dists = torch.cdist(self._xyz[None], cam_origins[None].to(device))[0] + dists_percentile = torch.quantile( + dists, torch.tensor(nearfar_percent).to(device), dim=0, + ) + reject_mask = (dists < dists_percentile[0:1, :]) | ( + dists > dists_percentile[1:2, :] + ) + reject_mask = reject_mask.any(dim=1) + valid_mask = ~reject_mask + + return self.filter(valid_mask) + + def apply_all_filters( + self, + opacity_thres: float = 0.05, + crop_bbx: list[float] | None = None, + cam_origins: torch.Tensor | None = None, + nearfar_percent: tuple[float, float] = (0.005, 1.0), + ) -> 'GaussianModel': + """Apply opacity pruning, bounding-box crop, and near/far pruning. + + Parameters + ---------- + opacity_thres: opacity threshold for pruning. + crop_bbx: optional bounding box [x_min, x_max, y_min, y_max, z_min, z_max]. + cam_origins: optional camera origins for near/far pruning. + nearfar_percent: (near_pct, far_pct) quantile thresholds. + + Returns + ------- + self (for chaining). + + """ + if crop_bbx is None: + crop_bbx = [-1, 1, -1, 1, -1, 1] + self.prune(opacity_thres) + if crop_bbx is not None: + self.crop(crop_bbx) + if cam_origins is not None: + self.prune_by_nearfar(cam_origins, nearfar_percent) + return self + + def shrink_bbx(self, drop_ratio: float = 0.05) -> 'GaussianModel': + """Shrink the bounding box by dropping outlier Gaussians. + + Parameters + ---------- + drop_ratio: fraction of extreme values to drop from each side. + + Returns + ------- + self (for chaining). + + """ + xyz = self._xyz + xyz_min, xyz_max = torch.quantile( + xyz, + torch.tensor([drop_ratio, 1 - drop_ratio]).float().to(xyz.device), + dim=0, + ) + xyz_min = xyz_min.detach().cpu().numpy() + xyz_max = xyz_max.detach().cpu().numpy() + crop_bbx = [ + xyz_min[0], + xyz_max[0], + xyz_min[1], + xyz_max[1], + xyz_min[2], + xyz_max[2], + ] + _logger.info(f'Shrinking bbx to {crop_bbx}') + return self.crop(crop_bbx) + + def report_stats(self) -> None: + """Log statistics for all Gaussian attributes.""" + _logger.info( + f'xyz: {self._xyz.shape}, {self._xyz.min().item()}, {self._xyz.max().item()}' + ) + _logger.info( + f'features_dc: {self._features_dc.shape}, ' + f'{self._features_dc.min().item()}, {self._features_dc.max().item()}' + ) + if self.sh_degree > 0: + assert self._features_rest is not None + _logger.info( + f'features_rest: {self._features_rest.shape}, ' + f'{self._features_rest.min().item()}, {self._features_rest.max().item()}' + ) + _logger.info( + f'scaling: {self._scaling.shape}, ' + f'{self._scaling.min().item()}, {self._scaling.max().item()}' + ) + _logger.info( + f'rotation: {self._rotation.shape}, ' + f'{self._rotation.min().item()}, {self._rotation.max().item()}' + ) + _logger.info( + f'opacity: {self._opacity.shape}, ' + f'{self._opacity.min().item()}, {self._opacity.max().item()}' + ) + _logger.info( + f'after activation, xyz: {self.get_xyz.shape}, ' + f'{self.get_xyz.min().item()}, {self.get_xyz.max().item()}' + ) + _logger.info( + f'after activation, features: {self.get_features.shape}, ' + f'{self.get_features.min().item()}, {self.get_features.max().item()}' + ) + _logger.info( + f'after activation, scaling: {self.get_scaling.shape}, ' + f'{self.get_scaling.min().item()}, {self.get_scaling.max().item()}' + ) + _logger.info( + f'after activation, rotation: {self.get_rotation.shape}, ' + f'{self.get_rotation.min().item()}, {self.get_rotation.max().item()}' + ) + _logger.info( + f'after activation, opacity: {self.get_opacity.shape}, ' + f'{self.get_opacity.min().item()}, {self.get_opacity.max().item()}' + ) + _logger.info( + f'after activation, covariance: {self.get_covariance().shape}, ' + f'{self.get_covariance().min().item()}, {self.get_covariance().max().item()}' + ) + + @property + def get_scaling(self) -> torch.Tensor: + """Return activated scaling values.""" + if self.scaling_modifier is not None: + return self.scaling_activation(self._scaling) * self.scaling_modifier + else: + return self.scaling_activation(self._scaling) + + @property + def get_rotation(self) -> torch.Tensor: + """Return normalised rotation quaternions.""" + return self.rotation_activation(self._rotation) + + @property + def get_xyz(self) -> torch.Tensor: + """Return xyz positions.""" + return self._xyz + + @property + def get_features(self) -> torch.Tensor: + """Return all SH feature coefficients.""" + if self.sh_degree > 0: + assert self._features_rest is not None + features_dc = self._features_dc + features_rest = self._features_rest + return torch.cat((features_dc, features_rest), dim=1) + else: + return self._features_dc + + @property + def get_opacity(self) -> torch.Tensor: + """Return activated opacity values.""" + return self.opacity_activation(self._opacity) + + def get_covariance(self, scaling_modifier: float = 1) -> torch.Tensor: + """Return covariance matrices from scaling and rotation. + + Parameters + ---------- + scaling_modifier: global scale multiplier. + + Returns + ------- + covariance tensor of shape [N, 6]. + + """ + return self.covariance_activation( + self.get_scaling, scaling_modifier, self._rotation, + ) + + def construct_dtypes( + self, + use_fp16: bool = False, + enable_gs_viewer: bool = True, + ) -> list[tuple[str, str]]: + """Build PLY element dtype list for save_ply. + + Parameters + ---------- + use_fp16: use float16 instead of float32 for PLY storage. + enable_gs_viewer: pad SH to degree 3 for GS viewer compatibility. + + Returns + ------- + list of (name, dtype_str) pairs. + + """ + if not use_fp16: + dtype_list = [ + ('x', 'f4'), + ('y', 'f4'), + ('z', 'f4'), + ('red', 'u1'), + ('green', 'u1'), + ('blue', 'u1'), + ] + for i in range(self._features_dc.shape[1] * self._features_dc.shape[2]): + dtype_list.append((f'f_dc_{i}', 'f4')) + + if enable_gs_viewer: + assert self.sh_degree <= 3, 'GS viewer only supports SH up to degree 3' + sh_degree = 3 + for i in range(((sh_degree + 1) ** 2 - 1) * 3): + dtype_list.append((f'f_rest_{i}', 'f4')) + else: + if self.sh_degree > 0: + assert self._features_rest is not None + for i in range( + self._features_rest.shape[1] * self._features_rest.shape[2] + ): + dtype_list.append((f'f_rest_{i}', 'f4')) + + dtype_list.append(('opacity', 'f4')) + for i in range(self._scaling.shape[1]): + dtype_list.append((f'scale_{i}', 'f4')) + for i in range(self._rotation.shape[1]): + dtype_list.append((f'rot_{i}', 'f4')) + else: + dtype_list = [ + ('x', 'f2'), + ('y', 'f2'), + ('z', 'f2'), + ('red', 'u1'), + ('green', 'u1'), + ('blue', 'u1'), + ] + for i in range(self._features_dc.shape[1] * self._features_dc.shape[2]): + dtype_list.append((f'f_dc_{i}', 'f2')) + + if self.sh_degree > 0: + assert self._features_rest is not None + for i in range( + self._features_rest.shape[1] * self._features_rest.shape[2] + ): + dtype_list.append((f'f_rest_{i}', 'f2')) + dtype_list.append(('opacity', 'f2')) + for i in range(self._scaling.shape[1]): + dtype_list.append((f'scale_{i}', 'f2')) + for i in range(self._rotation.shape[1]): + dtype_list.append((f'rot_{i}', 'f2')) + return dtype_list + + def save_ply( + self, + path: str | Path, + use_fp16: bool = False, + enable_gs_viewer: bool = True, + color_code: bool = False, + filter_mask: np.ndarray | None = None, + ) -> None: + """Save Gaussians to a PLY file. + + Parameters + ---------- + path: output file path. + use_fp16: use float16 storage. + enable_gs_viewer: pad SH to degree 3 for viewer compatibility. + color_code: replace RGB with viridis colormap based on index. + filter_mask: optional boolean mask to select a subset. + + """ + Path(path).parent.mkdir(parents=True, exist_ok=True) + + xyz = self._xyz.detach().cpu().numpy() + f_dc = ( + self._features_dc.detach() + .transpose(1, 2) + .flatten(start_dim=1) + .contiguous() + .cpu() + .numpy() + ) + if not color_code: + rgb = (SH2RGB(f_dc) * 255.0).clip(0.0, 255.0).astype(np.uint8) + else: + index = np.linspace(0, 1, xyz.shape[0]) + rgb = matplotlib.colormaps['viridis'](index)[..., :3] + rgb = (rgb * 255.0).clip(0.0, 255.0).astype(np.uint8) + + opacities = self._opacity.detach().cpu().numpy() + if self.scaling_modifier is not None: + scale = self.inv_scaling_activation(self.get_scaling).detach().cpu().numpy() + else: + scale = self._scaling.detach().cpu().numpy() + rotation = self._rotation.detach().cpu().numpy() + + dtype_full = self.construct_dtypes(use_fp16, enable_gs_viewer) + elements = np.empty(xyz.shape[0], dtype=dtype_full) + + f_rest = None + if self.sh_degree > 0: + assert self._features_rest is not None + f_rest = ( + self._features_rest.detach() + .transpose(1, 2) + .flatten(start_dim=1) + .contiguous() + .cpu() + .numpy() + ) + + if enable_gs_viewer: + sh_degree = 3 + if f_rest is None: + f_rest = np.zeros( + (xyz.shape[0], 3 * ((sh_degree + 1) ** 2 - 1)), dtype=np.float32, + ) + elif f_rest.shape[1] < 3 * ((sh_degree + 1) ** 2 - 1): + f_rest_pad = np.zeros( + (xyz.shape[0], 3 * ((sh_degree + 1) ** 2 - 1)), dtype=np.float32, + ) + f_rest_pad[:, : f_rest.shape[1]] = f_rest + f_rest = f_rest_pad + + if f_rest is not None: + attributes = np.concatenate( + (xyz, rgb, f_dc, f_rest, opacities, scale, rotation), axis=1, + ) + else: + attributes = np.concatenate( + (xyz, rgb, f_dc, opacities, scale, rotation), axis=1, + ) + + if filter_mask is not None: + attributes = attributes[filter_mask] + elements = elements[filter_mask] + + elements[:] = list(map(tuple, attributes)) + el = PlyElement.describe(elements, 'vertex') + PlyData([el]).write(path) + + def load_ply(self, path: str | Path) -> None: + """Load Gaussians from a PLY file. + + Parameters + ---------- + path: path to the PLY file. + + """ + plydata = PlyData.read(path) + + xyz = np.stack( + ( + np.asarray(plydata.elements[0]['x']), + np.asarray(plydata.elements[0]['y']), + np.asarray(plydata.elements[0]['z']), + ), + axis=1, + ) + opacities = np.asarray(plydata.elements[0]['opacity'])[..., np.newaxis] + + features_dc = np.zeros((xyz.shape[0], 3, 1)) + features_dc[:, 0, 0] = np.asarray(plydata.elements[0]['f_dc_0']) + features_dc[:, 1, 0] = np.asarray(plydata.elements[0]['f_dc_1']) + features_dc[:, 2, 0] = np.asarray(plydata.elements[0]['f_dc_2']) + + if self.sh_degree > 0: + extra_f_names = [ + p.name + for p in plydata.elements[0].properties + if p.name.startswith('f_rest_') + ] + extra_f_names = sorted(extra_f_names, key=lambda x: int(x.split('_')[-1])) + assert len(extra_f_names) == 3 * (self.sh_degree + 1) ** 2 - 3 + features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) + for idx, attr_name in enumerate(extra_f_names): + features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) + features_extra = features_extra.reshape( + (features_extra.shape[0], 3, (self.sh_degree + 1) ** 2 - 1) + ) + + scale_names = [ + p.name + for p in plydata.elements[0].properties + if p.name.startswith('scale_') + ] + scale_names = sorted(scale_names, key=lambda x: int(x.split('_')[-1])) + scales = np.zeros((xyz.shape[0], len(scale_names))) + for idx, attr_name in enumerate(scale_names): + scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + rot_names = [ + p.name for p in plydata.elements[0].properties if p.name.startswith('rot') + ] + rot_names = sorted(rot_names, key=lambda x: int(x.split('_')[-1])) + rots = np.zeros((xyz.shape[0], len(rot_names))) + for idx, attr_name in enumerate(rot_names): + rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) + + self._xyz = torch.from_numpy(xyz.astype(np.float32)) + self._features_dc = ( + torch.from_numpy(features_dc.astype(np.float32)) + .transpose(1, 2) + .contiguous() + ) + if self.sh_degree > 0: + self._features_rest = ( + torch.from_numpy(features_extra.astype(np.float32)) + .transpose(1, 2) + .contiguous() + ) + self._opacity = torch.from_numpy( + np.copy(opacities).astype(np.float32) + ).contiguous() + self._scaling = torch.from_numpy(scales.astype(np.float32)).contiguous() + self._rotation = torch.from_numpy(rots.astype(np.float32)).contiguous() + + +def render_opencv_cam_gsplat( + pc: GaussianModel, + height: int, + width: int, + C2W: torch.Tensor, + fxfycxcy: torch.Tensor, + sh_degree: int | None = None, + near_plane: float = 0.2, + bg_color: tuple[float, float, float] | torch.Tensor = (1.0, 1.0, 1.0), + render_depth: bool = False, +) -> dict[str, torch.Tensor | None]: + """Render a batch of views using gsplat rasterisation. + + Parameters + ---------- + pc: Gaussian point cloud model. + height: render height in pixels. + width: render width in pixels. + C2W: camera-to-world matrices of shape [V, 4, 4]. + fxfycxcy: intrinsics of shape [V, 4]. + sh_degree: SH degree to evaluate; defaults to pc.sh_degree. + near_plane: near clipping plane distance. + bg_color: background colour (1-D or 2-D tensor or tuple). + render_depth: whether to also render a depth map. + + Returns + ------- + dict with keys 'render' ([V, 3, H, W]) and 'depth' ([V, 1, H, W] or None). + + """ + means3D = pc.get_xyz + opacity = pc.get_opacity + scales = pc.get_scaling + rotations = pc.get_rotation + shs = pc.get_features + for name, t in [('xyz', means3D), ('scales', scales), ('rotation', rotations)]: + if torch.isnan(t).any() or torch.isinf(t).any(): + raise ValueError( + f'Gaussian {name} contains NaN or Inf. ' + 'Check model outputs and consider using float32 or clamping scales.' + ) + num_cams = C2W.size(0) + if torch.is_tensor(bg_color): + bg_color = bg_color.to(device=C2W.device, dtype=torch.float32) + else: + bg_color = torch.tensor(list(bg_color), dtype=torch.float32, device=C2W.device) + if bg_color.ndim == 1: + bg_color = bg_color.unsqueeze(0).expand(num_cams, -1) + elif bg_color.ndim == 2 and bg_color.size(0) == 1: + bg_color = bg_color.expand(num_cams, -1) + elif bg_color.ndim == 2 and bg_color.size(0) == num_cams: + pass + else: + raise ValueError( + f'Invalid bg_color shape {tuple(bg_color.shape)}; ' + f'expected [3], [1,3], or [{num_cams},3].' + ) + W2C = C2W.inverse() + + width = int(width) + height = int(height) + if width <= 0 or height <= 0: + raise ValueError( + f'Invalid render dimensions: width={width}, height={height}. ' + 'Check that input images have valid (H, W) from data.' + ) + + intr = torch.zeros(fxfycxcy.size(0), 3, 3, device=fxfycxcy.device, dtype=fxfycxcy.dtype) + intr[:, 0, 0] = fxfycxcy[:, 0] + intr[:, 1, 1] = fxfycxcy[:, 1] + intr[:, 0, 2] = fxfycxcy[:, 2] + intr[:, 1, 2] = fxfycxcy[:, 3] + intr[:, 2, 2] = 1.0 + + render_mode = 'RGB+ED' if render_depth else 'RGB' + render_colors, _, _ = rasterization( + means3D, rotations, scales, opacity.squeeze(), + shs, W2C, intr, width, height, + near_plane=near_plane, + sh_degree=sh_degree, + backgrounds=bg_color, + render_mode=render_mode, + packed=False, + ) + if render_mode != 'RGB': + render_colors, render_depth_map = render_colors[..., :3], render_colors[..., 3:] + else: + render_depth_map = None + return { + 'render': render_colors.permute(0, 3, 1, 2), + 'depth': ( + render_depth_map.permute(0, 3, 1, 2) if torch.is_tensor(render_depth_map) else None + ), + } + + +def render_opencv_cam( + pc: GaussianModel, + height: int, + width: int, + C2W: torch.Tensor, + fxfycxcy: torch.Tensor, + sh_degree: int | None = None, + near_plane: float = 0.2, + bg_color: tuple[float, float, float] | torch.Tensor = (1.0, 1.0, 1.0), + render_depth: bool = False, +) -> dict[str, torch.Tensor | None]: + """Compatibility wrapper for legacy single-camera render paths. + + Routes calls through the batched gsplat implementation. + + Parameters + ---------- + pc: Gaussian point cloud model. + height: render height in pixels. + width: render width in pixels. + C2W: camera-to-world matrix of shape [4, 4] or [1, 4, 4]. + fxfycxcy: intrinsics of shape [4] or [1, 4]. + sh_degree: SH degree to evaluate; defaults to pc.sh_degree. + near_plane: near clipping plane distance. + bg_color: background colour. + render_depth: whether to also render a depth map. + + Returns + ------- + dict with keys 'render', 'depth', 'alpha'. + + """ + squeeze_view = C2W.ndim == 2 + if squeeze_view: + C2W = C2W.unsqueeze(0) + if fxfycxcy.ndim == 1: + fxfycxcy = fxfycxcy.unsqueeze(0) + + if sh_degree is None: + sh_degree = pc.sh_degree + + buffers = render_opencv_cam_gsplat( + pc, + height, + width, + C2W, + fxfycxcy, + sh_degree=sh_degree, + near_plane=near_plane, + bg_color=bg_color, + render_depth=render_depth, + ) + if squeeze_view: + render = buffers['render'] + assert render is not None + buffers['render'] = render[0] + depth = buffers.get('depth') + if depth is not None: + buffers['depth'] = depth[0] + buffers.setdefault('alpha', None) + return buffers + + +class DeferredGaussianRender(torch.autograd.Function): + """Custom autograd function for deferred Gaussian rendering.""" + + @staticmethod + def forward( + ctx, + xyz, + features, + scaling, + rotation, + opacity, + height, + width, + C2W, + fxfycxcy, + scaling_modifier=None, + ): + """Forward pass: render all views and return image tensor. + + Parameters + ---------- + ctx: autograd context. + xyz: Gaussian positions of shape [b, n_gaussians, 3]. + features: SH features of shape [b, n_gaussians, (sh_degree+1)^2, 3]. + scaling: log-scale of shape [b, n_gaussians, 3]. + rotation: quaternions of shape [b, n_gaussians, 4]. + opacity: logit-opacity of shape [b, n_gaussians, 1]. + height: render height. + width: render width. + C2W: camera-to-world matrices of shape [b, v, 4, 4]. + fxfycxcy: intrinsics of shape [b, v, 4]. + scaling_modifier: optional global scaling factor. + + Returns + ------- + renders of shape [b, v, 3, height, width]. + + """ + ctx.scaling_modifier = scaling_modifier + + sh_degree = int(math.sqrt(features.shape[-2])) - 1 + + gaussians_model = GaussianModel(sh_degree, scaling_modifier) + + with torch.no_grad(): + b, v = C2W.size(0), C2W.size(1) + renders = [] + for i in range(b): + pc = gaussians_model.set_data( + xyz[i], features[i], scaling[i], rotation[i], opacity[i], + ) + for j in range(v): + renders.append( + render_opencv_cam(pc, height, width, C2W[i, j], fxfycxcy[i, j])[ + 'render' + ] + ) + renders = torch.stack(renders, dim=0) + renders = renders.reshape(b, v, 3, height, width) + + renders = renders.requires_grad_() + + ctx.save_for_backward(xyz, features, scaling, rotation, opacity, C2W, fxfycxcy) + ctx.rendering_size = (height, width) + ctx.sh_degree = sh_degree + + del gaussians_model + + return renders + + @staticmethod + def backward(ctx, grad_output): + """Backward pass: compute gradients w.r.t. Gaussian parameters. + + Parameters + ---------- + ctx: autograd context. + grad_output: upstream gradient tensor. + + Returns + ------- + gradients w.r.t. xyz, features, scaling, rotation, opacity, and None for + height, width, C2W, fxfycxcy, scaling_modifier. + + """ + xyz, features, scaling, rotation, opacity, C2W, fxfycxcy = ctx.saved_tensors + height, width = ctx.rendering_size + sh_degree = ctx.sh_degree + + input_dict = OrderedDict( + [ + ('xyz', xyz), + ('features', features), + ('scaling', scaling), + ('rotation', rotation), + ('opacity', opacity), + ] + ) + input_dict = {k: v.detach().requires_grad_() for k, v in input_dict.items()} + + gaussians_model = GaussianModel(sh_degree, ctx.scaling_modifier) + + with torch.enable_grad(): + b, v = C2W.size(0), C2W.size(1) + for i in range(b): + for j in range(v): + pc = gaussians_model.set_data( + **{k: v[i] for k, v in input_dict.items()} + ) + + render = render_opencv_cam( + pc, height, width, C2W[i, j], fxfycxcy[i, j], + )['render'] + + assert render is not None + render.backward(grad_output[i, j]) + + del gaussians_model + + return *[var.grad for var in input_dict.values()], None, None, None, None, None + + +deferred_gaussian_render = DeferredGaussianRender.apply diff --git a/beast/train.py b/beast/train.py index 10e4a90..cad83b0 100644 --- a/beast/train.py +++ b/beast/train.py @@ -16,7 +16,7 @@ import beast from beast.data.augmentations import imgaug_pipeline -from beast.data.datamodules import BaseDataModule +from beast.data.datamodules import BaseDataModule, MultiViewDataModule from beast.data.datasets import BaseDataset from beast.logging import log_step from beast.models.base import BaseLightningModel @@ -96,59 +96,87 @@ def train(config: dict, model: BaseLightningModel, output_dir: str | Path) -> Ba # Set up data objects # ---------------------------------------------------------------------------------- - # imgaug transform - if rank_zero_only.rank == 0: - log_step("Setting up imgaug pipeline", level='debug') - pipe_params = config.get('training', {}).get('imgaug', 'none') - if isinstance(pipe_params, str): - from beast.data.augmentations import expand_imgaug_str_to_dict - pipe_params = expand_imgaug_str_to_dict(pipe_params) # type: ignore[arg-type] - imgaug_pipeline_ = imgaug_pipeline(pipe_params) - if rank_zero_only.rank == 0: - log_step("Imgaug pipeline created", level='debug') - - # dataset - dataset = BaseDataset( - data_dir=config['data']['data_dir'], - imgaug_pipeline=imgaug_pipeline_, - num_channels=config['model']['model_params'].get('num_channels', 3), - ) - - # datamodule; breaks up dataset into train/val/test - if rank_zero_only.rank == 0: - log_step("Creating BaseDataModule", level='debug') - datamodule = BaseDataModule( - dataset=dataset, - train_batch_size=config['training']['train_batch_size'], - val_batch_size=config['training']['val_batch_size'], - test_batch_size=config['training']['test_batch_size'], - use_sampler=config['model']['model_params'].get('use_infoNCE', False), - num_workers=config['training']['num_workers'], - train_probability=config['training'].get('train_probability', 0.95), - val_probability=config['training'].get('val_probability', 0.05), - seed=config['training']['seed'], - ) - if rank_zero_only.rank == 0: - log_step("BaseDataModule created", level='debug') + model_class = config['model'].get('model_class', '').lower() + + if model_class == 'erayzer': + if rank_zero_only.rank == 0: + log_step('Creating MultiViewDataModule', level='debug') + datamodule = MultiViewDataModule( + data_dir=config['data']['data_dir'], + image_size=config['model']['image_tokenizer']['image_size'], + train_batch_size=config['training']['train_batch_size'], + val_batch_size=config['training']['val_batch_size'], + train_fraction=config['training'].get('train_fraction', 0.9), + num_workers=config['training']['num_workers'], + seed=config['training'].get('seed', 0), + ) + datamodule.setup() + if rank_zero_only.rank == 0: + log_step('MultiViewDataModule created', level='debug') + # ERayZer schedules by steps; max_fwdbwd_passes drives OneCycleLR directly + trainer_epoch_kwargs: dict = { + 'max_steps': config['training']['max_fwdbwd_passes'], + } + use_distributed_sampler = True + else: + # imgaug transform + if rank_zero_only.rank == 0: + log_step('Setting up imgaug pipeline', level='debug') + pipe_params = config.get('training', {}).get('imgaug', 'none') + if isinstance(pipe_params, str): + from beast.data.augmentations import expand_imgaug_str_to_dict + pipe_params = expand_imgaug_str_to_dict(pipe_params) # type: ignore[arg-type] + imgaug_pipeline_ = imgaug_pipeline(pipe_params) + if rank_zero_only.rank == 0: + log_step('Imgaug pipeline created', level='debug') + + # dataset + dataset = BaseDataset( + data_dir=config['data']['data_dir'], + imgaug_pipeline=imgaug_pipeline_, + num_channels=config['model']['model_params'].get('num_channels', 3), + ) - # update number of training steps (for learning rate scheduler with step information) - if rank_zero_only.rank == 0: - log_step("Calculating training steps", level='debug') - num_epochs = config['training']['num_epochs'] - train_size = int(np.floor(config['training'].get('train_probability', 0.95) * len(dataset))) - steps_per_epoch = int(np.ceil( - train_size - / config['training']['train_batch_size'] - / config['training']['num_gpus'] - / config['training']['num_nodes'] - )) - model.config['optimizer']['steps_per_epoch'] = steps_per_epoch - model.config['optimizer']['total_steps'] = steps_per_epoch * num_epochs - if rank_zero_only.rank == 0: - log_step( - f"Training steps calculated: {steps_per_epoch} steps/epoch, {num_epochs} epochs", - level='debug', + # datamodule; breaks up dataset into train/val/test + if rank_zero_only.rank == 0: + log_step('Creating BaseDataModule', level='debug') + datamodule = BaseDataModule( + dataset=dataset, + train_batch_size=config['training']['train_batch_size'], + val_batch_size=config['training']['val_batch_size'], + test_batch_size=config['training']['test_batch_size'], + use_sampler=config['model']['model_params'].get('use_infoNCE', False), + num_workers=config['training']['num_workers'], + train_probability=config['training'].get('train_probability', 0.95), + val_probability=config['training'].get('val_probability', 0.05), + seed=config['training']['seed'], ) + if rank_zero_only.rank == 0: + log_step('BaseDataModule created', level='debug') + + # update number of training steps (for learning rate scheduler) + if rank_zero_only.rank == 0: + log_step('Calculating training steps', level='debug') + num_epochs = config['training']['num_epochs'] + train_size = int( + np.floor(config['training'].get('train_probability', 0.95) * len(dataset)) + ) + steps_per_epoch = int(np.ceil( + train_size + / config['training']['train_batch_size'] + / config['training']['num_gpus'] + / config['training']['num_nodes'] + )) + model.config['optimizer']['steps_per_epoch'] = steps_per_epoch + model.config['optimizer']['total_steps'] = steps_per_epoch * num_epochs + if rank_zero_only.rank == 0: + log_step( + f'Training steps calculated: {steps_per_epoch} steps/epoch, {num_epochs} epochs', + level='debug', + ) + + trainer_epoch_kwargs = {'max_epochs': num_epochs, 'min_epochs': num_epochs} + use_distributed_sampler = not config['model']['model_params'].get('use_infoNCE', False) # ---------------------------------------------------------------------------------- # Save configuration in output directory @@ -157,13 +185,13 @@ def train(config: dict, model: BaseLightningModel, output_dir: str | Path) -> Ba # save config file if rank_zero_only.rank == 0: - log_step(f"Saving config to {output_dir}", level='debug') + log_step(f'Saving config to {output_dir}', level='debug') output_dir.mkdir(parents=True, exist_ok=True) dest_config_file = Path(output_dir) / 'config.yaml' with open(dest_config_file, 'w') as file: yaml.dump(config, file) if rank_zero_only.rank == 0: - log_step("Config saved", level='debug') + log_step('Config saved', level='debug') # ---------------------------------------------------------------------------------- # Set up and run training @@ -171,43 +199,32 @@ def train(config: dict, model: BaseLightningModel, output_dir: str | Path) -> Ba # logger if rank_zero_only.rank == 0: - log_step("Creating TensorBoardLogger", level='debug') + log_step('Creating TensorBoardLogger', level='debug') logger = pl_loggers.TensorBoardLogger('tb_logs', name='') if rank_zero_only.rank == 0: - log_step("TensorBoardLogger created", level='debug') + log_step('TensorBoardLogger created', level='debug') # early stopping, learning rate monitoring, model checkpointing, backbone unfreezing if rank_zero_only.rank == 0: - log_step("Setting up callbacks", level='debug') + log_step('Setting up callbacks', level='debug') callbacks = get_callbacks( lr_monitor=True, ckpt_every_n_epochs=config['training'].get('ckpt_every_n_epochs', None), ) if rank_zero_only.rank == 0: - log_step(f"Callbacks created: {len(callbacks)} callbacks", level='debug') - - # initialize to Trainer defaults. Note max_steps defaults to -1. - min_epochs = config['training']['num_epochs'] - max_epochs = min_epochs - - # our custom sampler does not play nice with DDP - if config['model']['model_params'].get('use_infoNCE', False): - use_distributed_sampler = False - else: - use_distributed_sampler = True + log_step(f'Callbacks created: {len(callbacks)} callbacks', level='debug') if rank_zero_only.rank == 0: - log_step("Creating PyTorch Lightning Trainer", level='debug') - log_step(" - accelerator: gpu", level='debug') + log_step('Creating PyTorch Lightning Trainer', level='debug') + log_step(' - accelerator: gpu', level='debug') log_step(f" - devices: {config['training']['num_gpus']}", level='debug') log_step(f" - num_nodes: {config['training']['num_nodes']}", level='debug') - log_step(f" - max_epochs: {max_epochs}", level='debug') + log_step(f' - trainer_epoch_kwargs: {trainer_epoch_kwargs}', level='debug') trainer = pl.Trainer( accelerator='gpu', devices=config['training']['num_gpus'], num_nodes=config['training']['num_nodes'], - max_epochs=max_epochs, - min_epochs=min_epochs, + precision=config['training'].get('precision', '32-true'), check_val_every_n_epoch=config['training'].get('check_val_every_n_epoch', 1), log_every_n_steps=config['training'].get('log_every_n_steps', 10), callbacks=callbacks, @@ -215,6 +232,7 @@ def train(config: dict, model: BaseLightningModel, output_dir: str | Path) -> Ba accumulate_grad_batches=config['optimizer'].get('accumulate_grad_batches', 1), sync_batchnorm=True, use_distributed_sampler=use_distributed_sampler, + **trainer_epoch_kwargs, ) if rank_zero_only.rank == 0: log_step("Trainer created", level='debug') diff --git a/configs/multiview/erayzer.yaml b/configs/multiview/erayzer.yaml new file mode 100644 index 0000000..cc2d37b --- /dev/null +++ b/configs/multiview/erayzer.yaml @@ -0,0 +1,118 @@ +# ERayZer: feed-forward multi-view 3D Gaussian splatting model. +# Encodes posed (or to-be-posed) multi-view images into a 3D Gaussian field +# and renders novel views via Gaussian splatting. +# +# Camera source: +# base ERayZer — learns to predict cameras from image context +# BEAST3D — reads GT cameras from the dataset (override _resolve_cameras) +# +# Run with: beast train --config configs/multiview/erayzer.yaml + +# --------------------------------------------------------------------------- +# Model architecture +# --------------------------------------------------------------------------- +model: + model_class: erayzer + seed: 0 + checkpoint: null # path to .ckpt for fine-tuning or inference + + image_tokenizer: + image_size: 256 # square input resolution (H = W) + patch_size: 16 # spatial tokens per view: (image_size / patch_size)^2 + in_channels: 3 # RGB + + target_image: + height: 256 # novel-view render resolution + width: 256 + + transformer: + d: 768 # token embedding dimension + d_head: 64 # per-head dimension (d must be divisible by d_head) + encoder_n_layer: 12 # camera-prediction encoder depth (must be even; 0 to disable) + encoder_geom_n_layer: 12 # geometry encoder depth (must be even) + use_qk_norm: true # QK normalization inside attention + special_init: true # apply layerwise weight scaling at init + depth_init: true # scale init std by layer depth rather than total depth + fix_decoder: false # freeze Gaussian decoder (encoder-only fine-tuning) + + pose_latent: + canonical: first # canonical anchor view: first | middle | unordered + mode: pairwise # pairwise (anchor+relative) | global (all-views direct) + representation: 6d # rotation parameterization: 6d (continuous) | quat + + gaussians: + sh_degree: 3 # spherical-harmonics degree (0 = RGB only, 3 = full SH) + range_setting: + type: object_centric_depth # depth mapping: object_centric_depth | linear_depth | + # log_depth | disparity + + hard_pixelalign: false # project each Gaussian onto its source pixel ray + input_with_pe: true # add 2-D sinusoidal positional encoding to image tokens + mask_ratio: 0.0 # random token dropout during training (0 = disabled) + + # Gaussian activation biases (applied before exp/sigmoid in GaussiansUpsampler) + scaling_bias: -2.3 + scaling_max: -1.2 + opacity_bias: -2.0 + + near_plane: 0.2 # near clipping plane passed to gsplat rasterizer + +# --------------------------------------------------------------------------- +# Training +# --------------------------------------------------------------------------- +training: + seed: 0 + num_views: 3 # cameras per sample in the dataset + num_input_views: 2 # views fed to the geometry encoder + num_target_views: 1 # novel views to render and supervise + random_num_input_views: false # train with a random number of input views in [min,max] + min_input_views: 2 # min reference views when random_num_input_views + max_input_views: 5 # max reference views when random_num_input_views + freeze_focal_steps: 200 # detach predicted intrinsics for this many steps, then learn + + train_batch_size: 8 # per GPU + val_batch_size: 4 + num_epochs: 200 + num_workers: 8 + num_gpus: 1 + num_nodes: 1 + + max_fwdbwd_passes: 200000 # total gradient steps (controls OneCycleLR span) + grad_checkpoint_every: 1 # gradient-checkpoint granularity (1 = per layer) + + log_every_n_steps: 10 + check_val_every_n_epoch: 1 + + train_fraction: 0.9 # fraction of samples used for training (rest = validation) + random_split: false # randomize input/target partition order per batch + random_inputs: false # randomize which views are used as input + render_interpolate: false # render an interpolated video trajectory at eval time + + # Loss weights (0 = disabled) + l2_loss_weight: 1.0 + gs_reg_loss_weight: 0.0 + perceptual_loss_weight: 0.0 + pose_consistency_reg_weight: 0.0 + +# --------------------------------------------------------------------------- +# Optimizer (AdamW + OneCycleLR) +# --------------------------------------------------------------------------- +optimizer: + lr: 4.0e-4 + beta1: 0.9 + beta2: 0.95 + wd: 0.05 + warmup: 3000 # warmup steps for OneCycleLR pct_start + div_factor: 1.0 # initial_lr = lr / div_factor + final_div_factor: 1.0 # min_lr = initial_lr / final_div_factor + +# --------------------------------------------------------------------------- +# Data +# --------------------------------------------------------------------------- +data: + data_dir: /PATH/TO/MULTIVIEW/DATA # root of a beast extract_3d output directory + +# --------------------------------------------------------------------------- +# Runtime flags +# --------------------------------------------------------------------------- +inference: false # set to true at inference time; disables view split and enables video render diff --git a/docs/beast3d.md b/docs/beast3d.md index 4295584..9d21355 100644 --- a/docs/beast3d.md +++ b/docs/beast3d.md @@ -5,3 +5,96 @@ Documentation for BEAST3D multi-view self-supervised pretraining. ## Guides - [Extracting data for BEAST3D](extract_data_for_3d.md) + +--- + +## Environment setup + +BEAST3D depends on a [custom fork of gsplat](https://github.com/QitaoZhao/gsplat) for +differentiable Gaussian splatting with intrinsics gradient support. This fork is compiled +from source during installation because it is not available as a pre-built wheel. + +### Prerequisites + +gsplat's `setup.py` imports `torch` at build time to find CUDA paths, so `torch` must be +installed before gsplat. The build backend (`poetry-core`) must also be present. The +standard beast install handles this: + +```bash +pip install lightning poetry-core +pip install -e . --no-build-isolation +``` + +`--no-build-isolation` tells pip to use the existing environment (where torch is already +installed) rather than creating a temporary venv without it. + +### GPU architecture + +gsplat's CUDA kernels are compiled for the GPU architectures visible during the build. +The compiled kernels only run on matching hardware. If you need to target a specific +architecture (e.g. building on a machine without a GPU, or targeting a different GPU than +the one present), set `TORCH_CUDA_ARCH_LIST`: + +```bash +TORCH_CUDA_ARCH_LIST="8.6" pip install -e . --no-build-isolation +``` + +Common values: `7.5` (RTX 2080), `8.0` (A100), `8.6` (A40, RTX 3090), `8.9` (L40, RTX +4090), `12.0` (RTX 5090). + +### Verifying the installation + +```python +from gsplat.cuda._backend import _C +print(_C) # should be a compiled module, not None +``` + +If the output is `None`, gsplat's CUDA extension did not load. Common causes: + +- **torch was not available during the build** — install lightning first, then reinstall + with `--no-build-isolation`. +- **CUDA toolkit not found** — ensure `CUDA_HOME` is set, or that `nvcc` is on your PATH. +- **Incompatible gcc** — nvcc requires gcc <= 12 for CUDA 12.x. Check with + `nvcc --version` and `gcc --version`. +- **Architecture mismatch** — the kernels were compiled for a different GPU. Rebuild with + the correct `TORCH_CUDA_ARCH_LIST`. + +### Source build on Blackwell GPUs (SM 12.0) or PyTorch >= 2.7 + +Building on newer hardware may require additional CUDA development headers not included in +a standard conda/pip PyTorch install. These must be installed right after creating the conda +env, before running `pip install`: + +```bash +conda create --yes --name beast python=3.10 +conda activate beast + +# install CUDA dev packages before anything else +conda install -c nvidia \ + cuda-nvcc= \ + cuda-cudart-dev= \ + libcusparse-dev \ + libcublas-dev \ + libcurand-dev \ + libcufft-dev \ + libcusolver-dev +``` + +Match `` to the CUDA version you plan to use with PyTorch (e.g. `13.0`). + +If conda installs CUDA headers under `targets/x86_64-linux/include/` rather than the +standard `include/`, set these environment variables before installing beast: + +```bash +conda env config vars set -n beast \ + CPATH=$CONDA_PREFIX/targets/x86_64-linux/include \ + CUDA_HOME=$CONDA_PREFIX +conda activate beast # re-activate so the variables take effect +``` + +Then proceed with the standard install: + +```bash +pip install lightning poetry-core +pip install -e . --no-build-isolation +``` diff --git a/docs/developer_guide.md b/docs/developer_guide.md new file mode 100644 index 0000000..247348c --- /dev/null +++ b/docs/developer_guide.md @@ -0,0 +1,386 @@ +# BEAST Developer Guide + +This guide is aimed at contributors (human or AI agent) adding a new model to the BEAST +codebase. It covers the overall architecture, the building blocks available for reuse, and a +concrete step-by-step checklist for registering a new model end-to-end. + +--- + +## Package layout + +``` +beast/ + api/ + model.py # High-level Model wrapper (from_config, from_dir, train, predict_*) + cli/ + commands/ # One file per CLI subcommand (train, predict, extract, extract_3d) + main.py # Entry point; dispatches to COMMANDS registry + config.py # Shared Pydantic schemas (TrainingConfig, OptimizerConfig, DataConfig, + # BeastConfig) + per-model dispatch via get_beast_config_class() + data/ + datasets.py # BaseDataset, MultiViewDataset + datamodules.py # BaseDataModule, MultiViewDataModule + samplers.py # ContrastBatchSampler (contrastive learning) + augmentations.py # imgaug pipeline builders + types.py # ExampleDict, MultiViewExampleDict (TypedDicts) + video.py # Video frame extraction utilities + geometry/ + camera.py # c2w/w2c conversions, intrinsic normalisation, SE(3) helpers + rotations.py # 6D rotation ↔ matrix, quaternion utilities + positional_encoding.py # 2-D sinusoidal PE used by ERayZer + inference.py # predict_images(), predict_video() — used by Model.predict_* + io.py # load_config(), apply_config_overrides(), calibration loaders + logging.py # log_step() helper + models/ + base.py # BaseLightningModel (abstract; all models inherit from this) + registry.py # MODEL_REGISTRY, TRAIN_REGISTRY, CONFIG_REGISTRY + _register_all() + beast_resnet/ # Per-model package — see "Model package layout" below + beast_vit/ + erayzer/ + preprocess/ # extract pipeline (extraction.py, extraction_3d.py) + SAM segmenter + rendering/ + gaussians_renderer.py # gsplat-based Gaussian splatting renderer + transformer.py # shared transformer building blocks (used by ERayZer) + losses.py # rendering losses (L2, perceptual) + dino.py # DINOv2 perceptual feature extractor + train.py # Shared training loop (data setup, Lightning Trainer, callbacks) +configs/ + resnet_ae.yaml + vit.yaml + multiview/ + erayzer.yaml # Reference config for ERayZer / BEAST3D +tests/ # Mirror of beast/ package tree; see Testing section +``` + +--- + +## Data flow + +``` +CLI / user script + │ + ▼ +beast.io.load_config(path) + │ reads YAML → peeks at model.model_class + │ → get_beast_config_class(model_class) → appropriate Pydantic config class + │ → model_validate(raw) → model_dump() → plain dict + ▼ +beast.api.model.Model.from_config(config) + │ validates dict if not already validated + │ → MODEL_REGISTRY[model_class](config) → LightningModule instance + ▼ +Model.train(output_dir) + │ → TRAIN_REGISTRY[model_class](config, model, output_dir) + │ (currently all three models delegate to beast.train.train) + ▼ +beast.train.train(config, model, output_dir) + │ branches on model_class to select DataModule: + │ 'erayzer' / beast3d → MultiViewDataModule + │ 'resnet' / 'vit' → BaseDataModule (with BaseDataset + imgaug) + │ saves config.yaml to output_dir + │ → pl.Trainer.fit(model, datamodule) + ▼ +BaseLightningModel subclass + training_step / validation_step + → get_model_outputs() → compute_loss() (abstract; implemented per-model) +``` + +--- + +## Config architecture + +Each model owns its complete Pydantic schemas in its `_config.py` file: + +| Schema | Location | Purpose | +|---|---|---| +| `ResnetModelConfig` | `beast_resnet/beast_resnet_config.py` | model section | +| `VitModelConfig` | `beast_vit/beast_vit_config.py` | model section | +| `ERayZerModelConfig` | `erayzer/erayzer_config.py` | model section | +| `ERayZerTrainingConfig` | `erayzer/erayzer_config.py` | training section | +| `ERayZerOptimizerConfig` | `erayzer/erayzer_config.py` | optimizer section | + +Models that share the standard training loop (resnet, vit) reuse `TrainingConfig` and +`OptimizerConfig` from `beast/config.py`. Models with divergent training behaviour (ERayZer, +BEAST3D) define their own training/optimizer schemas. + +Top-level full-config classes live in `beast/config.py` where they assemble model + training + +optimizer + data schemas: + +- `BeastConfig` — used for `resnet` and `vit` +- `ERayZerBeastConfig` — used for `erayzer` (and subclasses like BEAST3D if they share the + same training schema) + +`get_beast_config_class(model_class)` in `config.py` is the single dispatch point. It returns +`ERayZerBeastConfig` for `'erayzer'` and `BeastConfig` for everything else. Both `load_config` +and `Model.from_config` use it. + +**Adding a model with a distinct training schema** requires: +1. Define `TrainingConfig` and `OptimizerConfig` in `_config.py` +2. Define `BeastConfig` in `beast/config.py` +3. Add `'': BeastConfig` to `_MODEL_CONFIG_CLASSES` in `config.py` + +--- + +## Model package layout + +Every model lives in `beast/models//` and contains exactly: + +``` +beast/models// + __init__.py # re-exports model class, all config classes, train + _config.py # Pydantic schemas for model (+ training/optimizer if divergent) + _model.py # LightningModule subclass + _train.py # training entry point; either delegates to beast.train.train or + # provides a custom loop +``` + +The model's `model_class` string identifier (e.g. `'resnet'`, `'vit'`, `'erayzer'`) is a +stable logical name set in `Literal['']` inside the config class. It does **not** have +to match the directory name (just as HuggingFace's `model_type` strings don't match directory +names). + +--- + +## Registry + +`beast/models/registry.py` holds three dicts populated by `_register_all()`: + +```python +MODEL_REGISTRY: dict[str, type] # model_class → LightningModule subclass +TRAIN_REGISTRY: dict[str, Callable] # model_class → train function +CONFIG_REGISTRY: dict[str, type[BaseModel]] # model_class → model-section config class +``` + +`_register_all()` runs at import time (bottom of `registry.py`). To add a new model, add +three lines to `_register_all()` — see step 6 of the checklist below. + +--- + +## Reusable building blocks + +### Data + +| Class | When to use | +|---|---| +| `BaseDataset` | Single-view image folder; handles JPEG/PNG, imgaug transforms, grayscale→RGB | +| `MultiViewDataset` | Multi-view scenes from `beast extract_3d` output; loads images + optional cameras (c2w, fxfycxcy) + optional masks | +| `BaseDataModule` | Wraps `BaseDataset`; random train/val/test split by probability | +| `MultiViewDataModule` | Wraps `MultiViewDataset`; random split by `train_fraction`; used by ERayZer and BEAST3D | + +`MultiViewExampleDict` (from `beast.data.types`) is the typed batch dict returned by +`MultiViewDataset.__getitem__`. Keys: `image`, `view_names`, `video_id`, `frame_id`, and +optionally `c2w`, `fxfycxcy`, `input_mask`. + +Camera tensors use the following conventions throughout: +- **c2w**: camera-to-world transform, shape `(views, 4, 4)` +- **fxfycxcy**: intrinsics in absolute pixels `(fx, fy, cx, cy)`, shape `(views, 4)` +- `normalized_intrinsics=False` must be set on the model when passing absolute-pixel intrinsics + +### Model base class + +`BaseLightningModel` (`beast/models/base.py`) handles: +- `configure_optimizers` — AdamW or Adam + `MultiStepLR` or `OneCycleLR` +- `training_step`, `validation_step`, `test_step` — delegate to `evaluate_batch` +- `evaluate_batch` — calls `get_model_outputs` then `compute_loss`; handles logging + +Subclasses **must** implement three abstract methods: + +```python +def get_model_outputs(self, batch_dict: dict) -> dict: ... +def compute_loss(self, stage: str | None, **kwargs) -> tuple[torch.Tensor, list[dict]]: ... +def predict_step(self, batch_dict: dict, batch_idx: int) -> dict: ... +``` + +`BaseLightningModel` does **not** implement an optimizer for ERayZer-style models (ERayZer +overrides `configure_optimizers` entirely). If your model uses a non-standard optimizer, +override `configure_optimizers` in the subclass. + +### Rendering and geometry + +| Module | Contents | +|---|---| +| `beast.rendering.gaussians_renderer` | Gaussian splatting renderer (wraps gsplat) | +| `beast.rendering.transformer` | Transformer building blocks used by ERayZer | +| `beast.rendering.losses` | L2, perceptual (DINO), GS regularisation losses | +| `beast.geometry.camera` | `intrinsics_to_fxfycxcy`, `normalize_camera_sequence`, `scale_intrinsics`, `w2c_to_c2w`, SE(3) helpers | +| `beast.geometry.rotations` | `rotation_6d_to_matrix`, `matrix_to_rotation_6d`, quaternion ↔ matrix | +| `beast.geometry.positional_encoding` | 2-D sinusoidal PE | + +--- + +## Step-by-step: adding a new model + +This checklist uses `beast3d` as an example. Adjust names accordingly. + +### 1. Create the model package directory + +``` +beast/models/beast3d/ + __init__.py + beast3d_config.py + beast3d_model.py + beast3d_train.py +``` + +### 2. Write `beast3d_config.py` + +Define all Pydantic schemas the model needs. If the model shares the standard training loop +(BaseDataModule + epoch-based training), import and reuse `TrainingConfig` and `OptimizerConfig` +from `beast.config`. If it needs different training fields, define `Beast3DTrainingConfig` and +`Beast3DOptimizerConfig` here. + +```python +from typing import Literal +from pydantic import BaseModel + +class Beast3DModelConfig(BaseModel): + model_class: Literal['beast3d'] + seed: int = 0 + checkpoint: str | None = None + # ... model-specific fields ... +``` + +Follow the ERayZer configs as a reference for a fully typed example. + +### 3. Write `beast3d_model.py` + +Subclass `BaseLightningModel` (or another model if the architecture is closely related): + +```python +from beast.models.base import BaseLightningModel + +class Beast3D(BaseLightningModel): + def __init__(self, config: dict) -> None: + super().__init__(config) + # build architecture from config['model'] + + def get_model_outputs(self, batch_dict: dict) -> dict: ... + def compute_loss(self, stage, **kwargs) -> tuple[Tensor, list[dict]]: ... + def predict_step(self, batch_dict, batch_idx) -> dict: ... +``` + +If BEAST3D inherits from ERayZer, override only the methods that differ (e.g. +`_resolve_cameras` to return GT cameras from `data['c2w']` and `data['fxfycxcy']` instead of +predicting them). + +### 4. Write `beast3d_train.py` + +If the model uses the shared training loop unchanged: + +```python +from beast.train import train + +__all__ = ['train'] +``` + +If it needs a custom loop, implement `train(config, model, output_dir)` here and make sure it +has the same signature. + +### 5. Write `__init__.py` + +Re-export the public API of the package: + +```python +from beast.models.beast3d.beast3d_config import Beast3DModelConfig +from beast.models.beast3d.beast3d_model import Beast3D +from beast.models.beast3d.beast3d_train import train + +__all__ = ['Beast3D', 'Beast3DModelConfig', 'train'] +``` + +### 6. Register the model in `beast/models/registry.py` + +Add three imports and three dict assignments inside `_register_all()`: + +```python +from beast.models.beast3d.beast3d_config import Beast3DModelConfig +from beast.models.beast3d.beast3d_model import Beast3D +from beast.models.beast3d.beast3d_train import train as beast3d_train + +MODEL_REGISTRY['beast3d'] = Beast3D +TRAIN_REGISTRY['beast3d'] = beast3d_train +CONFIG_REGISTRY['beast3d'] = Beast3DModelConfig +``` + +### 7. Update `beast/config.py` + +**If the model uses `BeastConfig`** (shared `TrainingConfig` + `OptimizerConfig`): add +`Beast3DModelConfig` to the `ModelConfig` union so `BeastConfig` can validate it: + +```python +ModelConfig = Annotated[ + ResnetModelConfig | VitModelConfig | Beast3DModelConfig, + Field(discriminator='model_class'), +] +``` + +**If the model needs its own training/optimizer schemas**: import them and assemble a new +full-config class, then add it to `_MODEL_CONFIG_CLASSES`: + +```python +from beast.models.beast3d.beast3d_config import Beast3DModelConfig # (already imported above) + +class Beast3DBeastConfig(BaseModel): + model: Beast3DModelConfig + training: Beast3DTrainingConfig + optimizer: Beast3DOptimizerConfig + data: DataConfig + inference: bool = False + +_MODEL_CONFIG_CLASSES['beast3d'] = Beast3DBeastConfig +``` + +### 8. Add a YAML config file + +Create `configs/multiview/beast3d.yaml` (or `configs/beast3d.yaml` for non-multiview models). +Keep it alongside its `model_class: beast3d` sibling configs. It will be automatically picked +up by `TestConfigFiles` in `tests/test_config.py`. + +### 9. Write tests + +Mirror the source tree under `tests/`: + +``` +tests/models/beast3d/ + __init__.py + test_beast3d_model.py +``` + +Fixtures live in `tests/conftest.py` (shared) or a `conftest.py` in the same directory. +Test assets (sample tensors, etc.) go in `tests/models/beast3d/assets/`. + +The test for the model config typically looks like: + +```python +class TestBeast3DModelConfig: + def test_valid_config(self) -> None: + Beast3DModelConfig.model_validate({'model_class': 'beast3d', ...}) + + def test_missing_required_field_raises(self) -> None: + with pytest.raises(ValidationError): + Beast3DModelConfig.model_validate({}) +``` + +--- + +## CLI + +The CLI is built from `beast/cli/commands/`. Each command module exposes two functions: + +- `register_parser(subparsers)` — adds the subcommand and its arguments +- `handle(args)` — executes the command + +`beast/cli/commands/__init__.py` maintains a `COMMANDS` dict that the main parser iterates. +To add a CLI subcommand for a new model operation, add a new module there. + +--- + +## Testing conventions + +- Test files mirror the source tree: `beast/a/b/c.py` → `tests/a/b/test_c.py` +- Test class names: `Test` +- Fixture shared across a directory: `conftest.py` in that directory +- External test data is downloaded by `tests/conftest.py` at import time via `fetch_test_data_if_needed` +- All YAML files in `configs/` are validated automatically by `TestConfigFiles` in `tests/test_config.py` + (files listed in `_NON_BEAST_CONFIG_NAMES` are excluded, e.g. `extraction_pipeline.yaml` + which uses `Beast3DConfig` rather than `BeastConfig`) diff --git a/pyproject.toml b/pyproject.toml index 2a927aa..ac0178b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,12 +29,16 @@ classifiers = [ dependencies = [ "aniposelib (>=0.8.0)", + "einops", + "gsplat @ git+https://github.com/QitaoZhao/gsplat.git@daad91d9e667cf49bab815d872aa65a5cda0a77e", "imaug", "jaxtyping", "lightning (~=2.5)", + "matplotlib", "numpy (>=2.0.0)", "opencv-python-headless", "pillow", + "plyfile", "pydantic", "scikit-learn", "tensorboard", @@ -42,6 +46,7 @@ dependencies = [ "torchvision", "tqdm", "transformers", + "typing_extensions", ] [project.urls] diff --git a/scripts/buildbot/README.md b/scripts/buildbot/README.md index 65adf31..ea0812f 100644 --- a/scripts/buildbot/README.md +++ b/scripts/buildbot/README.md @@ -1,64 +1,178 @@ -These scripts are used to run tests on axon. +# Buildbot CI Scripts + +These scripts run the test suite on the Axon SLURM cluster. The GitHub Actions +workflow (`.github/workflows/tests.yaml`) SSHs into Axon and calls +`build_srun.sh`, which submits `build.sh` as a SLURM job. + +``` +.github/workflows/tests.yaml + → ssh axon 'sh buildbot/build_srun.sh ' + → srun ... buildbot/build.sh +``` + +## How it works -.github/workflows/run-tests.yml calls +### `build_srun.sh` - ssh axon 'sh buildbot/build_srun.sh ' +Submits `build.sh` via `srun` with resource requests: -from a self-hosted runner. You can run this locally too. +- **GPU type: `a40`** — gsplat's CUDA kernels are compiled for a specific GPU + architecture (sm_86 for A40). The job must run on the same GPU type the kernels + were built for, otherwise you get `no kernel image is available for execution + on the device` errors at runtime. See "gsplat and GPU architecture" below. +- **Memory: 32GB** — the test suite loads large models (ViT ~450MB, DINOv2, etc.) + plus DataLoader workers. 8GB is not enough and causes OOM kills. +- **`-X`** — forwards SIGINT so GitHub Actions cancellation works. +- **`-u`** — unbuffered output for real-time log streaming. -I. Setup of self-hosted runner: +### `build.sh` - A self-hosted runner is required that has ssh access to axon. +1. **Environment setup** — loads modules (gcc, CUDA), activates the conda env, + and sets compiler variables. +2. **Fetches PR code** — shallow-clones just the PR merge ref. +3. **Installs beast** — `pip install -e ".[dev]" --no-build-isolation`. +4. **Runs pytest** — with coverage reporting for Codecov. - 1. You need a server. I.e. AWS, a physical PC, - but the axon cluster does not work for this as axon nodes do not support outbound internet. (to be verified!!!) +## gsplat and GPU architecture - 2. The server should support `ssh axon` alias. - I.e. ensure ~/.ssh/config contains: - Host axon - HostName axon-remote.rc.zi.columbia.edu - User - Port 55 +`beast` depends on a [custom gsplat fork](https://github.com/QitaoZhao/gsplat) +that compiles CUDA kernels from source. These kernels are compiled for a specific +GPU architecture (e.g. sm_86 for A40) and will only run on matching hardware. - 3. ssh must be setup for passwordless authentication via public/private keys. - This is achieved by generating keys using ssh-keygen, and copying the public key over - to your authorized key file in axon (ssh-copy-id). +Key constraints: - 4. Install and run Github Self-hosted Runner software +- **gsplat requires `torch` at build time** — its `setup.py` imports torch to + find CUDA paths. This is why `--no-build-isolation` is needed (so pip uses the + existing torch instead of creating a clean venv without it). +- **gsplat requires a CUDA-compatible gcc** — CUDA 12.4 requires gcc <= 12. + The `gcc/10.4` module provides gcc 12.3 (the module name doesn't match the + actual version). The `cuda/12.4.0` module sets `NVCC_CCBIN` to this compiler, + and `build.sh` exports it as `CC`/`CXX` so torch's build system passes it + to nvcc. +- **gsplat requires compute capability >= 7.0** — it uses + `cg::labeled_partition` from CUDA cooperative groups, which is not available + on Pascal GPUs (sm_61, e.g. GTX 1080). This rules out ax[03-05]. +- **The compiled kernels are cached in the conda env** — gsplat takes ~30 + minutes to build. Once installed, it persists across CI runs as long as the + conda env and GPU type stay the same. The `srun` GPU type constraint + (`--gres=gpu:a40:2`) ensures the job always lands on a matching node. - Follow the steps at Repository settings > Actions > Runners > New self hosted runner +### Checking the installed gsplat architecture -II. Setup on axon (or other SLURM cluster): +```bash +/share/apps/cuda/12.4/bin/cuobjdump -lelf \ + $(python -c "import gsplat.csrc as m; print(m.__file__)") | grep sm_ +``` - Setup the following directory structure using the files herein: - ~/buildbot/build.sh - ~/buildbot/build_srun.sh +### Rebuilding gsplat for a different GPU type - Ensure build.sh has executable permissions, as required by `srun`: - chmod +x ~/buildbot/build.sh +If you need to switch GPU types (e.g. from A40 to L40), you must rebuild gsplat +on a node with the target GPU. Run interactively: - Setup your conda environment, and configure the variables in build.sh. - Note that it was necessary to install packages via conda rather than pip on Axon, because - ffmpeg needs to be install via `conda install -c conda-forge opencv`. +```bash +srun --gres=gpu:a40:1 --mem=16g -c4 -t 1:00:00 bash -c ' + ml gcc/10.4 cuda/12.4.0 + export CC="$NVCC_CCBIN" + export CXX="$(dirname "$NVCC_CCBIN")/g++" + conda activate beast_build + pip uninstall -y gsplat + pip install --no-cache-dir \ + "gsplat @ git+https://github.com/QitaoZhao/gsplat.git@daad91d9e667cf49bab815d872aa65a5cda0a77e" \ + --no-build-isolation +' +``` -III. Test that this is working. +Then update `--gres=gpu::2` in `build_srun.sh` to match. -From a server that meets the requirements of the self-hosted runner in I, run: +**Important:** use `--no-cache-dir` when rebuilding — pip caches wheels by +package version, so without this flag it will reuse the old wheel compiled for +the wrong architecture. - ssh axon 'sh buildbot/build_srun.sh ' +## Compiler toolchain -It should successfully run tests. If not, debug. +The build environment juggles two gcc versions: -IV. A note on security and self-hosted runners. +| Purpose | Version | Source | +|---|---|---| +| CUDA host compiler (nvcc `-ccbin`) | gcc 12.3 | `ml gcc/10.4` + `ml cuda/12.4.0` sets `NVCC_CCBIN` | +| Runtime `libstdc++` (NumPy 2 requirement) | gcc 14.1 | `LD_PRELOAD` + `LD_LIBRARY_PATH` in `build.sh` | - Ensure external contributors cannot just run whatever workflow they - want on your self-hosted runner: they would be able to execute - arbitrary code on axon with the full permissions of your UNI. +The `gcc/10.4` module name is misleading — it provides gcc 12.3 via spack. The +CUDA module depends on it and sets `NVCC_CCBIN` to the correct binary path. +`build.sh` then exports `CC` and `CXX` from `NVCC_CCBIN` so that torch's +`cpp_extension.py` passes the right compiler to nvcc. - The only way I've found to be able to control this is via the following setting in Github: +The gcc 14.1 `libstdc++` is only needed at runtime (via `LD_PRELOAD`) for NumPy +2 compatibility — it is not used as a compiler. - Settings > Actions > General > - Approval for running fork pull request workflows from contributors +## Initial setup - The default is "Require approval for first-time contributors", but for maximum security, - consider further restricting to repo owners organization members. +### I. Self-hosted runner + +A self-hosted GitHub Actions runner is required with SSH access to Axon. + +1. You need a server (e.g. AWS, a physical PC). Axon nodes do not support + outbound internet so they cannot serve as runners directly. + +2. Configure SSH access to Axon. Ensure `~/.ssh/config` contains: + ``` + Host axon + HostName axon-remote.rc.zi.columbia.edu + User + Port 55 + ``` + +3. Set up passwordless SSH authentication via `ssh-keygen` and `ssh-copy-id`. + +4. Install the GitHub self-hosted runner software. Follow the steps at: + Repository settings > Actions > Runners > New self-hosted runner. + +### II. Axon setup + +Create the buildbot directory and scripts: + +``` +~/buildbot/build.sh +~/buildbot/build_srun.sh +``` + +Make `build.sh` executable: + +```bash +chmod +x ~/buildbot/build.sh +``` + +Set up the conda environment: + +```bash +conda create -n beast_build python=3.11 +conda activate beast_build +conda install -c conda-forge opencv # needed for ffmpeg +pip install lightning poetry-core +``` + +Then install gsplat (see "Rebuilding gsplat" above) and beast: + +```bash +pip install -e ".[dev]" --no-build-isolation +``` + +### III. Test the setup + +From the self-hosted runner server: + +```bash +ssh axon 'sh buildbot/build_srun.sh ' +``` + +### IV. Security + +Ensure external contributors cannot run arbitrary workflows on your self-hosted +runner — they would be able to execute code on Axon with your UNI's permissions. + +Configure this in GitHub: + +Settings > Actions > General > Approval for running fork pull request workflows + +The default is "Require approval for first-time contributors". For maximum +security, restrict to repository owners or organization members. diff --git a/scripts/buildbot/build.sh b/scripts/buildbot/build.sh index d5eab51..0045d76 100644 --- a/scripts/buildbot/build.sh +++ b/scripts/buildbot/build.sh @@ -24,13 +24,23 @@ echo "Running from $(hostname)" echo "Setting up environment..." source ~/.bashrc ml Miniforge-24.7.1-2 -ml gcc/14.1 # satisfies NumPy >= 2 requirement +ml gcc/10.4 # nvcc requires gcc <= 12 for CUDA 12.4 +ml cuda/12.4.0 +export CC="$NVCC_CCBIN" +export CXX="$(dirname "$NVCC_CCBIN")/g++" export LD_PRELOAD=/home/$(whoami)/.conda/envs/$CONDA_ENV/lib/libstdc++.so.6 export LD_LIBRARY_PATH=/share/apps/spack/gcc/14.1/lib64:$LD_LIBRARY_PATH conda activate $CONDA_ENV echo "Active conda environment: $CONDA_ENV" echo "Python location: $(which python)" -echo "Pip location $(which pip)" +echo "Pip location: $(which pip)" +echo "CC=$CC" +echo "CXX=$CXX" +echo "CUDA_HOME=$CUDA_HOME" +echo "NVCC_CCBIN=$NVCC_CCBIN" +echo "nvcc: $(which nvcc)" +echo "gcc: $(which gcc)" +$CC --version 2>&1 | head -1 # Remove builds older than 24 hours find "$BASE_DIR" -maxdepth 1 -type d -mtime +0 -print0 | while IFS= read -r -d $'\0' directory; do @@ -57,9 +67,17 @@ else fi # Install with checks -pip install -e ".[dev]" -echo "Pip install exit code: $?" -pip show beast +pip install -e ".[dev]" --no-build-isolation 2>&1 | tee /tmp/pip-install.log +PIP_EXIT=${PIPESTATUS[0]} +echo "Pip install exit code: $PIP_EXIT" +if [ "$PIP_EXIT" -ne 0 ]; then + echo "=== COMPILATION ERRORS ===" + grep -E '(^error:|nvcc error|fatal error|cannot find|undefined reference|exit code)' /tmp/pip-install.log || echo "No specific error pattern found" + echo "=== LAST 50 LINES ===" + tail -50 /tmp/pip-install.log + exit 1 +fi +pip show beast-backbones python -c "import beast; print('Beast location:', beast.__file__); print('Beast import successful')" # Run with html reporting (for github actions) and codecov reporting diff --git a/scripts/buildbot/build_srun.sh b/scripts/buildbot/build_srun.sh index 509182f..e3ae392 100644 --- a/scripts/buildbot/build_srun.sh +++ b/scripts/buildbot/build_srun.sh @@ -2,4 +2,4 @@ # -X is needed for github cancel to work. # Otherwise srun needs two successive ctrl-C / SIGINT to cancel # -u was found to enable colors. --pty would also work. -srun -X -u -t 1:00:00 -c4 --mem=8g --gres=gpu:2 buildbot/build.sh $@ \ No newline at end of file +srun -X -u -t 1:00:00 -c4 --mem=32g --gres=gpu:a40:1 buildbot/build.sh $@ \ No newline at end of file diff --git a/tests/api/test_model.py b/tests/api/test_model.py index 1c338ea..91a726b 100644 --- a/tests/api/test_model.py +++ b/tests/api/test_model.py @@ -21,6 +21,26 @@ def _make_valid_config(model_class: str) -> dict: } +def _make_valid_erayzer_config() -> dict: + """Minimal dict config accepted by ERayZerBeastConfig.""" + return { + 'model': { + 'model_class': 'erayzer', + 'transformer': {'d': 768, 'd_head': 64, 'encoder_geom_n_layer': 12}, + }, + 'training': { + 'train_batch_size': 4, + 'val_batch_size': 2, + 'num_views': 3, + 'num_input_views': 2, + 'num_target_views': 1, + 'max_fwdbwd_passes': 100000, + }, + 'optimizer': {'lr': 1e-4}, + 'data': {'data_dir': '/path/to/data'}, + } + + class TestChdir: """Test the chdir context manager.""" @@ -70,14 +90,14 @@ def _mock_class(self) -> MagicMock: def test_dict_config_vit(self) -> None: config = _make_valid_config('vit') mock_class = self._mock_class() - with patch.dict(Model.MODEL_REGISTRY, {'vit': mock_class}): + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'vit': mock_class}): m = Model.from_config(config) assert isinstance(m, Model) def test_dict_config_resnet(self) -> None: config = _make_valid_config('resnet') mock_class = self._mock_class() - with patch.dict(Model.MODEL_REGISTRY, {'resnet': mock_class}): + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'resnet': mock_class}): m = Model.from_config(config) assert isinstance(m, Model) @@ -85,7 +105,7 @@ def test_file_config(self, tmp_path: Path) -> None: config_path = tmp_path / 'config.yaml' config_path.write_text(yaml.dump(_make_valid_config('vit'))) mock_class = self._mock_class() - with patch.dict(Model.MODEL_REGISTRY, {'vit': mock_class}): + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'vit': mock_class}): m = Model.from_config(config_path) assert isinstance(m, Model) @@ -94,10 +114,33 @@ def test_unknown_model_type_raises(self) -> None: with pytest.raises(ValidationError): Model.from_config(config) + def test_dict_config_erayzer(self) -> None: + config = _make_valid_erayzer_config() + mock_class = self._mock_class() + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'erayzer': mock_class}): + m = Model.from_config(config) + assert isinstance(m, Model) + mock_class.assert_called_once() + + def test_erayzer_typed_fields_pass_through(self) -> None: + # ERayZer-specific fields are fully typed in the schema and must survive model_dump() + config = _make_valid_erayzer_config() + config['model']['transformer'] = { + 'd': 768, 'd_head': 64, 'encoder_n_layer': 12, 'encoder_geom_n_layer': 12, + } + config['optimizer']['beta1'] = 0.95 + mock_class = self._mock_class() + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'erayzer': mock_class}): + Model.from_config(config) + passed_config = mock_class.call_args[0][0] + assert passed_config['model']['transformer']['d'] == 768 + assert passed_config['training']['num_views'] == 3 + assert passed_config['optimizer']['beta1'] == 0.95 + def test_model_dir_is_none_after_from_config(self) -> None: config = _make_valid_config('vit') mock_class = self._mock_class() - with patch.dict(Model.MODEL_REGISTRY, {'vit': mock_class}): + with patch.dict('beast.models.registry.MODEL_REGISTRY', {'vit': mock_class}): m = Model.from_config(config) assert m.model_dir is None @@ -125,7 +168,7 @@ def test_loads_config_and_checkpoint(self, tmp_path: Path) -> None: mock_instance = mock_class.return_value mock_state = {'state_dict': {'layer': 'weights'}} with ( - patch.dict(Model.MODEL_REGISTRY, {'vit': mock_class}), + patch.dict('beast.models.registry.MODEL_REGISTRY', {'vit': mock_class}), patch('beast.api.model.torch.load', return_value=mock_state), ): m = Model.from_dir(tmp_path) @@ -146,7 +189,7 @@ def test_accepts_string_path(self, tmp_path: Path) -> None: mock_class = self._mock_class() mock_state = {'state_dict': {}} with ( - patch.dict(Model.MODEL_REGISTRY, {'resnet': mock_class}), + patch.dict('beast.models.registry.MODEL_REGISTRY', {'resnet': mock_class}), patch('beast.api.model.torch.load', return_value=mock_state), ): m = Model.from_dir(str(tmp_path)) @@ -157,13 +200,18 @@ class TestModelTrain: """Test the Model.train method.""" def _make_model(self, model_dir: Path | None = None) -> Model: - return Model(MagicMock(), {'model': {}, 'training': {}}, model_dir=model_dir) + return Model( + MagicMock(), + {'model': {'model_class': 'vit'}, 'training': {}}, + model_dir=model_dir, + ) def test_sets_model_dir(self, tmp_path: Path) -> None: m = self._make_model() output = tmp_path / 'run1' output.mkdir() - with patch('beast.api.model.train', return_value=MagicMock()): + mock_train = MagicMock(return_value=MagicMock()) + with patch.dict('beast.models.registry.TRAIN_REGISTRY', {'vit': mock_train}): m.train(output_dir=output) assert m.model_dir == output @@ -172,7 +220,8 @@ def test_calls_train_with_correct_args(self, tmp_path: Path) -> None: original_model = m.model # train() reassigns m.model to its return value output = tmp_path / 'run1' output.mkdir() - with patch('beast.api.model.train', return_value=MagicMock()) as mock_train: + mock_train = MagicMock(return_value=MagicMock()) + with patch.dict('beast.models.registry.TRAIN_REGISTRY', {'vit': mock_train}): m.train(output_dir=output) mock_train.assert_called_once_with(m.config, original_model, output_dir=output) @@ -181,7 +230,8 @@ def test_updates_model_after_train(self, tmp_path: Path) -> None: output = tmp_path / 'run1' output.mkdir() trained_model = MagicMock() - with patch('beast.api.model.train', return_value=trained_model): + mock_train = MagicMock(return_value=trained_model) + with patch.dict('beast.models.registry.TRAIN_REGISTRY', {'vit': mock_train}): m.train(output_dir=output) assert m.model is trained_model diff --git a/tests/conftest.py b/tests/conftest.py index 02f2c23..8cd87b9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,4 @@ +import copy import gc import io import json @@ -12,10 +13,12 @@ from beast.api.model import Model from beast.data.augmentations import expand_imgaug_str_to_dict, imgaug_pipeline -from beast.data.datamodules import BaseDataModule -from beast.data.datasets import BaseDataset +from beast.data.datamodules import BaseDataModule, MultiViewDataModule +from beast.data.datasets import BaseDataset, MultiViewDataset from beast.io import load_config +_MULTIVIEW_IMAGE_SIZE = 64 + ROOT = Path(__file__).parent.parent @@ -58,6 +61,7 @@ def fetch_test_data_if_needed(save_dir: str | Path, dataset_name: str = 'testing datasets_url_dict = { # 'testing_data': 'https://figshare.com/ndownloader/articles/29207330/versions/1', 'testing_data': 'https://github.com/paninski-lab/beast-test-fixtures/releases/download/v1/test_models.zip', # noqa + 'multiview_testing_data': 'https://github.com/paninski-lab/beast-test-fixtures/releases/download/v2/multiview_testing_data.zip', # noqa } dst_dir = Path(save_dir) / dataset_name @@ -90,6 +94,7 @@ def fetch_test_data_if_needed(save_dir: str | Path, dataset_name: str = 'testing fetch_test_data_if_needed(save_dir=Path(__file__).parent) +fetch_test_data_if_needed(save_dir=Path(__file__).parent, dataset_name='multiview_testing_data') # --------------------------------------------- # pytest fixtures @@ -101,6 +106,11 @@ def data_dir() -> Path: return ROOT.joinpath('tests/testing_data/data_dir') +@pytest.fixture +def multiview_data_dir() -> Path: + return ROOT.joinpath('tests/multiview_testing_data/multiview_testing_data') + + @pytest.fixture def video_file() -> Path: return ROOT.joinpath('tests/testing_data/test_vid.mp4') @@ -137,6 +147,35 @@ def config_vit(config_vit_path, data_dir) -> dict: return config +@pytest.fixture +def config_erayzer_path() -> Path: + return ROOT.joinpath('configs/multiview/erayzer.yaml') + + +@pytest.fixture +def config_erayzer(config_erayzer_path) -> dict: + config = load_config(config_erayzer_path) + config['model']['image_tokenizer']['image_size'] = 32 + config['model']['image_tokenizer']['patch_size'] = 8 # 4×4 = 16 tokens per view + config['model']['target_image']['height'] = 32 + config['model']['target_image']['width'] = 32 + config['model']['transformer']['d'] = 32 + config['model']['transformer']['d_head'] = 8 + config['model']['transformer']['encoder_n_layer'] = 2 + config['model']['transformer']['encoder_geom_n_layer'] = 2 + config['model']['transformer']['use_qk_norm'] = False + config['model']['transformer']['special_init'] = False + config['model']['transformer']['depth_init'] = False + config['model']['gaussians']['sh_degree'] = 0 + config['training']['num_views'] = 3 + config['training']['num_input_views'] = 2 + config['training']['num_target_views'] = 1 + config['training']['grad_checkpoint_every'] = 1 + config['training']['max_fwdbwd_passes'] = 100 + config['optimizer']['warmup'] = 10 # must be < max_fwdbwd_passes + return copy.deepcopy(config) + + @pytest.fixture def aug_pipeline() -> Callable: params_dict = expand_imgaug_str_to_dict('default') @@ -182,6 +221,31 @@ def base_datamodule_contrastive(base_dataset) -> BaseDataModule: return datamodule +@pytest.fixture +def multiview_dataset(multiview_data_dir) -> MultiViewDataset: + """MultiViewDataset in test mode (stable view order, no mask).""" + return MultiViewDataset( + data_dir=multiview_data_dir, + image_size=_MULTIVIEW_IMAGE_SIZE, + mode='test', + ) + + +@pytest.fixture +def multiview_datamodule(multiview_data_dir) -> MultiViewDataModule: + """MultiViewDataModule with 80/20 split, batch size 2, already set up.""" + dm = MultiViewDataModule( + data_dir=multiview_data_dir, + image_size=_MULTIVIEW_IMAGE_SIZE, + train_batch_size=2, + val_batch_size=2, + train_fraction=0.8, + num_workers=0, + ) + dm.setup() + return dm + + @pytest.fixture def run_model_test(tmp_path, data_dir) -> Callable: diff --git a/tests/data/test_datamodules.py b/tests/data/test_datamodules.py index 156c34f..ed31ea5 100644 --- a/tests/data/test_datamodules.py +++ b/tests/data/test_datamodules.py @@ -5,7 +5,11 @@ import torch from torch.utils.data import DataLoader, RandomSampler -from beast.data.datamodules import BaseDataModule, split_sizes_from_probabilities +from beast.data.datamodules import ( + BaseDataModule, + MultiViewDataModule, + split_sizes_from_probabilities, +) from beast.data.datasets import BaseDataset from beast.data.samplers import ContrastBatchSampler @@ -203,3 +207,84 @@ def test_probabilities_not_summing_to_one_raises(self) -> None: val_probability=0.1, test_probability=0.2, ) + + +# --------------------------------------------------------------------------- +# MultiViewDataModule tests +# --------------------------------------------------------------------------- + +_IMAGE_SIZE = 64 +_N_VIEWS = 3 +_N_TOTAL_FRAMES = 20 # 2 sessions × 10 frames + + +class TestMultiViewDataModule: + """Test the MultiViewDataModule class.""" + + def test_setup_creates_splits(self, multiview_datamodule) -> None: + assert multiview_datamodule.train_dataset is not None + assert multiview_datamodule.val_dataset is not None + assert ( + len(multiview_datamodule.train_dataset) + len(multiview_datamodule.val_dataset) + == _N_TOTAL_FRAMES + ) + + def test_split_sizes(self, multiview_datamodule) -> None: + assert len(multiview_datamodule.train_dataset) == 16 + assert len(multiview_datamodule.val_dataset) == 4 + + def test_train_batch_shape(self, multiview_datamodule) -> None: + batch = next(iter(multiview_datamodule.train_dataloader())) + assert batch['image'].shape == (2, _N_VIEWS, 3, _IMAGE_SIZE, _IMAGE_SIZE) + assert batch['c2w'].shape == (2, _N_VIEWS, 4, 4) + assert batch['fxfycxcy'].shape == (2, _N_VIEWS, 4) + + def test_val_batch_shape(self, multiview_datamodule) -> None: + batch = next(iter(multiview_datamodule.val_dataloader())) + assert batch['image'].shape == (2, _N_VIEWS, 3, _IMAGE_SIZE, _IMAGE_SIZE) + + def test_train_mode_shuffles_views_val_does_not(self, multiview_datamodule) -> None: + assert multiview_datamodule.train_dataset.mode == 'train' + assert multiview_datamodule.val_dataset.mode == 'test' + + def test_sequential_split_no_overlap(self, multiview_datamodule) -> None: + train_ids = set(multiview_datamodule.train_dataset.unique_frame_ids) + val_ids = set(multiview_datamodule.val_dataset.unique_frame_ids) + assert train_ids.isdisjoint(val_ids) + + def test_mask_propagated_to_batch(self, multiview_data_dir) -> None: + dm = MultiViewDataModule( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, + train_batch_size=2, use_mask=True, num_workers=0, + ) + dm.setup() + batch = next(iter(dm.train_dataloader())) + assert 'input_mask' in batch + assert batch['input_mask'].shape == (2, _N_VIEWS, 1, _IMAGE_SIZE, _IMAGE_SIZE) + + def test_val_at_least_one_frame(self, multiview_data_dir) -> None: + # train_fraction=1.0 should still give at least one val frame + dm = MultiViewDataModule( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, train_fraction=1.0, + ) + dm.setup() + assert dm.val_dataset is not None + assert len(dm.val_dataset) >= 1 + + def test_invalid_train_fraction_raises(self, multiview_data_dir) -> None: + with pytest.raises(ValueError, match='train_fraction'): + MultiViewDataModule( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, train_fraction=0.0, + ) + + def test_setup_required_before_dataloader(self, multiview_data_dir) -> None: + dm = MultiViewDataModule(data_dir=multiview_data_dir, image_size=_IMAGE_SIZE) + with pytest.raises(RuntimeError, match='call setup()'): + dm.train_dataloader() + with pytest.raises(RuntimeError, match='call setup()'): + dm.val_dataloader() + + def test_slurm_env_var_sets_num_workers(self, multiview_data_dir, monkeypatch) -> None: + monkeypatch.setenv('SLURM_CPUS_PER_TASK', '4') + dm = MultiViewDataModule(data_dir=multiview_data_dir, image_size=_IMAGE_SIZE) + assert dm.num_workers == 4 diff --git a/tests/data/test_datasets.py b/tests/data/test_datasets.py index d5af048..3f6aeca 100644 --- a/tests/data/test_datasets.py +++ b/tests/data/test_datasets.py @@ -3,7 +3,7 @@ import pytest import torch -from beast.data.datasets import _IMAGENET_MEAN, _IMAGENET_STD, BaseDataset +from beast.data.datasets import _IMAGENET_MEAN, _IMAGENET_STD, BaseDataset, MultiViewDataset class TestBaseDataset: @@ -60,3 +60,112 @@ def test_list_indexing_returns_list(self, base_dataset) -> None: for item in result: assert 'image' in item assert item['image'].shape == (3, 224, 224) + + +# --------------------------------------------------------------------------- +# MultiViewDataset tests +# --------------------------------------------------------------------------- + +_IMAGE_SIZE = 64 +_N_VIEWS = 3 # cameras in fixture: camera1, camera2, camera3 +_N_FRAMES = 20 # 2 sessions × 10 frames each + + +class TestMultiViewDataset: + """Test the MultiViewDataset class.""" + + def test_length(self, multiview_dataset) -> None: + assert len(multiview_dataset) == _N_FRAMES + + def test_item_shapes(self, multiview_dataset) -> None: + item = multiview_dataset[0] + assert item['image'].shape == (_N_VIEWS, 3, _IMAGE_SIZE, _IMAGE_SIZE) + assert item['c2w'].shape == (_N_VIEWS, 4, 4) + assert item['fxfycxcy'].shape == (_N_VIEWS, 4) + assert len(item['view_names']) == _N_VIEWS + assert isinstance(item['video_id'], str) + assert isinstance(item['frame_id'], str) + + def test_image_range(self, multiview_dataset) -> None: + # images should be float32 in [0, 1] + item = multiview_dataset[0] + assert item['image'].dtype == torch.float32 + assert item['image'].min() >= 0.0 + assert item['image'].max() <= 1.0 + + def test_c2w_is_valid_se3(self, multiview_dataset) -> None: + # last row must be [0, 0, 0, 1] and rotation block must be orthogonal + c2w = multiview_dataset[0]['c2w'] + expected_bottom = torch.tensor([0.0, 0.0, 0.0, 1.0]) + bottom_rows = expected_bottom.unsqueeze(0).expand(_N_VIEWS, -1) + assert torch.allclose(c2w[:, 3, :], bottom_rows, atol=1e-5) + R = c2w[:, :3, :3] + RRt = torch.bmm(R, R.transpose(1, 2)) + assert torch.allclose(RRt, torch.eye(3).unsqueeze(0).expand(_N_VIEWS, -1, -1), atol=1e-5) + + def test_camera_normalization(self, multiview_dataset) -> None: + # fixture is mode='test' (stable order) and normalize_cameras=True (default); + # camera-0 is sorted-first (camera1), which is re-centered to the origin + c2w = multiview_dataset[0]['c2w'] + assert torch.allclose(c2w[0, :3, 3], torch.zeros(3), atol=1e-5) + + def test_no_camera_normalization(self, multiview_data_dir) -> None: + # with normalize_cameras=False translations should be large (original anipose scale) + ds = MultiViewDataset( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, normalize_cameras=False, + ) + item = ds[0] + assert 'c2w' in item + assert item['c2w'][:, :3, 3].abs().max() > 1.0 + + def test_mask_loading(self, multiview_data_dir) -> None: + ds = MultiViewDataset( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, use_mask=True, + ) + item = ds[0] + assert 'input_mask' in item + assert item['input_mask'].shape == (_N_VIEWS, 1, _IMAGE_SIZE, _IMAGE_SIZE) + assert item['input_mask'].dtype == torch.float32 + assert set(item['input_mask'].unique().tolist()).issubset({0.0, 1.0}) + + def test_no_mask_by_default(self, multiview_dataset) -> None: + assert 'input_mask' not in multiview_dataset[0] + + def test_train_mode_shuffles_views(self, multiview_data_dir) -> None: + # over many draws, view order should not always be the same + ds = MultiViewDataset( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, mode='train', + ) + orders = [tuple(ds[0]['view_names']) for _ in range(20)] + assert len(set(orders)) > 1 + + def test_test_mode_stable_view_order(self, multiview_dataset) -> None: + orders = [tuple(multiview_dataset[0]['view_names']) for _ in range(5)] + assert len(set(orders)) == 1 + + def test_frame_ids_subset(self, multiview_dataset, multiview_data_dir) -> None: + # passing explicit frame_ids restricts the dataset length + subset_ids = multiview_dataset.unique_frame_ids[:5] + ds_sub = MultiViewDataset( + data_dir=multiview_data_dir, image_size=_IMAGE_SIZE, frame_ids=subset_ids, + ) + assert len(ds_sub) == 5 + + def test_missing_npy_raises(self, multiview_data_dir, tmp_path) -> None: + import shutil + shutil.copytree(multiview_data_dir, tmp_path / 'mv') + npy = next((tmp_path / 'mv').rglob('*.npy')) + npy.unlink() + ds = MultiViewDataset(data_dir=tmp_path / 'mv', image_size=_IMAGE_SIZE) + with pytest.raises(FileNotFoundError, match='camera file missing'): + for i in range(len(ds)): + ds[i] + + def test_invalid_data_dir_raises(self, tmp_path) -> None: + with pytest.raises(ValueError, match='is not a directory'): + MultiViewDataset(data_dir=tmp_path / 'nonexistent', image_size=_IMAGE_SIZE) + + def test_missing_info_json_raises(self, tmp_path) -> None: + (tmp_path / 'empty').mkdir() + with pytest.raises(ValueError, match='info.json not found'): + MultiViewDataset(data_dir=tmp_path / 'empty', image_size=_IMAGE_SIZE) diff --git a/tests/data/test_samplers.py b/tests/data/test_samplers.py index 1b8d0b9..176d53e 100644 --- a/tests/data/test_samplers.py +++ b/tests/data/test_samplers.py @@ -8,7 +8,7 @@ contrastive_collate_fn, extract_anchor_indices, ) -from beast.models.vits import batch_wise_contrastive_loss, topk +from beast.models.beast_vit.beast_vit_model import batch_wise_contrastive_loss, topk class TestExtractAnchorIndices: diff --git a/tests/data/test_types.py b/tests/data/test_types.py index f8db64b..40ab2c5 100644 --- a/tests/data/test_types.py +++ b/tests/data/test_types.py @@ -2,7 +2,7 @@ import torch -from beast.data.types import ExampleDict +from beast.data.types import ExampleDict, MultiViewExampleDict class TestExampleDict: @@ -30,3 +30,46 @@ def test_valid_batched_item(self) -> None: } assert isinstance(item['video'], list) assert isinstance(item['idx'], list) + + +class TestMultiViewExampleDict: + """Test the MultiViewExampleDict TypedDict.""" + + def test_required_keys(self) -> None: + annotations = MultiViewExampleDict.__annotations__ + assert {'image', 'view_names', 'video_id', 'frame_id'}.issubset(set(annotations)) + + def test_optional_keys(self) -> None: + # c2w, fxfycxcy, input_mask must be present as annotations (even if NotRequired) + annotations = MultiViewExampleDict.__annotations__ + assert 'c2w' in annotations + assert 'fxfycxcy' in annotations + assert 'input_mask' in annotations + + def test_valid_minimal_item(self) -> None: + # Arrange / Act — only required fields + item: MultiViewExampleDict = { + 'image': torch.zeros(3, 3, 64, 64), + 'view_names': ['camera1', 'camera2', 'camera3'], + 'video_id': 's1-d1', + 'frame_id': 'img00000001.png', + } + # Assert + assert item['image'].shape == (3, 3, 64, 64) + assert len(item['view_names']) == 3 + + def test_valid_full_item(self) -> None: + # Arrange / Act — all fields including optional ones + item: MultiViewExampleDict = { + 'image': torch.zeros(3, 3, 64, 64), + 'view_names': ['camera1', 'camera2', 'camera3'], + 'video_id': 's1-d1', + 'frame_id': 'img00000001.png', + 'c2w': torch.eye(4).unsqueeze(0).expand(3, -1, -1), + 'fxfycxcy': torch.ones(3, 4), + 'input_mask': torch.zeros(3, 1, 64, 64), + } + # Assert + assert item['c2w'].shape == (3, 4, 4) + assert item['fxfycxcy'].shape == (3, 4) + assert item['input_mask'].shape == (3, 1, 64, 64) diff --git a/tests/geometry/__init__.py b/tests/geometry/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/geometry/test_camera.py b/tests/geometry/test_camera.py new file mode 100644 index 0000000..e65f162 --- /dev/null +++ b/tests/geometry/test_camera.py @@ -0,0 +1,570 @@ +"""Tests for beast.geometry.camera utility functions.""" + +import math + +import numpy as np +import pytest +import torch + +from beast.geometry.camera import ( + cam_info_to_plucker, + get_interpolated_k, + get_interpolated_poses, + get_interpolated_poses_many, + get_ordered_poses_and_k, + intrinsics_to_fxfycxcy, + normalize_camera_sequence, + quaternion_from_matrix, + quaternion_matrix, + quaternion_slerp, + scale_intrinsics, + unit_vector, + w2c_to_c2w, +) + + +def _make_w2c(R: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Build a 4×4 w2c from a 3×3 rotation and 3-vector translation.""" + m = torch.eye(4) + m[:3, :3] = R + m[:3, 3] = t + return m + + +def _rot_z(theta: float) -> torch.Tensor: + """3×3 rotation matrix around z-axis by theta radians.""" + c, s = torch.cos(torch.tensor(theta)), torch.sin(torch.tensor(theta)) + return torch.tensor([[c, -s, 0.], [s, c, 0.], [0., 0., 1.]]) + + +class TestW2cToC2w: + """Test the w2c_to_c2w function.""" + + def test_identity_maps_to_identity(self) -> None: + w2c = torch.eye(4) + assert torch.allclose(w2c_to_c2w(w2c), torch.eye(4), atol=1e-6) + + def test_pure_translation_inverted(self) -> None: + # w2c with t=[1,2,3] and R=I → c2w with t=[-1,-2,-3] + w2c = torch.eye(4) + w2c[:3, 3] = torch.tensor([1., 2., 3.]) + c2w = w2c_to_c2w(w2c) + assert torch.allclose(c2w[:3, 3], torch.tensor([-1., -2., -3.]), atol=1e-6) + assert torch.allclose(c2w[:3, :3], torch.eye(3), atol=1e-6) + + def test_rotation_transposed(self) -> None: + R = _rot_z(torch.pi / 4) + w2c = _make_w2c(R, torch.zeros(3)) + c2w = w2c_to_c2w(w2c) + assert torch.allclose(c2w[:3, :3], R.T, atol=1e-6) + + def test_round_trip(self) -> None: + R = _rot_z(1.2) + w2c = _make_w2c(R, torch.tensor([3., -1., 2.])) + assert torch.allclose(w2c_to_c2w(w2c_to_c2w(w2c)), w2c, atol=1e-6) + + def test_bottom_row_preserved(self) -> None: + w2c = _make_w2c(_rot_z(0.5), torch.tensor([1., 2., 3.])) + c2w = w2c_to_c2w(w2c) + assert torch.allclose(c2w[3], torch.tensor([0., 0., 0., 1.]), atol=1e-6) + + def test_output_rotation_is_orthogonal(self) -> None: + w2c = _make_w2c(_rot_z(0.7), torch.tensor([1., 0., -2.])) + R = w2c_to_c2w(w2c)[:3, :3] + assert torch.allclose(R @ R.T, torch.eye(3), atol=1e-6) + + def test_batched_input(self) -> None: + # [B, 4, 4] input + B = 5 + w2c = torch.eye(4).unsqueeze(0).expand(B, -1, -1).clone() + w2c[:, :3, 3] = torch.randn(B, 3) + c2w = w2c_to_c2w(w2c) + assert c2w.shape == (B, 4, 4) + # each should round-trip + assert torch.allclose(w2c_to_c2w(c2w), w2c, atol=1e-6) + + def test_higher_batch_dims(self) -> None: + # [B, V, 4, 4] input + w2c = torch.eye(4).reshape(1, 1, 4, 4).expand(3, 4, -1, -1) + c2w = w2c_to_c2w(w2c) + assert c2w.shape == (3, 4, 4, 4) + + +class TestIntrinsicsToFxfycxcy: + """Test the intrinsics_to_fxfycxcy function.""" + + def test_known_values(self) -> None: + K = torch.tensor([[500., 0., 320.], [0., 400., 240.], [0., 0., 1.]]) + out = intrinsics_to_fxfycxcy(K) + assert torch.allclose(out, torch.tensor([500., 400., 320., 240.])) + + def test_output_shape_unbatched(self) -> None: + K = torch.eye(3) + assert intrinsics_to_fxfycxcy(K).shape == (4,) + + def test_output_shape_batched(self) -> None: + K = torch.eye(3).unsqueeze(0).expand(7, -1, -1) + assert intrinsics_to_fxfycxcy(K).shape == (7, 4) + + def test_batched_values(self) -> None: + fx, fy, cx, cy = 800., 600., 400., 300. + K = torch.tensor([[fx, 0., cx], [0., fy, cy], [0., 0., 1.]]) + K_batch = K.unsqueeze(0).expand(3, -1, -1) + out = intrinsics_to_fxfycxcy(K_batch) + expected = torch.tensor([fx, fy, cx, cy]).unsqueeze(0).expand(3, -1) + assert torch.allclose(out, expected) + + +class TestScaleIntrinsics: + """Test the scale_intrinsics function.""" + + def test_unit_scale_unchanged(self) -> None: + K = torch.tensor([[500., 0., 320.], [0., 400., 240.], [0., 0., 1.]]) + assert torch.allclose(scale_intrinsics(K, 1.0, 1.0), K) + + def test_uniform_scale(self) -> None: + K = torch.tensor([[500., 0., 320.], [0., 400., 240.], [0., 0., 1.]]) + K2 = scale_intrinsics(K, 2.0, 2.0) + assert torch.allclose(K2[0, 0], torch.tensor(1000.)) + assert torch.allclose(K2[1, 1], torch.tensor(800.)) + assert torch.allclose(K2[0, 2], torch.tensor(640.)) + assert torch.allclose(K2[1, 2], torch.tensor(480.)) + + def test_anisotropic_scale(self) -> None: + K = torch.tensor([[400., 0., 200.], [0., 400., 200.], [0., 0., 1.]]) + K2 = scale_intrinsics(K, 0.5, 2.0) + # fx, cx scaled by 0.5 + assert torch.allclose(K2[0, 0], torch.tensor(200.)) + assert torch.allclose(K2[0, 2], torch.tensor(100.)) + # fy, cy scaled by 2.0 + assert torch.allclose(K2[1, 1], torch.tensor(800.)) + assert torch.allclose(K2[1, 2], torch.tensor(400.)) + + def test_does_not_modify_input(self) -> None: + K = torch.tensor([[500., 0., 320.], [0., 400., 240.], [0., 0., 1.]]) + K_orig = K.clone() + scale_intrinsics(K, 2.0, 2.0) + assert torch.allclose(K, K_orig) + + def test_skew_and_bottom_row_untouched(self) -> None: + K = torch.tensor([[500., 1., 320.], [0., 400., 240.], [0., 0., 1.]]) + K2 = scale_intrinsics(K, 2.0, 2.0) + assert K2[0, 1].item() == 1.0 # skew unchanged + assert torch.allclose(K2[2], torch.tensor([0., 0., 1.])) + + def test_batched_input(self) -> None: + K = torch.tensor([[500., 0., 320.], [0., 400., 240.], [0., 0., 1.]]) + K_batch = K.unsqueeze(0).expand(4, -1, -1).clone() + K2 = scale_intrinsics(K_batch, 0.5, 0.5) + assert K2.shape == (4, 3, 3) + assert torch.allclose(K2[:, 0, 0], torch.full((4,), 250.)) + + +class TestNormalizeCameraSequence: + """Test the normalize_camera_sequence function.""" + + def _simple_w2c(self, t: list[float]) -> torch.Tensor: + """w2c with identity rotation and given translation.""" + m = torch.eye(4) + m[:3, 3] = torch.tensor(t) + return m + + def test_camera0_at_origin(self) -> None: + # after normalization camera 0 should be at the world origin + w2c = torch.stack([ + self._simple_w2c([0., 0., 5.]), + self._simple_w2c([0., 0., 10.]), + self._simple_w2c([0., 0., 15.]), + ]) + c2w = normalize_camera_sequence(w2c) + assert torch.allclose(c2w[0, :3, 3], torch.zeros(3), atol=1e-6) + + def test_mean_distance_is_one(self) -> None: + w2c = torch.stack([ + self._simple_w2c([0., 0., 0.]), + self._simple_w2c([0., 0., 10.]), + self._simple_w2c([0., 0., 20.]), + ]) + c2w = normalize_camera_sequence(w2c) + mean_dist = c2w[:, :3, 3].norm(dim=-1).mean() + assert torch.allclose(mean_dist, torch.tensor(1.0), atol=1e-5) + + def test_output_is_valid_se3(self) -> None: + w2c = torch.stack([ + _make_w2c(_rot_z(0.0), torch.tensor([0., 0., 5.])), + _make_w2c(_rot_z(1.0), torch.tensor([3., 0., 5.])), + _make_w2c(_rot_z(2.0), torch.tensor([-3., 0., 5.])), + ]) + c2w = normalize_camera_sequence(w2c) + # bottom row + expected_bottom = torch.tensor([0., 0., 0., 1.]).expand(3, -1) + assert torch.allclose(c2w[:, 3, :], expected_bottom, atol=1e-6) + # orthogonal rotation blocks + R = c2w[:, :3, :3] + RRt = torch.bmm(R, R.transpose(1, 2)) + assert torch.allclose(RRt, torch.eye(3).unsqueeze(0).expand(3, -1, -1), atol=1e-5) + + def test_single_camera_is_identity(self) -> None: + # one camera at origin → c2w should be identity + w2c = self._simple_w2c([0., 0., 0.]).unsqueeze(0) + c2w = normalize_camera_sequence(w2c) + assert torch.allclose(c2w[0], torch.eye(4), atol=1e-6) + + def test_output_shape(self) -> None: + w2c = torch.eye(4).unsqueeze(0).expand(6, -1, -1) + assert normalize_camera_sequence(w2c).shape == (6, 4, 4) + + +# --------------------------------------------------------------------------- +# numpy-based camera utilities +# --------------------------------------------------------------------------- + + +def _rot_z_np(theta: float) -> np.ndarray: + """3×3 rotation matrix around z-axis by theta radians (numpy).""" + c, s = math.cos(theta), math.sin(theta) + return np.array([[c, -s, 0.], [s, c, 0.], [0., 0., 1.]]) + + +def _pose_np(R: np.ndarray, t: np.ndarray) -> np.ndarray: + """Build a 4×4 pose matrix from a 3×3 rotation and 3-vector translation.""" + m = np.eye(4) + m[:3, :3] = R + m[:3, 3] = t + return m + + +class TestUnitVector: + """Test the unit_vector function.""" + + def test_1d_unit_length(self) -> None: + v = np.array([3., 4., 0.]) + uv = unit_vector(v) + assert abs(np.linalg.norm(uv) - 1.0) < 1e-10 + + def test_1d_known_values(self) -> None: + v = np.array([1., 0., 0.]) + np.testing.assert_allclose(unit_vector(v), v, atol=1e-10) + + def test_2d_axis0(self) -> None: + v = np.array([[3., 1.], [4., 0.]]) + uv = unit_vector(v, axis=0) + assert uv.shape == (2, 2) + norms = np.linalg.norm(uv, axis=0) + np.testing.assert_allclose(norms, [1.0, 1.0], atol=1e-10) + + def test_2d_axis1(self) -> None: + v = np.array([[3., 4.], [0., 5.]]) + uv = unit_vector(v, axis=1) + # each row should have norm 1 (or 0 if zero row) + np.testing.assert_allclose(np.linalg.norm(uv[0]), 1.0, atol=1e-10) + np.testing.assert_allclose(np.linalg.norm(uv[1]), 1.0, atol=1e-10) + + def test_does_not_modify_original(self) -> None: + v = np.array([3., 4., 0.]) + v_orig = v.copy() + unit_vector(v) + np.testing.assert_array_equal(v, v_orig) + + +class TestQuaternionFromMatrix: + """Test the quaternion_from_matrix function.""" + + def test_identity_matrix(self) -> None: + q = quaternion_from_matrix(np.eye(3)) + # identity → (1, 0, 0, 0); sign may flip but |q|=1 + np.testing.assert_allclose(abs(q[0]), 1.0, atol=1e-6) + np.testing.assert_allclose(q[1:], [0., 0., 0.], atol=1e-6) + + def test_output_is_unit_quaternion(self) -> None: + R = _rot_z_np(1.0) + q = quaternion_from_matrix(R) + np.testing.assert_allclose(np.linalg.norm(q), 1.0, atol=1e-9) + + def test_nonnegative_scalar_part(self) -> None: + # function normalises so q[0] >= 0 + q = quaternion_from_matrix(_rot_z_np(math.pi)) + assert q[0] >= 0.0 + + def test_isprecise_matches_default(self) -> None: + R = _pose_np(_rot_z_np(0.7), np.zeros(3)) # 4×4 required for isprecise=True + q_default = quaternion_from_matrix(R) + q_precise = quaternion_from_matrix(R, isprecise=True) + # same rotation, so either q or -q; compare absolute values + np.testing.assert_allclose(np.abs(q_default), np.abs(q_precise), atol=1e-6) + + def test_4x4_input_accepted(self) -> None: + M = np.eye(4) + q = quaternion_from_matrix(M) + assert q.shape == (4,) + + +class TestQuaternionSlerp: + """Test the quaternion_slerp function.""" + + def _quat_z(self, theta: float) -> np.ndarray: + """Quaternion for rotation around z by theta.""" + return np.array([math.cos(theta / 2), 0., 0., math.sin(theta / 2)]) + + def test_fraction_zero_returns_q0(self) -> None: + q0 = self._quat_z(0.0) + q1 = self._quat_z(1.0) + np.testing.assert_allclose(quaternion_slerp(q0, q1, 0.0), q0, atol=1e-9) + + def test_fraction_one_returns_q1(self) -> None: + q0 = self._quat_z(0.0) + q1 = self._quat_z(1.0) + np.testing.assert_allclose(quaternion_slerp(q0, q1, 1.0), q1, atol=1e-9) + + def test_midpoint_is_unit_quaternion(self) -> None: + q0 = self._quat_z(0.0) + q1 = self._quat_z(1.0) + qm = quaternion_slerp(q0, q1, 0.5) + np.testing.assert_allclose(np.linalg.norm(qm), 1.0, atol=1e-9) + + def test_midpoint_angle(self) -> None: + # midpoint between 0° and 90° rotation should be 45° + q0 = self._quat_z(0.0) + q1 = self._quat_z(math.pi / 2) + qm = quaternion_slerp(q0, q1, 0.5) + expected = self._quat_z(math.pi / 4) + np.testing.assert_allclose(np.abs(qm), np.abs(expected), atol=1e-6) + + def test_same_quaternions_returns_input(self) -> None: + q = self._quat_z(0.3) + np.testing.assert_allclose(quaternion_slerp(q, q, 0.5), q, atol=1e-9) + + +class TestQuaternionMatrix: + """Test the quaternion_matrix function.""" + + def test_identity_quaternion(self) -> None: + M = quaternion_matrix([1., 0., 0., 0.]) + np.testing.assert_allclose(M[:3, :3], np.eye(3), atol=1e-9) + + def test_output_shape(self) -> None: + M = quaternion_matrix([1., 0., 0., 0.]) + assert M.shape == (4, 4) + + def test_bottom_row_and_column(self) -> None: + M = quaternion_matrix([1., 0., 0., 0.]) + np.testing.assert_allclose(M[3], [0., 0., 0., 1.], atol=1e-9) + np.testing.assert_allclose(M[:, 3], [0., 0., 0., 1.], atol=1e-9) + + def test_rotation_block_is_orthogonal(self) -> None: + q = np.array([math.cos(0.3), 0., 0., math.sin(0.3)]) + M = quaternion_matrix(q) + R = M[:3, :3] + np.testing.assert_allclose(R @ R.T, np.eye(3), atol=1e-9) + + def test_known_rot_z_90(self) -> None: + # 90° around z: (w=cos45°, x=0, y=0, z=sin45°) + half = math.pi / 4 + q = [math.cos(half), 0., 0., math.sin(half)] + M = quaternion_matrix(q) + expected = np.array([[0., -1., 0.], [1., 0., 0.], [0., 0., 1.]]) + np.testing.assert_allclose(M[:3, :3], expected, atol=1e-7) + + def test_near_zero_quaternion_returns_identity(self) -> None: + M = quaternion_matrix([0., 0., 0., 0.]) + np.testing.assert_allclose(M, np.eye(4), atol=1e-9) + + +class TestGetInterpolatedPoses: + """Test the get_interpolated_poses function.""" + + def _pose3(self, t: list[float]) -> np.ndarray: + """4×4 pose with identity rotation and given translation.""" + m = np.eye(4) + m[:3, 3] = t + return m + + def test_output_length(self) -> None: + a = self._pose3([0., 0., 0.]) + b = self._pose3([1., 0., 0.]) + out = get_interpolated_poses(a, b, steps=7) + assert len(out) == 7 + + def test_each_pose_is_3x4(self) -> None: + a = self._pose3([0., 0., 0.]) + b = self._pose3([1., 0., 0.]) + for pose in get_interpolated_poses(a, b, steps=5): + assert np.array(pose).shape == (3, 4) + + def test_endpoints_match(self) -> None: + a = self._pose3([0., 0., 0.]) + b = self._pose3([3., 0., 0.]) + out = get_interpolated_poses(a, b, steps=5) + np.testing.assert_allclose(out[0][:, 3], a[:3, 3], atol=1e-6) + np.testing.assert_allclose(out[-1][:, 3], b[:3, 3], atol=1e-6) + + def test_default_steps(self) -> None: + a = self._pose3([0., 0., 0.]) + b = self._pose3([1., 0., 0.]) + assert len(get_interpolated_poses(a, b)) == 10 + + +class TestGetInterpolatedK: + """Test the get_interpolated_k function.""" + + def test_output_length(self) -> None: + K = torch.eye(3) + out = get_interpolated_k(K, K * 2, steps=6) + assert len(out) == 6 + + def test_endpoints_match(self) -> None: + k_a = torch.tensor([[400., 0., 200.], [0., 400., 200.], [0., 0., 1.]]) + k_b = torch.tensor([[800., 0., 400.], [0., 800., 400.], [0., 0., 1.]]) + out = get_interpolated_k(k_a, k_b, steps=5) + assert torch.allclose(out[0], k_a, atol=1e-6) + assert torch.allclose(out[-1], k_b, atol=1e-6) + + def test_midpoint_is_average(self) -> None: + k_a = torch.tensor([[400., 0., 200.], [0., 400., 200.], [0., 0., 1.]]) + k_b = torch.tensor([[800., 0., 400.], [0., 800., 400.], [0., 0., 1.]]) + out = get_interpolated_k(k_a, k_b, steps=3) + expected_mid = (k_a + k_b) / 2 + assert torch.allclose(out[1], expected_mid, atol=1e-6) + + def test_default_steps(self) -> None: + K = torch.eye(3) + assert len(get_interpolated_k(K, K)) == 10 + + +class TestGetOrderedPosesAndK: + """Test the get_ordered_poses_and_k function.""" + + def _make_poses(self, translations: list[list[float]]) -> torch.Tensor: + """Build [N, 3, 4] poses with identity rotations.""" + poses = [] + for t in translations: + p = torch.zeros(3, 4) + p[:3, :3] = torch.eye(3) + p[:, 3] = torch.tensor(t) + poses.append(p) + return torch.stack(poses) + + def test_output_shapes_preserved(self) -> None: + poses = self._make_poses([[0., 0., 0.], [1., 0., 0.], [2., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0).expand(3, -1, -1).clone() + ordered_poses, ordered_Ks = get_ordered_poses_and_k(poses, Ks) + assert ordered_poses.shape == poses.shape + assert ordered_Ks.shape == Ks.shape + + def test_first_pose_unchanged(self) -> None: + # greedy nearest-neighbor always starts from poses[0] + poses = self._make_poses([[0., 0., 0.], [10., 0., 0.], [1., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0).expand(3, -1, -1).clone() + ordered_poses, _ = get_ordered_poses_and_k(poses, Ks) + assert torch.allclose(ordered_poses[0], poses[0]) + + def test_nearest_neighbor_ordering(self) -> None: + # poses[0]→[1]→[3]→[2] by distance; [1] is at x=1, [3] at x=2, [2] at x=10 + translations = [[0., 0., 0.], [1., 0., 0.], [10., 0., 0.], [2., 0., 0.]] + poses = self._make_poses(translations) + Ks = torch.eye(3).unsqueeze(0).expand(4, -1, -1).clone() + ordered_poses, _ = get_ordered_poses_and_k(poses, Ks) + # second should be the x=1 pose (nearest to x=0) + assert torch.allclose(ordered_poses[1, :, 3], torch.tensor([1., 0., 0.])) + + +class TestGetInterpolatedPosesMany: + """Test the get_interpolated_poses_many function.""" + + def _make_poses(self, translations: list[list[float]]) -> torch.Tensor: + """Build [N, 3, 4] poses with identity rotations.""" + poses = [] + for t in translations: + p = torch.zeros(3, 4) + p[:3, :3] = torch.eye(3) + p[:, 3] = torch.tensor(t) + poses.append(p) + return torch.stack(poses) + + def test_empty_poses_raises(self) -> None: + poses = torch.zeros(0, 3, 4) + Ks = torch.zeros(0, 3, 3) + with pytest.raises(ValueError): + get_interpolated_poses_many(poses, Ks) + + def test_single_pose_returns_one_step(self) -> None: + poses = self._make_poses([[0., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0) + out_poses, out_Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=5) + assert out_poses.shape[0] == 1 + assert out_Ks.shape[0] == 1 + + def test_two_poses_output_length(self) -> None: + poses = self._make_poses([[0., 0., 0.], [1., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0).expand(2, -1, -1).clone() + out_poses, out_Ks = get_interpolated_poses_many(poses, Ks, steps_per_transition=4) + # one transition → 4 steps + assert out_poses.shape[0] == 4 + assert out_Ks.shape[0] == 4 + + def test_output_tensor_types(self) -> None: + poses = self._make_poses([[0., 0., 0.], [1., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0).expand(2, -1, -1).clone() + out_poses, out_Ks = get_interpolated_poses_many(poses, Ks) + assert isinstance(out_poses, torch.Tensor) + assert isinstance(out_Ks, torch.Tensor) + + def test_output_pose_shape(self) -> None: + poses = self._make_poses([[0., 0., 0.], [1., 0., 0.]]) + Ks = torch.eye(3).unsqueeze(0).expand(2, -1, -1).clone() + out_poses, _ = get_interpolated_poses_many(poses, Ks, steps_per_transition=3) + # each pose is 3×4 + assert out_poses.shape[1:] == (3, 4) + + +class TestCamInfoToPlucker: + """Test the cam_info_to_plucker function.""" + + def _identity_cam(self, b: int = 1) -> tuple[torch.Tensor, torch.Tensor]: + """Identity c2w and simple fxfycxcy for b cameras at 32×32.""" + c2w = torch.eye(4).unsqueeze(0).expand(b, -1, -1).clone() + fxfycxcy = torch.tensor([0.5, 0.5, 0.5, 0.5]).unsqueeze(0).expand(b, -1).clone() + return c2w, fxfycxcy + + def test_output_shape_single_batch(self) -> None: + c2w, fxfycxcy = self._identity_cam(b=2) + imgs_info = {'height': 8, 'width': 8} + out = cam_info_to_plucker(c2w, fxfycxcy, imgs_info, normalized=True) + assert out.shape == (2, 6, 8, 8) + + def test_output_shape_bv_input(self) -> None: + # [b, v, 4, 4] input → [b*v, 6, h, w] + b, v = 2, 3 + c2w = torch.eye(4).reshape(1, 1, 4, 4).expand(b, v, -1, -1).clone() + fxfycxcy = torch.tensor([0.5, 0.5, 0.5, 0.5]).reshape(1, 1, 4).expand(b, v, -1).clone() + imgs_info = {'height': 4, 'width': 4} + out = cam_info_to_plucker(c2w, fxfycxcy, imgs_info, normalized=True) + assert out.shape == (b * v, 6, 4, 4) + + def test_direction_vectors_are_unit(self) -> None: + c2w, fxfycxcy = self._identity_cam(b=1) + imgs_info = {'height': 8, 'width': 8} + out = cam_info_to_plucker(c2w, fxfycxcy, imgs_info, normalized=True) + # last 3 channels are direction; each pixel should be unit + directions = out[0, 3:] # [3, h, w] + norms = directions.norm(dim=0) # [h, w] + assert torch.allclose(norms, torch.ones_like(norms), atol=1e-5) + + def test_return_moment_false_gives_origin_direction(self) -> None: + c2w, fxfycxcy = self._identity_cam(b=1) + imgs_info = {'height': 4, 'width': 4} + out = cam_info_to_plucker(c2w, fxfycxcy, imgs_info, normalized=True, return_moment=False) + # origin channels: identity camera at origin → all zeros + origins = out[0, :3] # [3, h, w] + assert torch.allclose(origins, torch.zeros_like(origins), atol=1e-6) + + def test_normalized_false_uses_pixel_intrinsics(self) -> None: + # when normalized=False, fxfycxcy is already in pixels + b = 1 + h, w = 8, 8 + c2w = torch.eye(4).unsqueeze(0) + # pixel-space intrinsics: fx=fy=4, cx=cy=4 + fxfycxcy = torch.tensor([[4., 4., 4., 4.]]) + imgs_info = {'height': h, 'width': w} + out = cam_info_to_plucker(c2w, fxfycxcy, imgs_info, normalized=False) + assert out.shape == (b, 6, h, w) diff --git a/tests/geometry/test_positional_encoding.py b/tests/geometry/test_positional_encoding.py new file mode 100644 index 0000000..46e693e --- /dev/null +++ b/tests/geometry/test_positional_encoding.py @@ -0,0 +1,113 @@ +"""Tests for sinusoidal positional encoding utilities.""" + +import pytest +import torch + +from beast.geometry.positional_encoding import ( + get_1d_sincos_pos_emb_from_grid, + get_2d_sincos_pos_embed, + get_2d_sincos_pos_embed_from_grid, +) + + +class TestGet1dSincosPosEmbFromGrid: + """Test the get_1d_sincos_pos_emb_from_grid function.""" + + def test_output_shape(self) -> None: + pos = torch.arange(10, dtype=torch.float32) + emb = get_1d_sincos_pos_emb_from_grid(embed_dim=16, pos=pos) + assert emb.shape == (10, 16) + + def test_single_position(self) -> None: + pos = torch.tensor([5.0]) + emb = get_1d_sincos_pos_emb_from_grid(embed_dim=8, pos=pos) + assert emb.shape == (1, 8) + + def test_odd_embed_dim_raises(self) -> None: + pos = torch.arange(4, dtype=torch.float32) + with pytest.raises(AssertionError): + get_1d_sincos_pos_emb_from_grid(embed_dim=7, pos=pos) + + def test_values_bounded(self) -> None: + pos = torch.linspace(0, 100, 50) + emb = get_1d_sincos_pos_emb_from_grid(embed_dim=32, pos=pos) + assert emb.abs().max().item() <= 1.0 + 1e-6 + + def test_zero_position_sin_components_are_zero(self) -> None: + # sin(0 * freq) = 0 for all freqs; cos(0 * freq) = 1 + pos = torch.tensor([0.0]) + emb = get_1d_sincos_pos_emb_from_grid(embed_dim=16, pos=pos) + sin_half = emb[0, :8] + cos_half = emb[0, 8:] + assert torch.allclose(sin_half, torch.zeros(8), atol=1e-6) + assert torch.allclose(cos_half, torch.ones(8), atol=1e-6) + + def test_different_positions_give_different_embeddings(self) -> None: + pos = torch.tensor([0.0, 1.0, 2.0]) + emb = get_1d_sincos_pos_emb_from_grid(embed_dim=16, pos=pos) + assert not torch.allclose(emb[0], emb[1]) + assert not torch.allclose(emb[1], emb[2]) + + +class TestGet2dSincosPosEmbed: + """Test the get_2d_sincos_pos_embed function.""" + + def test_square_grid_output_shape(self) -> None: + emb = get_2d_sincos_pos_embed(embed_dim=32, grid_size=(4, 4)) + assert emb.shape == (16, 32) + + def test_non_square_grid_output_shape(self) -> None: + emb = get_2d_sincos_pos_embed(embed_dim=32, grid_size=(2, 8)) + assert emb.shape == (16, 32) + + def test_single_token_grid(self) -> None: + emb = get_2d_sincos_pos_embed(embed_dim=16, grid_size=(1, 1)) + assert emb.shape == (1, 16) + + def test_embed_dim_is_full_width(self) -> None: + embed_dim = 64 + emb = get_2d_sincos_pos_embed(embed_dim=embed_dim, grid_size=(3, 5)) + assert emb.shape[-1] == embed_dim + + def test_different_grid_positions_give_different_embeddings(self) -> None: + emb = get_2d_sincos_pos_embed(embed_dim=32, grid_size=(4, 4)) + # at least some rows should differ — adjacent spatial positions differ + assert not torch.allclose(emb[0], emb[1]) + + +class TestGet2dSincosPosEmbedFromGrid: + """Test the get_2d_sincos_pos_embed_from_grid function.""" + + def _make_grid(self, h: int, w: int) -> torch.Tensor: + grid_h = torch.arange(h, dtype=torch.float32) + grid_w = torch.arange(w, dtype=torch.float32) + grid = torch.stack(torch.meshgrid(grid_w, grid_h, indexing='ij'), dim=0) + return grid.view(2, 1, w, h) + + def test_output_shape(self) -> None: + grid = self._make_grid(4, 4) + emb = get_2d_sincos_pos_embed_from_grid(embed_dim=32, grid=grid) + assert emb.shape == (16, 32) + + def test_odd_embed_dim_raises(self) -> None: + grid = self._make_grid(2, 2) + with pytest.raises(AssertionError): + get_2d_sincos_pos_embed_from_grid(embed_dim=7, grid=grid) + + def test_matches_get_2d_sincos_pos_embed(self) -> None: + # get_2d_sincos_pos_embed is a thin wrapper around this function; + # results should be identical for the same logical grid + h, w, d = 3, 5, 32 + emb_direct = get_2d_sincos_pos_embed(embed_dim=d, grid_size=(h, w)) + grid = self._make_grid(h, w) + emb_from_grid = get_2d_sincos_pos_embed_from_grid(embed_dim=d, grid=grid) + assert torch.allclose(emb_direct, emb_from_grid, atol=1e-6) + + def test_output_is_concatenation_of_h_and_w_halves(self) -> None: + # first half of embedding dim comes from H coords, second from W + h, w, d = 4, 4, 32 + grid = self._make_grid(h, w) + emb = get_2d_sincos_pos_embed_from_grid(embed_dim=d, grid=grid) + # each half is itself a valid 1D sincos embedding (values in [-1, 1]) + assert emb[:, :d // 2].abs().max().item() <= 1.0 + 1e-6 + assert emb[:, d // 2:].abs().max().item() <= 1.0 + 1e-6 diff --git a/tests/geometry/test_rotations.py b/tests/geometry/test_rotations.py new file mode 100644 index 0000000..0c6c7b5 --- /dev/null +++ b/tests/geometry/test_rotations.py @@ -0,0 +1,109 @@ +"""Tests for beast.geometry.rotations utility functions.""" + +import torch + +from beast.geometry.rotations import quat2mat, rot6d2mat + + +def _identity_6d() -> torch.Tensor: + """6D representation of the identity rotation.""" + return torch.tensor([1., 0., 0., 0., 1., 0.]) + + +def _quat_identity() -> torch.Tensor: + """Unit quaternion (w,x,y,z) for the identity rotation.""" + return torch.tensor([1., 0., 0., 0.]) + + +def _quat_rot_z_90() -> torch.Tensor: + """Unit quaternion for 90° rotation around z-axis.""" + return torch.tensor([0.7071068, 0., 0., 0.7071068]) + + +class TestRot6d2Mat: + """Test the rot6d2mat function.""" + + def test_identity_input(self) -> None: + R = rot6d2mat(_identity_6d()) + assert torch.allclose(R, torch.eye(3), atol=1e-6) + + def test_output_shape_unbatched(self) -> None: + assert rot6d2mat(_identity_6d()).shape == (3, 3) + + def test_output_shape_batched(self) -> None: + inp = _identity_6d().unsqueeze(0).expand(5, -1) + assert rot6d2mat(inp).shape == (5, 3, 3) + + def test_output_is_orthogonal(self) -> None: + # random 6D input should still produce an orthogonal matrix + torch.manual_seed(0) + inp = torch.randn(8, 6) + R = rot6d2mat(inp) + RRt = torch.bmm(R, R.transpose(1, 2)) + assert torch.allclose(RRt, torch.eye(3).unsqueeze(0).expand(8, -1, -1), atol=1e-5) + + def test_determinant_is_positive_one(self) -> None: + torch.manual_seed(1) + inp = torch.randn(8, 6) + R = rot6d2mat(inp) + dets = torch.linalg.det(R) + assert torch.allclose(dets, torch.ones(8), atol=1e-5) + + def test_non_orthogonal_input_is_projected(self) -> None: + # even when the two input columns are not orthogonal, output must be orthogonal + inp = torch.tensor([1., 1., 0., 1., 0., 0.]) + R = rot6d2mat(inp) + assert torch.allclose(R @ R.T, torch.eye(3), atol=1e-5) + + def test_higher_batch_dims(self) -> None: + inp = _identity_6d().reshape(1, 1, 6).expand(3, 4, -1) + assert rot6d2mat(inp).shape == (3, 4, 3, 3) + + +class TestQuat2Mat: + """Test the quat2mat function.""" + + def test_identity_quaternion(self) -> None: + R = quat2mat(_quat_identity()) + assert torch.allclose(R, torch.eye(3), atol=1e-6) + + def test_output_shape_unbatched(self) -> None: + assert quat2mat(_quat_identity()).shape == (3, 3) + + def test_output_shape_batched(self) -> None: + q = _quat_identity().unsqueeze(0).expand(5, -1) + assert quat2mat(q).shape == (5, 3, 3) + + def test_rot_z_90(self) -> None: + # 90° around z: x→y, y→-x, z→z + R = quat2mat(_quat_rot_z_90()) + expected = torch.tensor([[0., -1., 0.], [1., 0., 0.], [0., 0., 1.]]) + assert torch.allclose(R, expected, atol=1e-5) + + def test_output_is_orthogonal(self) -> None: + torch.manual_seed(2) + q = torch.randn(8, 4) + R = quat2mat(q) + RRt = torch.bmm(R, R.transpose(1, 2)) + assert torch.allclose(RRt, torch.eye(3).unsqueeze(0).expand(8, -1, -1), atol=1e-5) + + def test_determinant_is_positive_one(self) -> None: + torch.manual_seed(3) + q = torch.randn(8, 4) + R = quat2mat(q) + dets = torch.linalg.det(R) + assert torch.allclose(dets, torch.ones(8), atol=1e-5) + + def test_unnormalized_quaternion(self) -> None: + # scaled quaternion should give the same rotation as the unit version + q_unit = _quat_rot_z_90() + q_scaled = q_unit * 3.7 + assert torch.allclose(quat2mat(q_unit), quat2mat(q_scaled), atol=1e-5) + + def test_antipodal_quaternions_give_same_rotation(self) -> None: + q = _quat_rot_z_90() + assert torch.allclose(quat2mat(q), quat2mat(-q), atol=1e-5) + + def test_higher_batch_dims(self) -> None: + q = _quat_identity().reshape(1, 1, 4).expand(3, 4, -1) + assert quat2mat(q).shape == (3, 4, 3, 3) diff --git a/tests/models/beast_resnet/__init__.py b/tests/models/beast_resnet/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/models/test_resnets.py b/tests/models/beast_resnet/test_beast_resnet_model.py similarity index 95% rename from tests/models/test_resnets.py rename to tests/models/beast_resnet/test_beast_resnet_model.py index 53e30ec..3479fb0 100644 --- a/tests/models/test_resnets.py +++ b/tests/models/beast_resnet/test_beast_resnet_model.py @@ -1,9 +1,15 @@ +"""Tests for the ResNet autoencoder model.""" + import copy import pytest import torch -from beast.models.resnets import _RESNET_CONFIGS, ResnetAutoencoder, get_configs +from beast.models.beast_resnet.beast_resnet_model import ( + _RESNET_CONFIGS, + ResnetAutoencoder, + get_configs, +) class TestGetConfigs: @@ -36,11 +42,11 @@ def test_registry_covers_expected_archs(self) -> None: class TestResnetAutoencoder: + """Test ResnetAutoencoder forward pass shapes.""" def test_forward_features(self, config_ae): - config = copy.deepcopy(config_ae) - config['model']['model_params']['num_latents'] = None # no latents, just 2D features + config['model']['model_params']['num_latents'] = None input = torch.randn((5, 3, 224, 224)) @@ -75,7 +81,6 @@ def test_forward_features(self, config_ae): assert latents.shape == (input.shape[0], 2048, 7, 7) def test_forward_latents(self, config_ae): - config = copy.deepcopy(config_ae) num_latents = 16 config['model']['model_params']['num_latents'] = num_latents @@ -113,7 +118,6 @@ def test_forward_latents(self, config_ae): assert latents.shape == (input.shape[0], num_latents) def test_predict_step_return_reconstructions(self, config_ae): - config = copy.deepcopy(config_ae) config['model']['model_params']['backbone'] = 'resnet18' config['model']['model_params']['num_latents'] = 16 diff --git a/tests/models/beast_vit/__init__.py b/tests/models/beast_vit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/models/test_vits.py b/tests/models/beast_vit/test_beast_vit_model.py similarity index 93% rename from tests/models/test_vits.py rename to tests/models/beast_vit/test_beast_vit_model.py index 949b165..6366594 100644 --- a/tests/models/test_vits.py +++ b/tests/models/beast_vit/test_beast_vit_model.py @@ -1,11 +1,14 @@ +"""Tests for the Vision Transformer autoencoder model.""" + import copy import torch -from beast.models.vits import VisionTransformer +from beast.models.beast_vit.beast_vit_model import VisionTransformer class TestVisionTransformer: + """Test VisionTransformer forward pass and output shapes.""" def test_forward(self, config_vit): config = copy.deepcopy(config_vit) diff --git a/tests/models/erayzer/__init__.py b/tests/models/erayzer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/models/erayzer/conftest.py b/tests/models/erayzer/conftest.py new file mode 100644 index 0000000..dc7bc3c --- /dev/null +++ b/tests/models/erayzer/conftest.py @@ -0,0 +1,84 @@ +"""Fixtures for ERayZer integration tests.""" + +import copy +import gc + +import lightning.pytorch as pl +import pytest +import torch + +from beast.api.model import Model +from beast.data.datamodules import MultiViewDataModule + + +def _gsplat_cuda_available() -> bool: + """Return True if gsplat's CUDA extension compiled and loaded successfully.""" + try: + import gsplat.cuda._backend as gsplat_backend + return getattr(gsplat_backend, '_C', None) is not None + except Exception: + return False + + +# Applied to every integration test; skips gracefully when gsplat CUDA is absent +# (e.g. missing CUDA toolkit, architecture not yet supported by available wheels). +requires_gsplat_cuda = pytest.mark.skipif( + not _gsplat_cuda_available(), + reason='gsplat CUDA extension not available (no nvcc / unsupported GPU arch)', +) + + +@pytest.fixture +def run_erayzer_model_test(tmp_path, multiview_data_dir): + """Build, train, and run inference on an ERayZer model. + + Mirrors run_model_test from tests/conftest.py but uses MultiViewDataModule + for both training and inference, since ERayZer requires multi-view input. + + The fixture sets up a minimal step count so that at least one full training + epoch (and therefore one validation run + best-checkpoint save) completes + within the multiview test fixtures (20 frames, batch_size=2 → 9 steps/epoch). + """ + def _run(config: dict) -> None: + config = copy.deepcopy(config) + config['data']['data_dir'] = str(multiview_data_dir) + config['training']['train_batch_size'] = 2 + config['training']['val_batch_size'] = 2 + config['training']['num_workers'] = 0 + config['training']['log_every_n_steps'] = 1 + config['training']['check_val_every_n_epoch'] = 1 + # 20 steps covers 2 full epochs (9 steps each) so val + checkpoint run twice + config['training']['max_fwdbwd_passes'] = 20 + # warmup must be strictly < max_fwdbwd_passes + config['optimizer']['warmup'] = 5 + + model = Model.from_config(config) + try: + model.train(tmp_path) + + # verify predict_step works on multiview batches + dm = MultiViewDataModule( + data_dir=multiview_data_dir, + image_size=config['model']['image_tokenizer']['image_size'], + train_batch_size=2, + val_batch_size=2, + train_fraction=0.8, + num_workers=0, + ) + dm.setup() + trainer = pl.Trainer(accelerator='gpu', devices=1, logger=False) + preds = trainer.predict( + model.model, + dataloaders=dm.val_dataloader(), + return_predictions=True, + ) + assert preds is not None and len(preds) > 0 + + assert model.model_dir is not None + assert len(list(model.model_dir.rglob('*.ckpt'))) == 1 + finally: + del model + gc.collect() + torch.cuda.empty_cache() + + return _run diff --git a/tests/models/erayzer/test_erayzer_model.py b/tests/models/erayzer/test_erayzer_model.py new file mode 100644 index 0000000..e996dc0 --- /dev/null +++ b/tests/models/erayzer/test_erayzer_model.py @@ -0,0 +1,974 @@ +"""Tests for beast.models.erayzer.erayzer_model components.""" + +import copy +from pathlib import Path +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +import pytest +import torch +import torch.nn as nn + +from beast.models.erayzer.erayzer_model import ( + ERayZer, + GaussiansUpsampler, + LossComputer, + PoseEstimator, + _filter_init_state_dict, + _parse_hf_checkpoint_url, + _resolve_init_checkpoint, + build_transformer_blocks, + get_cam_se3, + get_point_range_func, + sanitize, +) +from beast.nn.transformer import QK_Norm_TransformerBlock +from tests.models.erayzer.conftest import requires_gsplat_cuda + +# --------------------------------------------------------------------------- +# TestSanitize +# --------------------------------------------------------------------------- + + +class TestSanitize: + """Test the sanitize function.""" + + def test_nan_replaced_by_zero(self) -> None: + t = torch.tensor([float('nan'), 1.0]) + out = sanitize(t) + assert out[0].item() == 0.0 + + def test_posinf_replaced_by_1e6(self) -> None: + t = torch.tensor([float('inf')]) + out = sanitize(t) + assert out[0].item() == pytest.approx(1e6) + + def test_neginf_replaced_by_neg_1e6(self) -> None: + t = torch.tensor([float('-inf')]) + out = sanitize(t) + assert out[0].item() == pytest.approx(-1e6) + + def test_finite_values_pass_through(self) -> None: + t = torch.tensor([0.0, 1.0, -1.0, 100.0]) + assert torch.allclose(sanitize(t), t) + + def test_large_values_clamped(self) -> None: + t = torch.tensor([2e6, -2e6]) + out = sanitize(t) + assert out[0].item() == pytest.approx(1e6) + assert out[1].item() == pytest.approx(-1e6) + + def test_shape_preserved(self) -> None: + t = torch.randn(3, 4, 5) + assert sanitize(t).shape == t.shape + + +# --------------------------------------------------------------------------- +# TestBuildTransformerBlocks +# --------------------------------------------------------------------------- + + +class TestBuildTransformerBlocks: + """Test the build_transformer_blocks factory.""" + + def test_returns_correct_count(self) -> None: + blocks = build_transformer_blocks(4, d=32, d_head=8, use_qk_norm=False) + assert len(blocks) == 4 + + def test_each_element_is_correct_type(self) -> None: + blocks = build_transformer_blocks(2, d=32, d_head=8, use_qk_norm=False) + for b in blocks: + assert isinstance(b, QK_Norm_TransformerBlock) + + def test_returns_module_list(self) -> None: + blocks = build_transformer_blocks(2, d=32, d_head=8, use_qk_norm=False) + assert isinstance(blocks, nn.ModuleList) + + def test_special_init_no_error(self) -> None: + build_transformer_blocks(2, d=32, d_head=8, use_qk_norm=False, special_init=True) + + def test_depth_init_no_error(self) -> None: + build_transformer_blocks( + 2, d=32, d_head=8, use_qk_norm=False, special_init=True, depth_init=True, + ) + + +# --------------------------------------------------------------------------- +# TestGetCamSe3 +# --------------------------------------------------------------------------- + + +class TestGetCamSe3: + """Test the get_cam_se3 function.""" + + def test_6d_output_shapes(self) -> None: + cam_info = torch.randn(4, 13) + c2w, fxfycxcy = get_cam_se3(cam_info) + assert c2w.shape == (4, 4, 4) + assert fxfycxcy.shape == (4, 4) + + def test_quat_output_shapes(self) -> None: + cam_info = torch.randn(4, 11) + c2w, fxfycxcy = get_cam_se3(cam_info) + assert c2w.shape == (4, 4, 4) + assert fxfycxcy.shape == (4, 4) + + def test_unknown_width_raises(self) -> None: + cam_info = torch.randn(2, 10) + with pytest.raises(NotImplementedError): + get_cam_se3(cam_info) + + def test_bottom_row_is_0001(self) -> None: + cam_info = torch.randn(3, 13) + c2w, _ = get_cam_se3(cam_info) + expected = torch.tensor([0., 0., 0., 1.]) + assert torch.allclose(c2w[:, 3, :], expected.expand(3, -1), atol=1e-6) + + def test_identity_6d_gives_identity_rotation(self) -> None: + # 6D identity: first two columns of I₃ = [1,0,0,0,1,0] + base = torch.zeros(1, 13) + base[0, :6] = torch.tensor([1., 0., 0., 0., 1., 0.]) + c2w, _ = get_cam_se3(base) + assert torch.allclose(c2w[0, :3, :3], torch.eye(3), atol=1e-5) + + def test_fxfycxcy_matches_last_4(self) -> None: + cam_info = torch.zeros(2, 13) + cam_info[:, 9:] = torch.tensor([[1., 2., 3., 4.]]) + _, fxfycxcy = get_cam_se3(cam_info) + assert torch.allclose(fxfycxcy, torch.tensor([[1., 2., 3., 4.]]).expand(2, -1)) + + +# --------------------------------------------------------------------------- +# TestPoseEstimator +# --------------------------------------------------------------------------- + + +class TestPoseEstimator: + """Test the PoseEstimator module.""" + + def _make_config(self, canonical: str = 'first', mode: str = 'pairwise', + rep: str = '6d', per_view_focal: bool = False) -> dict: + return { + 'model': { + 'transformer': {'d': 32}, + 'pose_latent': { + 'canonical': canonical, + 'mode': mode, + 'representation': rep, + 'per_view_focal': per_view_focal, + }, + }, + 'training': {'pose_consistency_reg_weight': 0.0}, + } + + def test_output_shape_6d_pairwise_first(self) -> None: + config = self._make_config() + estimator = PoseEstimator(config) + b, v = 2, 4 + x = torch.randn(b * v, 32) + out = estimator(x, v) + assert out.shape == (b * v, 13) # 6 rot + 3 trans + 4 intrinsics + + def test_output_shape_quat(self) -> None: + config = self._make_config(rep='quat') + estimator = PoseEstimator(config) + b, v = 2, 3 + x = torch.randn(b * v, 32) + out = estimator(x, v) + assert out.shape == (b * v, 11) # 4 quat + 3 trans + 4 intrinsics + + def test_pairwise_middle_canonical(self) -> None: + config = self._make_config(canonical='middle') + estimator = PoseEstimator(config) + b, v = 2, 4 + x = torch.randn(b * v, 32) + out = estimator(x, v) + assert out.shape == (b * v, 13) + + def test_global_unordered(self) -> None: + config = self._make_config(canonical='unordered', mode='global') + estimator = PoseEstimator(config) + b, v = 2, 4 + x = torch.randn(b, v, 32) + out = estimator(x, v) + assert out.shape == (b, v, 13) + + def test_single_view_degenerate(self) -> None: + config = self._make_config() + estimator = PoseEstimator(config) + b, v = 2, 1 + x = torch.randn(b * v, 32) + out = estimator(x, v) + assert out.shape == (b * v, 13) + + def test_invalid_canonical_raises(self) -> None: + with pytest.raises(ValueError, match='Unknown canonical mode'): + PoseEstimator(self._make_config(canonical='bad')) + + def test_invalid_representation_raises(self) -> None: + with pytest.raises(NotImplementedError, match='Unknown pose representation'): + PoseEstimator(self._make_config(rep='euler')) + + def test_cxcy_columns_are_half(self) -> None: + # _append_cxcy sets cx=cy=0.5 for all views + config = self._make_config() + estimator = PoseEstimator(config) + x = torch.randn(6, 32) + out = estimator(x, v=3) + # last two columns are cx, cy + assert torch.allclose(out[:, 11], torch.full((6,), 0.5), atol=1e-5) + assert torch.allclose(out[:, 12], torch.full((6,), 0.5), atol=1e-5) + + def test_per_view_focal_false_broadcasts_canonical(self) -> None: + # default (checkpoint-matching): focal predicted from the canonical view + # and broadcast, so every view shares the same fx + config = self._make_config(per_view_focal=False) + estimator = PoseEstimator(config) + b, v = 2, 4 + x = torch.randn(b, v, 32) + out = estimator(x, v) # [b, v, 13] + fx = out[..., 9] # focal x is column 9 (after 6 rot + 3 trans) + assert torch.allclose(fx, fx[:, :1].expand(-1, v), atol=1e-6) + + def test_per_view_focal_true_varies_per_view(self) -> None: + # opt-in: focal head applied per view, so distinct view tokens give + # distinct per-view focals + config = self._make_config(per_view_focal=True) + estimator = PoseEstimator(config) + b, v = 2, 4 + x = torch.randn(b, v, 32) + out = estimator(x, v) + fx = out[..., 9] + assert not torch.allclose(fx, fx[:, :1].expand(-1, v), atol=1e-4) + + +# --------------------------------------------------------------------------- +# TestGaussiansUpsampler +# --------------------------------------------------------------------------- + + +class TestGaussiansUpsampler: + """Test the GaussiansUpsampler.to_gs method.""" + + def _make_config(self, sh_degree: int = 0) -> dict: + return { + 'model': { + 'gaussians': {'sh_degree': sh_degree}, + 'hard_pixelalign': False, + 'scaling_bias': -2.3, + 'scaling_max': -1.2, + 'opacity_bias': -2.0, + } + } + + def _make_input(self, b: int, n: int, sh_degree: int) -> torch.Tensor: + n_sh = (sh_degree + 1) ** 2 * 3 + # 3 xyz + n_sh features + 3 scaling + 4 rotation + 1 opacity + d_out = 3 + n_sh + 3 + 4 + 1 + return torch.zeros(b, n, d_out) + + def test_returns_five_tensors(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=0) + result = up.to_gs(inp) + assert len(result) == 5 + + def test_xyz_shape(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=0) + xyz, *_ = up.to_gs(inp) + assert xyz.shape == (2, 50, 3) + + def test_features_shape_sh0(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=0) + _, features, *_ = up.to_gs(inp) + # sh_degree=0: (0+1)^2 = 1, features shape [B, N, 1, 3] + assert features.shape == (2, 50, 1, 3) + + def test_features_shape_sh1(self) -> None: + config = self._make_config(sh_degree=1) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=1) + _, features, *_ = up.to_gs(inp) + # sh_degree=1: (1+1)^2 = 4, features shape [B, N, 4, 3] + assert features.shape == (2, 50, 4, 3) + + def test_scaling_bounded_by_max(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + # large raw scaling values should be clamped to scaling_max + inp = self._make_input(2, 50, sh_degree=0) + inp[:, :, 3 + 3:3 + 3 + 3] = 100.0 # raw scaling cols + _, _, scaling, *_ = up.to_gs(inp) + assert (scaling <= config['model']['scaling_max'] + 1e-5).all() + + def test_rotation_shape(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=0) + _, _, _, rotation, _ = up.to_gs(inp) + assert rotation.shape == (2, 50, 4) + + def test_opacity_shape(self) -> None: + config = self._make_config(sh_degree=0) + up = GaussiansUpsampler(config) + inp = self._make_input(2, 50, sh_degree=0) + *_, opacity = up.to_gs(inp) + assert opacity.shape == (2, 50, 1) + + +# --------------------------------------------------------------------------- +# TestGetPointRangeFunc +# --------------------------------------------------------------------------- + + +class TestGetPointRangeFunc: + """Test the get_point_range_func factory.""" + + def test_object_centric_depth_output_range(self) -> None: + fn = get_point_range_func({'range_setting': {'type': 'object_centric_depth'}}) + t = torch.zeros(1) + val = fn(t).item() + # sigmoid(0)=0.5 → 2*0.5-1=0 → 0*1.5+2.7=2.7 + assert abs(val - 2.7) < 1e-5 + + def test_linear_depth_output_in_range(self) -> None: + near, far = 1.0, 100.0 + fn = get_point_range_func({'range_setting': { + 'type': 'linear_depth', 'near': near, 'far': far, + }}) + t = torch.linspace(-5, 5, 20) + out = fn(t) + assert (out >= near - 1e-5).all() + assert (out <= far + 1e-5).all() + + def test_log_depth_output_positive(self) -> None: + fn = get_point_range_func({'range_setting': {'type': 'log_depth'}}) + t = torch.randn(20) + out = fn(t) + assert (out > 0).all() + + def test_disparity_output_in_range(self) -> None: + near, far = 0.5, 50.0 + fn = get_point_range_func({'range_setting': { + 'type': 'disparity', 'near': near, 'far': far, + }}) + t = torch.linspace(-5, 5, 20) + out = fn(t) + assert (out >= near - 1e-4).all() + assert (out <= far + 1e-4).all() + + def test_unknown_type_raises(self) -> None: + with pytest.raises(NotImplementedError, match='Unknown range_setting type'): + get_point_range_func({'range_setting': {'type': 'bad'}}) + + def test_output_shape_preserved(self) -> None: + fn = get_point_range_func({'range_setting': {'type': 'object_centric_depth'}}) + t = torch.randn(3, 4, 5) + assert fn(t).shape == t.shape + + def test_default_type_is_object_centric(self) -> None: + # passing empty gaussians_config uses default range_setting + fn = get_point_range_func({}) + assert fn is not None + + +# --------------------------------------------------------------------------- +# TestLossComputer +# --------------------------------------------------------------------------- + + +class TestLossComputer: + """Test the LossComputer module.""" + + def _make_config( + self, l2_weight: float = 1.0, gs_reg_weight: float = 0.0, + ) -> dict: + return { + 'training': { + 'l2_loss_weight': l2_weight, + 'gs_reg_loss_weight': gs_reg_weight, + 'perceptual_loss_weight': 0.0, + } + } + + def _make_images(self, b: int, v: int, h: int = 8, w: int = 8, + ) -> tuple[torch.Tensor, torch.Tensor]: + rendering = torch.rand(b, v, 3, h, w) + target = torch.rand(b, v, 3, h, w) + return rendering, target + + def test_basic_loss_is_finite(self) -> None: + lc = LossComputer(self._make_config()) + rendering, target = self._make_images(2, 3) + result = lc(rendering, target, None, None) + assert torch.isfinite(result.loss) + + def test_psnr_is_positive(self) -> None: + lc = LossComputer(self._make_config()) + rendering, target = self._make_images(2, 3) + result = lc(rendering, target, None, None) + assert result.psnr.item() > 0.0 + + def test_gs_reg_zero_when_no_xyz(self) -> None: + lc = LossComputer(self._make_config(gs_reg_weight=1.0)) + rendering, target = self._make_images(2, 3) + result = lc(rendering, target, None, None) + assert result.gs_reg_loss.item() == 0.0 + + def test_gs_reg_nonzero_when_xyz_provided(self) -> None: + lc = LossComputer(self._make_config(gs_reg_weight=1.0)) + rendering, target = self._make_images(2, 3) + xyz_norm = torch.randn(2, 100, 3) + xyz_init = xyz_norm + 0.1 + result = lc(rendering, target, xyz_norm, xyz_init) + assert result.gs_reg_loss.item() > 0.0 + + def test_masked_loss_differs_from_unmasked(self) -> None: + torch.manual_seed(0) + lc = LossComputer(self._make_config()) + b, v, h, w = 2, 3, 8, 8 + rendering = torch.rand(b, v, 3, h, w) + target = torch.rand(b, v, 3, h, w) + # mask selects only half the pixels + mask = torch.zeros(b, v, 1, h, w) + mask[:, :, :, :h // 2, :] = 1.0 + result_unmasked = lc(rendering, target, None, None) + result_masked = lc(rendering, target, None, None, pixel_mask=mask) + assert not torch.allclose(result_unmasked.l2_loss, result_masked.l2_loss) + + def test_returns_expected_fields(self) -> None: + lc = LossComputer(self._make_config()) + rendering, target = self._make_images(2, 3) + result = lc(rendering, target, None, None) + assert hasattr(result, 'loss') + assert hasattr(result, 'l2_loss') + assert hasattr(result, 'psnr') + assert hasattr(result, 'gs_reg_loss') + assert hasattr(result, 'perceptual_loss') + + def test_target_with_alpha_channel(self) -> None: + # target with 4 channels (RGBA) should strip alpha + lc = LossComputer(self._make_config()) + rendering = torch.rand(2, 3, 3, 8, 8) + target_rgba = torch.rand(2, 3, 4, 8, 8) + result = lc(rendering, target_rgba, None, None) + assert torch.isfinite(result.loss) + + +# --------------------------------------------------------------------------- +# TestERayZer +# --------------------------------------------------------------------------- + + +class TestERayZer: + """Test the ERayZer model class.""" + + def test_construction_with_encoder(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + assert hasattr(model, 'transformer_encoder') + assert hasattr(model, 'pose_predictor') + assert hasattr(model, 'transformer_encoder_geom') + assert hasattr(model, 'image_tokenizer') + assert hasattr(model, 'renderer') + + def test_construction_without_encoder(self, config_erayzer) -> None: + config = copy.deepcopy(config_erayzer) + config['model']['transformer']['encoder_n_layer'] = 0 + model = ERayZer(config) + assert not hasattr(model, 'transformer_encoder') + assert not hasattr(model, 'pose_predictor') + + def test_odd_geom_layers_raises(self, config_erayzer) -> None: + config = copy.deepcopy(config_erayzer) + config['model']['transformer']['encoder_geom_n_layer'] = 3 + with pytest.raises(ValueError, match='encoder_geom_n_layer must be even'): + ERayZer(config) + + def test_odd_encoder_layers_raises(self, config_erayzer) -> None: + config = copy.deepcopy(config_erayzer) + config['model']['transformer']['encoder_n_layer'] = 3 + with pytest.raises(ValueError, match='encoder_n_layer must be even'): + ERayZer(config) + + def test_resolve_cameras_raises_without_encoder(self, config_erayzer) -> None: + config = copy.deepcopy(config_erayzer) + config['model']['transformer']['encoder_n_layer'] = 0 + model = ERayZer(config) + with pytest.raises(RuntimeError, match='camera prediction modules are not initialized'): + model._resolve_cameras( + torch.randn(2, 16, 32), b=1, v_all=2, n=16, data={}, device=torch.device('cpu'), + ) + + def test_resolve_view_indices_from_data(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + data = { + 'image': torch.zeros(2, 3, 3, 32, 32), + 'input_indices': torch.tensor([[0, 1], [0, 1]]), + 'target_indices': torch.tensor([[2], [2]]), + } + in_idx, tgt_idx = model.resolve_view_indices(data, num_real_views=3, device='cpu') + assert in_idx.tolist() == [[0, 1], [0, 1]] + assert tgt_idx.tolist() == [[2], [2]] + + def test_resolve_view_indices_from_config_split(self, config_erayzer) -> None: + # num_views=3, num_input_views=2, num_target_views=1, inference=False + model = ERayZer(config_erayzer) + model.random_index = False + data = {'image': torch.zeros(2, 3, 3, 32, 32)} + in_idx, tgt_idx = model.resolve_view_indices(data, num_real_views=3, device='cpu') + assert in_idx.tolist() == [[0, 1], [0, 1]] + assert tgt_idx.tolist() == [[2], [2]] + + def test_resolve_view_indices_inference_mode(self, config_erayzer) -> None: + # in inference mode, all views are both input and target + config = copy.deepcopy(config_erayzer) + config['inference'] = True + model = ERayZer(config) + data = {'image': torch.zeros(2, 3, 3, 32, 32)} + in_idx, tgt_idx = model.resolve_view_indices(data, num_real_views=3, device='cpu') + assert in_idx.shape[1] <= 3 + assert tgt_idx.shape[1] == 3 + + def test_maybe_randomize_single_view(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + torch.manual_seed(0) + input_idx = torch.tensor([[0]]) + target_idx = torch.tensor([[1]]) + # run many times; with p=0.5 swap both orderings should appear + got_swap = False + got_no_swap = False + for _ in range(30): + in_i, tgt_i = model.maybe_randomize_view_indices(input_idx, target_idx, 'cpu') + if in_i.tolist() == [[1]]: + got_swap = True + else: + got_no_swap = True + if got_swap and got_no_swap: + break + assert got_swap + assert got_no_swap + + def test_add_spatial_pe_shape_preserved(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + b, v, hh, ww = 2, 3, 4, 4 + n = hh * ww + d = config_erayzer['model']['transformer']['d'] + tokens = torch.randn(b * v, n, d) + out = model.add_spatial_pe(tokens, b, v, hh, ww, model.pe_embedder) + assert out.shape == (b * v, n, d) + + def test_add_spatial_pe_mismatched_tokens_raises(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + tokens = torch.randn(6, 15, 32) # 15 != 4*4 + with pytest.raises(ValueError, match='Token count'): + model.add_spatial_pe(tokens, b=2, v=3, h_tokens=4, w_tokens=4, + embedder=model.pe_embedder) + + def test_compute_loss_returns_loss_and_log_list(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + b, v, h, w = 2, 3, 32, 32 + kwargs = { + 'render': torch.rand(b, v, 3, h, w), + 'target_image': torch.rand(b, v, 3, h, w), + } + loss, log_list = model.compute_loss('train', **kwargs) + assert loss.ndim == 0 + assert isinstance(log_list, list) + assert any(d['name'] == 'train_psnr' for d in log_list) + + def test_forward_output_keys(self, config_erayzer) -> None: + """Forward pass with mocked renderer to avoid gsplat.""" + model = ERayZer(config_erayzer) + model.eval() + + b, v, h, w = 1, 3, 32, 32 + data = {'image': torch.rand(b, v, 3, h, w)} + + def _fake_render(xyz, features, scaling, rotation, opacity, + height, width, C2W, fxfycxcy): + return SimpleNamespace( + render=torch.zeros(C2W.shape[0], C2W.shape[1], 3, height, width), + depth=torch.zeros(C2W.shape[0], C2W.shape[1], 1, height, width), + alpha=torch.zeros(C2W.shape[0], C2W.shape[1], 1, height, width), + ) + + with patch.object(model.renderer, 'forward', side_effect=_fake_render): + result = model(data) + + assert hasattr(result, 'render') + assert hasattr(result, 'gaussians') + assert hasattr(result, 'c2w_input') + assert hasattr(result, 'fxfycxcy_input') + assert hasattr(result, 'target_image') + assert hasattr(result, 'input_indices') + assert hasattr(result, 'target_indices') + + def test_forward_render_shape(self, config_erayzer) -> None: + """render output has shape [B, n_target, 3, H, W].""" + model = ERayZer(config_erayzer) + model.eval() + + b, v, h, w = 1, 3, 32, 32 + n_target = config_erayzer['training']['num_target_views'] + data = {'image': torch.rand(b, v, 3, h, w)} + + def _fake_render(xyz, features, scaling, rotation, opacity, + height, width, C2W, fxfycxcy): + return SimpleNamespace( + render=torch.zeros(C2W.shape[0], C2W.shape[1], 3, height, width), + depth=torch.zeros(C2W.shape[0], C2W.shape[1], 1, height, width), + alpha=torch.zeros(C2W.shape[0], C2W.shape[1], 1, height, width), + ) + + with patch.object(model.renderer, 'forward', side_effect=_fake_render): + result = model(data) + + assert result.render.shape == (b, n_target, 3, h, w) + + def test_configure_optimizers_keys(self, config_erayzer) -> None: + model = ERayZer(config_erayzer) + opt_cfg = model.configure_optimizers() + assert 'optimizer' in opt_cfg + assert 'lr_scheduler' in opt_cfg + + +# --------------------------------------------------------------------------- +# TestERayZerIntegration +# --------------------------------------------------------------------------- + + +@requires_gsplat_cuda +class TestERayZerIntegration: + """Integration tests that train and run inference on ERayZer. + + Skipped automatically when the gsplat CUDA extension is unavailable + (e.g. CUDA toolkit not installed, GPU architecture not yet supported by + available pre-compiled wheels). + """ + + def test_integration_basic(self, config_erayzer, run_erayzer_model_test) -> None: + """Test ERayZer trains to completion and runs inference with L2 loss only.""" + run_erayzer_model_test(config=config_erayzer) + + def test_integration_gs_reg(self, config_erayzer, run_erayzer_model_test) -> None: + """Test ERayZer trains with Gaussian splatting regularization loss enabled.""" + config = copy.deepcopy(config_erayzer) + config['training']['gs_reg_loss_weight'] = 0.1 + run_erayzer_model_test(config=config) + + def test_integration_perceptual_loss(self, config_erayzer, run_erayzer_model_test) -> None: + """Test ERayZer trains with VGG perceptual loss enabled.""" + config = copy.deepcopy(config_erayzer) + config['training']['perceptual_loss_weight'] = 0.1 + run_erayzer_model_test(config=config) + + +# --------------------------------------------------------------------------- +# TestFilterInitStateDict +# --------------------------------------------------------------------------- + + +class TestFilterInitStateDict: + """Test the _filter_init_state_dict checkpoint-extraction helper.""" + + def test_extracts_from_state_dict_key(self) -> None: + raw = {'state_dict': {'a': torch.zeros(1), 'b': torch.zeros(2)}} + out = _filter_init_state_dict(raw) + assert set(out) == {'a', 'b'} + + def test_extracts_from_model_key(self) -> None: + raw = {'model': {'a': torch.zeros(1)}} + out = _filter_init_state_dict(raw) + assert set(out) == {'a'} + + def test_handles_bare_top_level_state_dict(self) -> None: + raw = {'camera_token': torch.zeros(1), 'register_token': torch.zeros(2)} + out = _filter_init_state_dict(raw) + assert set(out) == {'camera_token', 'register_token'} + + def test_drops_loss_computer_keys(self) -> None: + raw = { + 'camera_token': torch.zeros(1), + 'loss_computer.perceptual_loss_module.vgg.features.0.weight': torch.zeros(3), + 'loss_computer.anything': torch.zeros(1), + } + out = _filter_init_state_dict(raw) + assert set(out) == {'camera_token'} + + def test_keeps_non_dropped_keys(self) -> None: + raw = {'state_dict': { + 'transformer_encoder.0.norm1.weight': torch.zeros(4), + 'loss_computer.x': torch.zeros(1), + }} + out = _filter_init_state_dict(raw) + assert set(out) == {'transformer_encoder.0.norm1.weight'} + + def test_custom_drop_prefix(self) -> None: + raw = {'keep.a': torch.zeros(1), 'drop.b': torch.zeros(1)} + out = _filter_init_state_dict(raw, drop_prefixes=('drop.',)) + assert set(out) == {'keep.a'} + + +# --------------------------------------------------------------------------- +# TestLoadInitCheckpoint +# --------------------------------------------------------------------------- + + +class TestResolveViewIndices: + """Test ERayZer.resolve_view_indices view-sampling (no model build needed).""" + + def _standin( + self, training: bool, training_cfg: dict, random_index: bool = False, + ) -> SimpleNamespace: + # resolve_view_indices only touches self.training, self.random_index, + # and self.config — a stand-in exercises it without building the model + return SimpleNamespace( + training=training, + random_index=random_index, + config={'training': training_cfg, 'inference': False}, + ) + + def _call(self, standin, num_real_views: int, batch_size: int = 2): + data = {'image': torch.zeros(batch_size, num_real_views, 3, 8, 8)} + return ERayZer.resolve_view_indices( + standin, data, num_real_views, torch.device('cpu'), + ) + + def test_target_view_range_holds_out_exactly_n_targets(self) -> None: + standin = self._standin(True, {'target_view_range': [1, 1], 'num_views': 6}) + for _ in range(30): + in_idx, tgt_idx = self._call(standin, 6) + assert tgt_idx.shape[1] == 1 # exactly 1 target + assert in_idx.shape[1] >= 1 # at least 1 input + assert in_idx.shape[1] + tgt_idx.shape[1] <= 6 + + def test_total_view_count_varies_across_batches(self) -> None: + standin = self._standin(True, {'target_view_range': [1, 1], 'num_views': 6}) + totals = {int(self._call(standin, 6)[0].shape[1] + self._call(standin, 6)[1].shape[1]) + for _ in range(60)} + # multi-view-style sampler must use a variable total, not a fixed count + assert len(totals) >= 2 + + def test_input_and_target_indices_are_disjoint(self) -> None: + standin = self._standin(True, {'target_view_range': [1, 1], 'num_views': 6}) + for _ in range(30): + in_idx, tgt_idx = self._call(standin, 6) + a = set(in_idx[0].tolist()) + b = set(tgt_idx[0].tolist()) + assert a.isdisjoint(b) + assert max(a | b) < 6 + + def test_target_range_two_allows_two_targets(self) -> None: + standin = self._standin(True, {'target_view_range': [1, 2], 'num_views': 6}) + seen = {int(self._call(standin, 6)[1].shape[1]) for _ in range(60)} + assert seen <= {1, 2} and len(seen) >= 1 + + def test_validation_uses_fixed_split_not_sampler(self) -> None: + # at val time the sampler branch is skipped: deterministic fixed split + standin = self._standin( + False, + {'target_view_range': [1, 1], 'num_views': 6, + 'num_input_views': 2, 'num_target_views': 4}, + ) + in_idx, tgt_idx = self._call(standin, 6) + assert in_idx.shape[1] == 2 and tgt_idx.shape[1] == 4 + assert in_idx[0].tolist() == [0, 1] # deterministic order + assert tgt_idx[0].tolist() == [2, 3, 4, 5] + + def test_validation_random_split_alternates_views(self) -> None: + # random_split (random_index=True) shuffles which cameras are input vs + # target at val, while keeping the 2/4 counts + standin = self._standin( + False, + {'target_view_range': [1, 1], 'num_views': 6, + 'num_input_views': 2, 'num_target_views': 4}, + random_index=True, + ) + input_sets = set() + for _ in range(40): + in_idx, tgt_idx = self._call(standin, 6) + assert in_idx.shape[1] == 2 and tgt_idx.shape[1] == 4 + assert set(in_idx[0].tolist()).isdisjoint(tgt_idx[0].tolist()) + input_sets.add(tuple(sorted(in_idx[0].tolist()))) + assert len(input_sets) >= 2 # the input/target assignment varies + + +class TestValidationVisuals: + """Test ERayZer validation-visualization guards (no GPU/model build needed).""" + + def test_no_logger_is_noop(self) -> None: + # logger=None -> no image-capable experiment -> returns without error + standin = SimpleNamespace(logger=None) + ERayZer._log_validation_visuals(cast(ERayZer, standin), {}) + + def test_logger_without_add_image_is_noop(self) -> None: + # an experiment lacking add_image is treated as not image-capable + standin = SimpleNamespace(logger=SimpleNamespace(experiment=object())) + ERayZer._log_validation_visuals(cast(ERayZer, standin), {}) + + def test_forward_failure_is_swallowed(self) -> None: + # a capable logger but a get_model_outputs that raises must not propagate + class _Exp: + def add_image(self, *a, **k) -> None: + pass + + def _boom(_batch) -> dict: + raise RuntimeError('forward exploded') + + standin = SimpleNamespace( + logger=SimpleNamespace(experiment=_Exp()), + get_model_outputs=_boom, + ) + ERayZer._log_validation_visuals(cast(ERayZer, standin), {}) + + +class TestConfigureOptimizers: + """Test ERayZer.configure_optimizers LR scaling + schedule selection.""" + + def _standin(self, opt_cfg: dict, train_cfg: dict) -> ERayZer: + # a duck-typed stand-in exercises configure_optimizers without a full build; + # cast so the unbound-method calls type-check against the real signature + param = nn.Parameter(torch.zeros(2)) + return cast(ERayZer, SimpleNamespace( + config={'optimizer': opt_cfg, 'training': train_cfg}, + parameters=lambda: iter([param]), + )) + + def _train_cfg(self, **kw) -> dict: + cfg = {'train_batch_size': 8, 'num_gpus': 1, 'num_nodes': 1, 'max_fwdbwd_passes': 1000} + cfg.update(kw) + return cfg + + def test_constant_schedule_is_lambdalr(self) -> None: + from torch.optim.lr_scheduler import LambdaLR + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.999, 'wd': 0.05, 'schedule': 'constant'}, + self._train_cfg(), + ) + out = ERayZer.configure_optimizers(m) + assert isinstance(out['lr_scheduler']['scheduler'], LambdaLR) + + def test_onecycle_schedule(self) -> None: + from torch.optim.lr_scheduler import OneCycleLR + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.95, 'wd': 0.05, 'warmup': 100, + 'schedule': 'onecycle', 'div_factor': 10.0, 'final_div_factor': 100.0}, + self._train_cfg(), + ) + out = ERayZer.configure_optimizers(m) + assert isinstance(out['lr_scheduler']['scheduler'], OneCycleLR) + + def test_unknown_schedule_raises(self) -> None: + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.95, 'wd': 0.05, 'schedule': 'bogus'}, + self._train_cfg(), + ) + with pytest.raises(ValueError, match='unknown optimizer.schedule'): + ERayZer.configure_optimizers(m) + + def test_no_scaling_by_default(self) -> None: + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.95, 'wd': 0.05, 'schedule': 'constant'}, + self._train_cfg(train_batch_size=8), + ) + out = ERayZer.configure_optimizers(m) + assert out['optimizer'].param_groups[0]['lr'] == pytest.approx(4e-4) + + def test_lr_batch_scaling_at_256_is_identity(self) -> None: + # global_batch_size = 32 * 8 * 1 = 256 -> scaled lr == base lr + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.999, 'wd': 0.05, + 'schedule': 'constant', 'scale_lr_by_batch': True}, + self._train_cfg(train_batch_size=32, num_gpus=8), + ) + out = ERayZer.configure_optimizers(m) + assert out['optimizer'].param_groups[0]['lr'] == pytest.approx(4e-4) + + def test_lr_batch_scaling_small_batch(self) -> None: + # global_batch_size = 8 -> lr scaled down by 8/256 + m = self._standin( + {'lr': 4e-4, 'beta1': 0.9, 'beta2': 0.999, 'wd': 0.05, + 'schedule': 'constant', 'scale_lr_by_batch': True}, + self._train_cfg(train_batch_size=8), + ) + out = ERayZer.configure_optimizers(m) + assert out['optimizer'].param_groups[0]['lr'] == pytest.approx(4e-4 * 8 / 256) + + +class TestParseHfCheckpointUrl: + """Test the _parse_hf_checkpoint_url helper.""" + + def test_blob_url(self) -> None: + repo, rev, fn = _parse_hf_checkpoint_url( + 'https://huggingface.co/qitaoz/E-RayZer/blob/main/checkpoints/erayzer_multi.pt', + ) + assert repo == 'qitaoz/E-RayZer' + assert rev == 'main' + assert fn == 'checkpoints/erayzer_multi.pt' + + def test_resolve_url_with_revision(self) -> None: + repo, rev, fn = _parse_hf_checkpoint_url( + 'https://huggingface.co/org/repo/resolve/v1.0/sub/dir/model.pt', + ) + assert repo == 'org/repo' + assert rev == 'v1.0' + assert fn == 'sub/dir/model.pt' + + def test_bare_path_defaults_to_main(self) -> None: + repo, rev, fn = _parse_hf_checkpoint_url('https://huggingface.co/org/repo/model.pt') + assert repo == 'org/repo' + assert rev == 'main' + assert fn == 'model.pt' + + +class TestResolveInitCheckpoint: + """Test the _resolve_init_checkpoint spec resolver.""" + + def test_local_path_passthrough(self, tmp_path) -> None: + p = tmp_path / 'weights.bin' + p.write_bytes(b'x') + assert _resolve_init_checkpoint(str(p)) == Path(p) + + def test_nonexistent_local_path_still_returns_path(self) -> None: + # resolution doesn't check existence (the loader does) + assert _resolve_init_checkpoint('/no/such/file.pt') == Path('/no/such/file.pt') + + +class TestLoadInitCheckpoint: + """Test ERayZer._load_init_checkpoint file handling and strict=False load.""" + + def test_missing_file_raises(self, tmp_path) -> None: + # the file-existence check runs before any use of self, so a stand-in + # object exercises the error path without building the full model + missing = tmp_path / 'nope.bin' + with pytest.raises(FileNotFoundError, match='init_checkpoint not found'): + ERayZer._load_init_checkpoint(cast(ERayZer, object()), str(missing)) + + def test_loads_matching_keys_strict_false(self, tmp_path) -> None: + # build a tiny module, save part of its state dict (plus a loss_computer + # key and an unexpected key), and confirm the loader maps what it can + module = nn.Sequential(nn.Linear(3, 4)) + ckpt = { + 'state_dict': { + '0.weight': torch.ones(4, 3), + '0.bias': torch.ones(4), + 'loss_computer.vgg.x': torch.zeros(2), # must be dropped + 'unexpected.param': torch.zeros(1), # tolerated by strict=False + }, + } + path = tmp_path / 'ckpt.bin' + torch.save(ckpt, path) + + result = ERayZer._load_init_checkpoint(cast(ERayZer, module), str(path)) + + linear = cast(nn.Linear, module[0]) + assert torch.allclose(linear.weight, torch.ones(4, 3)) + assert torch.allclose(linear.bias, torch.ones(4)) + # loss_computer.* dropped, so it is not reported as unexpected + assert result.unexpected_keys == ['unexpected.param'] diff --git a/tests/models/erayzer/test_visualize.py b/tests/models/erayzer/test_visualize.py new file mode 100644 index 0000000..a466c6b --- /dev/null +++ b/tests/models/erayzer/test_visualize.py @@ -0,0 +1,209 @@ +"""Tests for beast.models.erayzer.visualize.""" + +from types import SimpleNamespace + +import pytest +import torch + +from beast.models.erayzer.visualize import ( + camera_intrinsic_stats, + export_gaussian_glb, + make_camera_pose_image, + make_render_grid, + viz_is_due, +) + +try: + import trimesh # noqa: F401 # pyright: ignore[reportMissingImports] + _HAS_TRIMESH = True +except Exception: + _HAS_TRIMESH = False + +requires_trimesh = pytest.mark.skipif(not _HAS_TRIMESH, reason='trimesh not installed') + +# --------------------------------------------------------------------------- +# TestMakeRenderGrid +# --------------------------------------------------------------------------- + + +class TestMakeRenderGrid: + """Test the make_render_grid 2-row GT-vs-render grid.""" + + def test_returns_chw_float_in_unit_range(self) -> None: + inp = torch.rand(1, 2, 3, 8, 8) + tgt = torch.rand(1, 4, 3, 8, 8) + ren = torch.rand(1, 4, 3, 8, 8) + ren_in = torch.rand(1, 2, 3, 8, 8) + grid = make_render_grid(inp, tgt, ren, render_input=ren_in) + assert grid.ndim == 3 + assert grid.shape[0] == 3 + assert grid.dtype == torch.float32 + assert grid.min() >= 0.0 and grid.max() <= 1.0 + + def test_is_two_rows_tall(self) -> None: + # output must be exactly 2 image rows (GT, render) plus a thin separator + h = 8 + inp = torch.rand(1, 2, 3, h, h) + tgt = torch.rand(1, 3, 3, h, h) + ren = torch.rand(1, 3, 3, h, h) + ren_in = torch.rand(1, 2, 3, h, h) + grid = make_render_grid(inp, tgt, ren, render_input=ren_in) + assert 2 * h <= grid.shape[1] < 3 * h # two rows, not three + + def test_render_input_adds_columns_not_rows(self) -> None: + # input views add COLUMNS to the 2 rows, not extra rows + h = 8 + inp = torch.rand(1, 2, 3, h, h) + tgt = torch.rand(1, 3, 3, h, h) + ren = torch.rand(1, 3, 3, h, h) + ren_in = torch.rand(1, 2, 3, h, h) + base = make_render_grid(inp, tgt, ren) # target only + merged = make_render_grid(inp, tgt, ren, render_input=ren_in) # input + target + assert merged.shape[1] == base.shape[1] # same height (still 2 rows) + assert merged.shape[2] > base.shape[2] # wider (input columns added) + + def test_accepts_unbatched_inputs(self) -> None: + inp = torch.rand(2, 3, 8, 8) + tgt = torch.rand(3, 3, 8, 8) + ren = torch.rand(3, 3, 8, 8) + ren_in = torch.rand(2, 3, 8, 8) + grid = make_render_grid(inp, tgt, ren, render_input=ren_in) + assert grid.shape[0] == 3 + + def test_clamps_out_of_range_values(self) -> None: + inp = torch.full((1, 1, 3, 8, 8), 5.0) + tgt = torch.full((1, 1, 3, 8, 8), -5.0) + ren = torch.zeros(1, 1, 3, 8, 8) + grid = make_render_grid(inp, tgt, ren) + assert grid.min() >= 0.0 and grid.max() <= 1.0 + + +# --------------------------------------------------------------------------- +# TestExportGaussianGlb +# --------------------------------------------------------------------------- + + +class TestExportGaussianGlb: + """Test the export_gaussian_glb point-cloud writer.""" + + def _fake_gaussian(self, n: int = 16) -> SimpleNamespace: + # mimics GaussianModel: get_xyz [N,3], get_features [N,K,3], get_opacity [N,1] + return SimpleNamespace( + get_xyz=torch.rand(n, 3), + get_features=torch.rand(n, 1, 3), + get_opacity=torch.rand(n, 1), + ) + + @requires_trimesh + def test_writes_nonempty_glb(self, tmp_path) -> None: + path = tmp_path / 'cloud.glb' + n = export_gaussian_glb(self._fake_gaussian(16), path) + assert path.is_file() + assert path.stat().st_size > 0 + assert n == 16 + + @requires_trimesh + def test_opacity_threshold_filters_points(self, tmp_path) -> None: + g = self._fake_gaussian(20) + g.get_opacity = torch.cat([torch.ones(5, 1), torch.zeros(15, 1)]) + n = export_gaussian_glb(g, tmp_path / 'c.glb', opacity_threshold=0.5) + assert n == 5 + + @requires_trimesh + def test_returns_zero_and_skips_when_empty_after_filter(self, tmp_path) -> None: + g = self._fake_gaussian(8) + g.get_opacity = torch.zeros(8, 1) + path = tmp_path / 'empty.glb' + n = export_gaussian_glb(g, path, opacity_threshold=0.5) + assert n == 0 + assert not path.is_file() + + +# --------------------------------------------------------------------------- +# TestMakeCameraPoseImage +# --------------------------------------------------------------------------- + + +class TestMakeCameraPoseImage: + """Test the make_camera_pose_image function (input + target frustums).""" + + def _identity_cameras(self, v: int, offset: float = 0.0) -> torch.Tensor: + c2w = torch.eye(4).reshape(1, 1, 4, 4).repeat(1, v, 1, 1) + # spread the centers so the plot has extent + for i in range(v): + c2w[0, i, :3, 3] = torch.tensor([float(i) + offset, 0.0, 0.0]) + return c2w + + def test_input_and_target_returns_chw_float(self) -> None: + c2w_in = self._identity_cameras(2) + c2w_tg = self._identity_cameras(4, offset=0.5) + img = make_camera_pose_image(c2w_in, c2w_tg) + assert img.ndim == 3 + assert img.shape[0] == 3 + assert img.dtype == torch.float32 + assert img.min() >= 0.0 and img.max() <= 1.0 + + def test_input_only(self) -> None: + img = make_camera_pose_image(self._identity_cameras(4)) + assert img.shape[0] == 3 + + def test_accepts_unbatched_cameras(self) -> None: + c2w_in = self._identity_cameras(3)[0] # [V, 4, 4] + c2w_tg = self._identity_cameras(2)[0] + img = make_camera_pose_image(c2w_in, c2w_tg) + assert img.shape[0] == 3 + + +# --------------------------------------------------------------------------- +# TestVizIsDue +# --------------------------------------------------------------------------- + + +class TestVizIsDue: + """Test the viz_is_due cadence helper.""" + + def test_first_batch_on_due_epoch(self) -> None: + assert viz_is_due(batch_idx=0, current_epoch=0, every_n_epochs=1) is True + + def test_non_first_batch_is_skipped(self) -> None: + assert viz_is_due(batch_idx=2, current_epoch=0, every_n_epochs=1) is False + + def test_off_epoch_is_skipped(self) -> None: + assert viz_is_due(batch_idx=0, current_epoch=1, every_n_epochs=2) is False + + def test_on_epoch_cadence(self) -> None: + assert viz_is_due(batch_idx=0, current_epoch=4, every_n_epochs=2) is True + + def test_disabled_when_zero(self) -> None: + assert viz_is_due(batch_idx=0, current_epoch=0, every_n_epochs=0) is False + + +# --------------------------------------------------------------------------- +# TestCameraIntrinsicStats +# --------------------------------------------------------------------------- + + +class TestCameraIntrinsicStats: + """Test the camera_intrinsic_stats summary function.""" + + def test_normalizes_by_image_size(self) -> None: + # fxfycxcy in pixels at image_size=256: fx=fy=256, cx=cy=128 + fxfycxcy = torch.tensor([[256.0, 256.0, 128.0, 128.0]]) + stats = camera_intrinsic_stats(fxfycxcy, image_size=256) + assert stats['focal_fx_mean'] == 1.0 + assert stats['focal_fy_mean'] == 1.0 + assert stats['cx_mean'] == 0.5 + assert stats['cy_mean'] == 0.5 + + def test_reports_focal_spread(self) -> None: + fxfycxcy = torch.tensor([ + [256.0, 256.0, 128.0, 128.0], + [512.0, 512.0, 128.0, 128.0], + ]) + stats = camera_intrinsic_stats(fxfycxcy, image_size=256) + assert stats['focal_fx_std'] > 0.0 + + def test_returns_plain_floats(self) -> None: + fxfycxcy = torch.tensor([[256.0, 256.0, 128.0, 128.0]]) + stats = camera_intrinsic_stats(fxfycxcy, image_size=256) + assert all(isinstance(v, float) for v in stats.values()) diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py new file mode 100644 index 0000000..e323b14 --- /dev/null +++ b/tests/models/test_registry.py @@ -0,0 +1,64 @@ +"""Tests for the model registry.""" + +from collections.abc import Callable + +import pytest + +from beast.models.beast_resnet.beast_resnet_model import ResnetAutoencoder +from beast.models.beast_vit.beast_vit_model import VisionTransformer +from beast.models.erayzer.erayzer_model import ERayZer +from beast.models.registry import MODEL_REGISTRY, TRAIN_REGISTRY, get_model_class, get_train_fn + + +class TestGetModelClass: + """Test the get_model_class function.""" + + def test_resnet_returns_resnet_autoencoder(self) -> None: + assert get_model_class('resnet') is ResnetAutoencoder + + def test_vit_returns_vision_transformer(self) -> None: + assert get_model_class('vit') is VisionTransformer + + def test_erayzer_returns_erayzer(self) -> None: + assert get_model_class('erayzer') is ERayZer + + def test_unknown_raises_key_error(self) -> None: + with pytest.raises(KeyError): + get_model_class('nonexistent_model') + + def test_error_message_lists_registered_classes(self) -> None: + with pytest.raises(KeyError, match='resnet'): + get_model_class('nonexistent_model') + + def test_all_registered_keys_resolve(self) -> None: + for key in MODEL_REGISTRY: + assert get_model_class(key) is MODEL_REGISTRY[key] + + +class TestGetTrainFn: + """Test the get_train_fn function.""" + + def test_resnet_returns_callable(self) -> None: + fn = get_train_fn('resnet') + assert callable(fn) + + def test_vit_returns_callable(self) -> None: + fn = get_train_fn('vit') + assert callable(fn) + + def test_erayzer_returns_callable(self) -> None: + fn = get_train_fn('erayzer') + assert callable(fn) + + def test_unknown_raises_key_error(self) -> None: + with pytest.raises(KeyError): + get_train_fn('nonexistent_model') + + def test_error_message_lists_registered_classes(self) -> None: + with pytest.raises(KeyError, match='erayzer'): + get_train_fn('nonexistent_model') + + def test_all_registered_keys_resolve(self) -> None: + for key in TRAIN_REGISTRY: + fn = get_train_fn(key) + assert isinstance(fn, Callable) diff --git a/tests/nn/__init__.py b/tests/nn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/nn/test_dino.py b/tests/nn/test_dino.py new file mode 100644 index 0000000..978bcde --- /dev/null +++ b/tests/nn/test_dino.py @@ -0,0 +1,63 @@ +"""Tests for the DinoV3 feature extractor.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import torch + +from beast.nn.dino import DinoV3 + +_EMBED_DIM = 64 +_N_TOKENS = 20 # CLS + registers + patches; patch_tokens = hidden[:, 5:] + + +def _make_mock_automodel(embed_dim: int = _EMBED_DIM, n_tokens: int = _N_TOKENS): + """Return a mock that AutoModel.from_pretrained can be patched with.""" + mock_model = MagicMock() + mock_model.config = SimpleNamespace(hidden_size=embed_dim) + + def _forward(pixel_values): + bv = pixel_values.shape[0] + hidden = torch.zeros(bv, n_tokens, embed_dim) + return SimpleNamespace(last_hidden_state=hidden) + + mock_model.side_effect = _forward + mock_model.parameters.return_value = iter([torch.nn.Parameter(torch.zeros(1))]) + return mock_model + + +class TestDinoV3: + """Test the DinoV3 feature extractor.""" + + def test_embed_dim_set_from_model_config(self) -> None: + mock = _make_mock_automodel(embed_dim=128) + with patch('beast.nn.dino.AutoModel.from_pretrained', return_value=mock): + dino = DinoV3() + assert dino.embed_dim == 128 + + def test_freeze_true_freezes_parameters(self) -> None: + param = torch.nn.Parameter(torch.zeros(4)) + mock = _make_mock_automodel() + mock.parameters.return_value = iter([param]) + with patch('beast.nn.dino.AutoModel.from_pretrained', return_value=mock): + DinoV3(freeze=True) + assert not param.requires_grad + + def test_freeze_false_leaves_parameters_trainable(self) -> None: + param = torch.nn.Parameter(torch.zeros(4)) + mock = _make_mock_automodel() + mock.parameters.return_value = iter([param]) + with patch('beast.nn.dino.AutoModel.from_pretrained', return_value=mock): + DinoV3(freeze=False) + assert param.requires_grad + + def test_forward_output_shapes(self) -> None: + b, v, h, w = 2, 3, 16, 16 + n_patch_tokens = _N_TOKENS - 5 # hidden[:, 5:] + mock = _make_mock_automodel() + with patch('beast.nn.dino.AutoModel.from_pretrained', return_value=mock): + dino = DinoV3() + images = torch.rand(b, v, 3, h, w) + patch_tokens, cls_tokens = dino(images) + assert patch_tokens.shape == (b, v, n_patch_tokens, _EMBED_DIM) + assert cls_tokens.shape == (b, v, _EMBED_DIM) diff --git a/tests/models/test_perceptual.py b/tests/nn/test_perceptual.py similarity index 51% rename from tests/models/test_perceptual.py rename to tests/nn/test_perceptual.py index 0951eb0..0bfd65c 100644 --- a/tests/models/test_perceptual.py +++ b/tests/nn/test_perceptual.py @@ -2,30 +2,7 @@ import torch -from beast.models.perceptual import AlexPerceptual, Perceptual - - -class TestPerceptual: - """Test the base Perceptual class.""" - - def test_forward(self) -> None: - torch.manual_seed(0) - mock_net = torch.nn.Sequential( - torch.nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2), - torch.nn.ReLU(inplace=True), - ) - with torch.no_grad(): - mock_net[0].weight.fill_(0.01) # type: ignore[union-attr] - mock_net[0].bias.zero_() # type: ignore[union-attr] - criterion = torch.nn.MSELoss() - perceptual = Perceptual(network=mock_net, criterion=criterion) - torch.manual_seed(1) - x_hat = 0.01 * torch.randn((5, 3, 224, 224)) - x = 0.01 * torch.randn((5, 3, 224, 224)) - loss = perceptual(x_hat, x) - assert isinstance(loss, torch.Tensor) - assert loss.ndim == 0 - assert loss.item() >= 0 +from beast.nn.perceptual import AlexPerceptual, VGGPerceptual class TestAlexPerceptual: @@ -61,3 +38,47 @@ def test_different_inputs_produce_nonzero_loss(self) -> None: x_hat = torch.randn((2, 3, 224, 224), device=device) loss_diff = perceptual(x_hat, x) assert loss_diff.item() > 0 + + def test_no_gradients_through_alexnet(self) -> None: + device = 'cuda' if torch.cuda.is_available() else 'cpu' + perceptual = AlexPerceptual(device=device, criterion=torch.nn.MSELoss()) + for param in perceptual.net.parameters(): + assert not param.requires_grad + + +class TestVGGPerceptual: + """Test the VGGPerceptual class.""" + + def test_forward(self) -> None: + torch.manual_seed(0) + device = 'cuda' if torch.cuda.is_available() else 'cpu' + perceptual = VGGPerceptual(device=device) + x_hat = torch.rand((2, 3, 64, 64), device=device) + x = torch.rand((2, 3, 64, 64), device=device) + loss = perceptual(x_hat, x) + assert isinstance(loss, torch.Tensor) + assert loss.ndim == 0 + assert loss.item() >= 0 + + def test_identical_inputs_produce_near_zero_loss(self) -> None: + device = 'cuda' if torch.cuda.is_available() else 'cpu' + perceptual = VGGPerceptual(device=device) + torch.manual_seed(0) + x = torch.rand((2, 3, 64, 64), device=device) + loss = perceptual(x, x) + assert loss.item() < 1e-5 + + def test_different_inputs_produce_nonzero_loss(self) -> None: + torch.manual_seed(0) + device = 'cuda' if torch.cuda.is_available() else 'cpu' + perceptual = VGGPerceptual(device=device) + x = torch.rand((2, 3, 64, 64), device=device) + x_hat = torch.rand((2, 3, 64, 64), device=device) + loss = perceptual(x_hat, x) + assert loss.item() > 0 + + def test_no_gradients_through_vgg(self) -> None: + device = 'cuda' if torch.cuda.is_available() else 'cpu' + perceptual = VGGPerceptual(device=device) + for param in perceptual.parameters(): + assert not param.requires_grad diff --git a/tests/nn/test_transformer.py b/tests/nn/test_transformer.py new file mode 100644 index 0000000..2aea49c --- /dev/null +++ b/tests/nn/test_transformer.py @@ -0,0 +1,77 @@ +"""Tests for transformer building blocks.""" + +import pytest +import torch + +from beast.nn.transformer import MLP, QK_Norm_SelfAttention, QK_Norm_TransformerBlock, RMSNorm + +_B, _L, _D, _D_HEAD = 2, 8, 32, 8 # batch, seq len, dim, head dim + + +class TestMLP: + """Test the MLP module.""" + + def test_output_shape(self) -> None: + mlp = MLP(d=_D) + x = torch.randn(_B, _L, _D) + assert mlp(x).shape == (_B, _L, _D) + + def test_explicit_mlp_dim(self) -> None: + mlp = MLP(d=_D, mlp_dim=24) + x = torch.randn(_B, _L, _D) + assert mlp(x).shape == (_B, _L, _D) + + +class TestRMSNorm: + """Test the RMSNorm module.""" + + def test_output_shape_preserved(self) -> None: + norm = RMSNorm(dim=_D) + x = torch.randn(_B, _L, _D) + assert norm(x).shape == (_B, _L, _D) + + def test_reduces_rms_to_near_one(self) -> None: + norm = RMSNorm(dim=_D) + # with unit weight the output RMS should be close to 1 along the last dim + x = torch.randn(_B, _L, _D) * 5 + out = norm(x) + rms = out.pow(2).mean(dim=-1).sqrt() + assert torch.allclose(rms, torch.ones_like(rms), atol=1e-4) + + +class TestQK_Norm_SelfAttention: + """Test the QK_Norm_SelfAttention module.""" + + def test_output_shape(self) -> None: + attn = QK_Norm_SelfAttention(d=_D, d_head=_D_HEAD) + x = torch.randn(_B, _L, _D) + assert attn(x).shape == (_B, _L, _D) + + def test_output_shape_without_qk_norm(self) -> None: + attn = QK_Norm_SelfAttention(d=_D, d_head=_D_HEAD, use_qk_norm=False) + x = torch.randn(_B, _L, _D) + assert attn(x).shape == (_B, _L, _D) + + def test_invalid_head_dim_raises(self) -> None: + with pytest.raises(AssertionError): + QK_Norm_SelfAttention(d=_D, d_head=7) # 32 % 7 != 0 + + +class TestQK_Norm_TransformerBlock: + """Test the QK_Norm_TransformerBlock module.""" + + def test_output_shape(self) -> None: + block = QK_Norm_TransformerBlock(d=_D, d_head=_D_HEAD) + x = torch.randn(_B, _L, _D) + assert block(x).shape == (_B, _L, _D) + + def test_output_differs_from_input(self) -> None: + torch.manual_seed(0) + block = QK_Norm_TransformerBlock(d=_D, d_head=_D_HEAD) + x = torch.randn(_B, _L, _D) + assert not torch.allclose(block(x), x) + + def test_output_shape_without_qk_norm(self) -> None: + block = QK_Norm_TransformerBlock(d=_D, d_head=_D_HEAD, use_qk_norm=False) + x = torch.randn(_B, _L, _D) + assert block(x).shape == (_B, _L, _D) diff --git a/tests/rendering/__init__.py b/tests/rendering/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/rendering/test_gaussians_renderer.py b/tests/rendering/test_gaussians_renderer.py new file mode 100644 index 0000000..6c7a4f6 --- /dev/null +++ b/tests/rendering/test_gaussians_renderer.py @@ -0,0 +1,141 @@ +"""Tests for beast.rendering.gaussians_renderer utility functions.""" + +import numpy as np +import torch + +from beast.rendering.gaussians_renderer import ( + RGB2SH, + SH2RGB, + build_rotation, + build_scaling_rotation, + strip_lowerdiag, + strip_symmetric, +) + + +class TestStripLowerdiag: + """Test the strip_lowerdiag function.""" + + def test_output_shape(self) -> None: + L = torch.randn(5, 3, 3) + result = strip_lowerdiag(L) + assert result.shape == (5, 6) + + def test_extracts_correct_elements(self) -> None: + L = torch.arange(9).float().reshape(1, 3, 3) + result = strip_lowerdiag(L) + expected = torch.tensor([[0., 1., 2., 4., 5., 8.]]) + assert torch.allclose(result, expected) + + def test_single_element_batch(self) -> None: + L = torch.eye(3).unsqueeze(0) + result = strip_lowerdiag(L) + expected = torch.tensor([[1., 0., 0., 1., 0., 1.]]) + assert torch.allclose(result, expected) + + +class TestStripSymmetric: + """Test the strip_symmetric function.""" + + def test_delegates_to_strip_lowerdiag(self) -> None: + sym = torch.randn(4, 3, 3) + assert torch.allclose(strip_symmetric(sym), strip_lowerdiag(sym)) + + def test_output_shape(self) -> None: + sym = torch.randn(3, 3, 3) + assert strip_symmetric(sym).shape == (3, 6) + + +class TestBuildRotation: + """Test the build_rotation function.""" + + def test_output_shape(self) -> None: + quats = torch.randn(7, 4) + R = build_rotation(quats) + assert R.shape == (7, 3, 3) + + def test_identity_quaternion(self) -> None: + quat = torch.tensor([[1., 0., 0., 0.]]) + R = build_rotation(quat) + assert torch.allclose(R[0], torch.eye(3), atol=1e-6) + + def test_output_is_orthogonal(self) -> None: + torch.manual_seed(42) + quats = torch.randn(10, 4) + R = build_rotation(quats) + RRt = torch.bmm(R, R.transpose(1, 2)) + assert torch.allclose(RRt, torch.eye(3).unsqueeze(0).expand(10, -1, -1), atol=1e-5) + + def test_determinant_is_positive_one(self) -> None: + torch.manual_seed(42) + quats = torch.randn(10, 4) + R = build_rotation(quats) + dets = torch.linalg.det(R) + assert torch.allclose(dets, torch.ones(10), atol=1e-5) + + def test_90_degree_z_rotation(self) -> None: + quat = torch.tensor([[0.7071068, 0., 0., 0.7071068]]) + R = build_rotation(quat) + expected = torch.tensor([[0., -1., 0.], [1., 0., 0.], [0., 0., 1.]]) + assert torch.allclose(R[0], expected, atol=1e-5) + + +class TestBuildScalingRotation: + """Test the build_scaling_rotation function.""" + + def test_output_shape(self) -> None: + s = torch.ones(5, 3) + r = torch.tensor([[1., 0., 0., 0.]]).expand(5, -1) + result = build_scaling_rotation(s, r) + assert result.shape == (5, 3, 3) + + def test_identity_rotation_uniform_scale(self) -> None: + s = torch.tensor([[2., 2., 2.]]) + r = torch.tensor([[1., 0., 0., 0.]]) + result = build_scaling_rotation(s, r) + expected = 2.0 * torch.eye(3) + assert torch.allclose(result[0], expected, atol=1e-5) + + def test_identity_rotation_nonuniform_scale(self) -> None: + s = torch.tensor([[1., 2., 3.]]) + r = torch.tensor([[1., 0., 0., 0.]]) + result = build_scaling_rotation(s, r) + expected = torch.diag(torch.tensor([1., 2., 3.])) + assert torch.allclose(result[0], expected, atol=1e-5) + + +class TestRGB2SH: + """Test the RGB2SH function.""" + + def test_midgray_maps_to_zero(self) -> None: + rgb = torch.tensor([0.5, 0.5, 0.5]) + sh = RGB2SH(rgb) + assert torch.allclose(sh, torch.zeros(3), atol=1e-7) + + def test_output_shape_preserved(self) -> None: + rgb = torch.randn(8, 3) + assert RGB2SH(rgb).shape == (8, 3) + + +class TestSH2RGB: + """Test the SH2RGB function.""" + + def test_zero_maps_to_midgray(self) -> None: + sh = torch.tensor([0., 0., 0.]) + rgb = SH2RGB(sh) + assert torch.allclose(rgb, torch.tensor([0.5, 0.5, 0.5]), atol=1e-7) + + def test_roundtrip_with_rgb2sh(self) -> None: + torch.manual_seed(0) + rgb = torch.rand(10, 3) + assert torch.allclose(SH2RGB(RGB2SH(rgb)), rgb, atol=1e-6) + + def test_numpy_input(self) -> None: + sh = np.zeros(3, dtype=np.float32) + rgb = SH2RGB(sh) + assert isinstance(rgb, np.ndarray) + np.testing.assert_allclose(rgb, np.array([0.5, 0.5, 0.5]), atol=1e-7) + + def test_output_shape_preserved(self) -> None: + sh = torch.randn(8, 3) + assert SH2RGB(sh).shape == (8, 3) diff --git a/tests/test_config.py b/tests/test_config.py index 10c3fb7..1efb396 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -7,15 +7,21 @@ from beast.config import ( BeastConfig, + ERayZerBeastConfig, OptimizerConfig, - ResnetModelParams, TrainingConfig, - VitModelParams, + get_beast_config_class, ) from beast.io import load_config +from beast.models.beast_resnet.beast_resnet_config import ResnetModelParams +from beast.models.beast_vit.beast_vit_config import VitModelParams +from beast.models.erayzer.erayzer_config import ERayZerOptimizerConfig, ERayZerTrainingConfig _CONFIGS_DIR = Path(__file__).parent.parent / 'configs' -_CONFIG_FILES = list(_CONFIGS_DIR.glob('*.yaml')) +_NON_BEAST_CONFIG_NAMES = {'extraction_pipeline.yaml'} +_CONFIG_FILES = sorted( + p for p in _CONFIGS_DIR.rglob('*.yaml') if p.name not in _NON_BEAST_CONFIG_NAMES +) _MINIMAL_RESNET = { 'model': {'model_class': 'resnet', 'model_params': {}}, @@ -31,6 +37,23 @@ 'data': {'data_dir': '/path/to/data'}, } +_MINIMAL_ERAYZER = { + 'model': { + 'model_class': 'erayzer', + 'transformer': {'d': 256, 'd_head': 64, 'encoder_geom_n_layer': 4}, + }, + 'training': { + 'train_batch_size': 4, + 'val_batch_size': 2, + 'num_views': 3, + 'num_input_views': 2, + 'num_target_views': 1, + 'max_fwdbwd_passes': 50000, + }, + 'optimizer': {'lr': 4e-4}, + 'data': {'data_dir': '/path/to/data'}, +} + class TestBeastConfig: """Test the BeastConfig model.""" @@ -135,10 +158,93 @@ def test_defaults_applied(self) -> None: assert cfg.use_perceptual_loss is False +class TestERayZerTrainingConfig: + """Test the ERayZerTrainingConfig model.""" + + def test_defaults_applied(self) -> None: + cfg = ERayZerTrainingConfig( + train_batch_size=4, + val_batch_size=2, + num_views=3, + num_input_views=2, + num_target_views=1, + max_fwdbwd_passes=50000, + ) + assert cfg.num_epochs == 200 + assert cfg.seed == 0 + assert cfg.train_fraction == 0.9 + assert cfg.l2_loss_weight == 1.0 + assert cfg.random_split is False + assert cfg.ckpt_every_n_epochs is None + + def test_missing_required_fields_raises(self) -> None: + with pytest.raises(ValidationError): + ERayZerTrainingConfig(train_batch_size=4, val_batch_size=2) # type: ignore[call-arg] + + +class TestERayZerOptimizerConfig: + """Test the ERayZerOptimizerConfig model.""" + + def test_defaults_applied(self) -> None: + cfg = ERayZerOptimizerConfig(lr=4e-4) + assert cfg.beta1 == 0.9 + assert cfg.beta2 == 0.95 + assert cfg.wd == 0.05 + assert cfg.warmup == 3000 + assert cfg.div_factor == 1.0 + assert cfg.accumulate_grad_batches == 1 + + def test_missing_lr_raises(self) -> None: + with pytest.raises(ValidationError): + ERayZerOptimizerConfig() # type: ignore[call-arg] + + +class TestERayZerBeastConfig: + """Test the ERayZerBeastConfig model.""" + + def test_valid_config(self) -> None: + ERayZerBeastConfig.model_validate(_MINIMAL_ERAYZER) + + def test_missing_required_training_field_raises(self) -> None: + raw = {**_MINIMAL_ERAYZER, 'training': {'train_batch_size': 4, 'val_batch_size': 2}} + with pytest.raises(ValidationError): + ERayZerBeastConfig.model_validate(raw) + + def test_model_dump_contains_erayzer_fields(self) -> None: + config = ERayZerBeastConfig.model_validate(_MINIMAL_ERAYZER) + dumped = config.model_dump() + assert dumped['training']['num_views'] == 3 + assert dumped['training']['max_fwdbwd_passes'] == 50000 + assert 'beta1' in dumped['optimizer'] + + +class TestGetBeastConfigClass: + """Test the get_beast_config_class dispatcher.""" + + def test_erayzer_returns_erayzer_beast_config(self) -> None: + assert get_beast_config_class('erayzer') is ERayZerBeastConfig + + def test_resnet_falls_back_to_beast_config(self) -> None: + assert get_beast_config_class('resnet') is BeastConfig + + def test_vit_falls_back_to_beast_config(self) -> None: + assert get_beast_config_class('vit') is BeastConfig + + def test_unknown_falls_back_to_beast_config(self) -> None: + assert get_beast_config_class('totally_unknown_model') is BeastConfig + + def test_empty_string_falls_back_to_beast_config(self) -> None: + assert get_beast_config_class('') is BeastConfig + + class TestConfigFiles: """Validate all config files in the configs/ directory.""" - @pytest.mark.parametrize('config_path', _CONFIG_FILES, ids=[p.name for p in _CONFIG_FILES]) + @pytest.mark.parametrize( + 'config_path', + _CONFIG_FILES, + ids=[str(p.relative_to(_CONFIGS_DIR)) for p in _CONFIG_FILES], + ) def test_config_is_valid(self, config_path: Path) -> None: config = load_config(config_path) assert isinstance(config, dict) diff --git a/tests/test_losses.py b/tests/test_losses.py new file mode 100644 index 0000000..3c27ab3 --- /dev/null +++ b/tests/test_losses.py @@ -0,0 +1,48 @@ +"""Tests for beast.losses.""" + +import torch + +from beast.losses import masked_mse_loss + + +class TestMaskedMseLoss: + """Test the masked_mse_loss function.""" + + def test_returns_scalar(self) -> None: + rendering = torch.rand(2, 3, 4, 4) + target = torch.rand(2, 3, 4, 4) + mask = torch.ones(2, 1, 4, 4) + loss = masked_mse_loss(rendering, target, mask) + assert loss.ndim == 0 + + def test_identical_inputs_give_zero_loss(self) -> None: + x = torch.rand(2, 3, 4, 4) + mask = torch.ones(2, 1, 4, 4) + assert masked_mse_loss(x, x, mask).item() == 0.0 + + def test_full_mask_matches_standard_mse(self) -> None: + torch.manual_seed(0) + rendering = torch.rand(2, 3, 8, 8) + target = torch.rand(2, 3, 8, 8) + mask = torch.ones(2, 1, 8, 8) + loss = masked_mse_loss(rendering, target, mask) + expected = torch.nn.functional.mse_loss(rendering, target) + assert torch.allclose(loss, expected, atol=1e-6) + + def test_zero_mask_outside_region_ignores_those_pixels(self) -> None: + # only top-left 2x2 pixels are unmasked; error only in that region + rendering = torch.zeros(1, 3, 4, 4) + target = torch.zeros(1, 3, 4, 4) + target[:, :, :2, :2] = 1.0 # error only in the unmasked region + mask = torch.zeros(1, 1, 4, 4) + mask[:, :, :2, :2] = 1.0 + loss = masked_mse_loss(rendering, target, mask) + assert loss.item() == 1.0 + + def test_3d_mask_is_accepted(self) -> None: + rendering = torch.rand(2, 3, 4, 4) + target = torch.rand(2, 3, 4, 4) + mask_3d = torch.ones(2, 4, 4) # no channel dim + loss = masked_mse_loss(rendering, target, mask_3d) + assert loss.ndim == 0 + assert loss.item() >= 0 diff --git a/tests/test_train.py b/tests/test_train.py new file mode 100644 index 0000000..86a3323 --- /dev/null +++ b/tests/test_train.py @@ -0,0 +1,88 @@ +"""Tests for training helper functions.""" + +import os +import random + +import numpy as np +import torch +from lightning.pytorch import callbacks as pl_callbacks + +from beast.train import get_callbacks, reset_seeds + + +class TestResetSeeds: + """Test the reset_seeds function.""" + + def test_sets_pythonhashseed_env_var(self) -> None: + reset_seeds(42) + assert os.environ['PYTHONHASHSEED'] == '42' + + def test_torch_produces_same_values_after_reset(self) -> None: + reset_seeds(0) + a = torch.randn(5) + reset_seeds(0) + b = torch.randn(5) + assert torch.equal(a, b) + + def test_numpy_produces_same_values_after_reset(self) -> None: + reset_seeds(0) + a = np.random.rand(5) + reset_seeds(0) + b = np.random.rand(5) + assert np.array_equal(a, b) + + def test_python_random_produces_same_values_after_reset(self) -> None: + reset_seeds(0) + a = [random.random() for _ in range(5)] + reset_seeds(0) + b = [random.random() for _ in range(5)] + assert a == b + + def test_cudnn_flags_set(self) -> None: + reset_seeds(0) + assert torch.backends.cudnn.benchmark is False + assert torch.backends.cudnn.deterministic is True + + +class TestGetCallbacks: + """Test the get_callbacks function.""" + + def test_defaults_return_lr_monitor_and_best_checkpoint(self) -> None: + callbacks = get_callbacks() + assert len(callbacks) == 2 + types = {type(cb) for cb in callbacks} + assert pl_callbacks.LearningRateMonitor in types + assert pl_callbacks.ModelCheckpoint in types + + def test_no_lr_monitor(self) -> None: + callbacks = get_callbacks(lr_monitor=False) + assert not any(isinstance(cb, pl_callbacks.LearningRateMonitor) for cb in callbacks) + + def test_no_checkpointing(self) -> None: + callbacks = get_callbacks(checkpointing=False) + assert not any(isinstance(cb, pl_callbacks.ModelCheckpoint) for cb in callbacks) + + def test_all_disabled_returns_empty_list(self) -> None: + assert get_callbacks(checkpointing=False, lr_monitor=False) == [] + + def test_ckpt_every_n_epochs_adds_extra_checkpoint(self) -> None: + callbacks = get_callbacks(ckpt_every_n_epochs=5) + checkpoint_cbs = [cb for cb in callbacks if isinstance(cb, pl_callbacks.ModelCheckpoint)] + assert len(checkpoint_cbs) == 2 + + def test_best_checkpoint_monitors_val_loss(self) -> None: + callbacks = get_callbacks(lr_monitor=False) + ckpt = callbacks[0] + assert isinstance(ckpt, pl_callbacks.ModelCheckpoint) + assert ckpt.monitor == 'val_loss' + assert ckpt.mode == 'min' + + def test_periodic_checkpoint_saves_all(self) -> None: + callbacks = get_callbacks(lr_monitor=False, ckpt_every_n_epochs=3) + periodic = [ + cb for cb in callbacks + if isinstance(cb, pl_callbacks.ModelCheckpoint) and cb.monitor is None + ] + assert len(periodic) == 1 + assert periodic[0].every_n_epochs == 3 + assert periodic[0].save_top_k == -1 diff --git a/tests/test_video.py b/tests/test_video.py index bb5e494..e4685e9 100644 --- a/tests/test_video.py +++ b/tests/test_video.py @@ -18,6 +18,7 @@ downsample_video, get_frames_from_idxs, get_video_stats, + merge_videos, read_nth_frames, reencode_video, trim_video, @@ -386,3 +387,43 @@ def test_ffmpeg_failure_raises(self, tmp_path) -> None: with patch('beast.video.subprocess.run', return_value=mock_result): with pytest.raises(RuntimeError, match='ffmpeg downsample failed'): downsample_video(Path('in.mp4'), tmp_path / 'out.mp4', target_fps=5.0) + + +class TestMergeVideos: + """Test the merge_videos function.""" + + def test_empty_list_returns_without_writing(self, tmp_path) -> None: + output = tmp_path / 'out.mp4' + merge_videos([], str(output), fps=30) + assert not output.exists() + + def test_single_video_copies_without_ffmpeg(self, tmp_path) -> None: + src = tmp_path / 'src.mp4' + src.write_bytes(b'fake video bytes') + dst = tmp_path / 'out.mp4' + with patch('beast.video.subprocess.run') as mock_run: + merge_videos([str(src)], str(dst), fps=30) + mock_run.assert_not_called() + assert dst.read_bytes() == b'fake video bytes' + + def test_multiple_videos_calls_ffmpeg_concat(self, tmp_path) -> None: + mock_result = Mock(returncode=0) + with patch('beast.video.subprocess.run', return_value=mock_result) as mock_run: + merge_videos(['/a.mp4', '/b.mp4'], str(tmp_path / 'out.mp4'), fps=30) + cmd = mock_run.call_args[0][0] + assert 'ffmpeg' in cmd + assert '-f' in cmd + assert 'concat' in cmd + + def test_file_list_cleaned_up_on_success(self, tmp_path) -> None: + mock_result = Mock(returncode=0) + with patch('beast.video.subprocess.run', return_value=mock_result): + merge_videos(['/a.mp4', '/b.mp4'], str(tmp_path / 'out.mp4'), fps=30) + assert not (tmp_path / 'video_list.txt').exists() + + def test_ffmpeg_failure_raises_and_cleans_up(self, tmp_path) -> None: + mock_result = Mock(returncode=1, stderr='fatal error') + with patch('beast.video.subprocess.run', return_value=mock_result): + with pytest.raises(RuntimeError, match='ffmpeg merge failed'): + merge_videos(['/a.mp4', '/b.mp4'], str(tmp_path / 'out.mp4'), fps=30) + assert not (tmp_path / 'video_list.txt').exists()