| |
| """ |
| Entropy-coded variable-bit compressor for CC-4DGS point clouds. |
| This prototype quantises coordinates/features/rest attributes, then |
| encodes each dimension with canonical Huffman coding to produce a |
| compact bitstream. Decompression reconstructs the PLY exactly as |
| quantised (loss defined solely by the chosen bit depths). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import math |
| import heapq |
| from pathlib import Path |
| from typing import Dict, Iterable, List, Tuple |
|
|
| import numpy as np |
| import subprocess |
| import shutil |
| import sys |
|
|
| REPO_ROOT = Path(__file__).resolve().parents[1] |
| if str(REPO_ROOT) not in sys.path: |
| sys.path.insert(0, str(REPO_ROOT)) |
|
|
| from tools.compress_point_cloud import ( |
| _load_ply, |
| _write_ply, |
| _quantise_linear, |
| _dequantise_linear, |
| ) |
|
|
|
|
| class BitWriter: |
| def __init__(self) -> None: |
| self.buffer = bytearray() |
| self.current = 0 |
| self.nbits = 0 |
|
|
| def write(self, value: int, bits: int) -> None: |
| for shift in range(bits - 1, -1, -1): |
| bit = (value >> shift) & 1 |
| self.current = (self.current << 1) | bit |
| self.nbits += 1 |
| if self.nbits == 8: |
| self.buffer.append(self.current) |
| self.current = 0 |
| self.nbits = 0 |
|
|
| def finish(self) -> Tuple[bytes, int]: |
| pad = 0 |
| if self.nbits: |
| pad = 8 - self.nbits |
| self.current <<= pad |
| self.buffer.append(self.current) |
| self.current = 0 |
| self.nbits = 0 |
| return bytes(self.buffer), pad |
|
|
|
|
| class BitReader: |
| def __init__(self, data: bytes) -> None: |
| self.data = data |
| self.byte_pos = 0 |
| self.bit_pos = 0 |
|
|
| def read_bit(self) -> int: |
| if self.byte_pos >= len(self.data): |
| raise EOFError("BitReader exhausted.") |
| byte = self.data[self.byte_pos] |
| bit = (byte >> (7 - self.bit_pos)) & 1 |
| self.bit_pos += 1 |
| if self.bit_pos == 8: |
| self.bit_pos = 0 |
| self.byte_pos += 1 |
| return bit |
|
|
|
|
| def _build_huffman_lengths(values: np.ndarray, alphabet_size: int) -> List[int]: |
| counts = np.bincount(values, minlength=alphabet_size).astype(np.int64) |
| heap: List[Tuple[int, int, int | Tuple[int, int]]] = [] |
| counter = 0 |
| for symbol, freq in enumerate(counts): |
| if freq > 0: |
| heapq.heappush(heap, (freq, counter, symbol)) |
| counter += 1 |
| if not heap: |
| raise ValueError("Empty frequency table; cannot build Huffman codes.") |
| if len(heap) == 1: |
| freq, _, symbol = heap[0] |
| lengths = [0] * alphabet_size |
| lengths[symbol] = 1 |
| return lengths |
| while len(heap) > 1: |
| freq1, _, node1 = heapq.heappop(heap) |
| freq2, _, node2 = heapq.heappop(heap) |
| heapq.heappush(heap, (freq1 + freq2, counter, (node1, node2))) |
| counter += 1 |
| root = heap[0][2] |
| lengths = [0] * alphabet_size |
|
|
| def assign(node: int | Tuple[int, int], depth: int) -> None: |
| if isinstance(node, int): |
| lengths[node] = depth |
| return |
| left, right = node |
| assign(left, depth + 1) |
| assign(right, depth + 1) |
|
|
| assign(root, 0) |
| return lengths |
|
|
|
|
| class CanonicalHuffman: |
| def __init__(self, lengths: List[int]) -> None: |
| self.lengths = lengths |
| self.code_map: Dict[int, Tuple[int, int]] = {} |
| self.decode_tables: Dict[int, Dict[int, int]] = {} |
| self.max_len = 0 |
| self._build() |
|
|
| def _build(self) -> None: |
| entries = [(l, sym) for sym, l in enumerate(self.lengths) if l > 0] |
| if not entries: |
| raise ValueError("No symbols to encode.") |
| entries.sort() |
| code = 0 |
| prev_len = entries[0][0] |
| code <<= prev_len |
| for length, symbol in entries: |
| if length > prev_len: |
| code <<= length - prev_len |
| prev_len = length |
| self.code_map[symbol] = (code, length) |
| self.decode_tables.setdefault(length, {})[code] = symbol |
| code += 1 |
| self.max_len = prev_len |
|
|
| def encode(self, values: np.ndarray) -> Tuple[bytes, int]: |
| writer = BitWriter() |
| for val in values: |
| code, length = self.code_map[val] |
| writer.write(code, length) |
| return writer.finish() |
|
|
| def decode(self, data: bytes, count: int) -> np.ndarray: |
| reader = BitReader(data) |
| out = np.empty(count, dtype=np.int64) |
| for idx in range(count): |
| code = 0 |
| for length in range(1, self.max_len + 1): |
| bit = reader.read_bit() |
| code = (code << 1) | bit |
| table = self.decode_tables.get(length) |
| if table and code in table: |
| out[idx] = table[code] |
| break |
| else: |
| raise ValueError("Failed to decode symbol.") |
| return out |
|
|
|
|
| def _encode_dimension(values: np.ndarray, bits: int) -> Tuple[np.ndarray, np.ndarray, bytes, int]: |
| alphabet_size = 1 << bits |
| lengths = _build_huffman_lengths(values, alphabet_size) |
| codec = CanonicalHuffman(lengths) |
| stream, pad = codec.encode(values) |
| return np.array(lengths, dtype=np.uint8), np.array([values.size], dtype=np.int64), stream, pad |
|
|
|
|
| def _decode_dimension(stream: bytes, lengths: np.ndarray, count: int) -> np.ndarray: |
| codec = CanonicalHuffman(lengths.tolist()) |
| return codec.decode(stream, count).astype(np.int64) |
|
|
|
|
| def _stack_fields(arrays: Dict[str, np.ndarray], names: Iterable[str]) -> np.ndarray: |
| return np.stack([arrays[name] for name in names], axis=1) |
|
|
|
|
| def compress_point_cloud( |
| input_path: Path, |
| output_path: Path, |
| *, |
| pos_bits: int, |
| feature_bits: int, |
| rest_bits: int, |
| ) -> None: |
| print(f"[entropy] Loading {input_path} ...") |
| arrays, original_names = _load_ply(input_path) |
| pos_names = ["x", "y", "z"] |
| rest_names = sorted([n for n in arrays if n.startswith("f_rest_")], key=lambda x: int(x.split("_")[-1])) |
| other_names = [n for n in original_names if n not in pos_names + rest_names] |
|
|
| pos = _stack_fields(arrays, pos_names) |
| others = _stack_fields(arrays, other_names) if other_names else None |
| rest = _stack_fields(arrays, rest_names) if rest_names else None |
|
|
| meta: Dict[str, object] = { |
| "version": 1, |
| "original_names": original_names, |
| "pos_names": pos_names, |
| "rest_names": rest_names, |
| "other_names": other_names, |
| "pos_bits": pos_bits, |
| "feature_bits": feature_bits, |
| "other_bits": feature_bits, |
| "rest_bits": rest_bits, |
| "count": int(pos.shape[0]), |
| } |
|
|
| compressed: Dict[str, np.ndarray] = {} |
|
|
| print("[entropy] Quantising positions ...") |
| pos_q, pos_min, pos_max = _quantise_linear(pos, pos_bits) |
| compressed["pos_min"] = pos_min |
| compressed["pos_max"] = pos_max |
| for dim in range(pos_q.shape[1]): |
| lengths, counts, stream, pad = _encode_dimension(pos_q[:, dim].astype(np.int64), pos_bits) |
| compressed[f"pos_lengths_{dim}"] = lengths |
| compressed[f"pos_counts_{dim}"] = counts |
| compressed[f"pos_stream_{dim}"] = np.frombuffer(stream, dtype=np.uint8) |
| compressed[f"pos_pad_{dim}"] = np.array([pad], dtype=np.uint8) |
|
|
| if others is not None: |
| print("[entropy] Quantising other attributes ...") |
| other_q, other_min, other_max = _quantise_linear(others, feature_bits) |
| compressed["other_min"] = other_min |
| compressed["other_max"] = other_max |
| for dim in range(other_q.shape[1]): |
| lengths, counts, stream, pad = _encode_dimension(other_q[:, dim].astype(np.int64), feature_bits) |
| compressed[f"other_lengths_{dim}"] = lengths |
| compressed[f"other_counts_{dim}"] = counts |
| compressed[f"other_stream_{dim}"] = np.frombuffer(stream, dtype=np.uint8) |
| compressed[f"other_pad_{dim}"] = np.array([pad], dtype=np.uint8) |
| else: |
| compressed["other_min"] = np.empty(0, dtype=np.float32) |
| compressed["other_max"] = np.empty(0, dtype=np.float32) |
|
|
| if rest is not None: |
| print("[entropy] Quantising f_rest* ...") |
| rest_q, rest_min, rest_max = _quantise_linear(rest, rest_bits) |
| compressed["rest_min"] = rest_min |
| compressed["rest_max"] = rest_max |
| for dim in range(rest_q.shape[1]): |
| lengths, counts, stream, pad = _encode_dimension(rest_q[:, dim].astype(np.int64), rest_bits) |
| compressed[f"rest_lengths_{dim}"] = lengths |
| compressed[f"rest_counts_{dim}"] = counts |
| compressed[f"rest_stream_{dim}"] = np.frombuffer(stream, dtype=np.uint8) |
| compressed[f"rest_pad_{dim}"] = np.array([pad], dtype=np.uint8) |
| else: |
| compressed["rest_min"] = np.empty(0, dtype=np.float32) |
| compressed["rest_max"] = np.empty(0, dtype=np.float32) |
|
|
| compressed["meta"] = np.array(json.dumps(meta)) |
| print(f"[entropy] Saving archive to {output_path} ...") |
| np.savez_compressed(output_path, **compressed) |
| print("[entropy] Compression done.") |
|
|
|
|
| def _decode_group(prefix: str, meta: Dict[str, object], packed) -> np.ndarray | None: |
| names: List[str] = meta[prefix + "_names"] |
| if not names: |
| return None |
| bits = int(meta[f"{prefix}_bits"]) |
| mins = packed[f"{prefix}_min"] |
| maxs = packed[f"{prefix}_max"] |
| decoded_cols = [] |
| for dim in range(len(names)): |
| lengths = packed[f"{prefix}_lengths_{dim}"] |
| counts = int(packed[f"{prefix}_counts_{dim}"][0]) |
| stream = bytes(packed[f"{prefix}_stream_{dim}"]) |
| values = _decode_dimension(stream, lengths, counts).astype(np.int64) |
| decoded = _dequantise_linear(values.reshape(-1, 1), mins[dim : dim + 1], maxs[dim : dim + 1], bits) |
| decoded_cols.append(decoded[:, 0]) |
| return np.stack(decoded_cols, axis=1) |
|
|
|
|
| def decompress_point_cloud(input_path: Path, output_path: Path) -> None: |
| print(f"[entropy] Loading archive {input_path} ...") |
| packed = np.load(input_path, allow_pickle=False) |
| meta = json.loads(str(packed["meta"])) |
| arrays: Dict[str, np.ndarray] = {} |
|
|
| pos = _decode_group("pos", meta, packed) |
| for idx, name in enumerate(meta["pos_names"]): |
| arrays[name] = pos[:, idx] |
|
|
| others = _decode_group("other", meta, packed) |
| if others is not None: |
| for idx, name in enumerate(meta["other_names"]): |
| arrays[name] = others[:, idx] |
|
|
| rest = _decode_group("rest", meta, packed) |
| if rest is not None: |
| for idx, name in enumerate(meta["rest_names"]): |
| arrays[name] = rest[:, idx] |
|
|
| print(f"[entropy] Writing reconstructed PLY to {output_path} ...") |
| _write_ply(output_path, arrays, meta["original_names"]) |
| print("[entropy] Decompression done.") |
|
|
|
|
| def build_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Entropy-coded point cloud compressor.") |
| sub = parser.add_subparsers(dest="command", required=True) |
|
|
| comp = sub.add_parser("compress") |
| 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=10) |
| comp.add_argument("--rest-bits", type=int, default=12) |
| comp.add_argument("--ram", action="store_true", help="After compression, render & metric the reconstructed PLY.") |
| comp.add_argument("--model-path", type=Path) |
| comp.add_argument("--dataset-path", type=Path) |
| comp.add_argument("--config-path", type=Path) |
| comp.add_argument("--iteration", type=int) |
| comp.add_argument("--deformation-source", type=Path) |
| comp.add_argument( |
| "--keep-temp-ply", |
| action="store_true", |
| help="Keep the reconstructed PLY produced for --ram.", |
| ) |
|
|
| decomp = sub.add_parser("decompress") |
| decomp.add_argument("input", type=Path) |
| decomp.add_argument("output", type=Path) |
| decomp.add_argument("--ram", action="store_true", help="After decompression, render & metric.") |
| decomp.add_argument("--model-path", type=Path) |
| decomp.add_argument("--dataset-path", type=Path) |
| decomp.add_argument("--config-path", type=Path) |
| decomp.add_argument("--iteration", type=int) |
| decomp.add_argument("--deformation-source", type=Path) |
|
|
| return parser |
|
|
|
|
| def main() -> None: |
| parser = build_parser() |
| args = parser.parse_args() |
| if args.command == "compress": |
| input_path = args.input |
| default_deform = None |
| if input_path.is_dir(): |
| ply_candidate = input_path / "point_cloud.ply" |
| if not ply_candidate.exists(): |
| parser.error(f"Iteration folder {input_path} is missing point_cloud.ply") |
| input_path = ply_candidate |
| possible_def = input_path.parent / "deformation.pth" |
| if possible_def.exists(): |
| default_deform = possible_def |
| else: |
| possible_def = input_path.parent / "deformation.pth" |
| if possible_def.exists(): |
| default_deform = possible_def |
|
|
| compress_point_cloud( |
| input_path, |
| args.output, |
| pos_bits=args.pos_bits, |
| feature_bits=args.feature_bits, |
| rest_bits=args.rest_bits, |
| ) |
| if args.ram: |
| required = [args.model_path, args.dataset_path, args.config_path, args.iteration] |
| if any(v is None for v in required): |
| parser.error("--ram requires --model-path, --dataset-path, --config-path, --iteration") |
| temp_ply = args.output.with_suffix(".ply") |
| decompress_point_cloud(args.output, temp_ply) |
|
|
| iter_dir = args.model_path / "point_cloud" / f"iteration_{args.iteration}" |
| iter_dir.mkdir(parents=True, exist_ok=True) |
| shutil.copy2(temp_ply, iter_dir / "point_cloud.ply") |
| deform_src = args.deformation_source or default_deform or (input_path.parent / "deformation.pth") |
| if not deform_src.exists(): |
| raise FileNotFoundError(f"Missing deformation checkpoint: {deform_src}") |
| shutil.copy2(deform_src, iter_dir / "deformation.pth") |
| render_cmd = [ |
| sys.executable, |
| str(REPO_ROOT / "render.py"), |
| "-s", |
| str(args.dataset_path), |
| "--configs", |
| str(args.config_path), |
| "--model_path", |
| str(args.model_path), |
| "--iteration", |
| str(args.iteration), |
| "--skip_train", |
| "--skip_video", |
| ] |
| metrics_cmd = [ |
| sys.executable, |
| str(REPO_ROOT / "metrics.py"), |
| "-m", |
| str(args.model_path), |
| ] |
| subprocess.run(render_cmd, cwd=REPO_ROOT, check=True) |
| subprocess.run(metrics_cmd, cwd=REPO_ROOT, check=True) |
| if not args.keep_temp_ply: |
| temp_ply.unlink(missing_ok=True) |
| elif args.command == "decompress": |
| if args.ram: |
| required = [args.model_path, args.dataset_path, args.config_path, args.iteration] |
| if any(v is None for v in required): |
| parser.error("--ram requires --model-path, --dataset-path, --config-path, --iteration") |
| iter_dir = args.model_path / "point_cloud" / f"iteration_{args.iteration}" |
| iter_dir.mkdir(parents=True, exist_ok=True) |
| target_ply = iter_dir / "point_cloud.ply" |
| decompress_point_cloud(args.input, target_ply) |
| if args.output.resolve() != target_ply.resolve(): |
| shutil.copy2(target_ply, args.output) |
| deform_src = args.deformation_source |
| if deform_src and deform_src.exists(): |
| shutil.copy2(deform_src, iter_dir / "deformation.pth") |
| render_cmd = [ |
| sys.executable, |
| str(REPO_ROOT / "render.py"), |
| "-s", |
| str(args.dataset_path), |
| "--configs", |
| str(args.config_path), |
| "--model_path", |
| str(args.model_path), |
| "--iteration", |
| str(args.iteration), |
| "--skip_train", |
| "--skip_video", |
| ] |
| metrics_cmd = [ |
| sys.executable, |
| str(REPO_ROOT / "metrics.py"), |
| "-m", |
| str(args.model_path), |
| ] |
| subprocess.run(render_cmd, cwd=REPO_ROOT, check=True) |
| subprocess.run(metrics_cmd, cwd=REPO_ROOT, check=True) |
| else: |
| decompress_point_cloud(args.input, args.output) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
