gsvvcompressor 1.2.0__cp310-cp310-win_amd64.whl

gsvvcompressor/interframe/twopass.py

@@ -0,0 +1,226 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Iterator, List, Optional
+
+from gaussian_splatting import GaussianModel
+
+from ..encoder import AbstractEncoder
+from ..payload import Payload
+from ..serializer import AbstractSerializer
+from .interface import (
+    InterframeCodecContext,
+    InterframeCodecInterface,
+    InterframeEncoderInitConfig,
+)
+from .encoder import InterframeEncoder
+
+
+@dataclass
+class PassOneContext:
+    """
+    Context data from the first pass of two-pass encoding.
+
+    This dataclass holds the information gathered during the first pass
+    over all frames, which is then used to optimize the encoding in the
+    second pass. For example, this could contain statistics about the
+    entire video sequence, codebook data, or other global information.
+
+    Subclasses should define specific fields for their encoding scheme.
+    """
+    pass
+
+
+class TwoPassInterframeCodecInterface(InterframeCodecInterface):
+    """
+    Abstract interface for two-pass inter-frame encoding/decoding algorithms.
+
+    This interface extends InterframeCodecInterface to support two-pass encoding,
+    where the first pass gathers information about all frames, and the second
+    pass uses this information to optimize encoding.
+
+    Design Guidelines:
+        - PassOneContext: Store information gathered during the first pass
+          (e.g., global statistics, codebook data, quality parameters).
+        - The keyframe_to_context method receives PassOneContext instead of
+          InterframeEncoderInitConfig, allowing it to use first-pass data.
+    """
+
+    @abstractmethod
+    def keyframe_to_context_pass_one(
+        self,
+        frame: GaussianModel,
+        init_config: InterframeEncoderInitConfig,
+    ) -> PassOneContext:
+        """
+        Process the keyframe during the first pass to initialize PassOneContext.
+
+        This method is called by the encoder when processing the first frame
+        during the first pass. It initializes the PassOneContext with information
+        from the keyframe.
+
+        Args:
+            frame: The GaussianModel keyframe to process.
+            init_config: Encoder initialization configuration.
+
+        Returns:
+            PassOneContext initialized from the keyframe.
+        """
+        pass
+
+    @abstractmethod
+    def interframe_to_context_pass_one(
+        self,
+        frame: GaussianModel,
+        prev_pass_one_context: PassOneContext,
+    ) -> PassOneContext:
+        """
+        Process a frame during the first pass to update PassOneContext.
+
+        This method is called by the encoder when processing subsequent frames
+        during the first pass. It updates the PassOneContext with information
+        gathered from the current frame.
+
+        Args:
+            frame: The GaussianModel frame to process.
+            prev_pass_one_context: The PassOneContext from processing previous frames.
+
+        Returns:
+            Updated PassOneContext incorporating information from this frame.
+        """
+        pass
+
+    @abstractmethod
+    def keyframe_to_context(
+        self,
+        frame: GaussianModel,
+        pass_one_context: PassOneContext,
+    ) -> InterframeCodecContext:
+        """
+        Convert a keyframe to a Context using first-pass information.
+
+        This method is called by the encoder when processing the first frame
+        in the second pass. The pass_one_context provides information gathered
+        during the first pass for optimized encoding.
+
+        Args:
+            frame: The GaussianModel frame to convert.
+            pass_one_context: Context from the first pass containing
+                global encoding information.
+
+        Returns:
+            The corresponding Context representation.
+        """
+        pass
+
+
+class TwoPassInterframeEncoder(InterframeEncoder):
+    """
+    Encoder that uses two-pass inter-frame compression.
+
+    This encoder collects frames during pack() calls (pass one), then
+    performs the actual encoding in flush_pack() (pass two).
+
+    Pass one: Gather information from all frames to build PassOneContext.
+    Pass two: Encode all frames using the gathered information.
+    """
+
+    def __init__(
+        self,
+        serializer: AbstractSerializer,
+        interface: TwoPassInterframeCodecInterface,
+        init_config: InterframeEncoderInitConfig,
+        payload_device=None,
+    ):
+        """
+        Initialize the two-pass inter-frame encoder.
+
+        Args:
+            serializer: The serializer to use for converting Payload to bytes.
+            interface: The TwoPassInterframeCodecInterface instance that
+                provides encoding methods.
+            init_config: Configuration parameters for encoder initialization.
+        """
+        super().__init__(serializer=serializer, interface=interface, init_config=init_config, payload_device=payload_device)
+        self._interface = interface
+        self._frames: List[GaussianModel] = []
+        self._pass_one_context: Optional[PassOneContext] = None
+
+    def pack(self, frame: GaussianModel) -> Iterator[Payload]:
+        """
+        Perform pass one on the frame and store it for pass two.
+
+        During pass one, this method processes each frame to gather encoding
+        information and stores the frame for later encoding in flush_pack().
+
+        Args:
+            frame: A GaussianModel instance to process.
+
+        Yields:
+            No payloads during pass one (empty iterator).
+        """
+        # Pass one: gather information
+        if self._pass_one_context is None:
+            # First frame: keyframe
+            self._pass_one_context = self._interface.keyframe_to_context_pass_one(
+                frame, self._init_config
+            )
+        else:
+            # Subsequent frames: interframe
+            self._pass_one_context = self._interface.interframe_to_context_pass_one(
+                frame, self._pass_one_context
+            )
+
+        # Store frame for pass two
+        self._frames.append(frame)
+
+        # No payloads during pass one
+        return
+        yield  # Make this a generator
+
+    def flush_pack(self) -> Iterator[Payload]:
+        """
+        Perform pass two: encode all stored frames using pass one information.
+
+        This method encodes all frames that were collected during pack() calls,
+        using the PassOneContext gathered during pass one.
+
+        Yields:
+            Packed Payload instances for all frames.
+        """
+        if not self._frames or self._pass_one_context is None:
+            return
+            yield  # Make this a generator
+
+        prev_context: Optional[InterframeCodecContext] = None
+
+        for frame in self._frames:
+            if prev_context is None:
+                # First frame: convert and encode as keyframe using pass_one_context
+                current_context = self._interface.keyframe_to_context(
+                    frame, self._pass_one_context
+                )
+                payload = self._interface.encode_keyframe(current_context)
+                # Decode back to get reconstructed context (avoid error accumulation)
+                reconstructed_context = self._interface.decode_keyframe_for_encode(
+                    payload, current_context
+                )
+            else:
+                # Subsequent frames: convert and encode as delta from previous
+                current_context = self._interface.interframe_to_context(
+                    frame, prev_context
+                )
+                payload = self._interface.encode_interframe(
+                    prev_context, current_context
+                )
+                # Decode back to get reconstructed context (avoid error accumulation)
+                reconstructed_context = self._interface.decode_interframe_for_encode(
+                    payload, prev_context
+                )
+
+            # Use reconstructed context as previous for next frame
+            prev_context = reconstructed_context
+
+            yield payload
+
+        # Clear stored frames after encoding
+        self._frames.clear()

