Skip to content

mlx3d.structures

mlx3d.structures

Meshes

Batch of triangle meshes.

Parameters:

Name Type Description Default
verts list[array] | array

list of (V_i, 3) float arrays, or a padded (N, V, 3) array.

 required
faces list[array] | array

list of (F_i, 3) integer arrays, or a padded (N, F, 3) array where unused rows are filled with -1.

 required
Source code in src/mlx3d/structures/meshes.py 
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
class Meshes:
    """Batch of triangle meshes.

    Args:
        verts (list[mx.array] | mx.array): list of ``(V_i, 3)`` float arrays,
            or a padded ``(N, V, 3)`` array.
        faces (list[mx.array] | mx.array): list of ``(F_i, 3)`` integer
            arrays, or a padded ``(N, F, 3)`` array where unused rows are
            filled with ``-1``.
    """

    def __init__(self, verts, faces) -> None:
        if isinstance(verts, (list, tuple)) and isinstance(faces, (list, tuple)):
            if len(verts) != len(faces):
                raise ValueError("verts and faces must contain the same number of meshes.")
            self._verts_list = [_as_float_array(v) for v in verts]
            self._faces_list = [_as_int_array(f) for f in faces]
            for v, f in zip(self._verts_list, self._faces_list):
                if v.ndim != 2 or v.shape[-1] != 3:
                    raise ValueError("Each verts entry must have shape (V, 3).")
                if f.ndim != 2 or f.shape[-1] != 3:
                    raise ValueError("Each faces entry must have shape (F, 3).")
        elif isinstance(verts, mx.array) and isinstance(faces, mx.array):
            if verts.ndim != 3 or verts.shape[-1] != 3:
                raise ValueError("Padded verts must have shape (N, V, 3).")
            if faces.ndim != 3 or faces.shape[-1] != 3:
                raise ValueError("Padded faces must have shape (N, F, 3).")
            if verts.shape[0] != faces.shape[0]:
                raise ValueError("verts and faces must have the same batch dimension.")
            faces = _as_int_array(faces)
            verts = _as_float_array(verts)
            # Number of valid (non -1-padded) faces per mesh.
            valid = (faces > -1).all(axis=-1)
            num_faces = valid.sum(axis=-1)
            self._verts_list = [verts[i] for i in range(verts.shape[0])]
            nf = num_faces.tolist()
            self._faces_list = [faces[i, : nf[i]] for i in range(faces.shape[0])]
        else:
            raise ValueError(
                "verts and faces must both be lists of arrays or both padded mx.arrays."
            )

        self._N = len(self._verts_list)
        # Caches.
        self._verts_packed = None
        self._faces_packed = None
        self._verts_padded = None
        self._faces_padded = None
        self._verts_normals_packed = None
        self._faces_normals_packed = None
        self._faces_areas_packed = None
        self._edges_packed = None

    # ------------------------------------------------------------------ basics
    def __len__(self) -> int:
        return self._N

    def __getitem__(self, index) -> "Meshes":
        if isinstance(index, int):
            index = [index]
        if isinstance(index, slice):
            index = list(range(self._N))[index]
        return Meshes(
            [self._verts_list[i] for i in index],
            [self._faces_list[i] for i in index],
        )

    def isempty(self) -> bool:
        return self._N == 0 or all(v.shape[0] == 0 for v in self._verts_list)

    @property
    def num_verts_per_mesh(self) -> mx.array:
        return mx.array([v.shape[0] for v in self._verts_list], dtype=mx.int32)

    @property
    def num_faces_per_mesh(self) -> mx.array:
        return mx.array([f.shape[0] for f in self._faces_list], dtype=mx.int32)

    # -------------------------------------------------------------- list views
    def verts_list(self) -> list[mx.array]:
        return self._verts_list

    def faces_list(self) -> list[mx.array]:
        return self._faces_list

    # ------------------------------------------------------------ packed views
    def verts_packed(self) -> mx.array:
        if self._verts_packed is None:
            self._verts_packed = (
                mx.concatenate(self._verts_list, axis=0) if self._N > 0 else mx.zeros((0, 3))
            )
        return self._verts_packed

    def faces_packed(self) -> mx.array:
        """(sum(F_i), 3) faces with vertex indices into ``verts_packed``."""
        if self._faces_packed is None:
            faces = []
            offset = 0
            for v, f in zip(self._verts_list, self._faces_list):
                faces.append(f + offset)
                offset += v.shape[0]
            self._faces_packed = (
                mx.concatenate(faces, axis=0) if faces else mx.zeros((0, 3), dtype=mx.int32)
            )
        return self._faces_packed

    def mesh_to_verts_packed_first_idx(self) -> mx.array:
        counts = self.num_verts_per_mesh
        return mx.concatenate([mx.zeros((1,), dtype=mx.int32), mx.cumsum(counts)[:-1]])

    def mesh_to_faces_packed_first_idx(self) -> mx.array:
        counts = self.num_faces_per_mesh
        return mx.concatenate([mx.zeros((1,), dtype=mx.int32), mx.cumsum(counts)[:-1]])

    def verts_packed_to_mesh_idx(self) -> mx.array:
        return mx.concatenate(
            [mx.full((v.shape[0],), i, dtype=mx.int32) for i, v in enumerate(self._verts_list)]
            or [mx.zeros((0,), dtype=mx.int32)]
        )

    def faces_packed_to_mesh_idx(self) -> mx.array:
        return mx.concatenate(
            [mx.full((f.shape[0],), i, dtype=mx.int32) for i, f in enumerate(self._faces_list)]
            or [mx.zeros((0,), dtype=mx.int32)]
        )

    # ------------------------------------------------------------ padded views
    def verts_padded(self) -> mx.array:
        if self._verts_padded is None:
            V = max((v.shape[0] for v in self._verts_list), default=0)
            rows = []
            for v in self._verts_list:
                pad = V - v.shape[0]
                rows.append(mx.pad(v, ((0, pad), (0, 0))) if pad > 0 else v)
            self._verts_padded = mx.stack(rows, axis=0) if rows else mx.zeros((0, 0, 3))
        return self._verts_padded

    def faces_padded(self) -> mx.array:
        if self._faces_padded is None:
            F = max((f.shape[0] for f in self._faces_list), default=0)
            rows = []
            for f in self._faces_list:
                pad = F - f.shape[0]
                rows.append(mx.pad(f, ((0, pad), (0, 0)), constant_values=-1) if pad > 0 else f)
            self._faces_padded = (
                mx.stack(rows, axis=0) if rows else mx.zeros((0, 0, 3), dtype=mx.int32)
            )
        return self._faces_padded

    # -------------------------------------------------------- derived geometry
    def faces_normals_packed(self) -> mx.array:
        """(sum(F_i), 3) unnormalized face normals (length = 2 * face area)."""
        if self._faces_normals_packed is None:
            verts = self.verts_packed()
            faces = self.faces_packed()
            v0 = verts[faces[:, 0]]
            v1 = verts[faces[:, 1]]
            v2 = verts[faces[:, 2]]
            self._faces_normals_packed = mx.linalg.cross(v1 - v0, v2 - v0)
        return self._faces_normals_packed

    def faces_areas_packed(self) -> mx.array:
        """(sum(F_i),) face areas."""
        if self._faces_areas_packed is None:
            n = self.faces_normals_packed()
            self._faces_areas_packed = 0.5 * mx.linalg.norm(n, axis=-1)
        return self._faces_areas_packed

    def verts_normals_packed(self, eps: float = 1e-12) -> mx.array:
        """(sum(V_i), 3) area-weighted vertex normals."""
        if self._verts_normals_packed is None:
            verts = self.verts_packed()
            faces = self.faces_packed()
            fn = self.faces_normals_packed()
            normals = mx.zeros_like(verts)
            for k in range(3):
                normals = normals.at[faces[:, k]].add(fn)
            self._verts_normals_packed = normals / mx.maximum(
                mx.linalg.norm(normals, axis=-1, keepdims=True), eps
            )
        return self._verts_normals_packed

    def verts_normals_list(self) -> list[mx.array]:
        normals = self.verts_normals_packed()
        out = []
        offset = 0
        for v in self._verts_list:
            out.append(normals[offset : offset + v.shape[0]])
            offset += v.shape[0]
        return out

    def edges_packed(self) -> mx.array:
        """(E, 2) unique undirected edges with indices into ``verts_packed``.

        Edge connectivity carries no gradients, so this is computed once on
        CPU with NumPy and cached.
        """
        if self._edges_packed is None:
            faces = np.array(self.faces_packed())
            if faces.shape[0] == 0:
                self._edges_packed = mx.zeros((0, 2), dtype=mx.int32)
            else:
                e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0)
                e.sort(axis=1)
                e = np.unique(e, axis=0)
                self._edges_packed = mx.array(e.astype(np.int32))
        return self._edges_packed

    def bounding_boxes(self) -> mx.array:
        """(N, 3, 2) per-mesh min/max corners."""
        boxes = []
        for v in self._verts_list:
            boxes.append(mx.stack([v.min(axis=0), v.max(axis=0)], axis=-1))
        return mx.stack(boxes, axis=0)

    # ------------------------------------------------------------ modification
    def offset_verts(self, offsets: mx.array) -> "Meshes":
        """Return a new ``Meshes`` with packed-vertex ``offsets`` added.

        ``offsets`` has shape ``(sum(V_i), 3)`` (or broadcastable ``(3,)``).
        """
        new_verts = self.verts_packed() + offsets
        out = []
        offset = 0
        for v in self._verts_list:
            out.append(new_verts[offset : offset + v.shape[0]])
            offset += v.shape[0]
        return Meshes(out, self._faces_list)

    def update_padded(self, new_verts_padded: mx.array) -> "Meshes":
        """Return a new ``Meshes`` with the same topology and new padded vertices."""
        new_list = [new_verts_padded[i, : v.shape[0]] for i, v in enumerate(self._verts_list)]
        return Meshes(new_list, self._faces_list)

    def scale_verts(self, scale) -> "Meshes":
        return Meshes([v * scale for v in self._verts_list], self._faces_list)

