mlx3d.io

mlx3d.io

GltfData dataclass

Result of :func:load_gltf.

Attributes:

Name	Type	Description
verts	array	(V, 3) positions.
faces	array	(F, 3) triangle indices.
normals	array | None	(V, 3) vertex normals, or None.
uvs	array | None	(V, 2) texture coordinates, or None.
material_ids	array | None	(F,) material index per face, or None.
materials	list[GltfMaterial] | None	decoded material summaries.
texture_image	array | None	first base-color texture image as (H, W, 3) in [0, 1] when the scene uses one; None otherwise.

Source code in src/mlx3d/io/gltf_io.py

@dataclass
class GltfData:
    """Result of :func:`load_gltf`.

    Attributes:
        verts: ``(V, 3)`` positions.
        faces: ``(F, 3)`` triangle indices.
        normals: ``(V, 3)`` vertex normals, or ``None``.
        uvs: ``(V, 2)`` texture coordinates, or ``None``.
        material_ids: ``(F,)`` material index per face, or ``None``.
        materials: decoded material summaries.
        texture_image: first base-color texture image as ``(H, W, 3)`` in
            ``[0, 1]`` when the scene uses one; ``None`` otherwise.
    """

    verts: mx.array
    faces: mx.array
    normals: mx.array | None = None
    uvs: mx.array | None = None
    material_ids: mx.array | None = None
    materials: list[GltfMaterial] | None = None
    texture_image: mx.array | None = None

ObjData dataclass

Result of :func:load_obj.

Attributes:

Name	Type	Description
verts	array	(V, 3) vertex positions.
faces	array	(F, 3) triangle vertex indices.
normals	array | None	(VN, 3) normal vectors from vn lines, or None.
texcoords	array | None	(VT, 2) texture coordinates from vt lines, or None.
faces_normals_idx	array | None	(F, 3) per-corner indices into normals, or None.
faces_texcoords_idx	array | None	(F, 3) per-corner indices into texcoords, or None.
verts_colors	array | None	(V, 3) vertex colors if the file used the color extension.
texture_path	str | None	resolved diffuse texture path from .mtl / map_Kd, or None.
texture_image	array | None	(H, W, 3) diffuse texture in [0, 1] when loaded, or None.

Source code in src/mlx3d/io/obj_io.py

@dataclass
class ObjData:
    """Result of :func:`load_obj`.

    Attributes:
        verts: (V, 3) vertex positions.
        faces: (F, 3) triangle vertex indices.
        normals: (VN, 3) normal vectors from ``vn`` lines, or ``None``.
        texcoords: (VT, 2) texture coordinates from ``vt`` lines, or ``None``.
        faces_normals_idx: (F, 3) per-corner indices into ``normals``, or ``None``.
        faces_texcoords_idx: (F, 3) per-corner indices into ``texcoords``, or ``None``.
        verts_colors: (V, 3) vertex colors if the file used the color extension.
        texture_path: resolved diffuse texture path from ``.mtl`` / ``map_Kd``, or ``None``.
        texture_image: (H, W, 3) diffuse texture in [0, 1] when loaded, or ``None``.
    """

    verts: mx.array
    faces: mx.array
    normals: mx.array | None = None
    texcoords: mx.array | None = None
    faces_normals_idx: mx.array | None = None
    faces_texcoords_idx: mx.array | None = None
    verts_colors: mx.array | None = None
    texture_path: str | None = None
    texture_image: mx.array | None = None

PlyData dataclass

Result of :func:load_ply.

Attributes:

Name	Type	Description
verts	array	(V, 3) positions.
faces	array | None	(F, 3) triangle indices, or None for pure point clouds.
normals	array | None	(V, 3) if nx, ny, nz present.
colors	array | None	(V, 3) in [0, 1] if red, green, blue present.
extra	dict[str, array]	dict of any remaining per-vertex scalar properties, each (V,).

Source code in src/mlx3d/io/ply_io.py

