| |
| """ |
| Utility to pack Gaussian point clouds (PLY) into a quantised NPZ archive |
| and optionally restore them back. The compression pipeline is used by the |
| CC-4DGS post-training CCA tooling. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| from pathlib import Path |
| from typing import Dict, Iterable, List, Tuple |
|
|
| import numpy as np |
| from plyfile import PlyData, PlyElement |
|
|
|
|
| def _load_ply(path: Path) -> Tuple[Dict[str, np.ndarray], List[str]]: |
| ply = PlyData.read(path) |
| vertex = ply["vertex"].data |
| names = list(vertex.dtype.names) |
| arrays = {name: vertex[name].astype(np.float32).copy() for name in names} |
| return arrays, names |
|
|
|
|
| def _write_ply(path: Path, arrays: Dict[str, np.ndarray], order: Iterable[str]) -> None: |
| dtype = [(name, "<f4") for name in order] |
| count = len(next(iter(arrays.values()))) |
| structured = np.empty(count, dtype=dtype) |
| for name in order: |
| structured[name] = arrays[name].astype(np.float32) |
| vertex = PlyElement.describe(structured, "vertex") |
| ply = PlyData([vertex], text=False) |
| ply.write(str(path)) |
|
|
|
|
| def _select_uint_dtype(bits: int) -> np.dtype: |
| """Return the smallest unsigned dtype that can hold `bits` levels.""" |
| if bits <= 8: |
| return np.uint8 |
| if bits <= 16: |
| return np.uint16 |
| if bits <= 32: |
| return np.uint32 |
| raise ValueError(f"Quantisation with {bits} bits is not supported (max 32).") |
|
|
|
|
| def _quantise_linear(values: np.ndarray, bits: int) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """Linearly quantise along the last axis.""" |
| mins = values.min(axis=0) |
| maxs = values.max(axis=0) |
| spans = np.where(maxs > mins, maxs - mins, 1.0) |
| levels = (1 << bits) - 1 |
| scaled = np.round((values - mins) / spans * levels) |
| q_dtype = _select_uint_dtype(bits) |
| quantised = scaled.clip(0, levels).astype(q_dtype) |
| return quantised, mins.astype(np.float32), maxs.astype(np.float32) |
|
|
|
|
| def _dequantise_linear(quantised: np.ndarray, mins: np.ndarray, maxs: np.ndarray, bits: int) -> np.ndarray: |
| levels = (1 << bits) - 1 |
| spans = np.where(maxs > mins, maxs - mins, 1.0) |
| values = mins + (quantised.astype(np.float32) / levels) * spans |
| return values |
|
|
|
|
| def compress_point_cloud( |
| input_path: Path, |
| output_path: Path, |
| *, |
| pos_bits: int = 16, |
| feature_bits: int = 8, |
| coeff_bits: int = 12, |
| pca_components: int = 16, |
| drop_normals: bool = True, |
| keep_percent: float = 1.0, |
| top_k: int | None = None, |
| sort_by: str = "opacity", |
| min_field_value: float | None = None, |
| ) -> None: |
| arrays, original_names = _load_ply(input_path) |
| count = len(next(iter(arrays.values()))) |
|
|
| if sort_by not in arrays: |
| raise KeyError(f"Cannot sort by '{sort_by}' because the field is missing.") |
|
|
| |
| keep_mask = np.ones(count, dtype=bool) |
| if min_field_value is not None: |
| keep_mask &= arrays[sort_by] >= min_field_value |
|
|
| if top_k is not None and top_k < count: |
| top_indices = np.argsort(arrays[sort_by])[::-1][:top_k] |
| keep_mask = np.zeros(count, dtype=bool) |
| keep_mask[top_indices] = True |
| elif keep_percent < 1.0: |
| keep_n = int(math.ceil(count * keep_percent)) |
| if keep_n <= 0: |
| raise ValueError("keep_percent too small; nothing would be kept.") |
| sorted_idx = np.argsort(arrays[sort_by])[::-1][:keep_n] |
| keep_mask = np.zeros(count, dtype=bool) |
| keep_mask[sorted_idx] = True |
|
|
| if not keep_mask.all(): |
| for key in arrays: |
| arrays[key] = arrays[key][keep_mask] |
|
|
| pos_names = ["x", "y", "z"] |
| missing = [name for name in pos_names if name not in arrays] |
| if missing: |
| raise KeyError(f"The point cloud is missing mandatory coordinate(s): {missing}") |
|
|
| normal_names = [name for name in ("nx", "ny", "nz") if name in arrays] |
| dropped_normals = normal_names.copy() |
| rest_names = sorted(name for name in arrays if name.startswith("f_rest")) |
| dynamic_names = [name for name in arrays if name == "dynamic"] |
| other_names = [ |
| name |
| for name in arrays |
| if name not in pos_names + normal_names + rest_names + dynamic_names |
| ] |
|
|
| if drop_normals: |
| for name in normal_names: |
| arrays.pop(name, None) |
| normal_names = [] |
|
|
| pos = np.stack([arrays[name] for name in pos_names], axis=1) |
| rest = np.stack([arrays[name] for name in rest_names], axis=1) if rest_names else None |
| others = np.stack([arrays[name] for name in other_names], axis=1) if other_names else None |
| dynamic = ( |
| np.stack([arrays[name] for name in dynamic_names], axis=1) |
| if dynamic_names |
| else None |
| ) |
|
|
| compressed: Dict[str, np.ndarray] = {} |
| meta: Dict[str, object] = { |
| "version": 1, |
| "pos_names": pos_names, |
| "rest_names": rest_names, |
| "other_names": other_names, |
| "normal_names": dropped_normals if drop_normals else normal_names, |
| "original_names": original_names, |
| "pos_bits": pos_bits, |
| "feature_bits": feature_bits, |
| "coeff_bits": coeff_bits, |
| "pca_components": pca_components if rest is not None else 0, |
| "sorted_by": sort_by, |
| "kept_points": int(pos.shape[0]), |
| "original_points": count, |
| "keep_percent": keep_percent, |
| "top_k": top_k, |
| "min_field_value": min_field_value, |
| "dynamic_names": dynamic_names, |
| } |
|
|
| pos_q, pos_min, pos_max = _quantise_linear(pos, pos_bits) |
| compressed["pos"] = pos_q |
| compressed["pos_min"] = pos_min |
| compressed["pos_max"] = pos_max |
|
|
| if others is not None: |
| features_q, features_min, features_max = _quantise_linear(others, feature_bits) |
| compressed["features"] = features_q |
| compressed["features_min"] = features_min |
| compressed["features_max"] = features_max |
| else: |
| compressed["features"] = np.empty((pos.shape[0], 0), dtype=np.uint8) |
| compressed["features_min"] = np.empty(0, dtype=np.float32) |
| compressed["features_max"] = np.empty(0, dtype=np.float32) |
|
|
| rest_mode = "none" |
| if rest is not None and pca_components > 0: |
| rest_mode = "pca" |
| mean = rest.mean(axis=0) |
| centered = rest - mean |
| u, s, vt = np.linalg.svd(centered, full_matrices=False) |
| k = min(pca_components, vt.shape[0]) |
| coeff = u[:, :k] * s[:k] |
| coeff_q, coeff_min, coeff_max = _quantise_linear(coeff, coeff_bits) |
| compressed["coeff"] = coeff_q |
| compressed["coeff_min"] = coeff_min |
| compressed["coeff_max"] = coeff_max |
| compressed["rest_mean"] = mean.astype(np.float32) |
| compressed["rest_basis"] = vt[:k].astype(np.float32) |
| elif rest is not None: |
| rest_mode = "direct" |
| meta["pca_components"] = 0 |
| rest_q, rest_min, rest_max = _quantise_linear(rest, coeff_bits) |
| compressed["rest_direct"] = rest_q |
| compressed["rest_direct_min"] = rest_min |
| compressed["rest_direct_max"] = rest_max |
| compressed["coeff"] = np.empty((pos.shape[0], 0), dtype=np.uint8) |
| compressed["coeff_min"] = np.empty(0, dtype=np.float32) |
| compressed["coeff_max"] = np.empty(0, dtype=np.float32) |
| compressed["rest_mean"] = np.empty(0, dtype=np.float32) |
| compressed["rest_basis"] = np.empty((0, 0), dtype=np.float32) |
| else: |
| meta["pca_components"] = 0 |
| compressed["coeff"] = np.empty((pos.shape[0], 0), dtype=np.uint8) |
| compressed["coeff_min"] = np.empty(0, dtype=np.float32) |
| compressed["coeff_max"] = np.empty(0, dtype=np.float32) |
| compressed["rest_mean"] = np.empty(0, dtype=np.float32) |
| compressed["rest_basis"] = np.empty((0, 0), dtype=np.float32) |
|
|
| meta["rest_mode"] = rest_mode |
| meta["rest_bits"] = coeff_bits |
|
|
| if dynamic is not None: |
| compressed["dynamic"] = dynamic.astype(np.float32) |
|
|
| compressed["meta"] = np.array(json.dumps(meta)) |
| np.savez_compressed(output_path, **compressed) |
|
|
|
|
| def decompress_point_cloud(input_path: Path, output_path: Path) -> None: |
| packed = np.load(input_path, allow_pickle=False) |
| meta = json.loads(str(packed["meta"])) |
|
|
| pos_bits = int(meta["pos_bits"]) |
| feature_bits = int(meta["feature_bits"]) |
| coeff_bits = int(meta["coeff_bits"]) |
|
|
| pos = _dequantise_linear(packed["pos"], packed["pos_min"], packed["pos_max"], pos_bits) |
|
|
| features = _dequantise_linear(packed["features"], packed["features_min"], packed["features_max"], feature_bits) |
| coeff = packed["coeff"] |
| if coeff.size: |
| coeff = _dequantise_linear(coeff, packed["coeff_min"], packed["coeff_max"], coeff_bits) |
| else: |
| coeff = np.empty((pos.shape[0], 0), dtype=np.float32) |
|
|
| rest_names: List[str] = meta["rest_names"] |
| rest_mode = meta.get("rest_mode", "pca" if int(meta.get("pca_components", 0)) > 0 else "none") |
| rest_bits = int(meta.get("rest_bits", coeff_bits)) |
| rest = None |
| if rest_names and rest_mode == "pca": |
| basis = packed["rest_basis"] |
| mean = packed["rest_mean"] |
| if coeff.size == 0: |
| raise RuntimeError("Compressed file lacks PCA coefficients required to reconstruct f_rest.") |
| rest = coeff @ basis + mean |
| elif rest_names and rest_mode == "direct": |
| rest = _dequantise_linear( |
| packed["rest_direct"], |
| packed["rest_direct_min"], |
| packed["rest_direct_max"], |
| rest_bits, |
| ) |
|
|
| arrays: Dict[str, np.ndarray] = {} |
|
|
| pos_names: List[str] = meta["pos_names"] |
| for idx, name in enumerate(pos_names): |
| arrays[name] = pos[:, idx] |
|
|
| offset = 0 |
| other_names: List[str] = meta["other_names"] |
| for name in other_names: |
| arrays[name] = features[:, offset] |
| offset += 1 |
|
|
| if rest is not None: |
| for idx, name in enumerate(rest_names): |
| arrays[name] = rest[:, idx] |
|
|
| dynamic_names: List[str] = meta["dynamic_names"] |
| if dynamic_names: |
| dynamic_arr = packed["dynamic"] |
| for idx, name in enumerate(dynamic_names): |
| arrays[name] = dynamic_arr[:, idx] |
|
|
| normal_names: List[str] = meta["normal_names"] |
| for name in normal_names: |
| arrays[name] = np.zeros(pos.shape[0], dtype=np.float32) |
|
|
| original_names: List[str] = meta["original_names"] |
| missing = [name for name in original_names if name not in arrays] |
| if missing: |
| raise KeyError(f"Failed to rebuild fields {missing}; metadata may be incompatible.") |
|
|
| _write_ply(output_path, arrays, original_names) |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Compress or decompress CC-4DGS point clouds.") |
| subparsers = parser.add_subparsers(dest="command", required=True) |
|
|
| comp = subparsers.add_parser("compress", help="Compress a binary PLY into a quantised NPZ archive.") |
| comp.add_argument("input", type=Path) |
| comp.add_argument("output", type=Path) |
| comp.add_argument("--pos-bits", type=int, default=16) |
| comp.add_argument("--feature-bits", type=int, default=8) |
| comp.add_argument("--coeff-bits", type=int, default=12) |
| comp.add_argument("--pca-components", type=int, default=16) |
| comp.add_argument("--keep-percent", type=float, default=1.0) |
| comp.add_argument("--top-k", type=int, default=None) |
| comp.add_argument("--sort-by", type=str, default="opacity") |
| comp.add_argument("--min-field-value", type=float, default=None) |
| comp.add_argument("--keep-normals", action="store_false", dest="drop_normals") |
|
|
| decomp = subparsers.add_parser("decompress", help="Rebuild a binary PLY from a compressed archive.") |
| decomp.add_argument("input", type=Path) |
| decomp.add_argument("output", type=Path) |
|
|
| return parser |
|
|
|
|
| def main() -> None: |
| parser = build_parser() |
| args = parser.parse_args() |
| if args.command == "compress": |
| compress_point_cloud( |
| args.input, |
| args.output, |
| pos_bits=args.pos_bits, |
| feature_bits=args.feature_bits, |
| coeff_bits=args.coeff_bits, |
| pca_components=args.pca_components, |
| drop_normals=args.drop_normals, |
| keep_percent=args.keep_percent, |
| top_k=args.top_k, |
| sort_by=args.sort_by, |
| min_field_value=args.min_field_value, |
| ) |
| elif args.command == "decompress": |
| decompress_point_cloud(args.input, args.output) |
| else: |
| parser.error(f"Unsupported command: {args.command}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