faces_packed()

(sum(F_i), 3) faces with vertex indices into verts_packed.

Source code in src/mlx3d/structures/meshes.py 
127
128
129
130
131
132
133
134
135
136
137
138
def faces_packed(self) -> mx.array:
    """(sum(F_i), 3) faces with vertex indices into ``verts_packed``."""
    if self._faces_packed is None:
        faces = []
        offset = 0
        for v, f in zip(self._verts_list, self._faces_list):
            faces.append(f + offset)
            offset += v.shape[0]
        self._faces_packed = (
            mx.concatenate(faces, axis=0) if faces else mx.zeros((0, 3), dtype=mx.int32)
        )
    return self._faces_packed

faces_normals_packed()

(sum(F_i), 3) unnormalized face normals (length = 2 * face area).

Source code in src/mlx3d/structures/meshes.py 
184
185
186
187
188
189
190
191
192
193
def faces_normals_packed(self) -> mx.array:
    """(sum(F_i), 3) unnormalized face normals (length = 2 * face area)."""
    if self._faces_normals_packed is None:
        verts = self.verts_packed()
        faces = self.faces_packed()
        v0 = verts[faces[:, 0]]
        v1 = verts[faces[:, 1]]
        v2 = verts[faces[:, 2]]
        self._faces_normals_packed = mx.linalg.cross(v1 - v0, v2 - v0)
    return self._faces_normals_packed