gsvvcompressor/io/__init__.py

@@ -0,0 +1,31 @@
+"""
+IO module for reading and writing GaussianModel frames and bytes data.
+"""
+
+from .gaussian_model import FrameReader, FrameWriter
+from .bytes import BytesReader, BytesWriter
+from .config import (
+    FrameReaderConfig,
+    FrameWriterConfig,
+    build_frame_reader,
+    build_frame_writer,
+    BytesReaderConfig,
+    BytesWriterConfig,
+    build_bytes_reader,
+    build_bytes_writer,
+)
+
+__all__ = [
+    "FrameReader",
+    "FrameWriter",
+    "FrameReaderConfig",
+    "FrameWriterConfig",
+    "build_frame_reader",
+    "build_frame_writer",
+    "BytesReader",
+    "BytesWriter",
+    "BytesReaderConfig",
+    "BytesWriterConfig",
+    "build_bytes_reader",
+    "build_bytes_writer",
+]

gsvvcompressor/io/bytes.py

@@ -0,0 +1,103 @@
+"""
+Bytes data reader and writer classes.
+
+This module provides classes for reading and writing bytes data from/to files.
+"""
+
+import os
+from typing import Iterator
+
+
+class BytesReader:
+    """
+    Reader for bytes data from a file.
+
+    This class reads bytes data from a file and yields it as an iterator.
+    Each yielded item is a bytes chunk from the file.
+
+    Example:
+        reader = BytesReader("data/compressed.bin")
+        for chunk in reader.read():
+            process(chunk)
+    """
+
+    # Default chunk size for reading (64KB)
+    DEFAULT_CHUNK_SIZE = 64 * 1024
+
+    def __init__(self, path: str, chunk_size: int = DEFAULT_CHUNK_SIZE):
+        """
+        Initialize the bytes reader.
+
+        Args:
+            path: Path to the file to read.
+            chunk_size: Size of each chunk to yield in bytes.
+                Default is 64KB.
+        """
+        self._path = path
+        self._chunk_size = chunk_size
+
+    def read(self) -> Iterator[bytes]:
+        """
+        Read bytes data from the file.
+
+        Yields:
+            bytes chunks from the file.
+
+        Raises:
+            FileNotFoundError: If the file does not exist.
+        """
+        if not os.path.exists(self._path):
+            raise FileNotFoundError(f"File not found: {self._path}")
+
+        with open(self._path, "rb") as f:
+            while True:
+                chunk = f.read(self._chunk_size)
+                if not chunk:
+                    break
+                yield chunk
+
+
+class BytesWriter:
+    """
+    Writer for bytes data to a file.
+
+    This class writes bytes data from an iterator to a file.
+
+    Example:
+        writer = BytesWriter("output/compressed.bin")
+        writer.write(bytes_iterator)
+    """
+
+    def __init__(self, path: str, overwrite: bool = False):
+        """
+        Initialize the bytes writer.
+
+        Args:
+            path: Path to the file to write.
+            overwrite: If True, overwrite existing file. Default is False.
+        """
+        self._path = path
+        self._overwrite = overwrite
+
+    def write(self, data: Iterator[bytes]) -> None:
+        """
+        Write bytes data to the file.
+
+        Args:
+            data: An iterator yielding bytes chunks to write.
+
+        Raises:
+            FileExistsError: If the target file already exists and overwrite is False.
+        """
+        # Check if file exists before writing
+        if not self._overwrite and os.path.exists(self._path):
+            raise FileExistsError(f"File already exists: {self._path}")
+
+        # Create parent directory if needed
+        parent_dir = os.path.dirname(self._path)
+        if parent_dir and not os.path.exists(parent_dir):
+            os.makedirs(parent_dir, exist_ok=True)
+
+        with open(self._path, "wb") as f:
+            for chunk in data:
+                f.write(chunk)

