src/mcss/types.py
"""Typed tensor contracts shared by datasets, models, and measurements."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Self
import torch
from torch import Tensor
def _require_finite(name: str, value: Tensor) -> None:
if not torch.isfinite(value).all():
raise ValueError(f"{name} must contain only finite values")
@dataclass
class Cameras:
"""Calibrated OpenCV cameras with arbitrary leading batch/view dimensions."""
intrinsics: Tensor
c2w: Tensor
image_size: tuple[int, int]
def __post_init__(self) -> None:
height, width = self.image_size
if height <= 0 or width <= 0:
raise ValueError("image_size values must be positive")
if self.intrinsics.ndim < 2 or self.intrinsics.shape[-2:] != (3, 3):
raise ValueError("intrinsics must have shape [..., 3, 3]")
if self.c2w.ndim < 2 or self.c2w.shape[-2:] != (4, 4):
raise ValueError("c2w must have shape [..., 4, 4]")
if self.intrinsics.shape[:-2] != self.c2w.shape[:-2]:
raise ValueError("intrinsics and c2w must share leading dimensions")
if self.intrinsics.device != self.c2w.device:
raise ValueError("intrinsics and c2w must be on the same device")
_require_finite("intrinsics", self.intrinsics)
_require_finite("c2w", self.c2w)
expected_row = torch.tensor(
[0.0, 0.0, 0.0, 1.0], device=self.c2w.device, dtype=self.c2w.dtype
)
if not torch.allclose(self.c2w[..., 3, :], expected_row, atol=1e-5, rtol=1e-5):
raise ValueError("c2w must have a homogeneous final row [0, 0, 0, 1]")
if (self.intrinsics[..., 0, 0] <= 0).any() or (self.intrinsics[..., 1, 1] <= 0).any():
raise ValueError("intrinsics must have positive focal lengths")
@property
def leading_shape(self) -> tuple[int, ...]:
return tuple(self.intrinsics.shape[:-2])
@property
def device(self) -> torch.device:
return self.intrinsics.device
@property
def dtype(self) -> torch.dtype:
return self.intrinsics.dtype
def to(self, *args: object, **kwargs: object) -> Self:
return Cameras(
self.intrinsics.to(*args, **kwargs),
self.c2w.to(*args, **kwargs),
self.image_size,
)
def select_views(self, indices: Tensor) -> Self:
"""Select the final leading (view) dimension without changing camera convention."""
if not self.leading_shape:
raise ValueError("cameras require a view dimension for select_views")
if indices.ndim != 1:
raise ValueError("indices must be one-dimensional")
indices = indices.to(device=self.device, dtype=torch.long)
return Cameras(
self.intrinsics[..., indices, :, :],
self.c2w[..., indices, :, :],
self.image_size,
)
@dataclass
class StateEvidence:
"""Inspectable evidence/completion fields aligned with one typed 3D state."""
confidence: Tensor
unknown_probability: Tensor
completion_gate: Tensor
provenance: Tensor
base_density_logits: Tensor
base_color: Tensor
density_residual: Tensor
color_logit_residual: Tensor
def __post_init__(self) -> None:
scalar_fields = {
"confidence": self.confidence,
"unknown_probability": self.unknown_probability,
"completion_gate": self.completion_gate,
"base_density_logits": self.base_density_logits,
"density_residual": self.density_residual,
}
color_fields = {
"base_color": self.base_color,
"color_logit_residual": self.color_logit_residual,
}
for name, value in {**scalar_fields, **color_fields, "provenance": self.provenance}.items():
if value.ndim != 5:
raise ValueError(f"{name} must have shape [B, C, D, H, W]")
_require_finite(name, value)
if any(value.shape[1] != 1 for value in scalar_fields.values()):
raise ValueError("scalar evidence fields require one channel")
if any(value.shape[1] != 3 for value in color_fields.values()):
raise ValueError("color evidence fields require three channels")
if self.provenance.shape[1] < 1:
raise ValueError("provenance requires at least one source-view channel")
reference = self.confidence
for name, value in scalar_fields.items():
if value.shape != reference.shape:
raise ValueError(f"{name} must match confidence shape")
for name, value in color_fields.items():
if value.shape[0] != reference.shape[0] or value.shape[2:] != reference.shape[2:]:
raise ValueError(f"{name} must share confidence batch and spatial dimensions")
if (
self.provenance.shape[0] != reference.shape[0]
or self.provenance.shape[2:] != reference.shape[2:]
):
raise ValueError("provenance must share confidence batch and spatial dimensions")
devices = {value.device for value in (*scalar_fields.values(), *color_fields.values())}
devices.add(self.provenance.device)
if len(devices) != 1:
raise ValueError("evidence fields must be on one device")
for name, value in {
"confidence": self.confidence,
"unknown_probability": self.unknown_probability,
"completion_gate": self.completion_gate,
"base_color": self.base_color,
"provenance": self.provenance,
}.items():
if ((value < 0.0) | (value > 1.0)).any():
raise ValueError(f"{name} must be within [0, 1]")
tolerance = 5e-3 if self.confidence.dtype in {torch.float16, torch.bfloat16} else 1e-5
if not torch.allclose(
self.confidence + self.unknown_probability,
torch.ones_like(self.confidence),
atol=tolerance,
rtol=tolerance,
):
raise ValueError("unknown_probability must equal one minus confidence")
provenance_sum = self.provenance.sum(dim=1, keepdim=True)
normalized = torch.isclose(
provenance_sum,
torch.ones_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
empty = torch.isclose(
provenance_sum,
torch.zeros_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
if not (normalized | empty).all() or (empty & (self.confidence > tolerance)).any():
raise ValueError("provenance must sum to one, or zero where confidence is zero")
def to(self, *args: object, **kwargs: object) -> Self:
return StateEvidence(
self.confidence.to(*args, **kwargs),
self.unknown_probability.to(*args, **kwargs),
self.completion_gate.to(*args, **kwargs),
self.provenance.to(*args, **kwargs),
self.base_density_logits.to(*args, **kwargs),
self.base_color.to(*args, **kwargs),
self.density_residual.to(*args, **kwargs),
self.color_logit_residual.to(*args, **kwargs),
)
@dataclass
class DualEvidence:
"""Separate deterministic surface-localization and appearance evidence."""
surface_localization_support: Tensor
appearance_confidence: Tensor
surface_peakness: Tensor
localization_unknown_probability: Tensor
appearance_unknown_probability: Tensor
localization_completion_gate: Tensor
appearance_completion_gate: Tensor
provenance: Tensor
base_density_logits: Tensor
base_color: Tensor
density_residual: Tensor
color_logit_residual: Tensor
def __post_init__(self) -> None:
scalar_fields = {
"surface_localization_support": self.surface_localization_support,
"appearance_confidence": self.appearance_confidence,
"surface_peakness": self.surface_peakness,
"localization_unknown_probability": self.localization_unknown_probability,
"appearance_unknown_probability": self.appearance_unknown_probability,
"localization_completion_gate": self.localization_completion_gate,
"appearance_completion_gate": self.appearance_completion_gate,
"base_density_logits": self.base_density_logits,
"density_residual": self.density_residual,
}
color_fields = {
"base_color": self.base_color,
"color_logit_residual": self.color_logit_residual,
}
fields = {**scalar_fields, **color_fields, "provenance": self.provenance}
for name, value in fields.items():
if value.ndim != 5:
raise ValueError(f"{name} must have shape [B, C, D, H, W]")
_require_finite(name, value)
if any(value.shape[1] != 1 for value in scalar_fields.values()):
raise ValueError("scalar dual-evidence fields require one channel")
if any(value.shape[1] != 3 for value in color_fields.values()):
raise ValueError("color dual-evidence fields require three channels")
if self.provenance.shape[1] < 1:
raise ValueError("provenance requires at least one source-view channel")
reference = self.appearance_confidence
for name, value in scalar_fields.items():
if value.shape != reference.shape:
raise ValueError(f"{name} must match appearance_confidence shape")
for name, value in color_fields.items():
if value.shape[0] != reference.shape[0] or value.shape[2:] != reference.shape[2:]:
raise ValueError(
f"{name} must share appearance_confidence batch and spatial dimensions"
)
if (
self.provenance.shape[0] != reference.shape[0]
or self.provenance.shape[2:] != reference.shape[2:]
):
raise ValueError(
"provenance must share appearance_confidence batch and spatial dimensions"
)
if len({value.device for value in fields.values()}) != 1:
raise ValueError("dual-evidence fields must be on one device")
bounded = {
"surface_localization_support": self.surface_localization_support,
"appearance_confidence": self.appearance_confidence,
"surface_peakness": self.surface_peakness,
"localization_unknown_probability": self.localization_unknown_probability,
"appearance_unknown_probability": self.appearance_unknown_probability,
"localization_completion_gate": self.localization_completion_gate,
"appearance_completion_gate": self.appearance_completion_gate,
"base_color": self.base_color,
"provenance": self.provenance,
}
for name, value in bounded.items():
if ((value < 0.0) | (value > 1.0)).any():
raise ValueError(f"{name} must be within [0, 1]")
tolerance = 5e-3 if reference.dtype in {torch.float16, torch.bfloat16} else 1e-5
if (self.surface_localization_support > reference + tolerance).any():
raise ValueError(
"surface_localization_support cannot exceed appearance_confidence"
)
if not torch.allclose(
self.surface_localization_support,
reference * self.surface_peakness,
atol=tolerance,
rtol=tolerance,
):
raise ValueError(
"surface_localization_support must equal appearance_confidence "
"times surface_peakness"
)
for confidence, unknown, name in (
(
self.surface_localization_support,
self.localization_unknown_probability,
"localization_unknown_probability",
),
(
self.appearance_confidence,
self.appearance_unknown_probability,
"appearance_unknown_probability",
),
):
if not torch.allclose(
confidence + unknown,
torch.ones_like(confidence),
atol=tolerance,
rtol=tolerance,
):
raise ValueError(f"{name} must equal one minus its confidence")
provenance_sum = self.provenance.sum(dim=1, keepdim=True)
normalized = torch.isclose(
provenance_sum,
torch.ones_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
empty = torch.isclose(
provenance_sum,
torch.zeros_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
if not (normalized | empty).all() or (
empty & (self.appearance_confidence > tolerance)
).any():
raise ValueError(
"provenance must sum to one, or zero where appearance confidence is zero"
)
@property
def density_logits(self) -> Tensor:
return self.base_density_logits + (
self.localization_completion_gate * self.density_residual
)
@property
def color(self) -> Tensor:
base_logits = torch.logit(self.base_color.clamp(1e-4, 1.0 - 1e-4))
return torch.sigmoid(
base_logits + self.appearance_completion_gate * self.color_logit_residual
)
def to(self, *args: object, **kwargs: object) -> Self:
return DualEvidence(
self.surface_localization_support.to(*args, **kwargs),
self.appearance_confidence.to(*args, **kwargs),
self.surface_peakness.to(*args, **kwargs),
self.localization_unknown_probability.to(*args, **kwargs),
self.appearance_unknown_probability.to(*args, **kwargs),
self.localization_completion_gate.to(*args, **kwargs),
self.appearance_completion_gate.to(*args, **kwargs),
self.provenance.to(*args, **kwargs),
self.base_density_logits.to(*args, **kwargs),
self.base_color.to(*args, **kwargs),
self.density_residual.to(*args, **kwargs),
self.color_logit_residual.to(*args, **kwargs),
)
@dataclass
class TransportEvidence:
"""Fixed ray evidence transported from source voxels to candidate 3D positions."""
surface_localization_support: Tensor
transported_appearance_confidence: Tensor
surface_peakness: Tensor
transport_coverage: Tensor
appearance_confidence: Tensor
localization_unknown_probability: Tensor
appearance_unknown_probability: Tensor
localization_completion_gate: Tensor
appearance_completion_gate: Tensor
provenance: Tensor
base_density_logits: Tensor
base_color: Tensor
density_residual: Tensor
color_logit_residual: Tensor
def __post_init__(self) -> None:
scalar_fields = {
"surface_localization_support": self.surface_localization_support,
"transported_appearance_confidence": self.transported_appearance_confidence,
"surface_peakness": self.surface_peakness,
"transport_coverage": self.transport_coverage,
"appearance_confidence": self.appearance_confidence,
"localization_unknown_probability": self.localization_unknown_probability,
"appearance_unknown_probability": self.appearance_unknown_probability,
"localization_completion_gate": self.localization_completion_gate,
"appearance_completion_gate": self.appearance_completion_gate,
"base_density_logits": self.base_density_logits,
"density_residual": self.density_residual,
}
color_fields = {
"base_color": self.base_color,
"color_logit_residual": self.color_logit_residual,
}
fields = {**scalar_fields, **color_fields, "provenance": self.provenance}
for name, value in fields.items():
if value.ndim != 5:
raise ValueError(f"{name} must have shape [B, C, D, H, W]")
_require_finite(name, value)
if any(value.shape[1] != 1 for value in scalar_fields.values()):
raise ValueError("scalar transport-evidence fields require one channel")
if any(value.shape[1] != 3 for value in color_fields.values()):
raise ValueError("color transport-evidence fields require three channels")
if self.provenance.shape[1] < 1:
raise ValueError("provenance requires at least one source-view channel")
reference = self.transported_appearance_confidence
for name, value in scalar_fields.items():
if value.shape != reference.shape:
raise ValueError(f"{name} must match transported_appearance_confidence shape")
for name, value in color_fields.items():
if value.shape[0] != reference.shape[0] or value.shape[2:] != reference.shape[2:]:
raise ValueError(
f"{name} must share transported evidence batch and spatial dimensions"
)
if (
self.provenance.shape[0] != reference.shape[0]
or self.provenance.shape[2:] != reference.shape[2:]
):
raise ValueError("provenance must share transported evidence spatial dimensions")
if len({value.device for value in fields.values()}) != 1:
raise ValueError("transport-evidence fields must be on one device")
bounded = {
"surface_localization_support": self.surface_localization_support,
"transported_appearance_confidence": self.transported_appearance_confidence,
"surface_peakness": self.surface_peakness,
"appearance_confidence": self.appearance_confidence,
"localization_unknown_probability": self.localization_unknown_probability,
"appearance_unknown_probability": self.appearance_unknown_probability,
"localization_completion_gate": self.localization_completion_gate,
"appearance_completion_gate": self.appearance_completion_gate,
"base_color": self.base_color,
"provenance": self.provenance,
}
for name, value in bounded.items():
if ((value < 0.0) | (value > 1.0)).any():
raise ValueError(f"{name} must be within [0, 1]")
if (self.transport_coverage < 0.0).any():
raise ValueError("transport_coverage must be nonnegative")
tolerance = 5e-3 if reference.dtype in {torch.float16, torch.bfloat16} else 1e-5
if (self.surface_localization_support > reference + tolerance).any():
raise ValueError(
"surface_localization_support cannot exceed "
"transported_appearance_confidence"
)
if not torch.allclose(
self.surface_localization_support,
reference * self.surface_peakness,
atol=tolerance,
rtol=tolerance,
):
raise ValueError(
"surface_localization_support must equal transported_appearance_confidence "
"times surface_peakness"
)
for confidence, unknown, name in (
(
self.surface_localization_support,
self.localization_unknown_probability,
"localization_unknown_probability",
),
(
self.appearance_confidence,
self.appearance_unknown_probability,
"appearance_unknown_probability",
),
):
if not torch.allclose(
confidence + unknown,
torch.ones_like(confidence),
atol=tolerance,
rtol=tolerance,
):
raise ValueError(f"{name} must equal one minus its confidence")
provenance_sum = self.provenance.sum(dim=1, keepdim=True)
normalized = torch.isclose(
provenance_sum,
torch.ones_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
empty = torch.isclose(
provenance_sum,
torch.zeros_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
if not (normalized | empty).all() or (
empty & (self.appearance_confidence > tolerance)
).any():
raise ValueError(
"provenance must sum to one, or zero where appearance confidence is zero"
)
@property
def density_logits(self) -> Tensor:
return self.base_density_logits + (
self.localization_completion_gate * self.density_residual
)
@property
def color(self) -> Tensor:
base_logits = torch.logit(self.base_color.clamp(1e-4, 1.0 - 1e-4))
return torch.sigmoid(
base_logits + self.appearance_completion_gate * self.color_logit_residual
)
def to(self, *args: object, **kwargs: object) -> Self:
return TransportEvidence(
self.surface_localization_support.to(*args, **kwargs),
self.transported_appearance_confidence.to(*args, **kwargs),
self.surface_peakness.to(*args, **kwargs),
self.transport_coverage.to(*args, **kwargs),
self.appearance_confidence.to(*args, **kwargs),
self.localization_unknown_probability.to(*args, **kwargs),
self.appearance_unknown_probability.to(*args, **kwargs),
self.localization_completion_gate.to(*args, **kwargs),
self.appearance_completion_gate.to(*args, **kwargs),
self.provenance.to(*args, **kwargs),
self.base_density_logits.to(*args, **kwargs),
self.base_color.to(*args, **kwargs),
self.density_residual.to(*args, **kwargs),
self.color_logit_residual.to(*args, **kwargs),
)
@dataclass
class StateAppearance:
"""Inspectable high-resolution color evidence and bounded completion residual."""
confidence: Tensor
unknown_probability: Tensor
completion_gate: Tensor
provenance: Tensor
base_color: Tensor
color_logit_residual: Tensor
def __post_init__(self) -> None:
scalar_fields = {
"confidence": self.confidence,
"unknown_probability": self.unknown_probability,
"completion_gate": self.completion_gate,
}
color_fields = {
"base_color": self.base_color,
"color_logit_residual": self.color_logit_residual,
}
fields = {**scalar_fields, **color_fields, "provenance": self.provenance}
for name, value in fields.items():
if value.ndim != 5:
raise ValueError(f"{name} must have shape [B, C, D, H, W]")
_require_finite(name, value)
if any(value.shape[1] != 1 for value in scalar_fields.values()):
raise ValueError("scalar appearance fields require one channel")
if any(value.shape[1] != 3 for value in color_fields.values()):
raise ValueError("color appearance fields require three channels")
if self.provenance.shape[1] < 1:
raise ValueError("provenance requires at least one source-view channel")
reference = self.confidence
for name, value in scalar_fields.items():
if value.shape != reference.shape:
raise ValueError(f"{name} must match confidence shape")
for name, value in color_fields.items():
if value.shape[0] != reference.shape[0] or value.shape[2:] != reference.shape[2:]:
raise ValueError(f"{name} must share confidence batch and spatial dimensions")
if (
self.provenance.shape[0] != reference.shape[0]
or self.provenance.shape[2:] != reference.shape[2:]
):
raise ValueError("provenance must share confidence batch and spatial dimensions")
devices = {value.device for value in fields.values()}
if len(devices) != 1:
raise ValueError("appearance fields must be on one device")
for name, value in {
"confidence": self.confidence,
"unknown_probability": self.unknown_probability,
"completion_gate": self.completion_gate,
"base_color": self.base_color,
"provenance": self.provenance,
}.items():
if ((value < 0.0) | (value > 1.0)).any():
raise ValueError(f"{name} must be within [0, 1]")
tolerance = 5e-3 if self.confidence.dtype in {torch.float16, torch.bfloat16} else 1e-5
if not torch.allclose(
self.confidence + self.unknown_probability,
torch.ones_like(self.confidence),
atol=tolerance,
rtol=tolerance,
):
raise ValueError("unknown_probability must equal one minus confidence")
provenance_sum = self.provenance.sum(dim=1, keepdim=True)
normalized = torch.isclose(
provenance_sum,
torch.ones_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
empty = torch.isclose(
provenance_sum,
torch.zeros_like(provenance_sum),
atol=tolerance,
rtol=tolerance,
)
if not (normalized | empty).all() or (empty & (self.confidence > tolerance)).any():
raise ValueError("provenance must sum to one, or zero where confidence is zero")
@property
def color(self) -> Tensor:
base_logits = torch.logit(self.base_color.clamp(1e-4, 1.0 - 1e-4))
return torch.sigmoid(base_logits + self.completion_gate * self.color_logit_residual)
@property
def spatial_shape(self) -> tuple[int, int, int]:
return tuple(self.confidence.shape[-3:])
def to(self, *args: object, **kwargs: object) -> Self:
return StateAppearance(
self.confidence.to(*args, **kwargs),
self.unknown_probability.to(*args, **kwargs),
self.completion_gate.to(*args, **kwargs),
self.provenance.to(*args, **kwargs),
self.base_color.to(*args, **kwargs),
self.color_logit_residual.to(*args, **kwargs),
)
@dataclass
class SceneState:
"""Explicit fields read by the fixed measurement program."""
density_logits: Tensor
color: Tensor
log_variance: Tensor
bounds: Tensor
features: Tensor | None = None
evidence: StateEvidence | None = None
appearance: StateAppearance | None = None
dual_evidence: DualEvidence | None = None
transport_evidence: TransportEvidence | None = None
def __post_init__(self) -> None:
tensors = {
"density_logits": self.density_logits,
"color": self.color,
"log_variance": self.log_variance,
}
for name, value in tensors.items():
if value.ndim != 5:
raise ValueError(f"{name} must have shape [B, C, D, H, W]")
_require_finite(name, value)
if self.density_logits.shape[1] != 1 or self.log_variance.shape[1] != 1:
raise ValueError("density_logits and log_variance require one channel")
if self.color.shape[1] != 3:
raise ValueError("color requires three RGB channels")
if self.color.shape[0] != self.density_logits.shape[0]:
raise ValueError("typed fields must share batch size")
if self.color.shape[2:] != self.density_logits.shape[2:]:
raise ValueError("typed fields must share spatial dimensions")
if self.log_variance.shape != self.density_logits.shape:
raise ValueError("log_variance must match density_logits shape")
if self.bounds.shape != (self.density_logits.shape[0], 2, 3):
raise ValueError("bounds must have shape [B, 2, 3]")
_require_finite("bounds", self.bounds)
if not (self.bounds[:, 1] > self.bounds[:, 0]).all():
raise ValueError("bounds maximum must exceed minimum on every axis")
if self.features is not None:
if self.features.ndim != 5 or self.features.shape[0] != self.density_logits.shape[0]:
raise ValueError("features must have shape [B, C, D, H, W]")
if self.features.shape[2:] != self.density_logits.shape[2:]:
raise ValueError("features must share state spatial dimensions")
_require_finite("features", self.features)
native_evidence_count = sum(
value is not None
for value in (self.evidence, self.dual_evidence, self.transport_evidence)
)
if native_evidence_count > 1:
raise ValueError("native evidence payloads are mutually exclusive")
if self.evidence is not None:
if self.evidence.confidence.shape != self.density_logits.shape:
raise ValueError("evidence confidence must match density_logits shape")
if self.evidence.base_color.shape != self.color.shape:
raise ValueError("evidence base_color must match color shape")
if self.evidence.confidence.device != self.density_logits.device:
raise ValueError("evidence and typed state fields must be on the same device")
if self.appearance is not None:
if self.appearance.confidence.shape[0] != self.density_logits.shape[0]:
raise ValueError("appearance and typed state fields must share batch size")
if self.appearance.confidence.device != self.density_logits.device:
raise ValueError("appearance and typed state fields must be on the same device")
if self.dual_evidence is not None:
if self.dual_evidence.appearance_confidence.shape != self.density_logits.shape:
raise ValueError(
"dual evidence appearance_confidence must match density_logits shape"
)
if self.dual_evidence.base_color.shape != self.color.shape:
raise ValueError("dual evidence base_color must match color shape")
if self.dual_evidence.appearance_confidence.device != self.density_logits.device:
raise ValueError("dual evidence and typed state fields must be on the same device")
if self.transport_evidence is not None:
if (
self.transport_evidence.transported_appearance_confidence.shape
!= self.density_logits.shape
):
raise ValueError(
"transported appearance confidence must match density_logits shape"
)
if self.transport_evidence.base_color.shape != self.color.shape:
raise ValueError("transport evidence base_color must match color shape")
if (
self.transport_evidence.transported_appearance_confidence.device
!= self.density_logits.device
):
raise ValueError(
"transport evidence and typed state fields must be on the same device"
)
@property
def spatial_shape(self) -> tuple[int, int, int]:
return tuple(self.density_logits.shape[-3:])
@property
def appearance_shape(self) -> tuple[int, int, int]:
if self.appearance is None:
return self.spatial_shape
return self.appearance.spatial_shape
def to(self, *args: object, **kwargs: object) -> Self:
return SceneState(
self.density_logits.to(*args, **kwargs),
self.color.to(*args, **kwargs),
self.log_variance.to(*args, **kwargs),
self.bounds.to(*args, **kwargs),
None if self.features is None else self.features.to(*args, **kwargs),
None if self.evidence is None else self.evidence.to(*args, **kwargs),
None if self.appearance is None else self.appearance.to(*args, **kwargs),
None if self.dual_evidence is None else self.dual_evidence.to(*args, **kwargs),
(
None
if self.transport_evidence is None
else self.transport_evidence.to(*args, **kwargs)
),
)
@dataclass
class SceneExample:
"""One sparse-view scene item before DataLoader collation."""
context_rgb: Tensor
context_cameras: Cameras
target_rgb: Tensor
target_cameras: Cameras
bounds: Tensor
target_depth: Tensor | None = None
target_normal: Tensor | None = None
target_point: Tensor | None = None
target_visibility: Tensor | None = None
target_support: Tensor | None = None
scene_id: str | None = None
def __post_init__(self) -> None:
_validate_example_or_batch(self, batched=False)
def to(self, *args: object, **kwargs: object) -> Self:
return SceneExample(
self.context_rgb.to(*args, **kwargs),
self.context_cameras.to(*args, **kwargs),
self.target_rgb.to(*args, **kwargs),
self.target_cameras.to(*args, **kwargs),
self.bounds.to(*args, **kwargs),
_optional_to(self.target_depth, *args, **kwargs),
_optional_to(self.target_normal, *args, **kwargs),
_optional_to(self.target_point, *args, **kwargs),
_optional_to(self.target_visibility, *args, **kwargs),
_optional_to(self.target_support, *args, **kwargs),
self.scene_id,
)
@dataclass
class SceneBatch:
"""A collated batch of context observations and hidden target supervision."""
context_rgb: Tensor
context_cameras: Cameras
target_rgb: Tensor
target_cameras: Cameras
bounds: Tensor
target_depth: Tensor | None = None
target_normal: Tensor | None = None
target_point: Tensor | None = None
target_visibility: Tensor | None = None
target_support: Tensor | None = None
scene_ids: tuple[str | None, ...] = ()
def __post_init__(self) -> None:
_validate_example_or_batch(self, batched=True)
def to(self, *args: object, **kwargs: object) -> Self:
return SceneBatch(
self.context_rgb.to(*args, **kwargs),
self.context_cameras.to(*args, **kwargs),
self.target_rgb.to(*args, **kwargs),
self.target_cameras.to(*args, **kwargs),
self.bounds.to(*args, **kwargs),
_optional_to(self.target_depth, *args, **kwargs),
_optional_to(self.target_normal, *args, **kwargs),
_optional_to(self.target_point, *args, **kwargs),
_optional_to(self.target_visibility, *args, **kwargs),
_optional_to(self.target_support, *args, **kwargs),
self.scene_ids,
)
def targets(self) -> dict[str, Tensor]:
result = {"rgb": self.target_rgb}
optional = {
"depth": self.target_depth,
"normal": self.target_normal,
"point": self.target_point,
"visibility": self.target_visibility,
}
result.update({name: value for name, value in optional.items() if value is not None})
if self.target_support is not None:
result["support"] = self.target_support
return result
def _optional_to(value: Tensor | None, *args: object, **kwargs: object) -> Tensor | None:
return None if value is None else value.to(*args, **kwargs)
def _validate_example_or_batch(value: SceneExample | SceneBatch, *, batched: bool) -> None:
rgb_ndim = 5 if batched else 4
camera_shape = 2 if batched else 1
bounds_shape = (value.context_rgb.shape[0], 2, 3) if batched else (2, 3)
if value.context_rgb.ndim != rgb_ndim or value.context_rgb.shape[-3] != 3:
raise ValueError("context_rgb must have RGB channel-first image shape")
if value.target_rgb.ndim != rgb_ndim or value.target_rgb.shape[-3] != 3:
raise ValueError("target_rgb must have RGB channel-first image shape")
if value.context_rgb.shape[-2:] != value.context_cameras.image_size:
raise ValueError("context_rgb spatial size must match context cameras")
if value.target_rgb.shape[-2:] != value.target_cameras.image_size:
raise ValueError("target_rgb spatial size must match target cameras")
if len(value.context_cameras.leading_shape) != camera_shape:
raise ValueError("context camera leading dimensions do not match example/batch shape")
if len(value.target_cameras.leading_shape) != camera_shape:
raise ValueError("target camera leading dimensions do not match example/batch shape")
if batched:
if value.context_cameras.leading_shape != value.context_rgb.shape[:2]:
raise ValueError("context camera dimensions must match context_rgb")
if value.target_cameras.leading_shape != value.target_rgb.shape[:2]:
raise ValueError("target camera dimensions must match target_rgb")
else:
if value.context_cameras.leading_shape != value.context_rgb.shape[:1]:
raise ValueError("context camera dimensions must match context_rgb")
if value.target_cameras.leading_shape != value.target_rgb.shape[:1]:
raise ValueError("target camera dimensions must match target_rgb")
if value.bounds.shape != bounds_shape:
raise ValueError("bounds shape does not match example/batch")
_require_finite("context_rgb", value.context_rgb)
_require_finite("target_rgb", value.target_rgb)
_require_finite("bounds", value.bounds)
target_views = value.target_rgb.shape[:-3]
expected = {
"target_depth": (1,),
"target_normal": (3,),
"target_point": (3,),
"target_visibility": (1,),
"target_support": (1,),
}
for name, channels in expected.items():
tensor = getattr(value, name)
if tensor is None:
continue
if tensor.shape[: len(target_views)] != target_views or tensor.shape[-3] != channels[0]:
raise ValueError(f"{name} must align with target views and channel count")
if tensor.shape[-2:] != value.target_rgb.shape[-2:]:
raise ValueError(f"{name} spatial size must match target_rgb")
_require_finite(name, tensor)
if isinstance(value, SceneExample):
if value.scene_id is not None and not isinstance(value.scene_id, str):
raise ValueError("scene_id must be a string or null")
elif value.scene_ids and len(value.scene_ids) != value.context_rgb.shape[0]:
raise ValueError("scene_ids must align with the batch dimension")