faces_areas_packed()

(sum(F_i),) face areas.

Source code in src/mlx3d/structures/meshes.py 
195
196
197
198
199
200
def faces_areas_packed(self) -> mx.array:
    """(sum(F_i),) face areas."""
    if self._faces_areas_packed is None:
        n = self.faces_normals_packed()
        self._faces_areas_packed = 0.5 * mx.linalg.norm(n, axis=-1)
    return self._faces_areas_packed

verts_normals_packed(eps=1e-12)

(sum(V_i), 3) area-weighted vertex normals.

Source code in src/mlx3d/structures/meshes.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
def verts_normals_packed(self, eps: float = 1e-12) -> mx.array:
    """(sum(V_i), 3) area-weighted vertex normals."""
    if self._verts_normals_packed is None:
        verts = self.verts_packed()
        faces = self.faces_packed()
        fn = self.faces_normals_packed()
        normals = mx.zeros_like(verts)
        for k in range(3):
            normals = normals.at[faces[:, k]].add(fn)
        self._verts_normals_packed = normals / mx.maximum(
            mx.linalg.norm(normals, axis=-1, keepdims=True), eps
        )
    return self._verts_normals_packed

edges_packed()

(E, 2) unique undirected edges with indices into verts_packed.

Edge connectivity carries no gradients, so this is computed once on CPU with NumPy and cached.

