|
|
import torch
|
|
|
from torch import Tensor
|
|
|
from loguru import logger
|
|
|
from typing import Tuple
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
try:
|
|
|
# ipywidgets is available in interactive environments like Jupyter.
|
|
|
from ipywidgets import interact, IntSlider
|
|
|
|
|
|
HAS_IPYWIDGETS = True
|
|
|
except ImportError:
|
|
|
HAS_IPYWIDGETS = False
|
|
|
logger.warning(
|
|
|
"ipywidgets not installed. Interactive slicing will be disabled."
|
|
|
)
|
|
|
|
|
|
|
|
|
class GaussianSplat4DStateSpace:
|
|
|
"""
|
|
|
4D Gaussian splatting with a state space model in PyTorch.
|
|
|
|
|
|
Each Gaussian is defined by an 8D state vector:
|
|
|
[x, y, z, w, vx, vy, vz, vw],
|
|
|
where the first four dimensions are the spatial coordinates and the last
|
|
|
four are the velocities. Only the spatial (first four) dimensions are used
|
|
|
for the 4D Gaussian splat, with a corresponding 4×4 covariance matrix.
|
|
|
|
|
|
Attributes:
|
|
|
num_gaussians (int): Number of Gaussians.
|
|
|
state_dim (int): Dimension of the state vector (should be 8).
|
|
|
states (Tensor): Current state for each Gaussian of shape (num_gaussians, state_dim).
|
|
|
covariances (Tensor): Covariance matrices for the spatial dimensions, shape (num_gaussians, 4, 4).
|
|
|
A (Tensor): State transition matrix of shape (state_dim, state_dim).
|
|
|
dt (float): Time step for state updates.
|
|
|
"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
num_gaussians: int,
|
|
|
init_states: Tensor,
|
|
|
init_covariances: Tensor,
|
|
|
dt: float = 1.0,
|
|
|
) -> None:
|
|
|
"""
|
|
|
Initialize the 4D Gaussian splat model.
|
|
|
|
|
|
Args:
|
|
|
num_gaussians (int): Number of Gaussians.
|
|
|
init_states (Tensor): Initial states of shape (num_gaussians, 8).
|
|
|
Each state is assumed to be
|
|
|
[x, y, z, w, vx, vy, vz, vw].
|
|
|
init_covariances (Tensor): Initial covariance matrices for the spatial dimensions,
|
|
|
shape (num_gaussians, 4, 4).
|
|
|
dt (float): Time step for the state update.
|
|
|
"""
|
|
|
if init_states.shape[1] != 8:
|
|
|
raise ValueError(
|
|
|
"init_states should have shape (N, 8) where 8 = 4 position + 4 velocity."
|
|
|
)
|
|
|
if init_covariances.shape[1:] != (4, 4):
|
|
|
raise ValueError(
|
|
|
"init_covariances should have shape (N, 4, 4)."
|
|
|
)
|
|
|
|
|
|
self.num_gaussians = num_gaussians
|
|
|
self.states = init_states.clone() # shape: (N, 8)
|
|
|
self.covariances = (
|
|
|
init_covariances.clone()
|
|
|
) # shape: (N, 4, 4)
|
|
|
self.dt = dt
|
|
|
self.state_dim = init_states.shape[1]
|
|
|
|
|
|
# Create an 8x8 constant-velocity state transition matrix:
|
|
|
# New position = position + velocity*dt, velocity remains unchanged.
|
|
|
I4 = torch.eye(
|
|
|
4, dtype=init_states.dtype, device=init_states.device
|
|
|
)
|
|
|
zeros4 = torch.zeros(
|
|
|
(4, 4), dtype=init_states.dtype, device=init_states.device
|
|
|
)
|
|
|
top = torch.cat([I4, dt * I4], dim=1)
|
|
|
bottom = torch.cat([zeros4, I4], dim=1)
|
|
|
self.A = torch.cat([top, bottom], dim=0) # shape: (8, 8)
|
|
|
|
|
|
logger.info(
|
|
|
"Initialized 4D GaussianSplatStateSpace with {} Gaussians.",
|
|
|
num_gaussians,
|
|
|
)
|
|
|
|
|
|
def update_states(self) -> None:
|
|
|
"""
|
|
|
Update the state of each Gaussian using the constant-velocity state space model.
|
|
|
|
|
|
Applies:
|
|
|
state_next = A @ state_current.
|
|
|
"""
|
|
|
self.states = (
|
|
|
self.A @ self.states.t()
|
|
|
).t() # shape: (num_gaussians, 8)
|
|
|
logger.debug("States updated: {}", self.states)
|
|
|
|
|
|
def _compute_gaussian(
|
|
|
self, pos: Tensor, cov: Tensor, coords: Tensor
|
|
|
) -> Tensor:
|
|
|
"""
|
|
|
Compute the 4D Gaussian function over a grid of coordinates.
|
|
|
|
|
|
Args:
|
|
|
pos (Tensor): The center of the Gaussian (4,).
|
|
|
cov (Tensor): The 4×4 covariance matrix.
|
|
|
coords (Tensor): A grid of coordinates of shape (..., 4).
|
|
|
|
|
|
Returns:
|
|
|
Tensor: Evaluated Gaussian values on the grid with shape equal to coords.shape[:-1].
|
|
|
"""
|
|
|
try:
|
|
|
cov_inv = torch.linalg.inv(cov)
|
|
|
except RuntimeError as e:
|
|
|
logger.warning(
|
|
|
"Covariance inversion failed; using pseudo-inverse. Error: {}",
|
|
|
e,
|
|
|
)
|
|
|
cov_inv = torch.linalg.pinv(cov)
|
|
|
|
|
|
# Broadcast pos over the grid
|
|
|
diff = coords - pos.view(
|
|
|
*(1 for _ in range(coords.ndim - 1)), 4
|
|
|
)
|
|
|
mahal = torch.einsum("...i,ij,...j->...", diff, cov_inv, diff)
|
|
|
gaussian = torch.exp(-0.5 * mahal)
|
|
|
return gaussian
|
|
|
|
|
|
def render(
|
|
|
self,
|
|
|
canvas_size: Tuple[int, int, int, int],
|
|
|
sigma_scale: float = 1.0,
|
|
|
normalize: bool = False,
|
|
|
) -> Tensor:
|
|
|
"""
|
|
|
Render the current 4D Gaussian splats onto a 4D canvas.
|
|
|
|
|
|
Args:
|
|
|
canvas_size (Tuple[int, int, int, int]): The size of the canvas (d1, d2, d3, d4).
|
|
|
sigma_scale (float): Scaling factor for the covariance (affects spread).
|
|
|
normalize (bool): Whether to normalize the final canvas to [0, 1].
|
|
|
|
|
|
Returns:
|
|
|
Tensor: A 4D tensor (canvas) with the accumulated contributions from all Gaussians.
|
|
|
"""
|
|
|
d1, d2, d3, d4 = canvas_size
|
|
|
|
|
|
# Create coordinate grids for each dimension.
|
|
|
grid1 = torch.linspace(
|
|
|
0, d1 - 1, d1, device=self.states.device
|
|
|
)
|
|
|
grid2 = torch.linspace(
|
|
|
0, d2 - 1, d2, device=self.states.device
|
|
|
)
|
|
|
grid3 = torch.linspace(
|
|
|
0, d3 - 1, d3, device=self.states.device
|
|
|
)
|
|
|
grid4 = torch.linspace(
|
|
|
0, d4 - 1, d4, device=self.states.device
|
|
|
)
|
|
|
|
|
|
# Create a 4D meshgrid (using indexing "ij")
|
|
|
grid = torch.stack(
|
|
|
torch.meshgrid(grid1, grid2, grid3, grid4, indexing="ij"),
|
|
|
dim=-1,
|
|
|
) # shape: (d1, d2, d3, d4, 4)
|
|
|
|
|
|
# Initialize the canvas.
|
|
|
canvas = torch.zeros(
|
|
|
(d1, d2, d3, d4),
|
|
|
dtype=self.states.dtype,
|
|
|
device=self.states.device,
|
|
|
)
|
|
|
|
|
|
for i in range(self.num_gaussians):
|
|
|
pos = self.states[i, :4] # spatial center (4,)
|
|
|
cov = (
|
|
|
self.covariances[i] * sigma_scale
|
|
|
) # scaled covariance
|
|
|
gaussian = self._compute_gaussian(pos, cov, grid)
|
|
|
canvas += gaussian
|
|
|
logger.debug(
|
|
|
"Rendered Gaussian {} at position {}", i, pos.tolist()
|
|
|
)
|
|
|
|
|
|
if normalize:
|
|
|
max_val = canvas.max()
|
|
|
if max_val > 0:
|
|
|
canvas = canvas / max_val
|
|
|
logger.debug("Canvas normalized.")
|
|
|
|
|
|
logger.info("4D Rendering complete.")
|
|
|
return canvas
|
|
|
|
|
|
|
|
|
def interactive_slice(canvas: Tensor) -> None:
|
|
|
"""
|
|
|
Display an interactive 2D slice of the 4D canvas using ipywidgets.
|
|
|
|
|
|
This function fixes two of the four dimensions (d3 and d4) via sliders and
|
|
|
displays the resulting 2D slice (over dimensions d1 and d2).
|
|
|
|
|
|
Args:
|
|
|
canvas (Tensor): A 4D tensor with shape (d1, d2, d3, d4).
|
|
|
"""
|
|
|
d1, d2, d3, d4 = canvas.shape
|
|
|
|
|
|
def display_slice(slice_d3: int, slice_d4: int):
|
|
|
slice_2d = canvas[:, :, slice_d3, slice_d4].cpu().numpy()
|
|
|
plt.figure(figsize=(6, 6))
|
|
|
plt.imshow(slice_2d, cmap="hot", origin="lower")
|
|
|
plt.title(f"2D Slice at d3={slice_d3}, d4={slice_d4}")
|
|
|
plt.colorbar()
|
|
|
plt.show()
|
|
|
|
|
|
interact(
|
|
|
display_slice,
|
|
|
slice_d3=IntSlider(min=0, max=d3 - 1, step=1, value=d3 // 2),
|
|
|
slice_d4=IntSlider(min=0, max=d4 - 1, step=1, value=d4 // 2),
|
|
|
)
|
|
|
|
|
|
|
|
|
def mip_projection(canvas: Tensor) -> None:
|
|
|
"""
|
|
|
Render a 2D view of the 4D canvas using maximum intensity projection (MIP)
|
|
|
along the 3rd and 4th dimensions.
|
|
|
|
|
|
Args:
|
|
|
canvas (Tensor): A 4D tensor with shape (d1, d2, d3, d4).
|
|
|
"""
|
|
|
# MIP along dimension 3
|
|
|
mip_3d = canvas.max(dim=2)[0] # shape: (d1, d2, d4)
|
|
|
# MIP along dimension 4
|
|
|
mip_2d = mip_3d.max(dim=2)[0] # shape: (d1, d2)
|
|
|
|
|
|
plt.figure(figsize=(6, 6))
|
|
|
plt.imshow(mip_2d.cpu().numpy(), cmap="hot", origin="lower")
|
|
|
plt.title("2D MIP (Projecting dimensions d3 and d4)")
|
|
|
plt.colorbar()
|
|
|
plt.show()
|
|
|
|
|
|
|
|
|
def main() -> None:
|
|
|
"""
|
|
|
Main function that:
|
|
|
- Creates a 4D Gaussian splat model.
|
|
|
- Updates the states to simulate motion.
|
|
|
- Renders a 4D canvas.
|
|
|
- Visualizes the 4D volume via interactive slicing (if available) or MIP.
|
|
|
"""
|
|
|
torch.manual_seed(42)
|
|
|
num_gaussians = 2
|
|
|
|
|
|
# Define initial states for each Gaussian:
|
|
|
# Each state is [x, y, z, w, vx, vy, vz, vw].
|
|
|
init_states = torch.tensor(
|
|
|
[
|
|
|
[10.0, 15.0, 20.0, 25.0, 0.5, -0.2, 0.3, 0.1],
|
|
|
[30.0, 35.0, 40.0, 45.0, -0.3, 0.4, -0.1, 0.2],
|
|
|
],
|
|
|
dtype=torch.float32,
|
|
|
)
|
|
|
|
|
|
# Define initial 4x4 covariance matrices for the spatial dimensions.
|
|
|
init_covariances = torch.stack(
|
|
|
[
|
|
|
torch.diag(
|
|
|
torch.tensor(
|
|
|
[5.0, 5.0, 5.0, 5.0], dtype=torch.float32
|
|
|
)
|
|
|
),
|
|
|
torch.diag(
|
|
|
torch.tensor(
|
|
|
[3.0, 3.0, 3.0, 3.0], dtype=torch.float32
|
|
|
)
|
|
|
),
|
|
|
]
|
|
|
)
|
|
|
|
|
|
# Create the 4D Gaussian splat model.
|
|
|
model = GaussianSplat4DStateSpace(
|
|
|
num_gaussians, init_states, init_covariances, dt=1.0
|
|
|
)
|
|
|
|
|
|
# Update states to simulate one time step.
|
|
|
model.update_states()
|
|
|
|
|
|
# Render the 4D canvas.
|
|
|
canvas_size = (20, 20, 20, 20)
|
|
|
canvas = model.render(
|
|
|
canvas_size, sigma_scale=1.0, normalize=True
|
|
|
)
|
|
|
|
|
|
# Visualize the 4D data.
|
|
|
if HAS_IPYWIDGETS:
|
|
|
logger.info("Launching interactive slicing tool for 4D data.")
|
|
|
interactive_slice(canvas)
|
|
|
else:
|
|
|
logger.info(
|
|
|
"ipywidgets not available; using maximum intensity projection instead."
|
|
|
)
|
|
|
mip_projection(canvas)
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
main()
|
