Merge pull request #473 from patricksnape/from_mask_fixes

Fix from_mask for Trimesh subclasses

menpo/shape/mesh/base.py

@@ -78,8 +78,8 @@ def tojson(self):
     def from_mask(self, mask):
         """
         A 1D boolean array with the same number of elements as the number of
-        points in the pointcloud. This is then broadcast across the dimensions
-        of the pointcloud and returns a new pointcloud containing only those
+        points in the TriMesh. This is then broadcast across the dimensions
+        of the mesh and returns a new mesh containing only those
         points that were ``True`` in the mask.
 
         Parameters
@@ -89,8 +89,8 @@ def from_mask(self, mask):
 
         Returns
         -------
-        pointcloud : :map:`PointCloud`
-            A new pointcloud that has been masked.
+        mesh : :map:`TriMesh`
+            A new mesh that has been masked.
         """
         if mask.shape[0] != self.n_points:
             raise ValueError('Mask must be a 1D boolean array of the same '
@@ -100,22 +100,27 @@ def from_mask(self, mask):
         if np.all(mask):  # Fast path for all true
             return tm
         else:
-            # Find the triangles we need to keep
-            masked_adj = mask_adjacency_array(mask, self.trilist)
-            # Find isolated vertices (vertices that don't exist in valid
-            # triangles)
-            isolated_indices = np.setdiff1d(np.nonzero(mask)[0], masked_adj)
-
-            # Create a 'new mask' that contains the points the use asked
-            # for MINUS the points that we can't create triangles for
-            new_mask = mask.copy()
-            new_mask[isolated_indices] = False
+            # Recalculate the mask to remove isolated vertices
+            isolated_mask = self._isolated_mask(mask)
             # Recreate the adjacency array with the updated mask
-            masked_adj = mask_adjacency_array(new_mask, self.trilist)
+            masked_adj = mask_adjacency_array(isolated_mask, self.trilist)
             tm.trilist = reindex_adjacency_array(masked_adj)
-            tm.points = tm.points[new_mask, :]
+            tm.points = tm.points[isolated_mask, :]
             return tm
 
+    def _isolated_mask(self, mask):
+        # Find the triangles we need to keep
+        masked_adj = mask_adjacency_array(mask, self.trilist)
+        # Find isolated vertices (vertices that don't exist in valid
+        # triangles)
+        isolated_indices = np.setdiff1d(np.nonzero(mask)[0], masked_adj)
+
+        # Create a 'new mask' that contains the points the use asked
+        # for MINUS the points that we can't create triangles for
+        new_mask = mask.copy()
+        new_mask[isolated_indices] = False
+        return new_mask
+
     def as_pointgraph(self, copy=True):
         from .. import PointGraph
         # Since we have triangles we need the last connection

menpo/shape/mesh/coloured.py

@@ -1,5 +1,7 @@
 import numpy as np
 from menpo.visualize.base import ColouredTriMeshViewer3d
+
+from ..adjacency import mask_adjacency_array, reindex_adjacency_array
 from .base import TriMesh
 
 
@@ -50,6 +52,41 @@ def __init__(self, points, trilist=None, colours=None, copy=True):
             raise ValueError('Must provide a colour per-vertex.')
         self.colours = colours_handle
 
+    def from_mask(self, mask):
+        """
+        A 1D boolean array with the same number of elements as the number of
+        points in the ColouredTriMesh. This is then broadcast across the
+        dimensions of the mesh and returns a new mesh containing only those
+        points that were ``True`` in the mask.
+
+        Parameters
+        ----------
+        mask : ``(n_points,)`` `ndarray`
+            1D array of booleans
+
+        Returns
+        -------
+        mesh : :map:`ColouredTriMesh`
+            A new mesh that has been masked.
+        """
+        if mask.shape[0] != self.n_points:
+            raise ValueError('Mask must be a 1D boolean array of the same '
+                             'number of entries as points in this '
+                             'ColouredTriMesh.')
+
+        ctm = self.copy()
+        if np.all(mask):  # Fast path for all true
+            return ctm
+        else:
+            # Recalculate the mask to remove isolated vertices
+            isolated_mask = self._isolated_mask(mask)
+            # Recreate the adjacency array with the updated mask
+            masked_adj = mask_adjacency_array(isolated_mask, self.trilist)
+            ctm.trilist = reindex_adjacency_array(masked_adj)
+            ctm.points = ctm.points[isolated_mask, :]
+            ctm.colours = ctm.colours[isolated_mask, :]
+            return ctm
+
     def _view(self, figure_id=None, new_figure=False, coloured=True, **kwargs):
         r"""
         Visualize the :class:`ColouredTriMesh`. Only 3D objects are currently

menpo/shape/mesh/textured.py

@@ -4,6 +4,7 @@
 from menpo.transform import Scale
 from menpo.visualize import TexturedTriMeshViewer3d
 
+from ..adjacency import mask_adjacency_array, reindex_adjacency_array
 from .base import TriMesh
 
 
@@ -95,6 +96,41 @@ def from_vector(self, flattened):
                                self.tcoords.points, self.texture,
                                trilist=self.trilist)
 
+    def from_mask(self, mask):
+        """
+        A 1D boolean array with the same number of elements as the number of
+        points in the TexturedTriMesh. This is then broadcast across the
+        dimensions of the mesh and returns a new mesh containing only those
+        points that were ``True`` in the mask.
+
+        Parameters
+        ----------
+        mask : ``(n_points,)`` `ndarray`
+            1D array of booleans
+
+        Returns
+        -------
+        mesh : :map:`TexturedTriMesh`
+            A new mesh that has been masked.
+        """
+        if mask.shape[0] != self.n_points:
+            raise ValueError('Mask must be a 1D boolean array of the same '
+                             'number of entries as points in this '
+                             'TexturedTriMesh.')
+
+        ttm = self.copy()
+        if np.all(mask):  # Fast path for all true
+            return ttm
+        else:
+            # Recalculate the mask to remove isolated vertices
+            isolated_mask = self._isolated_mask(mask)
+            # Recreate the adjacency array with the updated mask
+            masked_adj = mask_adjacency_array(isolated_mask, self.trilist)
+            ttm.trilist = reindex_adjacency_array(masked_adj)
+            ttm.points = ttm.points[isolated_mask, :]
+            ttm.tcoords.points = ttm.tcoords.points[isolated_mask, :]
+            return ttm
+
     def tojson(self):
         r"""
         Convert this `TriMesh` to a dictionary JSON representation.