@dataclass
class PlyData:
    """Result of :func:`load_ply`.

    Attributes:
        verts: (V, 3) positions.
        faces: (F, 3) triangle indices, or ``None`` for pure point clouds.
        normals: (V, 3) if ``nx, ny, nz`` present.
        colors: (V, 3) in [0, 1] if ``red, green, blue`` present.
        extra: dict of any remaining per-vertex scalar properties, each (V,).
    """

    verts: mx.array
    faces: mx.array | None = None
    normals: mx.array | None = None
    colors: mx.array | None = None
    extra: dict[str, mx.array] = field(default_factory=dict)

load_gltf(path)

Load triangle mesh primitives from a .glb or .gltf file.

The loader merges all triangle primitives referenced by the default scene into one indexed mesh and applies node transforms. Material indices are preserved per face when present.

Source code in src/mlx3d/io/gltf_io.py

def load_gltf(path: str) -> GltfData:
    """Load triangle mesh primitives from a ``.glb`` or ``.gltf`` file.

    The loader merges all triangle primitives referenced by the default scene
    into one indexed mesh and applies node transforms. Material indices are
    preserved per face when present.
    """
    if path.lower().endswith(".glb"):
        gltf, glb_bin = _read_glb(path)
    else:
        with open(path) as f:
            gltf = json.load(f)
        glb_bin = b""
    root_dir = os.path.dirname(os.path.abspath(path))
    buffers = _buffer_bytes(gltf, root_dir, glb_bin)
    materials = _materials(gltf)
    texture_ids = [m.base_color_texture for m in materials if m.base_color_texture is not None]
    texture_image = None
    if len(set(texture_ids)) == 1:
        texture_image = _load_texture_image(gltf, buffers, root_dir, texture_ids[0])

    verts_parts: list[np.ndarray] = []
    faces_parts: list[np.ndarray] = []
    normal_parts: list[np.ndarray | None] = []
    uv_parts: list[np.ndarray | None] = []
    material_parts: list[np.ndarray] = []
    vert_offset = 0

    for node_idx, world in _scene_nodes(gltf):
        node = gltf["nodes"][node_idx]
        if "mesh" not in node:
            continue
        mesh = gltf["meshes"][node["mesh"]]
        normal_xform = np.linalg.inv(world[:3, :3]).T
        for prim in mesh.get("primitives", []):
            if prim.get("mode", 4) != 4:
                continue
            attrs = prim["attributes"]
            verts = _read_accessor(gltf, buffers, attrs["POSITION"]).astype(np.float32)
            vh = np.concatenate([verts, np.ones((verts.shape[0], 1), dtype=np.float32)], axis=1)
            verts = (vh @ world.T)[:, :3].astype(np.float32)
            if "indices" in prim:
                faces = (
                    _read_accessor(gltf, buffers, prim["indices"]).astype(np.int32).reshape(-1, 3)
                )
            else:
                faces = np.arange(verts.shape[0], dtype=np.int32).reshape(-1, 3)
            faces = faces + vert_offset

            normals = None
            if "NORMAL" in attrs:
                normals = _read_accessor(gltf, buffers, attrs["NORMAL"]).astype(np.float32)
                normals = (normals @ normal_xform.T).astype(np.float32)
                denom = np.maximum(np.linalg.norm(normals, axis=-1, keepdims=True), 1e-12)
                normals = normals / denom
            uvs = (
                _read_accessor(gltf, buffers, attrs["TEXCOORD_0"]).astype(np.float32)
                if "TEXCOORD_0" in attrs
                else None
            )

            verts_parts.append(verts)
            faces_parts.append(faces)
            normal_parts.append(normals)
            uv_parts.append(uvs)
            material_parts.append(
                np.full((faces.shape[0],), int(prim.get("material", -1)), dtype=np.int32)
            )
            vert_offset += verts.shape[0]

    if not verts_parts:
        raise ValueError("No triangle mesh primitives found in glTF scene.")

    verts = np.concatenate(verts_parts, axis=0)
    faces = np.concatenate(faces_parts, axis=0)
    normals = (
        np.concatenate(
            [
                n if n is not None else np.zeros_like(v)
                for n, v in zip(normal_parts, verts_parts, strict=True)
            ],
            axis=0,
        )
        if any(n is not None for n in normal_parts)
        else None
    )
    uvs = (
        np.concatenate(
            [
                uv if uv is not None else np.zeros((v.shape[0], 2), dtype=np.float32)
                for uv, v in zip(uv_parts, verts_parts, strict=True)
            ],
            axis=0,
        )
        if any(uv is not None for uv in uv_parts)
        else None
    )
    material_ids = np.concatenate(material_parts, axis=0)
    return GltfData(
        verts=mx.array(verts),
        faces=mx.array(faces),
        normals=mx.array(normals) if normals is not None else None,
        uvs=mx.array(uvs) if uvs is not None else None,
        material_ids=mx.array(material_ids) if material_ids.size else None,
        materials=materials,
        texture_image=texture_image,
    )

