add I/O utilities and remeshing helpers

point2mesh/io_utils.py

@@ -0,0 +1,443 @@
+"""
+I/O utilities — load point clouds (PCD, PLY, XYZ, OBJ), export meshes,
+build initial meshes (convex hull / Poisson), and remesh between levels.
+"""
+
+from __future__ import annotations
+
+import os
+import struct
+import numpy as np
+from typing import Tuple, Optional
+from pathlib import Path
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Point cloud loaders (no Open3D dependency — pure Python/numpy)
+# ──────────────────────────────────────────────────────────────────────
+def load_pointcloud(path: str) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+    """
+    Load a point cloud from .ply, .pcd, .xyz, or .obj.
+
+    Returns
+    -------
+    points  : (N, 3)  float64
+    normals : (N, 3)  float64 or None
+    """
+    ext = Path(path).suffix.lower()
+    if ext == ".ply":
+        return _load_ply(path)
+    elif ext == ".pcd":
+        return _load_pcd(path)
+    elif ext == ".xyz":
+        data = np.loadtxt(path)
+        if data.shape[1] >= 6:
+            return data[:, :3], data[:, 3:6]
+        return data[:, :3], None
+    elif ext == ".obj":
+        return _load_obj_points(path)
+    else:
+        raise ValueError(f"Unsupported format: {ext}")
+
+
+# ── PLY loader ────────────────────────────────────────────────────────
+def _load_ply(path: str):
+    with open(path, "rb") as f:
+        header = b""
+        while True:
+            line = f.readline()
+            header += line
+            if b"end_header" in line:
+                break
+        header_str = header.decode("ascii", errors="ignore")
+
+    lines = header_str.split("\n")
+    n_verts = 0
+    props = []
+    is_binary_le = False
+    is_binary_be = False
+    for line in lines:
+        line = line.strip()
+        if line.startswith("element vertex"):
+            n_verts = int(line.split()[-1])
+        elif line.startswith("property"):
+            parts = line.split()
+            props.append((parts[1], parts[2]))
+        elif line.startswith("format binary_little_endian"):
+            is_binary_le = True
+        elif line.startswith("format binary_big_endian"):
+            is_binary_be = True
+
+    # Get column indices for x y z nx ny nz
+    prop_names = [p[1] for p in props]
+    x_i = prop_names.index("x") if "x" in prop_names else 0
+    y_i = prop_names.index("y") if "y" in prop_names else 1
+    z_i = prop_names.index("z") if "z" in prop_names else 2
+    has_normals = "nx" in prop_names
+    if has_normals:
+        nx_i = prop_names.index("nx")
+        ny_i = prop_names.index("ny")
+        nz_i = prop_names.index("nz")
+
+    if is_binary_le or is_binary_be:
+        endian = "<" if is_binary_le else ">"
+        dtype_map = {
+            "float": f"{endian}f4", "double": f"{endian}f8",
+            "uchar": "u1", "uint8": "u1",
+            "int": f"{endian}i4", "uint": f"{endian}u4",
+            "short": f"{endian}i2", "ushort": f"{endian}u2",
+        }
+        dt = np.dtype([(p[1], dtype_map.get(p[0], f"{endian}f4")) for p in props])
+        header_end = header_str.index("end_header") + len("end_header") + 1
+        with open(path, "rb") as f:
+            f.seek(len(header))
+            raw = np.frombuffer(f.read(n_verts * dt.itemsize), dtype=dt, count=n_verts)
+        pts = np.column_stack([raw[prop_names[x_i]].astype(np.float64),
+                                raw[prop_names[y_i]].astype(np.float64),
+                                raw[prop_names[z_i]].astype(np.float64)])
+        normals = None
+        if has_normals:
+            normals = np.column_stack([raw[prop_names[nx_i]].astype(np.float64),
+                                        raw[prop_names[ny_i]].astype(np.float64),
+                                        raw[prop_names[nz_i]].astype(np.float64)])
+        return pts, normals
+    else:
+        # ASCII
+        data = np.loadtxt(path, skiprows=len(lines), max_rows=n_verts)
+        pts = data[:, [x_i, y_i, z_i]]
+        normals = None
+        if has_normals:
+            normals = data[:, [nx_i, ny_i, nz_i]]
+        return pts, normals
+
+
+# ── PCD loader ────────────────────────────────────────────────────────
+def _load_pcd(path: str):
+    with open(path, "rb") as f:
+        meta = {}
+        while True:
+            line = f.readline().decode("ascii", errors="ignore").strip()
+            if line.startswith("DATA"):
+                meta["data"] = line.split()[-1]
+                break
+            if " " in line:
+                key = line.split()[0]
+                val = line[len(key):].strip()
+                meta[key] = val
+        fields = meta.get("FIELDS", "x y z").split()
+        n_pts = int(meta.get("POINTS", meta.get("WIDTH", "0")))
+        sizes = list(map(int, meta.get("SIZE", "4 4 4").split()))
+        types = meta.get("TYPE", "F F F").split()
+
+        x_i = fields.index("x") if "x" in fields else 0
+        y_i = fields.index("y") if "y" in fields else 1
+        z_i = fields.index("z") if "z" in fields else 2
+        has_n = "normal_x" in fields
+        if has_n:
+            nx_i = fields.index("normal_x")
+            ny_i = fields.index("normal_y")
+            nz_i = fields.index("normal_z")
+
+        if meta["data"].lower() == "ascii":
+            data = np.loadtxt(f, max_rows=n_pts)
+            pts = data[:, [x_i, y_i, z_i]]
+            normals = data[:, [nx_i, ny_i, nz_i]] if has_n else None
+            return pts, normals
+        else:
+            # binary
+            row_size = sum(sizes)
+            raw = f.read(n_pts * row_size)
+            dt_map = {"F": "f", "U": "u", "I": "i"}
+            dtype = np.dtype(
+                [(fn, f"<{dt_map.get(t, 'f')}{s}") for fn, t, s in zip(fields, types, sizes)]
+            )
+            arr = np.frombuffer(raw, dtype=dtype, count=n_pts)
+            pts = np.column_stack([arr[fields[x_i]].astype(np.float64),
+                                    arr[fields[y_i]].astype(np.float64),
+                                    arr[fields[z_i]].astype(np.float64)])
+            normals = None
+            if has_n:
+                normals = np.column_stack([arr[fields[nx_i]].astype(np.float64),
+                                            arr[fields[ny_i]].astype(np.float64),
+                                            arr[fields[nz_i]].astype(np.float64)])
+            return pts, normals
+
+
+# ── OBJ point loader ──────────────────────────────────────────────────
+def _load_obj_points(path: str):
+    pts = []
+    normals = []
+    with open(path) as f:
+        for line in f:
+            if line.startswith("v "):
+                pts.append([float(x) for x in line.split()[1:4]])
+            elif line.startswith("vn "):
+                normals.append([float(x) for x in line.split()[1:4]])
+    pts = np.array(pts, dtype=np.float64)
+    normals = np.array(normals, dtype=np.float64) if normals else None
+    return pts, normals
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Mesh exporters
+# ──────────────────────────────────────────────────────────────────────
+def save_mesh_obj(path: str, vertices: np.ndarray, faces: np.ndarray):
+    """Write an OBJ file."""
+    with open(path, "w") as f:
+        for v in vertices:
+            f.write(f"v {v[0]:.8f} {v[1]:.8f} {v[2]:.8f}\n")
+        for face in faces:
+            f.write(f"f {face[0]+1} {face[1]+1} {face[2]+1}\n")
+
+
+def save_mesh_ply(path: str, vertices: np.ndarray, faces: np.ndarray):
+    """Write a binary-little-endian PLY mesh."""
+    nv, nf = len(vertices), len(faces)
+    header = (
+        "ply\n"
+        "format binary_little_endian 1.0\n"
+        f"element vertex {nv}\n"
+        "property float x\nproperty float y\nproperty float z\n"
+        f"element face {nf}\n"
+        "property list uchar int vertex_indices\n"
+        "end_header\n"
+    )
+    with open(path, "wb") as f:
+        f.write(header.encode("ascii"))
+        f.write(vertices.astype(np.float32).tobytes())
+        for face in faces:
+            f.write(struct.pack("<B3i", 3, *face.astype(np.int32)))
+
+
+def save_mesh(path: str, vertices: np.ndarray, faces: np.ndarray):
+    ext = Path(path).suffix.lower()
+    if ext == ".obj":
+        save_mesh_obj(path, vertices, faces)
+    elif ext == ".ply":
+        save_mesh_ply(path, vertices, faces)
+    elif ext == ".stl":
+        _save_stl(path, vertices, faces)
+    else:
+        save_mesh_obj(path, vertices, faces)
+
+
+def _save_stl(path: str, vertices: np.ndarray, faces: np.ndarray):
+    """Write a binary STL."""
+    nf = len(faces)
+    with open(path, "wb") as f:
+        f.write(b"\0" * 80)  # header
+        f.write(struct.pack("<I", nf))
+        for face in faces:
+            v0, v1, v2 = vertices[face[0]], vertices[face[1]], vertices[face[2]]
+            n = np.cross(v1 - v0, v2 - v0)
+            nl = np.linalg.norm(n)
+            if nl > 0:
+                n /= nl
+            f.write(struct.pack("<3f", *n.astype(np.float32)))
+            f.write(struct.pack("<3f", *v0.astype(np.float32)))
+            f.write(struct.pack("<3f", *v1.astype(np.float32)))
+            f.write(struct.pack("<3f", *v2.astype(np.float32)))
+            f.write(struct.pack("<H", 0))
+
+
+# ─────────────────���────────────────────────────────────────────────────
+#  Initial mesh construction
+# ──────────────────────────────────────────────────────────────────────
+def build_initial_mesh(
+    points: np.ndarray,
+    target_faces: int = 2000,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Build a coarse initial mesh from a point cloud.
+
+    Strategy: convex hull → subdivide to reach target face count →
+    Laplacian smooth → decimate back to target.
+
+    Returns (vertices, faces).
+    """
+    from scipy.spatial import ConvexHull
+
+    hull = ConvexHull(points)
+    verts = hull.points[hull.vertices].copy()
+
+    # Re-index faces so they reference the hull-vertex subset
+    full_to_hull = {v: i for i, v in enumerate(hull.vertices)}
+    faces = np.array(
+        [[full_to_hull[v] for v in f] for f in hull.simplices], dtype=np.int64
+    )
+
+    # Subdivide until we have at least target_faces
+    verts, faces = _subdivide_to_target(verts, faces, target_faces)
+
+    # Laplacian smoothing (5 iterations)
+    verts = _laplacian_smooth(verts, faces, iterations=5, lam=0.5)
+
+    # Decimate if we overshot
+    if len(faces) > target_faces * 1.5:
+        verts, faces = _decimate(verts, faces, target_faces)
+
+    return verts.astype(np.float64), faces.astype(np.int64)
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Remeshing between coarse-to-fine levels
+# ──────────────────────────────────────────────────────────────────────
+def remesh(
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    target_faces: int,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Subdivide the mesh then decimate to *target_faces*.
+    Applies one pass of Loop-like midpoint subdivision, Laplacian smooth,
+    and greedy edge-collapse decimation.
+    """
+    # 1. Subdivide (midpoint)
+    verts, faces_new = _subdivide_midpoint(vertices, faces)
+    # 2. Smooth
+    verts = _laplacian_smooth(verts, faces_new, iterations=3, lam=0.3)
+    # 3. Decimate to target
+    if len(faces_new) > target_faces:
+        verts, faces_new = _decimate(verts, faces_new, target_faces)
+    return verts.astype(np.float64), faces_new.astype(np.int64)
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Internal helpers
+# ──────────────────────────────────────────────────────────────────────
+def _subdivide_to_target(verts, faces, target):
+    while len(faces) < target:
+        verts, faces = _subdivide_midpoint(verts, faces)
+    return verts, faces
+
+
+def _subdivide_midpoint(verts, faces):
+    """One pass of midpoint subdivision (each tri → 4 tris)."""
+    edge_mid = {}
+    new_verts = list(verts)
+
+    def get_mid(a, b):
+        key = (min(a, b), max(a, b))
+        if key not in edge_mid:
+            mid = (verts[a] + verts[b]) / 2.0
+            idx = len(new_verts)
+            new_verts.append(mid)
+            edge_mid[key] = idx
+        return edge_mid[key]
+
+    new_faces = []
+    for f in faces:
+        a, b, c = int(f[0]), int(f[1]), int(f[2])
+        ab = get_mid(a, b)
+        bc = get_mid(b, c)
+        ca = get_mid(c, a)
+        new_faces.extend([
+            [a, ab, ca],
+            [ab, b, bc],
+            [ca, bc, c],
+            [ab, bc, ca],
+        ])
+    return np.array(new_verts), np.array(new_faces, dtype=np.int64)
+
+
+def _laplacian_smooth(verts, faces, iterations=3, lam=0.5):
+    from collections import defaultdict
+    n = len(verts)
+    adj = defaultdict(set)
+    for f in faces:
+        for i in range(3):
+            a, b = int(f[i]), int(f[(i + 1) % 3])
+            adj[a].add(b)
+            adj[b].add(a)
+    verts = verts.copy()
+    for _ in range(iterations):
+        new_v = verts.copy()
+        for i in range(n):
+            if adj[i]:
+                nbrs = np.array(list(adj[i]))
+                centroid = verts[nbrs].mean(axis=0)
+                new_v[i] = verts[i] + lam * (centroid - verts[i])
+        verts = new_v
+    return verts
+
+
+def _decimate(verts, faces, target):
+    """
+    Greedy edge-collapse decimation down to *target* faces.
+    Uses quadric error metric (simplified).
+    """
+    from collections import defaultdict
+    verts = verts.copy().astype(np.float64)
+    faces = faces.copy().astype(np.int64)
+
+    while len(faces) > target:
+        # Build edges and compute collapse cost (edge length)
+        edge_cost = {}
+        for fi, f in enumerate(faces):
+            for k in range(3):
+                a, b = int(f[k]), int(f[(k + 1) % 3])
+                key = (min(a, b), max(a, b))
+                if key not in edge_cost:
+                    edge_cost[key] = np.linalg.norm(verts[a] - verts[b])
+
+        if not edge_cost:
+            break
+
+        # Collapse cheapest edge
+        best_edge = min(edge_cost, key=edge_cost.get)
+        va, vb = best_edge
+
+        # Move va to midpoint
+        verts[va] = (verts[va] + verts[vb]) / 2.0
+
+        # Redirect vb → va in all faces
+        for fi in range(len(faces)):
+            for k in range(3):
+                if faces[fi][k] == vb:
+                    faces[fi][k] = va
+
+        # Remove degenerate faces (two or more identical vertex indices)
+        keep = []
+        for f in faces:
+            if f[0] != f[1] and f[1] != f[2] and f[0] != f[2]:
+                keep.append(f)
+        faces = np.array(keep, dtype=np.int64) if keep else np.zeros((0, 3), dtype=np.int64)
+
+        if len(faces) == 0:
+            break
+
+    # Compact vertex array (remove unreferenced verts)
+    used = np.unique(faces.ravel())
+    remap = np.full(len(verts), -1, dtype=np.int64)
+    remap[used] = np.arange(len(used))
+    verts = verts[used]
+    faces = remap[faces]
+
+    return verts, faces
+
+
+# ──────────────────────────────────────────────────────────────────────
+#  Normal estimation for point clouds without normals
+# ──────────────────────────────────────────────────────────────────────
+def estimate_normals(
+    points: np.ndarray, k: int = 20
+) -> np.ndarray:
+    """
+    PCA-based normal estimation using k nearest neighbors.
+    """
+    from scipy.spatial import cKDTree
+    tree = cKDTree(points)
+    _, idx = tree.query(points, k=k)
+    normals = np.zeros_like(points)
+    for i in range(len(points)):
+        nbrs = points[idx[i]]
+        centroid = nbrs.mean(axis=0)
+        cov = (nbrs - centroid).T @ (nbrs - centroid) / k
+        _, _, Vt = np.linalg.svd(cov)
+        normals[i] = Vt[-1]
+    # Orient normals consistently (towards centroid-outward)
+    centroid = points.mean(axis=0)
+    for i in range(len(normals)):
+        if np.dot(normals[i], points[i] - centroid) < 0:
+            normals[i] *= -1
+    return normals