gsvvcompressor/io/config.py

@@ -0,0 +1,78 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from omegaconf import MISSING
+from .gaussian_model import FrameReader, FrameWriter
+from .bytes import BytesReader, BytesWriter
+
+
+@dataclass
+class FrameReaderConfig:
+    """Configuration for reading GaussianModel frame sequences."""
+    first_frame_path: str = MISSING
+    subsequent_format: str = MISSING
+    start_index: int = 2
+    sh_degree: int = 3
+    max_frames: Optional[int] = None
+    device: Optional[str] = None
+
+
+@dataclass
+class FrameWriterConfig:
+    """Configuration for writing GaussianModel frame sequences."""
+    first_frame_path: str = MISSING
+    subsequent_format: str = MISSING
+    start_index: int = 2
+    overwrite: bool = False
+
+
+def build_frame_reader(config: FrameReaderConfig) -> FrameReader:
+    """Build a FrameReader from configuration."""
+    return FrameReader(
+        first_frame_path=config.first_frame_path,
+        subsequent_format=config.subsequent_format,
+        start_index=config.start_index,
+        sh_degree=config.sh_degree,
+        max_frames=config.max_frames,
+        device=config.device,
+    )
+
+
+def build_frame_writer(config: FrameWriterConfig) -> FrameWriter:
+    """Build a FrameWriter from configuration."""
+    return FrameWriter(
+        first_frame_path=config.first_frame_path,
+        subsequent_format=config.subsequent_format,
+        start_index=config.start_index,
+        overwrite=config.overwrite,
+    )
+
+
+@dataclass
+class BytesReaderConfig:
+    """Configuration for reading bytes data from a file."""
+    path: str = MISSING
+    chunk_size: int = BytesReader.DEFAULT_CHUNK_SIZE
+
+
+@dataclass
+class BytesWriterConfig:
+    """Configuration for writing bytes data to a file."""
+    path: str = MISSING
+    overwrite: bool = False
+
+
+def build_bytes_reader(config: BytesReaderConfig) -> BytesReader:
+    """Build a BytesReader from configuration."""
+    return BytesReader(
+        path=config.path,
+        chunk_size=config.chunk_size,
+    )
+
+
+def build_bytes_writer(config: BytesWriterConfig) -> BytesWriter:
+    """Build a BytesWriter from configuration."""
+    return BytesWriter(
+        path=config.path,
+        overwrite=config.overwrite,
+    )