save_gltf(path, verts, faces, normals=None, uvs=None, material_base_color=None, texture_image=None)

Save a triangle mesh as a self-contained binary .glb file.

Source code in src/mlx3d/io/gltf_io.py

def save_gltf(
    path: str,
    verts: mx.array,
    faces: mx.array,
    normals: mx.array | None = None,
    uvs: mx.array | None = None,
    material_base_color: tuple[float, float, float, float] | None = None,
    texture_image: mx.array | None = None,
) -> None:
    """Save a triangle mesh as a self-contained binary ``.glb`` file."""
    v = np.asarray(verts, dtype=np.float32)
    f = np.asarray(faces, dtype=np.uint32).reshape(-1, 3)
    n = np.asarray(normals, dtype=np.float32) if normals is not None else None
    uv = np.asarray(uvs, dtype=np.float32) if uvs is not None else None
    if texture_image is not None and uv is None:
        raise ValueError("uvs are required when saving texture_image.")

    blob = b""
    views, accessors, attributes = [], [], {}

    def _add(arr: np.ndarray, target: int, comp: int, typ: str, with_minmax: bool) -> int:
        nonlocal blob
        offset = len(blob)
        raw = arr.tobytes()
        blob += _pad4(raw)
        views.append({"buffer": 0, "byteOffset": offset, "byteLength": len(raw), "target": target})
        acc = {
            "bufferView": len(views) - 1,
            "componentType": comp,
            "count": int(arr.shape[0]),
            "type": typ,
        }
        if with_minmax:
            acc["min"] = arr.min(axis=0).tolist()
            acc["max"] = arr.max(axis=0).tolist()
        accessors.append(acc)
        return len(accessors) - 1

    attributes["POSITION"] = _add(v, 34962, 5126, "VEC3", with_minmax=True)
    if n is not None:
        attributes["NORMAL"] = _add(n, 34962, 5126, "VEC3", with_minmax=False)
    if uv is not None:
        attributes["TEXCOORD_0"] = _add(uv, 34962, 5126, "VEC2", with_minmax=False)
    idx_accessor = _add(f.reshape(-1), 34963, 5125, "SCALAR", with_minmax=False)
    texture_view = None
    if texture_image is not None:
        raw = _texture_png_bytes(texture_image)
        offset = len(blob)
        blob += _pad4(raw)
        views.append({"buffer": 0, "byteOffset": offset, "byteLength": len(raw)})
        texture_view = len(views) - 1

    prim = {"attributes": attributes, "indices": idx_accessor, "mode": 4}
    materials = []
    if material_base_color is not None or texture_view is not None:
        pbr: dict[str, object] = {}
        if material_base_color is not None:
            pbr["baseColorFactor"] = [float(v) for v in material_base_color]
        if texture_view is not None:
            pbr["baseColorTexture"] = {"index": 0}
        materials.append({"pbrMetallicRoughness": pbr})
        prim["material"] = 0

    gltf = {
        "asset": {"version": "2.0", "generator": "mlx3d"},
        "scene": 0,
        "scenes": [{"nodes": [0]}],
        "nodes": [{"mesh": 0}],
        "meshes": [{"primitives": [prim]}],
        "buffers": [{"byteLength": len(blob)}],
        "bufferViews": views,
        "accessors": accessors,
    }
    if materials:
        gltf["materials"] = materials
    if texture_view is not None:
        gltf["images"] = [{"bufferView": texture_view, "mimeType": "image/png"}]
        gltf["textures"] = [{"source": 0}]

    json_chunk = _pad4(json.dumps(gltf, separators=(",", ":")).encode("utf-8"), b" ")
    bin_chunk = _pad4(blob)
    total = 12 + 8 + len(json_chunk) + 8 + len(bin_chunk)

    parent = os.path.dirname(os.path.abspath(path))
    os.makedirs(parent, exist_ok=True)
    with open(path, "wb") as out:
        out.write(struct.pack("<III", _GLB_MAGIC, 2, total))
        out.write(struct.pack("<II", len(json_chunk), 0x4E4F534A))
        out.write(json_chunk)
        out.write(struct.pack("<II", len(bin_chunk), 0x004E4942))
        out.write(bin_chunk)

