decimate tool

decimate.py

@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+# filepath: decimate_splat.py
+"""
+Decimate a Gaussian Splat PLY file and output as PLY or SPLAT format.
+"""
+
+import argparse
+import numpy as np
+from io import BytesIO
+from pathlib import Path
+from plyfile import PlyData, PlyElement
+
+
+def load_gaussian_ply(ply_file_path: str) -> PlyData:
+    """Load a Gaussian splat PLY file."""
+    return PlyData.read(ply_file_path)
+
+
+def compute_importance_scores(vert) -> np.ndarray:
+    """
+    Compute importance scores for each Gaussian.
+    Higher scores = more important (larger and more opaque).
+    """
+    scales = np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+    opacities = 1 / (1 + np.exp(-vert["opacity"]))
+    return scales * opacities
+
+
+def decimate_ply(plydata: PlyData, keep_ratio: float) -> PlyData:
+    """
+    Decimate the PLY data by keeping only a fraction of the Gaussians.
+    Keeps the most important Gaussians based on scale and opacity.
+    """
+    vert = plydata["vertex"]
+    total_points = len(vert.data)
+    keep_count = max(1, int(total_points * keep_ratio))
+    
+    # Compute importance and get indices of top Gaussians
+    importance = compute_importance_scores(vert)
+    sorted_indices = np.argsort(-importance)[:keep_count]
+    
+    # Sort indices to maintain some spatial coherence
+    sorted_indices = np.sort(sorted_indices)
+    
+    # Create new vertex data with only kept points
+    new_vertex_data = vert.data[sorted_indices]
+    
+    # Create new PlyElement and PlyData
+    new_vertex_element = PlyElement.describe(new_vertex_data, "vertex")
+    new_plydata = PlyData([new_vertex_element])
+    
+    return new_plydata
+
+
+def convert_ply_to_splat(plydata: PlyData) -> bytes:
+    """
+    Convert PLY data to SPLAT format for the antimatter15 viewer.
+    Returns the splat data as bytes.
+    """
+    vert = plydata["vertex"]
+    
+    sorted_indices = np.argsort(
+        -np.exp(vert["scale_0"] + vert["scale_1"] + vert["scale_2"])
+        / (1 + np.exp(-vert["opacity"]))
+    )
+    
+    buffer = BytesIO()
+    for idx in sorted_indices:
+        v = plydata["vertex"][idx]
+        position = np.array([v["x"], v["y"], v["z"]], dtype=np.float32)
+        scales = np.exp(
+            np.array([v["scale_0"], v["scale_1"], v["scale_2"]], dtype=np.float32)
+        )
+        color = np.array([
+            0.5 + 0.28209479177387814 * v["f_dc_0"],
+            0.5 + 0.28209479177387814 * v["f_dc_1"],
+            0.5 + 0.28209479177387814 * v["f_dc_2"],
+            1 / (1 + np.exp(-v["opacity"])),
+        ])
+        rot = np.array([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], dtype=np.float32)
+        buffer.write(position.tobytes())
+        buffer.write(scales.tobytes())
+        buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+        buffer.write(
+            ((rot / np.linalg.norm(rot)) * 128 + 128).clip(0, 255).astype(np.uint8).tobytes()
+        )
+    
+    return buffer.getvalue()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Decimate a Gaussian Splat PLY file and output as PLY or SPLAT format."
+    )
+    parser.add_argument(
+        "input",
+        type=str,
+        help="Input PLY file path"
+    )
+    parser.add_argument(
+        "-o", "--output",
+        type=str,
+        help="Output file path (default: input_decimated.ply or .splat)"
+    )
+    parser.add_argument(
+        "-r", "--ratio",
+        type=float,
+        default=0.5,
+        help="Ratio of points to keep (0.0-1.0, default: 0.5)"
+    )
+    parser.add_argument(
+        "-f", "--format",
+        type=str,
+        choices=["ply", "splat"],
+        default="ply",
+        help="Output format: 'ply' or 'splat' (default: ply)"
+    )
+    parser.add_argument(
+        "-v", "--verbose",
+        action="store_true",
+        help="Print verbose output"
+    )
+    
+    args = parser.parse_args()
+    
+    # Validate ratio
+    if not 0.0 < args.ratio <= 1.0:
+        parser.error("Ratio must be between 0.0 (exclusive) and 1.0 (inclusive)")
+    
+    # Determine output path
+    input_path = Path(args.input)
+    if args.output:
+        output_path = Path(args.output)
+    else:
+        suffix = ".splat" if args.format == "splat" else ".ply"
+        output_path = input_path.with_stem(f"{input_path.stem}_decimated").with_suffix(suffix)
+    
+    if args.verbose:
+        print(f"Loading: {args.input}")
+    
+    # Load PLY file
+    plydata = load_gaussian_ply(args.input)
+    original_count = len(plydata["vertex"].data)
+    
+    if args.verbose:
+        print(f"Original point count: {original_count:,}")
+    
+    # Decimate
+    decimated_plydata = decimate_ply(plydata, args.ratio)
+    new_count = len(decimated_plydata["vertex"].data)
+    
+    if args.verbose:
+        print(f"Decimated point count: {new_count:,} ({args.ratio * 100:.1f}%)")
+    
+    # Write output
+    if args.format == "splat":
+        splat_data = convert_ply_to_splat(decimated_plydata)
+        with open(output_path, "wb") as f:
+            f.write(splat_data)
+    else:
+        decimated_plydata.write(str(output_path))
+    
+    if args.verbose:
+        print(f"Saved to: {output_path}")
+    else:
+        print(f"Decimated {original_count:,} → {new_count:,} points, saved to {output_path}")
+
+
+if __name__ == "__main__":
+    main()
+