gsvvcompressor 1.2.0__cp311-cp311-win_amd64.whl

gsvvcompressor/xyz/interface.py

@@ -0,0 +1,382 @@
+"""
+XYZ quantization-based inter-frame codec interface.
+
+This module provides an inter-frame codec that only operates on the xyz
+coordinates of a GaussianModel using quantization.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Self
+
+import torch
+
+from gaussian_splatting import GaussianModel
+
+from ..payload import Payload
+from ..interframe import InterframeEncoderInitConfig, InterframeCodecContext, InterframeCodecInterface
+from .quant import XYZQuantConfig, compute_quant_config, quantize_xyz, dequantize_xyz
+
+
+@dataclass
+class XYZQuantInterframeCodecConfig(InterframeEncoderInitConfig):
+    """
+    Configuration parameters for XYZ quantization-based inter-frame codec.
+
+    This dataclass holds the initialization settings for xyz coordinate
+    quantization.
+
+    Attributes:
+        k: Which nearest neighbor to use for step size estimation (1 = nearest).
+        sample_size: Number of points to sample for NN estimation.
+        seed: Random seed for reproducible sampling.
+        quantile: Quantile of NN distances to use for dense scale estimation.
+        alpha: Scaling factor for step size.
+        min_step: Optional minimum step size.
+        max_step: Optional maximum step size.
+        tolerance: Tolerance for inter-frame change detection. Only coordinates
+            with absolute difference > tolerance are considered changed.
+    """
+    k: int = 1
+    sample_size: Optional[int] = 10000
+    seed: Optional[int] = 42
+    quantile: float = 0.05
+    alpha: float = 0.2
+    min_step: Optional[float] = None
+    max_step: Optional[float] = None
+    tolerance: int = 0
+
+
+@dataclass
+class XYZQuantInterframeCodecContext(InterframeCodecContext):
+    """
+    Context data for XYZ quantization-based inter-frame encoding/decoding.
+
+    This dataclass holds the quantization state including the quantization
+    configuration and quantized xyz coordinates.
+
+    Attributes:
+        quant_config: The quantization configuration (step_size, origin).
+        quantized_xyz: The quantized xyz coordinates, shape (N, 3), dtype int32.
+        tolerance: Tolerance for inter-frame change detection (encoder-only,
+            not used during decoding).
+    """
+    quant_config: XYZQuantConfig
+    quantized_xyz: torch.Tensor  # shape (N, 3), dtype int32
+    tolerance: int = 0  # encoder-only
+
+
+@dataclass
+class XYZQuantKeyframePayload(Payload):
+    """
+    Payload for XYZ quantization keyframe data.
+
+    Contains the quantization configuration and full quantized xyz coordinates.
+
+    Attributes:
+        quant_config: The quantization configuration.
+        quantized_xyz: The quantized xyz coordinates.
+    """
+    quant_config: XYZQuantConfig
+    quantized_xyz: torch.Tensor
+
+    def to(self, device) -> Self:
+        """
+        Move the Payload to the specified device.
+
+        Args:
+            device: The target device (e.g., 'cpu', 'cuda', torch.device).
+
+        Returns:
+            A new XYZQuantKeyframePayload instance on the target device.
+        """
+        return XYZQuantKeyframePayload(
+            quant_config=XYZQuantConfig(
+                step_size=self.quant_config.step_size,
+                origin=self.quant_config.origin.to(device),
+            ),
+            quantized_xyz=self.quantized_xyz.to(device),
+        )
+
+
+@dataclass
+class XYZQuantInterframePayload(Payload):
+    """
+    Payload for XYZ quantization interframe data.
+
+    Contains only the changed quantized xyz coordinates for subsequent frames.
+    The quantization configuration is inherited from the keyframe context.
+
+    Attributes:
+        xyz_mask: Boolean tensor indicating which xyz values changed, shape (N,).
+        quantized_xyz: Only the changed quantized xyz values (sparse), shape (M, 3).
+    """
+    xyz_mask: torch.Tensor
+    quantized_xyz: torch.Tensor
+
+    def to(self, device) -> Self:
+        """
+        Move the Payload to the specified device.
+
+        Args:
+            device: The target device (e.g., 'cpu', 'cuda', torch.device).
+
+        Returns:
+            A new XYZQuantInterframePayload instance on the target device.
+        """
+        return XYZQuantInterframePayload(
+            xyz_mask=self.xyz_mask.to(device),
+            quantized_xyz=self.quantized_xyz.to(device),
+        )
+
+
+class XYZQuantInterframeCodecInterface(InterframeCodecInterface):
+    """
+    XYZ quantization-based inter-frame encoding/decoding interface.
+
+    This interface uses coordinate quantization to compress the xyz coordinates
+    of GaussianModel. The keyframe computes the quantization configuration,
+    and subsequent frames use the same configuration to quantize their coordinates.
+
+    Only operates on xyz coordinates; other GaussianModel attributes are not modified.
+    """
+
+    def decode_interframe(
+        self,
+        payload: XYZQuantInterframePayload,
+        prev_context: XYZQuantInterframeCodecContext,
+    ) -> XYZQuantInterframeCodecContext:
+        """
+        Decode a delta payload to reconstruct the next frame's context.
+
+        Applies the changed xyz values from the payload to the previous context.
+
+        Args:
+            payload: The delta payload containing changed quantized xyz with mask.
+            prev_context: The context of the previous frame (contains quant_config).
+
+        Returns:
+            The reconstructed context for the current frame.
+        """
+        # Clone previous quantized_xyz and apply changes
+        new_quantized_xyz = prev_context.quantized_xyz.clone()
+        new_quantized_xyz[payload.xyz_mask] = payload.quantized_xyz
+
+        return XYZQuantInterframeCodecContext(
+            quant_config=prev_context.quant_config,
+            quantized_xyz=new_quantized_xyz,
+            # tolerance is encoder-only, decoder doesn't use it
+        )
+
+    def encode_interframe(
+        self,
+        prev_context: XYZQuantInterframeCodecContext,
+        next_context: XYZQuantInterframeCodecContext,
+    ) -> XYZQuantInterframePayload:
+        """
+        Encode the difference between two consecutive frames.
+
+        Compares prev and next contexts to find changed xyz coordinates and stores
+        only the changed values with their corresponding mask.
+
+        Args:
+            prev_context: The context of the previous frame.
+            next_context: The context of the next frame.
+
+        Returns:
+            A payload containing only changed quantized xyz with mask.
+        """
+        prev_xyz = prev_context.quantized_xyz
+        next_xyz = next_context.quantized_xyz
+        tolerance = prev_context.tolerance
+
+        # Find changed xyz coordinates (any coordinate difference > tolerance)
+        diff = (prev_xyz - next_xyz).abs()
+        mask = (diff > tolerance).any(dim=-1)
+        changed_xyz = next_xyz[mask]
+
+        return XYZQuantInterframePayload(
+            xyz_mask=mask,
+            quantized_xyz=changed_xyz,
+        )
+
+    def decode_keyframe(self, payload: XYZQuantKeyframePayload) -> XYZQuantInterframeCodecContext:
+        """
+        Decode a keyframe payload to create initial context.
+
+        Args:
+            payload: The keyframe payload containing quant_config and quantized xyz.
+
+        Returns:
+            The context for the first/key frame.
+        """
+        return XYZQuantInterframeCodecContext(
+            quant_config=payload.quant_config,
+            quantized_xyz=payload.quantized_xyz,
+            # tolerance is encoder-only, decoder doesn't use it
+        )
+
+    def decode_keyframe_for_encode(
+        self,
+        payload: XYZQuantKeyframePayload,
+        context: XYZQuantInterframeCodecContext,
+    ) -> XYZQuantInterframeCodecContext:
+        """
+        Decode a keyframe payload during encoding to avoid error accumulation.
+
+        Since the encode/decode round-trip is lossless for XYZ quantization
+        (the quantized data is simply copied), we can reuse the original context.
+        This preserves context.tolerance which is needed by the encoder.
+
+        Args:
+            payload: The keyframe payload that was just encoded.
+            context: The original context used for encoding this keyframe.
+
+        Returns:
+            The reconstructed context (same as original for this codec).
+        """
+        # Round-trip is lossless, reuse original context (including tolerance)
+        return context
+
+    def decode_interframe_for_encode(
+        self,
+        payload: XYZQuantInterframePayload,
+        prev_context: XYZQuantInterframeCodecContext,
+    ) -> XYZQuantInterframeCodecContext:
+        """
+        Decode an interframe payload during encoding to avoid error accumulation.
+
+        Similar to decode_interframe, but preserves tolerance from prev_context
+        since it's needed by the encoder for subsequent frames.
+
+        Args:
+            payload: The interframe payload that was just encoded.
+            prev_context: The previous frame's context (reconstructed version).
+
+        Returns:
+            The reconstructed context with tolerance preserved.
+        """
+        # Clone previous quantized_xyz and apply changes
+        new_quantized_xyz = prev_context.quantized_xyz.clone()
+        new_quantized_xyz[payload.xyz_mask] = payload.quantized_xyz
+
+        return XYZQuantInterframeCodecContext(
+            quant_config=prev_context.quant_config,
+            quantized_xyz=new_quantized_xyz,
+            tolerance=prev_context.tolerance,  # preserve for encoder
+        )
+
+    def encode_keyframe(self, context: XYZQuantInterframeCodecContext) -> XYZQuantKeyframePayload:
+        """
+        Encode the first frame as a keyframe.
+
+        Args:
+            context: The context of the first frame.
+
+        Returns:
+            A payload containing the quant_config and quantized xyz.
+        """
+        return XYZQuantKeyframePayload(
+            quant_config=context.quant_config,
+            quantized_xyz=context.quantized_xyz,
+            # tolerance is encoder-only, not needed in payload
+        )
+
+    def keyframe_to_context(
+        self,
+        frame: GaussianModel,
+        init_config: XYZQuantInterframeCodecConfig,
+    ) -> XYZQuantInterframeCodecContext:
+        """
+        Convert a keyframe to a XYZQuantInterframeCodecContext.
+
+        Computes quantization configuration from the frame's xyz coordinates
+        and quantizes them.
+
+        Args:
+            frame: The GaussianModel frame to convert.
+            init_config: Configuration parameters for quantization.
+
+        Returns:
+            The corresponding XYZQuantInterframeCodecContext representation.
+        """
+        xyz = frame.get_xyz
+
+        # Compute quantization configuration from xyz coordinates
+        quant_config = compute_quant_config(
+            points=xyz,
+            k=init_config.k,
+            sample_size=init_config.sample_size,
+            seed=init_config.seed,
+            quantile=init_config.quantile,
+            alpha=init_config.alpha,
+            min_step=init_config.min_step,
+            max_step=init_config.max_step,
+        )
+
+        # Quantize xyz coordinates
+        quantized_xyz = quantize_xyz(xyz, quant_config)
+
+        return XYZQuantInterframeCodecContext(
+            quant_config=quant_config,
+            quantized_xyz=quantized_xyz,
+            tolerance=init_config.tolerance,
+        )
+
+    def interframe_to_context(
+        self,
+        frame: GaussianModel,
+        prev_context: XYZQuantInterframeCodecContext,
+    ) -> XYZQuantInterframeCodecContext:
+        """
+        Convert a frame to a XYZQuantInterframeCodecContext using the previous context's config.
+
+        Uses the quantization config from the previous context to quantize xyz coordinates.
+
+        Args:
+            frame: The GaussianModel frame to convert.
+            prev_context: The context from the previous frame.
+
+        Returns:
+            The corresponding XYZQuantInterframeCodecContext representation.
+        """
+        xyz = frame.get_xyz
+
+        # Quantize xyz using the existing quant_config from keyframe
+        quantized_xyz = quantize_xyz(xyz, prev_context.quant_config)
+
+        return XYZQuantInterframeCodecContext(
+            quant_config=prev_context.quant_config,
+            quantized_xyz=quantized_xyz,
+            tolerance=prev_context.tolerance,
+        )
+
+    def context_to_frame(
+        self,
+        context: XYZQuantInterframeCodecContext,
+        frame: GaussianModel,
+    ) -> GaussianModel:
+        """
+        Convert a XYZQuantInterframeCodecContext back to a GaussianModel frame.
+
+        Dequantizes the xyz coordinates and sets them on the frame.
+        Only modifies xyz coordinates; other attributes are not touched.
+
+        Args:
+            context: The XYZQuantInterframeCodecContext to convert.
+            frame: An empty GaussianModel or one from previous pipeline steps.
+                This frame will be modified in-place with the xyz data.
+
+        Returns:
+            The modified GaussianModel with the xyz data.
+        """
+        # Dequantize xyz coordinates
+        xyz = dequantize_xyz(
+            context.quantized_xyz,
+            context.quant_config,
+            dtype=frame.get_xyz.dtype,
+        )
+
+        # Set xyz on the frame
+        frame._xyz = torch.nn.Parameter(xyz.to(frame.get_xyz.device))
+
+        return frame

gsvvcompressor/xyz/knn.py

@@ -0,0 +1,141 @@
+"""
+K-nearest neighbor distance computation for 3D point clouds.
+"""
+
+from typing import Optional
+
+import numpy as np
+import torch
+from scipy.spatial import cKDTree
+
+
+def compute_nn_distances(
+    points: torch.Tensor,
+    k: int = 1,
+    sample_size: Optional[int] = None,
+    seed: Optional[int] = None,
+) -> torch.Tensor:
+    """
+    Compute k-th nearest neighbor distances for points.
+
+    Uses scipy's cKDTree for efficient O(N log N) computation.
+    Optionally samples a subset of points to reduce computation for very large point clouds.
+
+    Args:
+        points: Point coordinates, shape (N, 3).
+        k: Which nearest neighbor distance to compute (1 = nearest, 2 = second nearest, etc.).
+        sample_size: If provided, randomly sample this many points for distance estimation.
+            If None or >= N, use all points.
+        seed: Random seed for reproducible sampling.
+
+    Returns:
+        Tensor of k-th nearest neighbor distances, shape (M,) where M = min(sample_size, N).
+
+    Raises:
+        ValueError: If there are not enough points to compute k-th neighbor (n_points <= k).
+    """
+    device = points.device
+    points_np = points.detach().cpu().numpy().astype(np.float64)
+    n_points = points_np.shape[0]
+
+    if n_points <= k:
+        raise ValueError(
+            f"Not enough points to compute {k}-th nearest neighbor. "
+            f"Got {n_points} points, need at least {k + 1}."
+        )
+
+    # Optionally subsample for large point clouds
+    if sample_size is not None and sample_size < n_points:
+        rng = np.random.default_rng(seed)
+        indices = rng.choice(n_points, size=sample_size, replace=False)
+        query_points = points_np[indices]
+    else:
+        query_points = points_np
+
+    tree = cKDTree(points_np)
+    # Query k+1 neighbors because the first neighbor is the point itself
+    distances, _ = tree.query(query_points, k=k + 1)
+
+    # Extract the k-th neighbor distance (index k, since index 0 is self with distance 0)
+    nn_distances = distances[:, k]
+
+    return torch.from_numpy(nn_distances).to(device=device, dtype=points.dtype)
+
+
+if __name__ == "__main__":
+    import time
+
+    print("=" * 60)
+    print("Testing compute_nn_distances")
+    print("=" * 60)
+
+    # Test 1: Simple grid points
+    print("\n[Test 1] Simple 2x2x2 grid with spacing 1.0")
+    grid_points = torch.tensor([
+        [0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0],
+        [0, 0, 1], [1, 0, 1], [0, 1, 1], [1, 1, 1],
+    ], dtype=torch.float32)
+    nn_dist = compute_nn_distances(grid_points, k=1)
+    print(f"  Points shape: {grid_points.shape}")
+    print(f"  NN distances: {nn_dist}")
+    print(f"  Expected: all 1.0 (grid spacing)")
+    assert torch.allclose(nn_dist, torch.ones(8)), "Test 1 failed!"
+    print("  PASSED")
+
+    # Test 2: Different k values
+    print("\n[Test 2] k=2 on grid (second nearest neighbor)")
+    nn_dist_k2 = compute_nn_distances(grid_points, k=2)
+    print(f"  NN distances (k=2): {nn_dist_k2}")
+    print(f"  Expected: all 1.0 (each point has multiple neighbors at distance 1)")
+    assert torch.allclose(nn_dist_k2, torch.ones(8)), "Test 2 failed!"
+    print("  PASSED")
+
+    # Test 3: k=4 should give sqrt(2) for corner points (face diagonal)
+    print("\n[Test 3] k=4 on grid (fourth nearest neighbor)")
+    nn_dist_k4 = compute_nn_distances(grid_points, k=4)
+    print(f"  NN distances (k=4): {nn_dist_k4}")
+    print(f"  Expected: all sqrt(2) ~ 1.414 (face diagonal neighbors)")
+    # Each corner has 3 edge neighbors at dist 1, then 3 face-diagonal neighbors at sqrt(2)
+    assert torch.allclose(nn_dist_k4, torch.full((8,), 2**0.5), atol=1e-5), "Test 3 failed!"
+    print("  PASSED")
+
+    # Test 4: Sampling
+    print("\n[Test 4] Sampling from larger point cloud")
+    large_points = torch.rand(1000, 3)
+    nn_dist_full = compute_nn_distances(large_points, k=1)
+    nn_dist_sampled = compute_nn_distances(large_points, k=1, sample_size=100, seed=42)
+    print(f"  Full cloud: {len(nn_dist_full)} distances")
+    print(f"  Sampled: {len(nn_dist_sampled)} distances")
+    print(f"  Full mean: {nn_dist_full.mean():.4f}, Sampled mean: {nn_dist_sampled.mean():.4f}")
+    assert len(nn_dist_sampled) == 100, "Test 4 failed!"
+    print("  PASSED")
+
+    # Test 5: Error on insufficient points
+    print("\n[Test 5] Error when n_points <= k")
+    try:
+        compute_nn_distances(torch.rand(2, 3), k=2)
+        print("  FAILED - should have raised ValueError")
+    except ValueError as e:
+        print(f"  Caught expected error: {e}")
+        print("  PASSED")
+
+    # Test 6: Reproducibility with seed
+    print("\n[Test 6] Reproducibility with seed")
+    dist1 = compute_nn_distances(large_points, k=1, sample_size=50, seed=123)
+    dist2 = compute_nn_distances(large_points, k=1, sample_size=50, seed=123)
+    assert torch.allclose(dist1, dist2), "Test 6 failed!"
+    print("  Same seed gives same results: PASSED")
+
+    # Test 7: Performance test
+    print("\n[Test 7] Performance test")
+    for n in [10_000, 100_000, 500_000]:
+        points = torch.rand(n, 3)
+
+        start = time.perf_counter()
+        _ = compute_nn_distances(points, k=1, sample_size=10000, seed=42)
+        elapsed = time.perf_counter() - start
+        print(f"  {n:>7,} points (sampled 10k): {elapsed:.3f}s")
+
+    print("\n" + "=" * 60)
+    print("All tests passed!")
+    print("=" * 60)

gsvvcompressor/xyz/quant.py

@@ -0,0 +1,143 @@
+"""
+XYZ coordinate quantization and dequantization.
+"""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from .knn import compute_nn_distances
+from .dense import compute_dense_scale
+from .size import compute_step_size
+
+
+@dataclass
+class XYZQuantConfig:
+    """
+    Configuration for XYZ coordinate quantization.
+
+    Attributes:
+        step_size: The quantization step size (delta).
+        origin: The origin offset for quantization (median of coordinates).
+    """
+    step_size: float
+    origin: torch.Tensor  # shape (3,)
+
+
+def compute_quant_config(
+    points: torch.Tensor,
+    k: int = 1,
+    sample_size: Optional[int] = 10000,
+    seed: Optional[int] = 42,
+    quantile: float = 0.05,
+    alpha: float = 0.2,
+    min_step: Optional[float] = None,
+    max_step: Optional[float] = None,
+) -> XYZQuantConfig:
+    """
+    Compute quantization configuration from points.
+
+    Internally computes step_size via:
+        1. compute_nn_distances() - k-th nearest neighbor distances
+        2. compute_dense_scale() - low quantile of NN distances
+        3. compute_step_size() - alpha * dense_scale
+
+    Args:
+        points: Point coordinates, shape (N, 3).
+        k: Which nearest neighbor to use (1 = nearest).
+        sample_size: Number of points to sample for NN estimation.
+            Set to None to use all points. Default 10000 balances speed and accuracy.
+        seed: Random seed for reproducible sampling.
+        quantile: Quantile of NN distances to use (e.g., 0.01 or 0.05).
+            Lower values target denser regions.
+        alpha: Scaling factor for step size (typical range: 0.1 to 0.25).
+            Smaller values preserve more detail but increase data size.
+        min_step: Optional minimum step size to prevent extremely fine quantization.
+        max_step: Optional maximum step size to prevent overly coarse quantization.
+
+    Returns:
+        XYZQuantConfig with step size and origin (median of coordinates).
+    """
+    nn_distances = compute_nn_distances(points, k=k, sample_size=sample_size, seed=seed)
+    dense_scale = compute_dense_scale(nn_distances, quantile=quantile)
+    step_size = compute_step_size(dense_scale, alpha=alpha, min_step=min_step, max_step=max_step)
+    origin = points.median(dim=0).values
+    return XYZQuantConfig(step_size=step_size, origin=origin)
+
+
+def quantize_xyz(
+    points: torch.Tensor,
+    config: XYZQuantConfig,
+) -> torch.Tensor:
+    """
+    Quantize XYZ coordinates using the given configuration.
+
+    Applies the formula:
+        quantized = round((points - origin) / step_size)
+
+    Args:
+        points: Point coordinates, shape (N, 3).
+        config: Quantization configuration (step size and origin).
+
+    Returns:
+        Quantized coordinates as integers, shape (N, 3), dtype torch.int32.
+    """
+    normalized = (points - config.origin) / config.step_size
+    quantized = torch.round(normalized).to(torch.int32)
+    return quantized
+
+
+def dequantize_xyz(
+    quantized: torch.Tensor,
+    config: XYZQuantConfig,
+    dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """
+    Dequantize XYZ coordinates back to floating point.
+
+    Applies the formula:
+        points = quantized * step_size + origin
+
+    Args:
+        quantized: Quantized coordinates, shape (N, 3), dtype torch.int32.
+        config: Quantization configuration (step size and origin).
+        dtype: Output dtype for the dequantized coordinates.
+
+    Returns:
+        Dequantized coordinates, shape (N, 3).
+    """
+    return quantized.to(dtype) * config.step_size + config.origin
+
+
+def estimate_quantization_error(
+    points: torch.Tensor,
+    config: XYZQuantConfig,
+) -> dict:
+    """
+    Estimate the quantization error statistics.
+
+    Useful for debugging and validating quantization parameters.
+
+    Args:
+        points: Original point coordinates, shape (N, 3).
+        config: Quantization configuration.
+
+    Returns:
+        Dictionary with error statistics:
+            - 'max_error': Maximum absolute error across all coordinates.
+            - 'mean_error': Mean absolute error.
+            - 'rmse': Root mean squared error.
+            - 'relative_max_error': Max error relative to step size.
+    """
+    quantized = quantize_xyz(points, config)
+    reconstructed = dequantize_xyz(quantized, config, dtype=points.dtype)
+
+    error = (points - reconstructed).abs()
+
+    return {
+        'max_error': error.max().item(),
+        'mean_error': error.mean().item(),
+        'rmse': torch.sqrt((error ** 2).mean()).item(),
+        'relative_max_error': error.max().item() / config.step_size,
+    }