load_image(path, *, normalize=True)

Load an image from disk as an MLX array.

Parameters:

Name	Type	Description	Default
path	str	Path to a PNG/JPEG/... file readable by Pillow.	required
normalize	bool	If True (default), return float32 in [0, 1]; otherwise return the raw uint8 values.	True

Returns:

Type	Description
array	(H, W, 3) for RGB images (alpha is dropped), or (H, W) for
array	single-channel images.

Source code in src/mlx3d/io/image_io.py

def load_image(path: str, *, normalize: bool = True) -> mx.array:
    """Load an image from disk as an MLX array.

    Args:
        path: Path to a PNG/JPEG/... file readable by Pillow.
        normalize: If ``True`` (default), return float32 in ``[0, 1]``;
            otherwise return the raw ``uint8`` values.

    Returns:
        ``(H, W, 3)`` for RGB images (alpha is dropped), or ``(H, W)`` for
        single-channel images.
    """
    img = Image.open(path)
    if img.mode in ("RGBA", "P", "LA"):
        img = img.convert("RGB")
    arr = np.asarray(img)
    if not normalize:
        return mx.array(arr)
    return mx.array(arr.astype(np.float32) / 255.0)

save_image(path, image)

Save an MLX (or NumPy) image to disk.

Accepts float images in [0, 1] (clipped) or uint8 images. Grayscale (H, W), RGB (H, W, 3) and RGBA (H, W, 4) are all supported. Parent directories are created automatically.

Parameters:

Name	Type	Description	Default
path	str	Destination file path; the format is inferred from the extension.	required
image	array | ndarray	The image to save.	required

Source code in src/mlx3d/io/image_io.py

def save_image(path: str, image: mx.array | np.ndarray) -> None:
    """Save an MLX (or NumPy) image to disk.

    Accepts float images in ``[0, 1]`` (clipped) or ``uint8`` images. Grayscale
    ``(H, W)``, RGB ``(H, W, 3)`` and RGBA ``(H, W, 4)`` are all supported.
    Parent directories are created automatically.

    Args:
        path: Destination file path; the format is inferred from the extension.
        image: The image to save.
    """
    arr = np.asarray(image)
    if arr.dtype != np.uint8:
        arr = (np.clip(arr, 0.0, 1.0) * 255.0).round().astype(np.uint8)
    parent = os.path.dirname(os.path.abspath(path))
    os.makedirs(parent, exist_ok=True)
    Image.fromarray(arr).save(path)

load_obj(path, load_texture=True)

Load a Wavefront OBJ file. Polygon faces are fan-triangulated.

Source code in src/mlx3d/io/obj_io.py

