optgs/model/decoder/decoder.py

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Generic, Literal, TypeVar

import torch
from jaxtyping import Float, Int32, Bool, UInt8
from torch import Tensor, nn

from ..types import Gaussians
from ...dataset import DatasetCfg
from ...dataset.data_types import BatchedViews, BatchedViewsDict, BatchedExample
from ...scene_trainer.gaussian_module import GaussiansModule


DepthRenderingMode = Literal[
    "depth",
    "log",
    "disparity",
    "relative_disparity",
]


@dataclass
class DecoderOutput:
    color: Float[Tensor, "batch view 3 height width"] | UInt8[Tensor, "batch view 3 height width"]
    depth: Float[Tensor, "batch view height width"] | None
    normal: Float[Tensor, "batch view 3 height width"] | None = None
    distortion_map: Float[Tensor, "batch view height width"] | None = None
    accumulated_alpha: Float[Tensor, "batch view height width"] | None = None
    radii: Int32[Tensor, "batch view n 2"] | None = None
    means2d: Float[Tensor, "batch view n 2"] | None = None
    # FastGS Abs-GS split signal: gradient of the *absolute* screen-space mean (cols [2:] of the
    # FastGS [N,4] screen tensor). Only the FastGS decoder populates it; reachable via autograd.
    means2d_abs: Float[Tensor, "batch view n 2"] | None = None
    visibility_filter: Bool[Tensor, "batch view n"] | None = None


T = TypeVar("T")


class Decoder(nn.Module, ABC, Generic[T]):
    cfg: T
    dataset_cfg: DatasetCfg

    def __init__(self, cfg: T, dataset_cfg: DatasetCfg) -> None:
        super().__init__()
        self.cfg = cfg
        self.dataset_cfg = dataset_cfg

    def means2d_grad_to_ndc(
        self,
        grad: Float[Tensor, "*batch n 2"],
        image_shape: tuple[int, int],
    ) -> Float[Tensor, "*batch n 2"]:
        """Normalize an autograd.grad(loss, decoder_output.means2d) result to the resolution-
        independent NDC ([-1, 1]) screen convention.

        This makes the ADC / densification strategy renderer-agnostic: every backend hands it a
        uniform NDC gradient, so a single threshold (the 3DGS / FastGS
        ``densify_grad_threshold = 0.0002``) is correct for all of them, and the strategy itself no
        longer needs to know each renderer's pixel scale.

        The 3DGS-family backends (inria/fastgs) already emit NDC gradients, so the default is
        identity. Only the gsplat backend (pixel-space, gradient ∝ image size) overrides this."""
        return grad

    @abstractmethod
    def forward(
            self,
            gaussians: Gaussians | GaussiansModule,
            extrinsics: Float[Tensor, "batch view 4 4"],
            intrinsics: Float[Tensor, "batch view 3 3"],
            near: Float[Tensor, "batch view"],
            far: Float[Tensor, "batch view"],
            image_shape: tuple[int, int],
            depth_mode: DepthRenderingMode | None = None,
            to_cpu: bool = False,
    ) -> DecoderOutput:
        pass

    def forward_batch(
            self,
            gaussians: Gaussians | GaussiansModule,
            batch: BatchedExample,
            image_shape: tuple[int, int] | None = None,
            input_str: Literal["context", "target"] | None = None,
            eval_context_views: bool | None = None,
            depth_mode: DepthRenderingMode | None = None,
            start=None, end=None,
            camera_poses=None,  # In case of manipulating camera poses (e.g. for stabilization)
            to_cpu: bool = False,  # move outputs to cpu as they are rendered
            iter_batch_size: int = -1,  # -1 to render all views at once
    ) -> DecoderOutput:

        assert input_str is not None or eval_context_views is not None
        if input_str is None:
            input_str = "context" if eval_context_views else "target"

        input = batch[input_str]

        if image_shape is None:
            image_shape = input["image_shape"].shape[-2:]
        if camera_poses is None:
            camera_poses = input["extrinsics"]
        return self.forward(
            gaussians,
            camera_poses[:, start:end],
            input["intrinsics"][:, start:end],
            input["near"][:, start:end],
            input["far"][:, start:end],
            image_shape,
            depth_mode=depth_mode,
            to_cpu=to_cpu,
            iter_batch_size=iter_batch_size,
        )

    def forward_batch_subset(self, gaussians: Gaussians | GaussiansModule,
                             batch_subset: BatchedViewsDict | BatchedViews,
                             image_shape: tuple[int, int] | None = None,
                             start: int | None = None,
                             end: int | None = None,
                             indices: torch.Tensor | list | None = None,
                             **kwargs) -> DecoderOutput:

        assert not ((start is not None and end is not None) and (
                indices is not None)), "Either start and end or indices must be provided."
        if start is not None:
            indices = list(range(start, end))

        if indices is None:
            indices = list(range(batch_subset["extrinsics"].shape[1]))

        if isinstance(indices, list):
            # Convert list to tensor for one flow handling
            indices = torch.tensor(indices, device=batch_subset["extrinsics"].device)
            indices = indices.unsqueeze(0).expand(batch_subset["extrinsics"].shape[0], -1)  # (batch, num_indices)

        if image_shape is None:
            image_shape = batch_subset["image"].shape[-2:]

        assert indices.dim() == 2, "Indices tensor must be 2D (scene_batch, num_indices)."
        scene_batch = indices.size(0)
        scene_batch_idx = torch.arange(scene_batch, device=indices.device)[:, None]  # (batch, 1)
        return self.forward(gaussians,
                            batch_subset["extrinsics"][scene_batch_idx, indices],
                            batch_subset["intrinsics"][scene_batch_idx, indices],
                            batch_subset["near"][scene_batch_idx, indices],
                            batch_subset["far"][scene_batch_idx, indices],
                            image_shape,
                            **kwargs)

    def forward_context(
            self,
            gaussians: Gaussians | GaussiansModule,
            batch: BatchedExample,
            image_shape: tuple[int, int] | None = None,
            depth_mode: DepthRenderingMode | None = None,
            **kwargs,
    ) -> DecoderOutput:
        return self.forward_batch(
            gaussians,
            batch,
            image_shape,
            "context",
            depth_mode=depth_mode,
            **kwargs,
        )

    def forward_target(
            self,
            gaussians: Gaussians | GaussiansModule,
            batch: BatchedExample,
            image_shape: tuple[int, int] | None = None,
            depth_mode: DepthRenderingMode | None = None,
            **kwargs,
    ) -> DecoderOutput:
        return self.forward_batch(
            gaussians,
            batch,
            image_shape,
            "target",
            depth_mode=depth_mode,
            **kwargs,
        )