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