gsvvcompressor/io/gaussian_model.py

@@ -0,0 +1,127 @@
+"""
+Frame reader and writer classes for GaussianModel sequences.
+"""
+
+import os
+from typing import Iterator
+
+import torch
+from gaussian_splatting import GaussianModel
+
+
+class FrameReader:
+    """
+    Reader for GaussianModel frames from a series of files.
+
+    Example:
+        reader = FrameReader(
+            first_frame_path="data/frame_0000.ply",
+            subsequent_format="data/frame_{:04d}.ply",
+            start_index=1,
+            device="cuda"
+        )
+        for frame in reader.read():
+            process(frame)
+    """
+
+    def __init__(
+        self,
+        first_frame_path: str,
+        subsequent_format: str,
+        start_index: int,
+        sh_degree: int = 3,
+        max_frames: int | None = None,
+        device: str | torch.device | None = None,
+    ):
+        self._first_frame_path = first_frame_path
+        self._subsequent_format = subsequent_format
+        self._start_index = start_index
+        self._sh_degree = sh_degree
+        self._max_frames = max_frames
+        self._device = device
+
+    def to(self, device: str | torch.device | None) -> "FrameReader":
+        """Set the device for loaded models.
+
+        Args:
+            device: The device to move models to after loading.
+                    Can be a string like "cuda" or "cpu", a torch.device, or None.
+
+        Returns:
+            self for method chaining.
+        """
+        self._device = device
+        return self
+
+    def read(self) -> Iterator[GaussianModel]:
+        """Read GaussianModel frames from the file sequence."""
+        model = GaussianModel(sh_degree=self._sh_degree)
+        model.load_ply(self._first_frame_path)
+        if self._device is not None:
+            model.to(self._device)
+        yield model
+
+        frame_count = 1
+        if self._max_frames is not None and frame_count >= self._max_frames:
+            return
+
+        index = self._start_index
+        while True:
+            path = self._subsequent_format.format(index)
+            model = GaussianModel(sh_degree=self._sh_degree)
+            model.load_ply(path)
+            if self._device is not None:
+                model.to(self._device)
+            yield model
+            frame_count += 1
+            if self._max_frames is not None and frame_count >= self._max_frames:
+                break
+            index += 1
+
+
+class FrameWriter:
+    """
+    Writer for GaussianModel frames to a series of files.
+
+    Example:
+        writer = FrameWriter(
+            first_frame_path="output/frame_0000.ply",
+            subsequent_format="output/frame_{:04d}.ply",
+            start_index=1
+        )
+        writer.write(frame_iterator)
+    """
+
+    def __init__(
+        self,
+        first_frame_path: str,
+        subsequent_format: str,
+        start_index: int,
+        overwrite: bool = False,
+    ):
+        self._first_frame_path = first_frame_path
+        self._subsequent_format = subsequent_format
+        self._start_index = start_index
+        self._overwrite = overwrite
+
+    def write(self, frames: Iterator[GaussianModel]) -> None:
+        """Write GaussianModel frames to the file sequence."""
+        is_first = True
+        index = self._start_index
+
+        for frame in frames:
+            if is_first:
+                path = self._first_frame_path
+                is_first = False
+            else:
+                path = self._subsequent_format.format(index)
+                index += 1
+
+            if not self._overwrite and os.path.exists(path):
+                raise FileExistsError(f"File already exists: {path}")
+
+            parent_dir = os.path.dirname(path)
+            if parent_dir and not os.path.exists(parent_dir):
+                os.makedirs(parent_dir, exist_ok=True)
+
+            frame.save_ply(path)

