vedo.assembly

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3import numpy as np
  4
  5try:
  6    import vedo.vtkclasses as vtk
  7except ImportError:
  8    import vtkmodules.all as vtk
  9
 10import vedo
 11
 12__docformat__ = "google"
 13
 14__doc__ = """
 15Submodule for managing groups of vedo objects
 16
 17![](https://vedo.embl.es/images/basic/align4.png)
 18"""
 19
 20__all__ = ["Group", "Assembly", "procrustes_alignment"]
 21
 22
 23#################################################
 24def procrustes_alignment(sources, rigid=False):
 25    """
 26    Return an ``Assembly`` of aligned source meshes with the `Procrustes` algorithm.
 27    The output ``Assembly`` is normalized in size.
 28
 29    The `Procrustes` algorithm takes N set of points and aligns them in a least-squares sense
 30    to their mutual mean. The algorithm is iterated until convergence,
 31    as the mean must be recomputed after each alignment.
 32
 33    The set of average points generated by the algorithm can be accessed with
 34    ``algoutput.info['mean']`` as a numpy array.
 35
 36    Arguments:
 37        rigid : bool
 38            if `True` scaling is disabled.
 39
 40    Examples:
 41        - [align4.py](https://github.com/marcomusy/vedo/tree/master/examples/basic/align4.py)
 42
 43        ![](https://vedo.embl.es/images/basic/align4.png)
 44    """
 45
 46    group = vtk.vtkMultiBlockDataGroupFilter()
 47    for source in sources:
 48        if sources[0].npoints != source.npoints:
 49            vedo.logger.error("sources have different nr of points")
 50            raise RuntimeError()
 51        group.AddInputData(source.polydata())
 52    procrustes = vtk.vtkProcrustesAlignmentFilter()
 53    procrustes.StartFromCentroidOn()
 54    procrustes.SetInputConnection(group.GetOutputPort())
 55    if rigid:
 56        procrustes.GetLandmarkTransform().SetModeToRigidBody()
 57    procrustes.Update()
 58
 59    acts = []
 60    for i, s in enumerate(sources):
 61        poly = procrustes.GetOutput().GetBlock(i)
 62        mesh = vedo.mesh.Mesh(poly)
 63        mesh.SetProperty(s.GetProperty())
 64        if hasattr(s, "name"):
 65            mesh.name = s.name
 66        acts.append(mesh)
 67    assem = Assembly(acts)
 68    assem.transform = procrustes.GetLandmarkTransform()
 69    assem.info["mean"] = vedo.utils.vtk2numpy(procrustes.GetMeanPoints().GetData())
 70    return assem
 71
 72
 73#################################################
 74class Group(vtk.vtkPropAssembly):
 75    """Form groups of generic objects (not necessarily meshes)."""
 76
 77    def __init__(self, objects=()):
 78        """Form groups of generic objects (not necessarily meshes)."""
 79
 80        vtk.vtkPropAssembly.__init__(self)
 81
 82        self.name = ""
 83        self.created = ""
 84        self.trail = None
 85        self.trail_points = []
 86        self.trail_segment_size = 0
 87        self.trail_offset = None
 88        self.shadows = []
 89        self.info = {}
 90        self.rendered_at = set()
 91        self.transform = None
 92        self.scalarbar = None
 93
 94        for a in vedo.utils.flatten(objects):
 95            if a:
 96                self.AddPart(a)
 97
 98        self.PickableOff()
 99