def load_obj(path: str, load_texture: bool = True) -> ObjData:
    """Load a Wavefront OBJ file. Polygon faces are fan-triangulated."""
    path_obj = Path(path)
    verts: list[list[float]] = []
    colors: list[list[float]] = []
    normals: list[list[float]] = []
    texcoords: list[list[float]] = []
    faces: list[list[int]] = []
    faces_vt: list[list[int]] = []
    faces_vn: list[list[int]] = []
    any_vt = False
    any_vn = False
    texture_path: str | None = None

    with open(path_obj, "r") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split()
            tag = parts[0]
            if tag == "mtllib" and len(parts) > 1:
                texture_path = _load_mtl_texture(path_obj.parent / parts[1])
            elif tag == "v":
                vals = [float(x) for x in parts[1:]]
                verts.append(vals[:3])
                if len(vals) >= 6:
                    colors.append(vals[3:6])
            elif tag == "vn":
                normals.append([float(x) for x in parts[1:4]])
            elif tag == "vt":
                texcoords.append([float(x) for x in parts[1:3]])
            elif tag == "f":
                corner_v, corner_vt, corner_vn = [], [], []
                for corner in parts[1:]:
                    fields = corner.split("/")
                    corner_v.append(_resolve_index(int(fields[0]), len(verts)))
                    if len(fields) > 1 and fields[1]:
                        corner_vt.append(_resolve_index(int(fields[1]), len(texcoords)))
                        any_vt = True
                    else:
                        corner_vt.append(-1)
                    if len(fields) > 2 and fields[2]:
                        corner_vn.append(_resolve_index(int(fields[2]), len(normals)))
                        any_vn = True
                    else:
                        corner_vn.append(-1)
                # Fan triangulation for polygons with > 3 corners.
                for i in range(1, len(corner_v) - 1):
                    faces.append([corner_v[0], corner_v[i], corner_v[i + 1]])
                    faces_vt.append([corner_vt[0], corner_vt[i], corner_vt[i + 1]])
                    faces_vn.append([corner_vn[0], corner_vn[i], corner_vn[i + 1]])

    if not verts:
        raise ValueError(f"No vertices found in {path!r}.")

    return ObjData(
        verts=mx.array(np.asarray(verts, dtype=np.float32)),
        faces=mx.array(np.asarray(faces, dtype=np.int32).reshape(-1, 3)),
        normals=mx.array(np.asarray(normals, dtype=np.float32)) if normals else None,
        texcoords=mx.array(np.asarray(texcoords, dtype=np.float32)) if texcoords else None,
        faces_normals_idx=(
            mx.array(np.asarray(faces_vn, dtype=np.int32).reshape(-1, 3)) if any_vn else None
        ),
        faces_texcoords_idx=(
            mx.array(np.asarray(faces_vt, dtype=np.int32).reshape(-1, 3)) if any_vt else None
        ),
        verts_colors=(
            mx.array(np.asarray(colors, dtype=np.float32))
            if len(colors) == len(verts) and colors
            else None
        ),
        texture_path=texture_path,
        texture_image=_read_texture(texture_path) if texture_path and load_texture else None,
    )

save_obj(path, verts, faces, verts_colors=None)

Save a triangle mesh as a Wavefront OBJ file.

Source code in src/mlx3d/io/obj_io.py

def save_obj(
    path: str,
    verts: mx.array,
    faces: mx.array,
    verts_colors: mx.array | None = None,
) -> None:
    """Save a triangle mesh as a Wavefront OBJ file."""
    v = np.array(verts, dtype=np.float64)
    f = np.array(faces, dtype=np.int64)
    c = np.array(verts_colors, dtype=np.float64) if verts_colors is not None else None
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w") as out:
        for i, p in enumerate(v):
            if c is not None:
                out.write(
                    f"v {p[0]:.8f} {p[1]:.8f} {p[2]:.8f} {c[i][0]:.6f} {c[i][1]:.6f} {c[i][2]:.6f}\n"
                )
            else:
                out.write(f"v {p[0]:.8f} {p[1]:.8f} {p[2]:.8f}\n")
        for tri in f:
            out.write(f"f {tri[0] + 1} {tri[1] + 1} {tri[2] + 1}\n")

load_ply(path)

Load a PLY file (ascii or binary_little_endian).

Source code in src/mlx3d/io/ply_io.py

