Skip to content

amirhossein-razlighi/mlx3D

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

117 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

MLX3D

Differentiable 3D computer vision on Apple Silicon, built on MLX.

MLX3D brings the PyTorch3D workflow to Macs: batched 3D data structures, cameras, differentiable rendering, and modern view synthesis — NeRF and 3D Gaussian Splatting with custom Metal kernels — running natively on the Apple GPU.

📖 Documentation & tutorials

Features

  • Structures — batched Meshes / Pointclouds with list, packed and padded views; differentiable normals, areas, edges.
  • Cameras & transforms — OpenCV/COLMAP-convention pinhole cameras (ray generation, projection, look-at) and batched rotation conversions (quaternion, axis-angle, Euler, 6D).
  • Ops & losses — GPU brute-force k-NN, chamfer distance, area-weighted surface sampling, Laplacian/edge/normal-consistency mesh losses, PSNR and differentiable SSIM.
  • NeRF — positional encoding, the NeRF MLP, stratified + hierarchical sampling, volume rendering, Blender-synthetic dataset loader.
  • Mesh rendering — differentiable soft triangle rasterization, UV texture sampling for OBJ/MTL assets, and scalar-field mesh extraction.
  • Gaussian Splatting — a Metal translation of the reference CUDA rasterizer (tile-based forward & backward kernels wrapped in mx.custom_function), EWA projection, spherical harmonics, anti-aliased and arbitrary feature rendering, adaptive density control, COLMAP loading, and standard 3DGS .ply checkpoints. ~30 FPS forward at 720p with 100k Gaussians on an M-series GPU.
  • Interactive viewermlx3d-view point_cloud.ply opens a browser viewer with orbit/pan/zoom; frames are rendered on the Apple GPU by the Metal rasterizer and streamed live. Works for NeRFs too.
  • IO — OBJ and PLY (ascii + binary, including Gaussian Splatting checkpoint layouts), plus one-line image save_image / load_image for any renderer output.
  • Composable & extensible — every image renderer is a plain callable (camera, scene) -> {"image", "alpha", "depth"} (the Renderer protocol), so you can drop in your own rasterizer, shader, or ray tracer and reuse the rest of the pipeline — no base classes to subclass.

Installation

pip install mlx3d

Requires an Apple Silicon Mac and Python ≥ 3.10.

Quick example

import mlx.core as mx
from mlx3d.cameras import Camera
from mlx3d.splatting import GaussianModel

model = GaussianModel.from_points(
    points=mx.random.normal((10_000, 3)) * 0.5,
    colors=mx.random.uniform(shape=(10_000, 3)),
)
camera = Camera.look_at(eye=(0, 0, -4), at=(0, 0, 0), width=1280, height=720)
out = model.render(camera)            # differentiable end to end
print(out["image"].shape)             # (720, 1280, 3)

Train Gaussian Splatting on any COLMAP scene (same inputs as the original 3DGS):

python examples/train_gaussian_splatting.py --data /path/to/scene --iters 7000
mlx3d-view outputs/gs/point_cloud.ply   # inspect the result interactively
mlx3d-render outputs/gs/point_cloud.ply --out render.png --antialias
mlx3d-eval outputs/gs/point_cloud.ply --data /path/to/scene --views 20 --json-out metrics.json
mlx3d-compact outputs/gs/point_cloud.ply --out point_cloud_small.ply --max-gaussians 500000

More in the docs: mesh optimization, point cloud fitting, NeRF, Gaussian Splatting.

Examples

The examples/ folder has runnable scripts for every core feature. The self-contained ones generate their own synthetic data — no downloads — and finish in seconds:

uv run python examples/render_mesh.py        # soft mesh rasterization
uv run python examples/raytrace_volume.py    # ray casting + volume rendering
uv run python examples/extract_mesh.py       # marching cubes from an SDF
uv run python examples/fit_pointcloud.py     # point-cloud optimization
uv run python examples/fit_mesh.py           # mesh fitting (chamfer + regularizers)
uv run python examples/fit_nerf.py           # train a small NeRF
uv run python examples/fit_gaussians.py      # fit 3D Gaussians
uv run python examples/extend_renderer.py    # plug in a custom renderer

See examples/README.md for the full list, including the COLMAP/Blender training scripts.

Development

Development uses uv:

git clone https://github.com/amirhossein-razlighi/mlx3D
cd mlx3D
uv sync               # creates .venv with all dev dependencies
uv run pytest tests/
uv run mkdocs serve   # docs at http://127.0.0.1:8000

Contributions are welcome — file an issue to get started.

License

MIT

Packages

 
 
 

Contributors