Source code in src/mlx3d/structures/meshes.py 
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def edges_packed(self) -> mx.array:
    """(E, 2) unique undirected edges with indices into ``verts_packed``.

    Edge connectivity carries no gradients, so this is computed once on
    CPU with NumPy and cached.
    """
    if self._edges_packed is None:
        faces = np.array(self.faces_packed())
        if faces.shape[0] == 0:
            self._edges_packed = mx.zeros((0, 2), dtype=mx.int32)
        else:
            e = np.concatenate([faces[:, [0, 1]], faces[:, [1, 2]], faces[:, [2, 0]]], axis=0)
            e.sort(axis=1)
            e = np.unique(e, axis=0)
            self._edges_packed = mx.array(e.astype(np.int32))
    return self._edges_packed

bounding_boxes()

(N, 3, 2) per-mesh min/max corners.

Source code in src/mlx3d/structures/meshes.py 
242
243
244
245
246
247
def bounding_boxes(self) -> mx.array:
    """(N, 3, 2) per-mesh min/max corners."""
    boxes = []
    for v in self._verts_list:
        boxes.append(mx.stack([v.min(axis=0), v.max(axis=0)], axis=-1))
    return mx.stack(boxes, axis=0)

offset_verts(offsets)

Return a new Meshes with packed-vertex offsets added.

offsets has shape (sum(V_i), 3) (or broadcastable (3,)).

Source code in src/mlx3d/structures/meshes.py 
250
251
252
253
254
255
256
257
258
259
260
261
def offset_verts(self, offsets: mx.array) -> "Meshes":
    """Return a new ``Meshes`` with packed-vertex ``offsets`` added.

    ``offsets`` has shape ``(sum(V_i), 3)`` (or broadcastable ``(3,)``).
    """
    new_verts = self.verts_packed() + offsets
    out = []
    offset = 0
    for v in self._verts_list:
        out.append(new_verts[offset : offset + v.shape[0]])
        offset += v.shape[0]
    return Meshes(out, self._faces_list)

update_padded(new_verts_padded)

Return a new Meshes with the same topology and new padded vertices.

Source code in src/mlx3d/structures/meshes.py 
263
264
265
266
def update_padded(self, new_verts_padded: mx.array) -> "Meshes":
    """Return a new ``Meshes`` with the same topology and new padded vertices."""
    new_list = [new_verts_padded[i, : v.shape[0]] for i, v in enumerate(self._verts_list)]
    return Meshes(new_list, self._faces_list)

Pointclouds

Batch of point clouds with optional per-point normals and features.

Parameters:

Name Type Description Default
points list[array] | array

list of (P_i, 3) arrays or a padded (N, P, 3) array.

 required
normals list[array] | array | None

optional, same layout as points.

 None
features list[array] | array | None

optional, list of (P_i, C) arrays or padded (N, P, C).

 None