def load_ply(path: str) -> PlyData:
    """Load a PLY file (ascii or binary_little_endian)."""
    with open(path, "rb") as f:
        data = f.read()

    header_end = data.find(b"end_header\n")
    if header_end < 0:
        raise ValueError(f"{path!r} is not a valid PLY file (no end_header).")
    header = data[:header_end].decode("ascii", errors="replace").splitlines()
    body = data[header_end + len(b"end_header\n") :]

    if not header or header[0].strip() != "ply":
        raise ValueError(f"{path!r} is not a PLY file.")

    fmt = None
    elements: list[tuple[str, int, list]] = []  # (name, count, [(prop_name, dtype) or list-prop])
    for line in header[1:]:
        parts = line.strip().split()
        if not parts or parts[0] == "comment":
            continue
        if parts[0] == "format":
            fmt = parts[1]
        elif parts[0] == "element":
            elements.append((parts[1], int(parts[2]), []))
        elif parts[0] == "property":
            if not elements:
                raise ValueError("property before element in PLY header.")
            if parts[1] == "list":
                elements[-1][2].append(
                    ("__list__", parts[4], _PLY_TO_NP[parts[2]], _PLY_TO_NP[parts[3]])
                )
            else:
                elements[-1][2].append((parts[2], _PLY_TO_NP[parts[1]]))

    if fmt not in ("ascii", "binary_little_endian"):
        raise ValueError(f"Unsupported PLY format {fmt!r}.")

    parsed: dict[str, dict[str, np.ndarray]] = {}
    offset = 0
    ascii_lines = body.decode("ascii", errors="replace").splitlines() if fmt == "ascii" else None
    ascii_pos = 0

    for name, count, props in elements:
        has_list = any(p[0] == "__list__" for p in props)
        if not has_list:
            dtype = np.dtype([(p[0], "<" + p[1]) for p in props])
            if fmt == "ascii":
                rows = []
                for _ in range(count):
                    rows.append(tuple(ascii_lines[ascii_pos].split()))
                    ascii_pos += 1
                arr = np.array(
                    [tuple(float(x) for x in r) for r in rows],
                    dtype=[(p[0], "f8") for p in props],
                ).astype(dtype)
            else:
                arr = np.frombuffer(body, dtype=dtype, count=count, offset=offset)
                offset += dtype.itemsize * count
            parsed[name] = {p[0]: arr[p[0]] for p in props}
        else:
            # Element with a list property (faces). Assume one list property.
            lists = []
            if fmt == "ascii":
                for _ in range(count):
                    vals = ascii_lines[ascii_pos].split()
                    ascii_pos += 1
                    n = int(vals[0])
                    lists.append([int(x) for x in vals[1 : 1 + n]])
            else:
                lp = next(p for p in props if p[0] == "__list__")
                count_dt = np.dtype("<" + lp[2])
                index_dt = np.dtype("<" + lp[3])
                for _ in range(count):
                    n = int(np.frombuffer(body, dtype=count_dt, count=1, offset=offset)[0])
                    offset += count_dt.itemsize
                    idx = np.frombuffer(body, dtype=index_dt, count=n, offset=offset)
                    offset += index_dt.itemsize * n
                    lists.append(idx.tolist())
            # Fan-triangulate.
            tris = []
            for poly in lists:
                for i in range(1, len(poly) - 1):
                    tris.append([poly[0], poly[i], poly[i + 1]])
            parsed[name] = {"__faces__": np.asarray(tris, dtype=np.int32).reshape(-1, 3)}

    if "vertex" not in parsed:
        raise ValueError(f"No vertex element in {path!r}.")
    vert_props = parsed["vertex"]
    for axis in ("x", "y", "z"):
        if axis not in vert_props:
            raise ValueError(f"Vertex element missing {axis!r} property.")

    verts = np.stack([vert_props["x"], vert_props["y"], vert_props["z"]], axis=-1).astype(
        np.float32
    )
    consumed = {"x", "y", "z"}

    normals = None
    if all(k in vert_props for k in ("nx", "ny", "nz")):
        normals = np.stack([vert_props["nx"], vert_props["ny"], vert_props["nz"]], axis=-1).astype(
            np.float32
        )
        consumed |= {"nx", "ny", "nz"}

    colors = None
    if all(k in vert_props for k in ("red", "green", "blue")):
        rgb = np.stack(
            [vert_props["red"], vert_props["green"], vert_props["blue"]], axis=-1
        ).astype(np.float32)
        if rgb.max() > 1.0:
            rgb = rgb / 255.0
        colors = rgb
        consumed |= {"red", "green", "blue", "alpha"}

    extra = {
        k: mx.array(np.ascontiguousarray(v.astype(np.float32)))
        for k, v in vert_props.items()
        if k not in consumed
    }

    faces = None
    if "face" in parsed and parsed["face"]["__faces__"].size > 0:
        faces = mx.array(parsed["face"]["__faces__"])

    return PlyData(
        verts=mx.array(verts),
        faces=faces,
        normals=mx.array(normals) if normals is not None else None,
        colors=mx.array(colors) if colors is not None else None,
        extra=extra,
    )