gsvvcompressor/movecameras.py

@@ -0,0 +1,33 @@
+"""复制 cameras.json 和 cfg_args 文件从输入路径到输出路径"""
+
+import shutil
+from pathlib import Path
+import argparse
+
+
+def copy_configs(src: Path, dst: Path):
+    dst.mkdir(parents=True, exist_ok=True)
+    for name in ["cameras.json", "cfg_args"]:
+        if (src / name).exists():
+            shutil.copy2(src / name, dst / name)
+            print(f"{src / name} -> {dst / name}")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--first-frame-in")
+    p.add_argument("--first-frame-out")
+    p.add_argument("--subsequent-format-in")
+    p.add_argument("--subsequent-format-out")
+    p.add_argument("--max-frames", type=int)
+    args = p.parse_args()
+
+    if args.first_frame_in:
+        copy_configs(Path(args.first_frame_in), Path(args.first_frame_out))
+
+    if args.subsequent_format_in:
+        for i in range(2, args.max_frames + 1):
+            copy_configs(
+                Path(args.subsequent_format_in.format(i)),
+                Path(args.subsequent_format_out.format(i)),
+            )

gsvvcompressor/payload.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Self
+
+
+@dataclass
+class Payload(ABC):
+    """
+    Abstract base class for frame payload data.
+
+    This represents the intermediate data structure after processing a frame,
+    before serialization. Subclasses should define specific fields for their
+    encoding scheme.
+
+    The name "Payload" reflects its role as the processed data that will be
+    passed to a serializer for conversion to bytes.
+    """
+
+    @abstractmethod
+    def to(self, device) -> Self:
+        """
+        Move the Payload to the specified device for further processing.
+
+        Subclasses must implement this method to handle device transfer
+        (e.g., moving tensors to GPU/CPU).
+
+        Args:
+            device: The target device (e.g., 'cpu', 'cuda', torch.device).
+
+        Returns:
+            A new Payload instance on the target device, or self if already
+            on the target device.
+        """
+        pass

gsvvcompressor/serializer.py

@@ -0,0 +1,42 @@
+from abc import ABC, abstractmethod
+from typing import Iterator
+
+from .payload import Payload
+
+
+class AbstractSerializer(ABC):
+    """
+    Abstract base class for serializing Payload objects to bytes.
+
+    Subclasses must implement `serialize_frame` and `flush` methods to define
+    the specific serialization logic for Payload objects.
+    """
+
+    @abstractmethod
+    def serialize_frame(self, payload: Payload) -> Iterator[bytes]:
+        """
+        Serialize a single Payload object to bytes.
+
+        Args:
+            payload: A Payload instance to serialize.
+
+        Yields:
+            Serialized byte chunks. May yield zero, one, or multiple chunks.
+            When the iterator is exhausted, all data for this payload has been
+            serialized.
+        """
+        pass
+
+    @abstractmethod
+    def flush(self) -> Iterator[bytes]:
+        """
+        Flush any remaining buffered data.
+
+        This method should be called after all payloads have been serialized
+        to ensure any remaining buffered data is output.
+
+        Yields:
+            Remaining buffered byte chunks. May yield zero, one, or multiple
+            chunks until all buffered data has been flushed.
+        """
+        pass

gsvvcompressor/vq/__init__.py

@@ -0,0 +1,15 @@
+from .interface import (
+    VQInterframeCodecConfig,
+    VQInterframeCodecContext,
+    VQKeyframePayload,
+    VQInterframePayload,
+    VQInterframeCodecInterface,
+)
+
+__all__ = [
+    'VQInterframeCodecConfig',
+    'VQInterframeCodecContext',
+    'VQKeyframePayload',
+    'VQInterframePayload',
+    'VQInterframeCodecInterface',
+]