Spaces:
Running on Zero
Running on Zero
| """Explicit camera conventions shared by raw, latent, train, and inference paths.""" | |
| from __future__ import annotations | |
| import io | |
| from dataclasses import dataclass | |
| from typing import Literal | |
| import numpy as np | |
| from solarwm.errors import DataContractError | |
| CameraStorage = Literal["absolute_c2w", "absolute_w2c", "relative_w2c"] | |
| TranslationTransform = Literal["linear", "logd4"] | |
| class CameraGuardError(DataContractError): | |
| """A valid camera trajectory rejected by configured magnitude guards.""" | |
| def camera_audit_prefix_frames(num_frames: int, fps: float, max_seconds: float = 10.0) -> int: | |
| """Apply the first-ten-second magnitude-guard window.""" | |
| count = max(0, int(num_frames)) | |
| if count == 0: | |
| return 0 | |
| rate = float(fps) | |
| seconds = float(max_seconds) | |
| if not np.isfinite(rate) or rate <= 0 or not np.isfinite(seconds) or seconds <= 0: | |
| return count | |
| return min(count, max(1, int(np.ceil(rate * seconds)))) | |
| def load_camera_npz(value: bytes, array_key: str) -> tuple[np.ndarray, CameraStorage]: | |
| """Load one explicitly declared camera array without changing its dtype.""" | |
| key_contract: dict[str, CameraStorage] = { | |
| "c2w": "absolute_c2w", | |
| "w2c": "absolute_w2c", | |
| "vipe_c2w": "absolute_c2w", | |
| "vipe_w2c": "absolute_w2c", | |
| "relative_w2c": "relative_w2c", | |
| } | |
| if array_key not in key_contract: | |
| raise DataContractError(f"unsupported camera array key {array_key!r}") | |
| try: | |
| with np.load(io.BytesIO(value), allow_pickle=False) as archive: | |
| if array_key.startswith("vipe_") and "vipe_status" in archive.files: | |
| status = str(np.asarray(archive["vipe_status"]).reshape(-1)[0]).lower() | |
| if status != "ok": | |
| raise DataContractError(f"VIPE camera status is {status!r}, expected 'ok'") | |
| if array_key not in archive.files: | |
| raise DataContractError(f"camera archive lacks declared array {array_key!r}") | |
| matrices = np.asarray(archive[array_key]) | |
| except (OSError, ValueError) as exc: | |
| raise DataContractError("camera member is not a valid non-pickled NPZ") from exc | |
| return _validate_matrices(matrices, name=array_key), key_contract[array_key] | |
| def _validate_matrices(value: np.ndarray, *, name: str) -> np.ndarray: | |
| matrices = np.asarray(value) | |
| if matrices.ndim != 3 or matrices.shape[1:] != (4, 4): | |
| raise DataContractError(f"{name} must have shape [T,4,4], got {matrices.shape}") | |
| if not np.isfinite(matrices).all(): | |
| raise DataContractError(f"{name} contains non-finite values") | |
| return matrices | |
| def invert_se3(matrices: np.ndarray) -> np.ndarray: | |
| matrices = _validate_matrices(matrices, name="SE3") | |
| result = np.zeros_like(matrices) | |
| rotation = matrices[:, :3, :3] | |
| translation = matrices[:, :3, 3] | |
| rotation_t = np.swapaxes(rotation, -1, -2) | |
| result[:, :3, :3] = rotation_t | |
| result[:, :3, 3] = -np.einsum("tij,tj->ti", rotation_t, translation) | |
| result[:, 3, 3] = 1 | |
| return result | |
| def absolute_c2w(value: np.ndarray, storage: CameraStorage) -> np.ndarray: | |
| matrices = _validate_matrices(value, name="camera") | |
| if storage == "absolute_c2w": | |
| return matrices | |
| if storage == "absolute_w2c": | |
| return invert_se3(matrices) | |
| raise DataContractError("relative_w2c cannot be converted back to authoritative absolute C2W") | |
| def relative_w2c(value: np.ndarray, storage: CameraStorage) -> np.ndarray: | |
| matrices = _validate_matrices(value, name="camera") | |
| if storage == "relative_w2c": | |
| relative = matrices.copy() | |
| else: | |
| c2w = absolute_c2w(matrices, storage) | |
| relative = np.matmul(invert_se3(c2w), c2w[0]) | |
| identity = np.eye(4, dtype=relative.dtype) | |
| if not np.allclose(relative[0], identity, rtol=1e-5, atol=1e-5): | |
| raise DataContractError("first relative W2C row is not identity") | |
| return relative | |
| def relative_c2w(value: np.ndarray, storage: CameraStorage) -> np.ndarray: | |
| return invert_se3(relative_w2c(value, storage)) | |
| def transform_translation( | |
| relative: np.ndarray, | |
| transform: TranslationTransform, | |
| ) -> np.ndarray: | |
| """Apply model-conditioning compression without mutating source camera data.""" | |
| matrices = _validate_matrices(relative, name="relative camera") | |
| if transform == "linear": | |
| return matrices | |
| if transform != "logd4": | |
| raise DataContractError(f"unknown camera translation transform {transform!r}") | |
| result = matrices.copy() | |
| working_dtype = np.float64 if matrices.dtype == np.float64 else np.float32 | |
| translation = matrices[:, :3, 3].astype(working_dtype, copy=False) | |
| norms = np.linalg.norm(translation, axis=-1) | |
| scale = np.zeros_like(norms) | |
| nonzero = norms != 0 | |
| scale[nonzero] = np.log1p(norms[nonzero]) / (4.0 * norms[nonzero]) | |
| result[:, :3, 3] = (translation * scale[:, None]).astype(result.dtype, copy=False) | |
| return result | |
| class CameraGuards: | |
| max_rel_translation: float | None = 20.0 | |
| max_camera_abs: float | None = 20.0 | |
| def apply(self, relative: np.ndarray, *, audit_rows: int | None = None) -> np.ndarray: | |
| matrices = _validate_matrices(relative, name="relative camera") | |
| count = matrices.shape[0] if audit_rows is None else int(audit_rows) | |
| if not 1 <= count <= matrices.shape[0]: | |
| raise DataContractError("audit_rows is outside the camera trajectory") | |
| prefix = matrices[:count] | |
| if self.max_rel_translation is not None: | |
| magnitude = np.linalg.norm(prefix[:, :3, 3], axis=-1) | |
| if float(magnitude.max(initial=0.0)) > self.max_rel_translation: | |
| raise CameraGuardError("relative translation exceeds configured guard") | |
| result = matrices.copy() | |
| if self.max_camera_abs is not None and ( | |
| float(np.abs(result[:count]).max(initial=0.0)) > self.max_camera_abs | |
| ): | |
| raise CameraGuardError("camera matrix exceeds configured absolute guard") | |
| return result | |
| def causal_pixel_indices(pixel_frames: int, temporal_stride: int) -> np.ndarray: | |
| if pixel_frames < 1 or temporal_stride < 1: | |
| raise DataContractError("pixel_frames and temporal_stride must be positive") | |
| if pixel_frames == 1: | |
| return np.asarray([0], dtype=np.int64) | |
| return np.asarray([0, *range(1, pixel_frames, temporal_stride)], dtype=np.int64) | |