Source code in src/mlx3d/structures/pointclouds.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
class Pointclouds:
    """Batch of point clouds with optional per-point normals and features.

    Args:
        points (list[mx.array] | mx.array): list of ``(P_i, 3)`` arrays or a
            padded ``(N, P, 3)`` array.
        normals (list[mx.array] | mx.array | None): optional, same layout as
            ``points``.
        features (list[mx.array] | mx.array | None): optional, list of
            ``(P_i, C)`` arrays or padded ``(N, P, C)``.
    """

    def __init__(self, points, normals=None, features=None) -> None:
        if isinstance(points, mx.array):
            if points.ndim != 3 or points.shape[-1] != 3:
                raise ValueError("Padded points must have shape (N, P, 3).")
            points = [points[i] for i in range(points.shape[0])]
            if isinstance(normals, mx.array):
                normals = [normals[i] for i in range(normals.shape[0])]
            if isinstance(features, mx.array):
                features = [features[i] for i in range(features.shape[0])]
        if not isinstance(points, (list, tuple)):
            raise ValueError("points must be a list of arrays or a padded mx.array.")

        self._points_list = [_as_float_array(p) for p in points]
        for p in self._points_list:
            if p.ndim != 2 or p.shape[-1] != 3:
                raise ValueError("Each points entry must have shape (P, 3).")

        def _check_aux(aux, name):
            if aux is None:
                return None
            aux = [_as_float_array(a) for a in aux]
            if len(aux) != len(self._points_list):
                raise ValueError(f"{name} must have one entry per cloud.")
            for a, p in zip(aux, self._points_list):
                if a.shape[0] != p.shape[0]:
                    raise ValueError(f"{name} entries must match points per cloud.")
            return aux

        self._normals_list = _check_aux(normals, "normals")
        self._features_list = _check_aux(features, "features")
        self._N = len(self._points_list)

    def __len__(self) -> int:
        return self._N

    def __getitem__(self, index) -> "Pointclouds":
        if isinstance(index, int):
            index = [index]
        if isinstance(index, slice):
            index = list(range(self._N))[index]
        return Pointclouds(
            [self._points_list[i] for i in index],
            normals=[self._normals_list[i] for i in index] if self._normals_list else None,
            features=[self._features_list[i] for i in index] if self._features_list else None,
        )

    @property
    def num_points_per_cloud(self) -> mx.array:
        return mx.array([p.shape[0] for p in self._points_list], dtype=mx.int32)

    # -------------------------------------------------------------------- views
    def points_list(self) -> list[mx.array]:
        return self._points_list

    def normals_list(self) -> list[mx.array] | None:
        return self._normals_list

    def features_list(self) -> list[mx.array] | None:
        return self._features_list

    def points_packed(self) -> mx.array:
        return mx.concatenate(self._points_list, axis=0) if self._N else mx.zeros((0, 3))

    def normals_packed(self) -> mx.array | None:
        if self._normals_list is None:
            return None
        return mx.concatenate(self._normals_list, axis=0)

    def features_packed(self) -> mx.array | None:
        if self._features_list is None:
            return None
        return mx.concatenate(self._features_list, axis=0)

    def points_packed_to_cloud_idx(self) -> mx.array:
        return mx.concatenate(
            [mx.full((p.shape[0],), i, dtype=mx.int32) for i, p in enumerate(self._points_list)]
            or [mx.zeros((0,), dtype=mx.int32)]
        )

    def points_padded(self) -> mx.array:
        P = max((p.shape[0] for p in self._points_list), default=0)
        rows = []
        for p in self._points_list:
            pad = P - p.shape[0]
            rows.append(mx.pad(p, ((0, pad), (0, 0))) if pad > 0 else p)
        return mx.stack(rows, axis=0) if rows else mx.zeros((0, 0, 3))

    def padded_mask(self) -> mx.array:
        """(N, max(P_i)) boolean mask of valid points in ``points_padded``."""
        P = max((p.shape[0] for p in self._points_list), default=0)
        idx = mx.arange(P)[None, :]
        return idx < self.num_points_per_cloud[:, None]

    # ------------------------------------------------------------- modification
    def offset_points(self, offsets: mx.array) -> "Pointclouds":
        """Add ``offsets`` (packed ``(sum(P_i), 3)`` or broadcastable) to all points."""
        new_points = self.points_packed() + offsets
        out = []
        offset = 0
        for p in self._points_list:
            out.append(new_points[offset : offset + p.shape[0]])
            offset += p.shape[0]
        return Pointclouds(out, normals=self._normals_list, features=self._features_list)

    def scale_points(self, scale) -> "Pointclouds":
        return Pointclouds(
            [p * scale for p in self._points_list],
            normals=self._normals_list,
            features=self._features_list,
        )

padded_mask()

(N, max(P_i)) boolean mask of valid points in points_padded.

Source code in src/mlx3d/structures/pointclouds.py 
112
113
114
115
116
def padded_mask(self) -> mx.array:
    """(N, max(P_i)) boolean mask of valid points in ``points_padded``."""
    P = max((p.shape[0] for p in self._points_list), default=0)
    idx = mx.arange(P)[None, :]
    return idx < self.num_points_per_cloud[:, None]

offset_points(offsets)

Add offsets (packed (sum(P_i), 3) or broadcastable) to all points.

Source code in src/mlx3d/structures/pointclouds.py 
119
120
121
122
123
124
125
126
127
def offset_points(self, offsets: mx.array) -> "Pointclouds":
    """Add ``offsets`` (packed ``(sum(P_i), 3)`` or broadcastable) to all points."""
    new_points = self.points_packed() + offsets
    out = []
    offset = 0
    for p in self._points_list:
        out.append(new_points[offset : offset + p.shape[0]])
        offset += p.shape[0]
    return Pointclouds(out, normals=self._normals_list, features=self._features_list)

join_meshes_as_batch(meshes)

Concatenate several Meshes batches into one.

Source code in src/mlx3d/structures/meshes.py 
272
273
274
275
276
def join_meshes_as_batch(meshes: list[Meshes]) -> Meshes:
    """Concatenate several ``Meshes`` batches into one."""
    verts = [v for m in meshes for v in m.verts_list()]
    faces = [f for m in meshes for f in m.faces_list()]
    return Meshes(verts, faces)