Upload learn_region_grow/io.py

learn_region_grow/io.py

@@ -0,0 +1,279 @@
+"""I/O utilities for PLY, PCD, and NumPy point clouds."""
+
+import numpy as np
+from pathlib import Path
+from typing import Tuple, Optional
+
+
+def load_ply(ply_path: str) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Load an ASCII PLY file and return points (N, 6: xyzrgb) and optionally normals.
+
+    Only supports 'vertex' elements with x,y,z and optionally r,g,b,nx,ny,nz.
+    Returns a (N, 3) xyz array and a (N, 6+) feature array (xyz + rgb + ...).
+    """
+    path = Path(ply_path)
+    if not path.is_file():
+        raise FileNotFoundError(f"PLY file not found: {ply_path}")
+
+    with open(path, 'r') as f:
+        header = []
+        while True:
+            line = f.readline().strip()
+            header.append(line)
+            if line.startswith("end_header"):
+                break
+            if not line:
+                raise ValueError("Invalid PLY header: missing end_header")
+
+    # Parse header to find vertex count and property names
+    n_vertices = 0
+    in_vertex = False
+    prop_names = []
+    for line in header:
+        if line.startswith("element vertex"):
+            parts = line.split()
+            n_vertices = int(parts[2])
+            in_vertex = True
+            prop_names = []
+        elif line.startswith("element"):
+            in_vertex = False
+        elif line.startswith("property") and in_vertex:
+            parts = line.split()
+            prop_names.append(parts[-1])
+
+    # Read body
+    data = np.loadtxt(path, skiprows=len(header), max_rows=n_vertices)
+    if data.ndim == 1:
+        data = data.reshape(1, -1)
+
+    # Map known properties
+    xyz = _get_props(data, prop_names, ["x", "y", "z"])
+    rgb = _get_props(data, prop_names, ["r", "g", "b"], dtype=np.uint8)
+    normals = _get_props(data, prop_names, ["nx", "ny", "nz"], optional=True)
+    return xyz, rgb, normals
+
+
+def _get_props(data, prop_names, keys, optional=False, dtype=np.float32):
+    indices = []
+    for k in keys:
+        if k in prop_names:
+            indices.append(prop_names.index(k))
+        elif k.capitalize() in prop_names:
+            indices.append(prop_names.index(k.capitalize()))
+        elif optional:
+            return None
+        else:
+            raise ValueError(f"Property {k} not found in PLY file. Available: {prop_names}")
+    if indices:
+        return data[:, indices].astype(dtype)
+    return None
+
+
+def load_pcd(pcd_path: str) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    """
+    Load a PCD (Point Cloud Data) file.
+
+    Supports ASCII PCD with fields: x y z rgb / rgba / r g b / normal_x normal_y normal_z.
+    Returns (xyz, rgb, normals) where rgb is uint8 (N,3) and normals is float32 (N,3) or None.
+    """
+    path = Path(pcd_path)
+    if not path.is_file():
+        raise FileNotFoundError(f"PCD file not found: {pcd_path}")
+
+    with open(path, 'r') as f:
+        lines = f.readlines()
+
+    # Find header end
+    header_end = 0
+    data_mode = "ascii"
+    fields = []
+    for i, line in enumerate(lines):
+        l = line.strip().split()
+        if not l:
+            continue
+        if l[0] == "DATA":
+            data_mode = l[1] if len(l) > 1 else "ascii"
+            header_end = i + 1
+            break
+        if l[0] == "FIELDS":
+            fields = l[1:]
+
+    if data_mode != "ascii":
+        raise NotImplementedError("Only ASCII PCD files are supported.")
+
+    data = np.loadtxt(lines[header_end:])
+    if data.ndim == 1:
+        data = data.reshape(1, -1)
+
+    xyz = _get_props(data, fields, ["x", "y", "z"])
+    rgb = None
+    normals = None
+
+    # Try RGB channels
+    if all(c in fields for c in ["r", "g", "b"]):
+        rgb = _get_props(data, fields, ["r", "g", "b"], dtype=np.uint8)
+    elif "rgb" in fields:
+        rgb_idx = fields.index("rgb")
+        rgb_packed = data[:, rgb_idx].astype(np.float32)
+        rgb = np.zeros((len(rgb_packed), 3), dtype=np.uint8)
+        # PCL packs RGB as a single float: int bits = 0x00RRGGBB
+        rgb_floats = rgb_packed.astype(np.float32)
+        rgb[:, 0] = ((rgb_floats.view(np.int32) >> 16) & 0xFF).astype(np.uint8)
+        rgb[:, 1] = ((rgb_floats.view(np.int32) >> 8) & 0xFF).astype(np.uint8)
+        rgb[:, 2] = ((rgb_floats.view(np.int32)) & 0xFF).astype(np.uint8)
+    elif "rgba" in fields:
+        rgba_idx = fields.index("rgba")
+        rgb_floats = data[:, rgba_idx].astype(np.float32)
+        rgb = np.zeros((len(rgb_floats), 3), dtype=np.uint8)
+        rgb[:, 0] = ((rgb_floats.view(np.int32) >> 16) & 0xFF).astype(np.uint8)
+        rgb[:, 1] = ((rgb_floats.view(np.int32) >> 8) & 0xFF).astype(np.uint8)
+        rgb[:, 2] = ((rgb_floats.view(np.int32)) & 0xFF).astype(np.uint8)
+
+    # Try normals
+    normals = _get_props(data, fields, ["normal_x", "normal_y", "normal_z"], optional=True)
+    if normals is None:
+        normals = _get_props(data, fields, ["nx", "ny", "nz"], optional=True)
+
+    return xyz, rgb, normals
+
+
+def save_ply(ply_path: str, xyz: np.ndarray, rgb: np.ndarray = None,
+             labels: np.ndarray = None, normals: np.ndarray = None):
+    """
+    Save a point cloud as ASCII PLY.
+
+    Parameters
+    ----------
+    xyz : np.ndarray, shape (N, 3)
+        XYZ coordinates.
+    rgb : np.ndarray, shape (N, 3), uint8, optional
+        Per-point colors. If labels is given instead, rgb is ignored and
+        colors are looked up from label_to_color.
+    labels : np.ndarray, shape (N,), optional
+        Integer instance labels. Used to colorize output.
+    normals : np.ndarray, shape (N, 3), optional
+        Surface normals.
+    """
+    n = len(xyz)
+    if labels is not None:
+        rgb = label_to_color(labels)
+    if rgb is None:
+        rgb = np.full((n, 3), 128, dtype=np.uint8)
+
+    props = ["property float x", "property float y", "property float z"]
+    data_cols = [xyz[:, 0:1], xyz[:, 1:2], xyz[:, 2:3]]
+
+    if normals is not None:
+        props += ["property float nx", "property float ny", "property float nz"]
+        data_cols += [normals[:, 0:1], normals[:, 1:2], normals[:, 2:3]]
+
+    props += ["property uchar red", "property uchar green", "property uchar blue"]
+    data_cols += [rgb[:, 0:1], rgb[:, 1:2], rgb[:, 2:3]]
+
+    if labels is not None:
+        props += ["property int label"]
+        data_cols += [labels.reshape(-1, 1)]
+
+    header = f"""ply
+format ascii 1.0
+element vertex {n}
+""" + "\n".join(props) + """
+end_header
+"""
+    data = np.hstack(data_cols)
+    if labels is not None:
+        # last column is int, rest float/uint8 — save manually
+        lines = [header]
+        for i in range(n):
+            parts = [f"{data[i,j]:.6f}" for j in range(data.shape[1]-1)]
+            parts.append(f"{int(data[i,-1])}")
+            lines.append(" ".join(parts) + "\n")
+        with open(ply_path, 'w') as f:
+            f.writelines(lines)
+    else:
+        np.savetxt(ply_path, data, header=header.strip(), comments='', fmt='%.6f')
+
+
+def save_pcd(pcd_path: str, xyz: np.ndarray, rgb: np.ndarray = None,
+             labels: np.ndarray = None, normals: np.ndarray = None):
+    """
+    Save a point cloud as ASCII PCD.
+    """
+    n = len(xyz)
+    if labels is not None:
+        rgb = label_to_color(labels)
+    if rgb is None:
+        rgb = np.full((n, 3), 128, dtype=np.uint8)
+
+    fields = ["x", "y", "z"]
+    types = ["F", "F", "F"]
+    sizes = [4, 4, 4]
+    counts = [1, 1, 1]
+
+    if normals is not None:
+        fields += ["normal_x", "normal_y", "normal_z"]
+        types += ["F", "F", "F"]
+        sizes += [4, 4, 4]
+        counts += [1, 1, 1]
+
+    fields += ["r", "g", "b"]
+    types += ["U", "U", "U"]
+    sizes += [1, 1, 1]
+    counts += [1, 1, 1]
+
+    data = np.hstack([xyz, rgb.astype(np.float32)])
+    if normals is not None:
+        data = np.hstack([xyz, normals, rgb.astype(np.float32)])
+
+    header = f"""# .PCD v0.7 - Point Cloud Data file format
+VERSION 0.7
+FIELDS {' '.join(fields)}
+SIZE {' '.join(str(s) for s in sizes)}
+TYPE {' '.join(types)}
+COUNT {' '.join(str(c) for c in counts)}
+WIDTH {n}
+HEIGHT 1
+VIEWPOINT 0 0 0 1 0 0 0
+POINTS {n}
+DATA ascii
+"""
+    np.savetxt(pcd_path, data, header=header.strip(), comments='', fmt='%.6f')
+
+
+def load_point_cloud(path: str) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    """
+    Load a point cloud from either PLY or PCD format.
+
+    Returns
+    -------
+    xyz : np.ndarray, shape (N, 3), float32
+    rgb : np.ndarray, shape (N, 3), uint8
+    normals : np.ndarray, shape (N, 3), float32 or None
+    """
+    p = Path(path)
+    suffix = p.suffix.lower()
+    if suffix == ".ply":
+        return load_ply(path)
+    elif suffix == ".pcd":
+        return load_pcd(path)
+    else:
+        raise ValueError(f"Unsupported point cloud format: {suffix}. Use .ply or .pcd")
+
+
+# ------------------------------------------------------------------
+# Color map for instance labels (cyclic, similar to S3DIS / ScanNet)
+# ------------------------------------------------------------------
+def label_to_color(labels: np.ndarray) -> np.ndarray:
+    """Map integer instance labels to RGB colors."""
+    colors = np.array([
+        [0.8, 0.2, 0.2], [0.2, 0.8, 0.2], [0.2, 0.2, 0.8],
+        [0.8, 0.8, 0.2], [0.8, 0.2, 0.8], [0.2, 0.8, 0.8],
+        [0.5, 0.3, 0.1], [0.1, 0.5, 0.3], [0.3, 0.1, 0.5],
+        [0.9, 0.6, 0.1], [0.1, 0.9, 0.6], [0.6, 0.1, 0.9],
+        [0.4, 0.4, 0.4], [0.7, 0.7, 0.7], [0.3, 0.7, 0.5],
+        [0.7, 0.3, 0.5], [0.5, 0.7, 0.3], [0.3, 0.5, 0.7],
+        [0.9, 0.3, 0.1], [0.1, 0.3, 0.9], [0.3, 0.9, 0.1],
+    ], dtype=np.float32)
+    rgb = colors[labels % len(colors)] * 255
+    return rgb.astype(np.uint8)