gsvvcompressor 1.2.0__cp311-cp311-macosx_10_9_universal2.whl

gsvvcompressor/combinations/vq_xyz_zstd.py

@@ -0,0 +1,106 @@
+"""
+VQ + XYZ quantization + Zstd compression encoder/decoder combination.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+from ..compress.zstd import ZstdSerializer, ZstdDeserializer
+from ..interframe.combine import (
+    CombinedInterframeCodecInterface,
+    CombinedInterframeEncoderInitConfig,
+)
+from ..interframe.encoder import InterframeEncoder
+from ..interframe.decoder import InterframeDecoder
+from ..vq.interface import VQInterframeCodecInterface, VQInterframeCodecConfig
+from ..xyz.interface import XYZQuantInterframeCodecInterface, XYZQuantInterframeCodecConfig
+from .registry import register_encoder, register_decoder
+
+
+def VQXYZZstdEncoder(
+    vq_config: VQInterframeCodecConfig,
+    xyz_config: XYZQuantInterframeCodecConfig,
+    zstd_level: int = 7,
+    payload_device: Optional[str] = None,
+) -> InterframeEncoder:
+    """Create an encoder combining VQ + XYZ quantization + Zstd compression."""
+    vq_interface = VQInterframeCodecInterface()
+    xyz_interface = XYZQuantInterframeCodecInterface()
+    combined_interface = CombinedInterframeCodecInterface([xyz_interface, vq_interface])
+    combined_config = CombinedInterframeEncoderInitConfig(
+        init_configs=[xyz_config, vq_config]
+    )
+    serializer = ZstdSerializer(level=zstd_level)
+    return InterframeEncoder(
+        serializer=serializer,
+        interface=combined_interface,
+        init_config=combined_config,
+        payload_device=payload_device,
+    )
+
+
+def VQXYZZstdDecoder(
+    payload_device: Optional[str] = None,
+    device: Optional[str] = None,
+) -> InterframeDecoder:
+    """Create a decoder for VQ + XYZ quantization + Zstd compressed data."""
+    vq_interface = VQInterframeCodecInterface()
+    xyz_interface = XYZQuantInterframeCodecInterface()
+    combined_interface = CombinedInterframeCodecInterface([xyz_interface, vq_interface])
+    deserializer = ZstdDeserializer()
+    return InterframeDecoder(
+        deserializer=deserializer,
+        interface=combined_interface,
+        payload_device=payload_device,
+        device=device,
+    )
+
+
+@dataclass
+class VQXYZZstdEncoderConfig:
+    """Configuration for VQ + XYZ + Zstd encoder."""
+    vq: VQInterframeCodecConfig = field(default_factory=VQInterframeCodecConfig)
+    xyz: XYZQuantInterframeCodecConfig = field(default_factory=XYZQuantInterframeCodecConfig)
+    zstd_level: int = 7
+    payload_device: Optional[str] = None
+
+
+@dataclass
+class VQXYZZstdDecoderConfig:
+    """Configuration for VQ + XYZ + Zstd decoder."""
+    payload_device: Optional[str] = None
+    device: Optional[str] = None
+
+
+def build_vqxyzzstd_encoder(config: VQXYZZstdEncoderConfig) -> InterframeEncoder:
+    """Build encoder from configuration."""
+    return VQXYZZstdEncoder(
+        vq_config=config.vq,
+        xyz_config=config.xyz,
+        zstd_level=config.zstd_level,
+        payload_device=config.payload_device,
+    )
+
+
+def build_vqxyzzstd_decoder(config: VQXYZZstdDecoderConfig) -> InterframeDecoder:
+    """Build decoder from configuration."""
+    return VQXYZZstdDecoder(
+        payload_device=config.payload_device,
+        device=config.device,
+    )
+
+
+# Register
+register_encoder(
+    "vqxyzzstd",
+    build_vqxyzzstd_encoder,
+    VQXYZZstdEncoderConfig,
+    "VQ + XYZ quantization + Zstd compression",
+)
+
+register_decoder(
+    "vqxyzzstd",
+    build_vqxyzzstd_decoder,
+    VQXYZZstdDecoderConfig,
+    "VQ + XYZ quantization + Zstd decompression",
+)

gsvvcompressor/compress/__init__.py

@@ -0,0 +1,5 @@
+"""Zstandard-based streaming serialization and deserialization."""
+
+from .zstd import ZstdDeserializer, ZstdSerializer
+
+__all__ = ["ZstdSerializer", "ZstdDeserializer"]

gsvvcompressor/compress/zstd.py

@@ -0,0 +1,144 @@
+"""
+Zstandard-based streaming serialization and deserialization.
+
+Uses cloudpickle + length-prefix framing + zstd compression for efficient
+streaming of Payload objects. This approach can serialize almost anything
+including nested custom classes, torch.Tensor, and numpy.ndarray.
+
+WARNING: cloudpickle uses pickle under the hood. Do NOT use with untrusted
+data sources. Only use with trusted data (local files, same-process, etc.).
+"""
+
+import struct
+from typing import Iterator
+
+import cloudpickle as pickle
+import zstandard as zstd
+
+from ..deserializer import AbstractDeserializer
+from ..payload import Payload
+from ..serializer import AbstractSerializer
+
+# 4-byte big-endian length prefix (max 4GB per object)
+_LEN = struct.Struct(">I")
+
+
+class ZstdSerializer(AbstractSerializer):
+    """
+    Streaming serializer using cloudpickle + length-prefix framing + zstd.
+
+    Each Payload is pickled, prefixed with its length, then compressed
+    incrementally using zstd streaming compression.
+    """
+
+    def __init__(self, level: int = 7):
+        """
+        Initialize the serializer.
+
+        Args:
+            level: Zstd compression level (1-22). Default is 7.
+        """
+        self._compressor = zstd.ZstdCompressor(level=level).compressobj()
+
+    def serialize_frame(self, payload: Payload) -> Iterator[bytes]:
+        """
+        Serialize and compress a Payload object.
+
+        Args:
+            payload: A Payload instance to serialize.
+
+        Yields:
+            Compressed byte chunks (may yield zero chunks if buffered).
+        """
+        # Pickle the payload
+        pickled = pickle.dumps(payload, protocol=pickle.DEFAULT_PROTOCOL)
+        # Add length prefix framing
+        framed = _LEN.pack(len(pickled)) + pickled
+        # Compress incrementally
+        out = self._compressor.compress(framed)
+        if out:
+            yield out
+
+    def flush(self) -> Iterator[bytes]:
+        """
+        Flush remaining compressed data.
+
+        Yields:
+            Final compressed byte chunks.
+        """
+        tail = self._compressor.flush()
+        if tail:
+            yield tail
+
+
+class ZstdDeserializer(AbstractDeserializer):
+    """
+    Streaming deserializer using zstd + length-prefix framing + cloudpickle.
+
+    Decompresses incoming bytes incrementally, buffers until a complete
+    length-prefixed frame is available, then unpickles and yields Payloads.
+    """
+
+    def __init__(self):
+        """Initialize the deserializer."""
+        self._decompressor = zstd.ZstdDecompressor().decompressobj()
+        self._buffer = bytearray()
+
+    def deserialize_frame(self, data: bytes) -> Iterator[Payload]:
+        """
+        Decompress and deserialize bytes to Payload objects.
+
+        Args:
+            data: Compressed bytes to deserialize.
+
+        Yields:
+            Complete Payload objects as they become available.
+        """
+        # Decompress and add to buffer
+        decompressed = self._decompressor.decompress(data)
+        if decompressed:
+            self._buffer.extend(decompressed)
+
+        # Yield complete frames
+        yield from self._extract_payloads()
+
+    def flush(self) -> Iterator[Payload]:
+        """
+        Flush any remaining buffered data.
+
+        Yields:
+            Any remaining Payload objects in the buffer.
+        """
+        # Try to extract any remaining complete payloads
+        yield from self._extract_payloads()
+
+        # If there's leftover data, it's incomplete/corrupted
+        if self._buffer:
+            raise ValueError(
+                f"Incomplete data in buffer: {len(self._buffer)} bytes remaining"
+            )
+
+    def _extract_payloads(self) -> Iterator[Payload]:
+        """
+        Extract complete payloads from the buffer.
+
+        Yields:
+            Complete Payload objects.
+        """
+        while True:
+            # Need at least length prefix
+            if len(self._buffer) < _LEN.size:
+                break
+
+            # Read length
+            (length,) = _LEN.unpack(self._buffer[: _LEN.size])
+
+            # Check if we have complete frame
+            if len(self._buffer) < _LEN.size + length:
+                break
+
+            # Extract and unpickle payload
+            pickled = bytes(self._buffer[_LEN.size: _LEN.size + length])
+            del self._buffer[: _LEN.size + length]
+
+            yield pickle.loads(pickled)

gsvvcompressor/decoder.py

@@ -0,0 +1,155 @@
+from abc import ABC, abstractmethod
+from typing import Iterator, Self
+
+import torch
+
+from gaussian_splatting import GaussianModel
+
+from .deserializer import AbstractDeserializer
+from .payload import Payload
+
+
+class AbstractDecoder(ABC):
+    """
+    Abstract base class for decompression algorithms.
+
+    This decoder uses a two-stage process:
+    1. Deserialize bytes to Payload objects (via the deserializer)
+    2. Unpack Payloads to frames (via `unpack`)
+
+    Subclasses must implement `unpack` and `flush_unpack`, and provide
+    a deserializer. This design separates frame unpacking logic from
+    deserialization format.
+    """
+
+    def __init__(
+        self,
+        deserializer: AbstractDeserializer,
+        payload_device: str | torch.device | None = None,
+        device: str | torch.device | None = None,
+    ):
+        """
+        Initialize the decoder.
+
+        Args:
+            deserializer: The deserializer to use for converting bytes to Payload.
+            payload_device: The target device for input Payloads before
+                unpacking (e.g., 'cpu', 'cuda'). If None, no device
+                transfer is performed.
+            device: The target device for output GaussianModel frames
+                (e.g., 'cpu', 'cuda'). If None, no device transfer is performed.
+        """
+        self._deserializer = deserializer
+        self._payload_device = payload_device
+        self._device = device
+
+    def to(self, device: str | torch.device | None) -> Self:
+        """Set the device for decoded models.
+
+        Args:
+            device: The device to move models to after decoding.
+                    Can be a string like "cuda" or "cpu", a torch.device, or None.
+
+        Returns:
+            self for method chaining.
+        """
+        self._device = device
+        return self
+
+    @abstractmethod
+    def unpack(self, payload: Payload) -> Iterator[GaussianModel]:
+        """
+        Unpack frame(s) from a Payload.
+
+        This method reconstructs the original frame(s) from the Payload object.
+
+        Args:
+            payload: A Payload instance to unpack.
+
+        Yields:
+            Unpacked GaussianModel instances. May yield zero, one, or multiple
+            models. When the iterator is exhausted, all frames for this payload
+            have been yielded.
+        """
+        pass
+
+    @abstractmethod
+    def flush_unpack(self) -> Iterator[GaussianModel]:
+        """
+        Flush any remaining buffered frames from the unpacking stage.
+
+        This method should be called after all payloads have been unpacked
+        to ensure any remaining buffered frames are output.
+
+        Yields:
+            Remaining buffered GaussianModel instances. May yield zero, one,
+            or multiple models until all buffered data has been flushed.
+        """
+        pass
+
+    def decode_frame(self, data: bytes) -> Iterator[GaussianModel]:
+        """
+        Decode a single chunk of compressed data.
+
+        This method deserializes the bytes to Payloads and then unpacks
+        the original frames.
+
+        Args:
+            data: A bytes object containing compressed data to decode.
+
+        Yields:
+            Decoded GaussianModel instances. May yield zero, one, or multiple
+            models. When the iterator is exhausted, all frames for this data
+            chunk have been decoded.
+        """
+        for payload in self._deserializer.deserialize_frame(data):
+            if self._payload_device is not None:
+                payload = payload.to(self._payload_device)
+            for model in self.unpack(payload):
+                if self._device is not None:
+                    model.to(self._device)
+                yield model
+
+    def flush(self) -> Iterator[GaussianModel]:
+        """
+        Flush any remaining buffered data from both deserialization and unpacking.
+
+        This method should be called after all data has been decoded
+        to ensure any remaining buffered frames are output.
+
+        Yields:
+            Remaining buffered GaussianModel instances. May yield zero, one,
+            or multiple models until all buffered data has been flushed.
+        """
+        # Flush deserialization stage and unpack any remaining payloads
+        for payload in self._deserializer.flush():
+            if self._payload_device is not None:
+                payload = payload.to(self._payload_device)
+            for model in self.unpack(payload):
+                if self._device is not None:
+                    model.to(self._device)
+                yield model
+
+        # Flush unpacking stage
+        for model in self.flush_unpack():
+            if self._device is not None:
+                model.to(self._device)
+            yield model
+
+    def decode_stream(self, stream: Iterator[bytes]) -> Iterator[GaussianModel]:
+        """
+        Decode a stream of compressed byte chunks into GaussianModel frames.
+
+        This method processes each chunk, deserializing and unpacking frames.
+        It handles the flush logic for both deserialization and unpacking stages.
+
+        Args:
+            stream: An iterator that yields bytes objects to decode.
+
+        Yields:
+            Decoded GaussianModel instances for each processed chunk or flush.
+        """
+        for data in stream:
+            yield from self.decode_frame(data)
+
+        yield from self.flush()

gsvvcompressor/deserializer.py

@@ -0,0 +1,42 @@
+from abc import ABC, abstractmethod
+from typing import Iterator
+
+from .payload import Payload
+
+
+class AbstractDeserializer(ABC):
+    """
+    Abstract base class for deserializing bytes to Payload objects.
+
+    Subclasses must implement `deserialize_frame` and `flush` methods to define
+    the specific deserialization logic for byte sequences into Payload.
+    """
+
+    @abstractmethod
+    def deserialize_frame(self, data: bytes) -> Iterator[Payload]:
+        """
+        Deserialize a chunk of bytes to Payload objects.
+
+        Args:
+            data: A bytes object containing serialized data to deserialize.
+
+        Yields:
+            Deserialized Payload instances. May yield zero, one, or multiple
+            payloads depending on available data. When the iterator is exhausted,
+            all complete payloads from this data chunk have been yielded.
+        """
+        pass
+
+    @abstractmethod
+    def flush(self) -> Iterator[Payload]:
+        """
+        Flush any remaining buffered data.
+
+        This method should be called after all data has been deserialized
+        to ensure any remaining buffered payloads are output.
+
+        Yields:
+            Remaining buffered Payload instances. May yield zero, one, or
+            multiple payloads until all buffered data has been flushed.
+        """
+        pass

gsvvcompressor/draco/__init__.py

@@ -0,0 +1,34 @@
+"""
+Draco-based compression for reduced 3DGS data.
+
+This module provides:
+- DracoPayload: Payload structure for Draco-compatible data
+- TwoPassDracoPayload: DracoPayload with num_points field for two-pass serialization
+- DracoInterframeCodecTranscodingInterface: Abstract interface for payload transcoding
+- DracoInterframeCodecInterface: Wrapper codec that transcodes to/from DracoPayload
+- DracoSerializer/DracoDeserializer: Streaming serialization using Draco encoding
+- TwoPassDracoSerializer/TwoPassDracoDeserializer: Two-pass serialization for batched compression
+"""
+
+from .interface import (
+    DracoInterframeCodecInterface,
+    DracoInterframeCodecTranscodingInterface,
+    DracoPayload,
+)
+from .serialize import DracoDeserializer, DracoSerializer
+from .twopass import (
+    TwoPassDracoDeserializer,
+    TwoPassDracoPayload,
+    TwoPassDracoSerializer,
+)
+
+__all__ = [
+    "DracoPayload",
+    "TwoPassDracoPayload",
+    "DracoInterframeCodecTranscodingInterface",
+    "DracoInterframeCodecInterface",
+    "DracoSerializer",
+    "DracoDeserializer",
+    "TwoPassDracoSerializer",
+    "TwoPassDracoDeserializer",
+]