100    def __iadd__(self, obj):
101        """
102        Add an object to the group
103        """
104        if not vedo.utils.is_sequence(obj):
105            obj = [obj]
106        for a in obj:
107            if a:
108                self.AddPart(a)
109        return self
110
111    def unpack(self):
112        """Unpack the group into its elements"""
113        elements = []
114        self.InitPathTraversal()
115        parts = self.GetParts()
116        parts.InitTraversal()
117        for i in range(parts.GetNumberOfItems()):
118            ele = parts.GetItemAsObject(i)
119            elements.append(ele)
120
121        # gr.InitPathTraversal()
122        # for _ in range(gr.GetNumberOfPaths()):
123        #     path  = gr.GetNextPath()
124        #     print([path])
125        #     path.InitTraversal()
126        #     for i in range(path.GetNumberOfItems()):
127        #         a = path.GetItemAsObject(i).GetViewProp()
128        #         print([a])
129
130        return elements
131
132    def clear(self):
133        """Remove all parts"""
134        for a in self.unpack():
135            self.RemovePart(a)
136        return self
137
138    def on(self):
139        """Switch on visibility"""
140        self.VisibilityOn()
141        return self
142
143    def off(self):
144        """Switch off visibility"""
145        self.VisibilityOff()
146        return self
147
148    def pickable(self, value=None):
149        """Set/get the pickability property of an object."""
150        if value is None:
151            return self.GetPickable()
152        self.SetPickable(value)
153        return self
154
155    def draggable(self, value=None):
156        """Set/get the draggability property of an object."""
157        if value is None:
158            return self.GetDragable()
159        self.SetDragable(value)
160        return self
161
162    def pos(self, x=None, y=None):
163        """Set/Get object position."""
164        if x is None:  # get functionality
165            return np.array(self.GetPosition())
166
167        if y is None:  # assume x is of the form (x,y)
168            x, y = x
169        self.SetPosition(x, y)
170        return self
171
172    def shift(self, ds):
173        """Add a shift to the current object position."""
174        p = np.array(self.GetPosition())
175
176        self.SetPosition(p + ds)
177        return self
178
179    def bounds(self):
180        """
181        Get the object bounds.
182        Returns a list in format [xmin,xmax, ymin,ymax].
183        """
184        return self.GetBounds()
185
186    def diagonal_size(self):
187        """Get the length of the diagonal"""
188        b = self.GetBounds()
189        return np.sqrt((b[1] - b[0]) ** 2 + (b[3] - b[2]) ** 2)
190
191    def show(self, **options):
192        """
193        Create on the fly an instance of class ``Plotter`` or use the last existing one to
194        show one single object.
195
196        This method is meant as a shortcut. If more than one object needs to be visualised
197        please use the syntax `show(mesh1, mesh2, volume, ..., options)`.
198
199        Returns the ``Plotter`` class instance.
200        """
201        return vedo.plotter.show(self, **options)
202
203
204#################################################
205class Assembly(vedo.base.Base3DProp, vtk.vtkAssembly):
206    """
207    Group many objects and treat them as a single new object.
208    """
209
210    def __init__(self, *meshs):
211        """
212        Group many objects and treat them as a single new object,
213        keeping track of internal transformations.
214
215        Examples:
216            - [gyroscope1.py](https://github.com/marcomusy/vedo/tree/master/examples/simulations/gyroscope1.py)
217
218            ![](https://vedo.embl.es/images/simulations/39766016-85c1c1d6-52e3-11e8-8575-d167b7ce5217.gif)
219        """
220        vtk.vtkAssembly.__init__(self)
221        vedo.base.Base3DProp.__init__(self)
222
223        if len(meshs) == 1:
224            meshs = meshs[0]
225        else:
226            meshs = vedo.utils.flatten(meshs)
227
228        self.actors = meshs
229
230        if meshs and hasattr(meshs[0], "top"):
231            self.base = meshs[0].base
232            self.top = meshs[0].top
233        else:
234            self.base = None
235            self.top = None
236
237        scalarbars = []
238        for a in meshs:
239            if isinstance(a, vtk.vtkProp3D):  # and a.GetNumberOfPoints():
240                self.AddPart(a)
241            if hasattr(a, "scalarbar") and a.scalarbar is not None:
242                scalarbars.append(a.scalarbar)
243
244        if len(scalarbars) > 1:
245            self.scalarbar = Group(scalarbars)
246        elif len(scalarbars) == 1:
247            self.scalarbar = scalarbars[0]
248
249        self.pipeline = vedo.utils.OperationNode(
250            "Assembly", parents=meshs, comment=f"#meshes {len(meshs)}", c="#f08080"
251        )
252        ###################################################################
253
254    def _repr_html_(self):
255        """
256        HTML representation of the Assembly object for Jupyter Notebooks.
257
258        Returns:
259            HTML text with the image and some properties.
260        """
261        import io
262        import base64
263        from PIL import Image
264
265        library_name = "vedo.assembly.Assembly"
266        help_url = "https://vedo.embl.es/docs/vedo/assembly.html"
267
268        arr = self.thumbnail(zoom=1.1, elevation=-60)
269
270        im = Image.fromarray(arr)
271        buffered = io.BytesIO()
272        im.save(buffered, format="PNG", quality=100)
273        encoded = base64.b64encode(buffered.getvalue()).decode("utf-8")
274        url = "data:image/png;base64," + encoded
275        image = f"<img src='{url}'></img>"
276
277        # statisitics
278        bounds = "<br/>".join(
279            [
280                vedo.utils.precision(min_x, 4) + " ... " + vedo.utils.precision(max_x, 4)
281                for min_x, max_x in zip(self.bounds()[::2], self.bounds()[1::2])
282            ]
283        )
284
285        help_text = ""
286        if self.name:
287            help_text += f"<b> {self.name}: &nbsp&nbsp</b>"
288        help_text += '<b><a href="' + help_url + '" target="_blank">' + library_name + "</a></b>"
289        if self.filename:
290            dots = ""
291            if len(self.filename) > 30:
292                dots = "..."
293            help_text += f"<br/><code><i>({dots}{self.filename[-30:]})</i></code>"
294
295        allt = [
296            "<table>",
297            "<tr>",
298            "<td>",
299            image,
300            "</td>",
301            "<td style='text-align: center; vertical-align: center;'><br/>",
302            help_text,
303            "<table>",
304            "<tr><td><b> nr. of objects </b></td><td>"
305            + str(self.GetNumberOfPaths())
306            + "</td></tr>",
307            "<tr><td><b> position </b></td><td>" + str(self.GetPosition()) + "</td></tr>",
308            "<tr><td><b> diagonal size </b></td><td>"
309            + vedo.utils.precision(self.diagonal_size(), 5)
310            + "</td></tr>",
311            "<tr><td><b> bounds </b> <br/> (x/y/z) </td><td>" + str(bounds) + "</td></tr>",
312            "</table>",
313            "</table>",
314        ]
315        return "\n".join(allt)
316
317    def __add__(self, obj):
318        """
319        Add an object to the assembly
320        """
321        if isinstance(obj, vtk.vtkProp3D):
322            self.AddPart(obj)
323
324        self.actors.append(obj)
325
326        if hasattr(obj, "scalarbar") and obj.scalarbar is not None:
327            if self.scalarbar is None:
328                self.scalarbar = obj.scalarbar
329                return self
330
331            def unpack_group(scalarbar):
332                if isinstance(scalarbar, Group):
333                    return scalarbar.unpack()
334                else:
335                    return scalarbar
336
337            if isinstance(self.scalarbar, Group):
338                self.scalarbar += unpack_group(obj.scalarbar)
339            else:
340                self.scalarbar = Group([unpack_group(self.scalarbar), unpack_group(obj.scalarbar)])
341        self.pipeline = vedo.utils.OperationNode("add mesh", parents=[self, obj], c="#f08080")
342        return self
343
344    def __contains__(self, obj):
345        """Allows to use ``in`` to check if an object is in the Assembly."""
346        return obj in self.actors
347
348    def clone(self):
349        """Make a clone copy of the object."""
350        newlist = []
351        for a in self.actors:
352            newlist.append(a.clone())
353        return Assembly(newlist)
354
355    def unpack(self, i=None, transformed=False):
356        """Unpack the list of objects from a ``Assembly``.
357
358        If `i` is given, get `i-th` object from a ``Assembly``.
359        Input can be a string, in this case returns the first object
360        whose name contains the given string.
361
362        Examples:
363            - [custom_axes4.py](https://github.com/marcomusy/vedo/tree/master/examples/pyplot/custom_axes4.py)
364        """
365        if transformed:
366            actors = []
367            for a in self.actors:
368                actors.append(a.clone(transformed=True))
369        else:
370            actors = self.actors
371
372        if i is None:
373            return actors
374        elif isinstance(i, int):
375            return actors[i]
376        elif isinstance(i, str):
377            for m in actors:
378                if i in m.name:
379                    return m
380
381    def recursive_unpack(self):
382        """Flatten out an Assembly."""
383
384        def _genflatten(lst):
385            if not lst:
386                return []
387            ##
388            if isinstance(lst[0], Assembly):
389                lst = lst[0].unpack()
390            ##
391            for elem in lst:
392                if isinstance(elem, Assembly):
393                    apos = elem.GetPosition()
394                    asum = np.sum(apos)
395                    for x in elem.unpack():
396                        if asum:
397                            yield x.clone().shift(apos)
398                        else:
399                            yield x
400                else:
401                    yield elem
402
403        return list(_genflatten([self]))