off2centroids.py

#!/usr/bin/env python3
"""Convert COFF file to centroids.ply + triangles.bin for Potree triangle splatting."""

import argparse
import struct
import numpy as np
from plyfile import PlyData, PlyElement

# SH constants
SH_C0 = 0.28209479177387814  # degree 0
SH_C1 = 0.4886025119029199   # degree 1


def parse_coff(filepath: str) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Parse COFF file, return vertices, faces (indices), and face colors."""
    with open(filepath, "r") as f:
        # Header
        header = f.readline().strip()
        if header not in ("OFF", "COFF"):
            raise ValueError(f"Expected OFF or COFF header, got: {header}")

        # Counts
        counts = f.readline().strip().split()
        num_vertices, num_faces = int(counts[0]), int(counts[1])
        print(f"Reading {num_vertices} vertices, {num_faces} faces...")

        # Vertices (x, y, z per line)
        vertices = np.zeros((num_vertices, 3), dtype=np.float32)
        for i in range(num_vertices):
            parts = f.readline().strip().split()
            vertices[i] = [float(parts[0]), float(parts[1]), float(parts[2])]

        # Faces (3 v0 v1 v2 [r g b a])
        faces = np.zeros((num_faces, 3), dtype=np.uint32)
        colors = np.zeros((num_faces, 3), dtype=np.uint8)
        for i in range(num_faces):
            parts = f.readline().strip().split()
            # parts[0] is vertex count (3 for triangles)
            faces[i] = [int(parts[1]), int(parts[2]), int(parts[3])]
            if len(parts) >= 7:
                colors[i] = [int(parts[4]), int(parts[5]), int(parts[6])]

    return vertices, faces, colors


def main():
    parser = argparse.ArgumentParser(description="Convert COFF to centroids.ply + triangles.bin + sh.bin")
    parser.add_argument("input_off", help="Input .off file")
    parser.add_argument("output_ply", help="Output centroids.ply file")
    parser.add_argument("output_bin", help="Output triangles.bin file (36 bytes/triangle)")
    parser.add_argument("output_sh", help="Output sh.bin file (12 bytes/triangle, SH degree 0)")
    args = parser.parse_args()

    # Parse input
    vertices, faces, colors = parse_coff(args.input_off)
    num_faces = len(faces)

    # Compute centroids and offsets
    print("Computing centroids and offsets...")
    centroids = np.zeros((num_faces, 3), dtype=np.float32)
    offsets = np.zeros((num_faces, 9), dtype=np.float32)

    for i in range(num_faces):
        v0, v1, v2 = vertices[faces[i, 0]], vertices[faces[i, 1]], vertices[faces[i, 2]]
        centroid = (v0 + v1 + v2) / 3.0
        centroids[i] = centroid
        offsets[i, 0:3] = v0 - centroid
        offsets[i, 3:6] = v1 - centroid
        offsets[i, 6:9] = v2 - centroid

    # Write PLY
    print(f"Writing {args.output_ply}...")
    dtype = [
        ("x", "f4"), ("y", "f4"), ("z", "f4"),
        ("red", "u1"), ("green", "u1"), ("blue", "u1"),
        ("original_index", "u4"),
    ]
    ply_data = np.zeros(num_faces, dtype=dtype)
    ply_data["x"] = centroids[:, 0]
    ply_data["y"] = centroids[:, 1]
    ply_data["z"] = centroids[:, 2]
    ply_data["red"] = colors[:, 0]
    ply_data["green"] = colors[:, 1]
    ply_data["blue"] = colors[:, 2]
    ply_data["original_index"] = np.arange(num_faces, dtype=np.uint32)

    el = PlyElement.describe(ply_data, "vertex")
    PlyData([el], byte_order="<").write(args.output_ply)

    # Convert RGB to SH degree 0 (DC component): sh_dc = (rgb/255 - 0.5) / SH_C0
    print("Converting RGB to SH degree 0...")
    rgb_norm = colors.astype(np.float32) / 255.0
    sh_dc = ((rgb_norm - 0.5) / SH_C0).astype(np.float32)

    # Write triangles.bin (geometry only, 36 bytes per triangle)
    print(f"Writing {args.output_bin}...")
    with open(args.output_bin, "wb") as f:
        for i in range(num_faces):
            offsets[i].tofile(f)

    # Write sh.bin (SH degree 0 coefficients, 12 bytes per triangle)
    print(f"Writing {args.output_sh}...")
    sh_dc.tofile(args.output_sh)

    print("Done.")
    print(f"  PLY: {num_faces} centroids")
    print(f"  BIN: {num_faces * 36} bytes (geometry only)")
    print(f"  SH:  {num_faces * 12} bytes (SH degree 0, 3 floats per triangle)")


if __name__ == "__main__":
    main()