save_ply(path, verts, faces=None, normals=None, colors=None, extra=None, binary=True)

Save points/mesh to PLY.

colors are expected in [0, 1] and stored as uchar. extra properties are stored as float32 in insertion order (useful for Gaussian Splatting checkpoints).

Source code in src/mlx3d/io/ply_io.py

def save_ply(
    path: str,
    verts: mx.array,
    faces: mx.array | None = None,
    normals: mx.array | None = None,
    colors: mx.array | None = None,
    extra: dict[str, mx.array] | None = None,
    binary: bool = True,
) -> None:
    """Save points/mesh to PLY.

    ``colors`` are expected in [0, 1] and stored as uchar. ``extra`` properties
    are stored as float32 in insertion order (useful for Gaussian Splatting
    checkpoints).
    """
    v = np.array(verts, dtype=np.float32)
    n = np.array(normals, dtype=np.float32) if normals is not None else None
    c = (
        np.clip(np.array(colors, dtype=np.float32) * 255.0, 0, 255).astype(np.uint8)
        if colors is not None
        else None
    )
    extra_np = {k: np.array(a, dtype=np.float32).reshape(len(v)) for k, a in (extra or {}).items()}

    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    fmt = "binary_little_endian" if binary else "ascii"
    header = ["ply", f"format {fmt} 1.0", f"element vertex {len(v)}"]
    header += [f"property float {a}" for a in ("x", "y", "z")]
    if n is not None:
        header += [f"property float n{a}" for a in ("x", "y", "z")]
    if c is not None:
        header += [f"property uchar {a}" for a in ("red", "green", "blue")]
    header += [f"property float {k}" for k in extra_np]
    if faces is not None:
        header.append(f"element face {len(faces)}")
        header.append("property list uchar int vertex_indices")
    header.append("end_header")

    cols: list[np.ndarray] = [v]
    if n is not None:
        cols.append(n)
    if extra_np:
        pass  # appended after colors in struct order below

    with open(path, "wb") as out:
        out.write(("\n".join(header) + "\n").encode("ascii"))
        if binary:
            fields = [("xyz", "<f4", 3)]
            if n is not None:
                fields.append(("n", "<f4", 3))
            if c is not None:
                fields.append(("rgb", "u1", 3))
            for k in extra_np:
                fields.append((k, "<f4", 1))
            dtype = np.dtype([(name, dt, (cnt,)) for name, dt, cnt in fields])
            rec = np.empty(len(v), dtype=dtype)
            rec["xyz"] = v
            if n is not None:
                rec["n"] = n
            if c is not None:
                rec["rgb"] = c
            for k, a in extra_np.items():
                rec[k] = a[:, None]
            out.write(rec.tobytes())
            if faces is not None:
                f_arr = np.array(faces, dtype=np.int32)
                face_dtype = np.dtype([("n", "u1"), ("idx", "<i4", (3,))])
                frec = np.empty(len(f_arr), dtype=face_dtype)
                frec["n"] = 3
                frec["idx"] = f_arr
                out.write(frec.tobytes())
        else:
            for i in range(len(v)):
                row = [f"{x:.8f}" for x in v[i]]
                if n is not None:
                    row += [f"{x:.8f}" for x in n[i]]
                if c is not None:
                    row += [str(int(x)) for x in c[i]]
                row += [f"{extra_np[k][i]:.8f}" for k in extra_np]
                out.write((" ".join(row) + "\n").encode("ascii"))
            if faces is not None:
                for tri in np.array(faces, dtype=np.int64):
                    out.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n".encode("ascii"))