3RScan/mesh_obj_texture_to_ply.py

#!/usr/bin/env python3
"""
Convert 3RScan mesh (mesh.refined.v2.obj + mesh.refined_0.png) to raw point cloud .ply with RGB.

Use when you have the official downloaded scan (e.g. from 3rscan_download.py) and want
raw geometry + texture colors as a point cloud for Sonata etc.

Usage:
  python mesh_obj_texture_to_ply.py -i data/19f1a892-.../mesh.refined.v2.obj -o converted/19f1a892.ply
  (Texture mesh.refined_0.png must be in the same dir as the .obj)
  Or:
  python mesh_obj_texture_to_ply.py -i data/19f1a892-a988-2bf8-8c91-0705cf396888 -o converted/
  (Treat -i as scan dir: uses mesh.refined.v2.obj + mesh.refined_0.png from that dir)
"""
import argparse
import os
import numpy as np

try:
    from PIL import Image
except ImportError:
    Image = None


def parse_obj(path: str):
    """Parse OBJ: return vertices (N,3), and vertex_index -> (vt_u, vt_v) from first face use."""
    vertices = []   # 1-based index -> (x,y,z)
    vts = []        # 1-based index -> (u,v)
    v_to_vt = {}    # vertex index (1-based) -> vt index (1-based), from first occurrence in f

    with open(path, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            if parts[0] == "v" and len(parts) >= 4:
                vertices.append((float(parts[1]), float(parts[2]), float(parts[3])))
            elif parts[0] == "vt" and len(parts) >= 3:
                vts.append((float(parts[1]), float(parts[2])))
            elif parts[0] == "f":
                for p in parts[1:]:
                    # v/vt/vn or v//vn or v
                    ids = p.split("/")
                    if len(ids) >= 2 and ids[1]:
                        v_idx = int(ids[0])
                        vt_idx = int(ids[1])
                        if v_idx not in v_to_vt:
                            v_to_vt[v_idx] = vt_idx
    return np.array(vertices, dtype=np.float32), vts, v_to_vt


def sample_texture(tex: np.ndarray, u: float, v: float) -> tuple:
    """Sample texture at UV [0,1]. tex shape (H,W,3). OBJ v often 0=bottom, image 0=top -> flip v."""
    h, w = tex.shape[0], tex.shape[1]
    u = np.clip(u, 0, 1)
    v = np.clip(v, 0, 1)
    # OBJ: v=0 usually bottom; image row 0 is top
    y = int((1 - v) * (h - 1e-6))
    x = int(u * (w - 1e-6))
    y = max(0, min(y, h - 1))
    x = max(0, min(x, w - 1))
    c = tex[y, x]
    if len(c) >= 3:
        return (c[0] / 255.0, c[1] / 255.0, c[2] / 255.0)
    return (c[0] / 255.0, c[0] / 255.0, c[0] / 255.0)


def write_ply(coord: np.ndarray, color: np.ndarray, out_path: str):
    n = coord.shape[0]
    if color.max() <= 1.1:
        c8 = (np.clip(color, 0, 1) * 255).astype(np.uint8)
    else:
        c8 = np.clip(color, 0, 255).astype(np.uint8)
    coord = coord.astype(np.float32)
    header = (
        "ply\nformat binary_little_endian 1.0\n"
        "element vertex {}\n"
        "property float x\nproperty float y\nproperty float z\n"
        "property uchar red\nproperty uchar green\nproperty uchar blue\n"
        "end_header\n"
    ).format(n)
    with open(out_path, "wb") as f:
        f.write(header.encode("ascii"))
        f.write(coord.tobytes())
        f.write(c8.tobytes())


def main():
    parser = argparse.ArgumentParser(description="Mesh OBJ + texture -> raw point cloud PLY.")
    parser.add_argument("-i", "--input", required=True,
                        help="Path to mesh.refined.v2.obj OR to scan dir containing it and mesh.refined_0.png")
    parser.add_argument("-o", "--output", required=True, help="Output .ply path (or dir if -i is scan dir)")
    args = parser.parse_args()

    obj_path = args.input
    if os.path.isdir(obj_path):
        obj_path = os.path.join(obj_path, "mesh.refined.v2.obj")
    if not os.path.isfile(obj_path):
        raise FileNotFoundError(f"OBJ not found: {obj_path}")

    base_dir = os.path.dirname(os.path.abspath(obj_path))
    tex_path = os.path.join(base_dir, "mesh.refined_0.png")
    if not os.path.isfile(tex_path):
        raise FileNotFoundError(f"Texture not found: {tex_path} (must be next to .obj)")

    vertices, vts, v_to_vt = parse_obj(obj_path)
    n_verts = len(vertices)
    if Image is None:
        raise RuntimeError("PIL/Pillow required: pip install Pillow")
    tex = np.array(Image.open(tex_path).convert("RGB"))

    colors = np.zeros((n_verts, 3), dtype=np.float32)
    for i in range(n_verts):
        v_idx_1based = i + 1
        if v_idx_1based in v_to_vt:
            vt_idx = v_to_vt[v_idx_1based]
            if 1 <= vt_idx <= len(vts):
                u, v = vts[vt_idx - 1]
                colors[i] = sample_texture(tex, u, v)
            else:
                colors[i] = (0.5, 0.5, 0.5)
        else:
            colors[i] = (0.5, 0.5, 0.5)

    out_path = args.output
    if os.path.isdir(args.output):
        if os.path.isdir(args.input):
            scan_id = os.path.basename(os.path.abspath(args.input))
        else:
            scan_id = os.path.basename(os.path.dirname(os.path.abspath(obj_path)))
        out_path = os.path.join(args.output, f"{scan_id}.ply")
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    write_ply(vertices, colors, out_path)
    print(f"Saved {n_verts} vertices -> {out_path}")


if __name__ == "__main__":
    main()