vedo.assembly

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3import numpy as np
  4from weakref import ref as weak_ref_to
  5
  6import vedo.vtkclasses as vtk
  7
  8import vedo
  9from vedo.transformations import LinearTransform
 10from vedo.visual import CommonVisual, Actor3DHelper
 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.new("MultiBlockDataGroupFilter")
 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.dataset)
 52    procrustes = vtk.new("ProcrustesAlignmentFilter")
 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.actor.SetProperty(s.actor.GetProperty())
 64        mesh.properties = s.actor.GetProperty()
 65        if hasattr(s, "name"):
 66            mesh.name = s.name
 67        acts.append(mesh)
 68    assem = Assembly(acts)
 69    assem.transform = procrustes.GetLandmarkTransform()
 70    assem.info["mean"] = vedo.utils.vtk2numpy(procrustes.GetMeanPoints().GetData())
 71    return assem
 72
 73
 74#################################################
 75class Group(CommonVisual, vtk.vtkPropAssembly):
 76    """Form groups of generic objects (not necessarily meshes)."""
 77
 78    def __init__(self, objects=()):
 79        """Form groups of generic objects (not necessarily meshes)."""
 80
 81        super().__init__()
 82
 83        self.actor = self
 84
 85        self.name = "Group"
 86        self.filename = ""
 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.scalarbar = None
 95
 96        for a in vedo.utils.flatten(objects):
 97            if a:
 98                self.AddPart(a.actor)
 99
100        self.PickableOff()
101
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    def pos(self, x=None, y=None):
166        """Set/Get object 2D position on the screen."""
167        if x is None:  # get functionality
168            return np.array(self.GetPosition())
169
170        if y is None:  # assume x is of the form (x,y)
171            x, y = x
172        self.SetPosition(x, y)
173        return self
174
175    def shift(self, ds):
176        """Add a shift to the current object position on the screen."""
177        p = np.array(self.GetPosition())
178        self.SetPosition(p + ds)
179        return self
180
181    def bounds(self):
182        """
183        Get the object 2D bounds.
184        Returns a list in format [xmin,xmax, ymin,ymax].
185        """
186        return self.GetBounds()
187
188
189#################################################
190class Assembly(CommonVisual, Actor3DHelper, vtk.vtkAssembly):
191    """
192    Group many objects and treat them as a single new object.
193    """
194
195    def __init__(self, *meshs):
196        """
197        Group many objects and treat them as a single new object,
198        keeping track of internal transformations.
199
200        Examples:
201            - [gyroscope1.py](https://github.com/marcomusy/vedo/tree/master/examples/simulations/gyroscope1.py)
202
203            ![](https://vedo.embl.es/images/simulations/39766016-85c1c1d6-52e3-11e8-8575-d167b7ce5217.gif)
204        """
205        super().__init__()
206
207        # Init by filename
208        if len(meshs) == 1 and isinstance(meshs[0], str):
209            filename = vedo.file_io.download(meshs[0], verbose=False)
210            data = np.load(filename, allow_pickle=True)
211            meshs = [vedo.file_io._from_numpy(dd) for dd in data]
212
213        if len(meshs) == 1:
214            meshs = meshs[0]
215        else:
216            meshs = vedo.utils.flatten(meshs)
217
218        self.actor = self
219        self.actor.retrieve_object = weak_ref_to(self)
220
221        self.name = "Assembly"
222        self.filename = ""
223        self.rendered_at = set()
224        self.scalarbar = None
225        self.info = {}
226        self.time = 0
227
228        self.transform = LinearTransform()
229
230        self.objects = [m for m in meshs if m]
231        self.actors  = [m.actor for m in self.objects]
232
233        scalarbars = []
234        for a in self.actors:
235            if isinstance(a, vtk.get_class("Prop3D")): # and a.GetNumberOfPoints():
236                self.AddPart(a)
237            if hasattr(a, "scalarbar") and a.scalarbar is not None:
238                scalarbars.append(a.scalarbar)
239
240        if len(scalarbars) > 1:
241            self.scalarbar = Group(scalarbars)
242        elif len(scalarbars) == 1:
243            self.scalarbar = scalarbars[0]
244
245        self.pipeline = vedo.utils.OperationNode(
246            "Assembly",
247            parents=self.objects,
248            comment=f"#meshes {len(self.objects)}",
249            c="#f08080",
250        )
251        ##########################################
252
253    def __str__(self):
254        """Print info about Assembly object."""
255        module = self.__class__.__module__
256        name = self.__class__.__name__
257        out = vedo.printc(
258            f"{module}.{name} at ({hex(id(self))})".ljust(75),
259            bold=True, invert=True, return_string=True,
260        )
261        out += "\x1b[0m"
262
263        if self.name:
264            out += "name".ljust(14) + ": " + self.name
265            if "legend" in self.info.keys() and self.info["legend"]:
266                out+= f", legend='{self.info['legend']}'"
267            out += "\n"
268
269        n = len(self.unpack())
270        out += "n. of objects".ljust(14) + ": " + str(n) + " "
271        names = [a.name for a in self.unpack() if a.name]
272        if names:
273            out += str(names).replace("'","")[:56]
274        out += "\n"
275
276        pos = self.GetPosition()
277        out += "position".ljust(14) + ": " + str(pos) + "\n"
278
279        bnds = self.GetBounds()
280        bx1, bx2 = vedo.utils.precision(bnds[0], 3), vedo.utils.precision(bnds[1], 3)
281        by1, by2 = vedo.utils.precision(bnds[2], 3), vedo.utils.precision(bnds[3], 3)
282        bz1, bz2 = vedo.utils.precision(bnds[4], 3), vedo.utils.precision(bnds[5], 3)
283        out+= "bounds".ljust(14) + ":"
284        out+= " x=(" + bx1 + ", " + bx2 + "),"
285        out+= " y=(" + by1 + ", " + by2 + "),"
286        out+= " z=(" + bz1 + ", " + bz2 + ")\n"
287        return out.rstrip() + "\x1b[0m"
288
289    def _repr_html_(self):
290        """
291        HTML representation of the Assembly object for Jupyter Notebooks.
292
293        Returns:
294            HTML text with the image and some properties.
295        """
296        import io
297        import base64
298        from PIL import Image
299
300        library_name = "vedo.assembly.Assembly"
301        help_url = "https://vedo.embl.es/docs/vedo/assembly.html"
302
303        arr = self.thumbnail(zoom=1.1, elevation=-60)
304
305        im = Image.fromarray(arr)
306        buffered = io.BytesIO()
307        im.save(buffered, format="PNG", quality=100)
308        encoded = base64.b64encode(buffered.getvalue()).decode("utf-8")
309        url = "data:image/png;base64," + encoded
310        image = f"<img src='{url}'></img>"
311
312        # statisitics
313        bounds = "<br/>".join(
314            [
315                vedo.utils.precision(min_x, 4) + " ... " + vedo.utils.precision(max_x, 4)
316                for min_x, max_x in zip(self.bounds()[::2], self.bounds()[1::2])
317            ]
318        )
319
320        help_text = ""
321        if self.name:
322            help_text += f"<b> {self.name}: &nbsp&nbsp</b>"
323        help_text += '<b><a href="' + help_url + '" target="_blank">' + library_name + "</a></b>"
324        if self.filename:
325            dots = ""
326            if len(self.filename) > 30:
327                dots = "..."
328            help_text += f"<br/><code><i>({dots}{self.filename[-30:]})</i></code>"
329
330        allt = [
331            "<table>",
332            "<tr>",
333            "<td>",
334            image,
335            "</td>",
336            "<td style='text-align: center; vertical-align: center;'><br/>",
337            help_text,
338            "<table>",
339            "<tr><td><b> nr. of objects </b></td><td>"
340            + str(self.GetNumberOfPaths())
341            + "</td></tr>",
342            "<tr><td><b> position </b></td><td>" + str(self.GetPosition()) + "</td></tr>",
343            "<tr><td><b> diagonal size </b></td><td>"
344            + vedo.utils.precision(self.diagonal_size(), 5)
345            + "</td></tr>",
346            "<tr><td><b> bounds </b> <br/> (x/y/z) </td><td>" + str(bounds) + "</td></tr>",
347            "</table>",
348            "</table>",
349        ]
350        return "\n".join(allt)
351
352    def __add__(self, obj):
353        """
354        Add an object to the assembly
355        """
356        if isinstance(obj, vtk.get_class("Prop3D")):
357
358            self.objects.append(obj)
359            self.actors.append(obj.actor)
360            self.AddPart(obj.actor)
361
362            if hasattr(obj, "scalarbar") and obj.scalarbar is not None:
363                if self.scalarbar is None:
364                    self.scalarbar = obj.scalarbar
365                    return self
366
367                def unpack_group(scalarbar):
368                    if isinstance(scalarbar, Group):
369                        return scalarbar.unpack()
370                    else:
371                        return scalarbar
372
373                if isinstance(self.scalarbar, Group):
374                    self.scalarbar += unpack_group(obj.scalarbar)
375                else:
376                    self.scalarbar = Group([unpack_group(self.scalarbar), unpack_group(obj.scalarbar)])
377            self.pipeline = vedo.utils.OperationNode("add mesh", parents=[self, obj], c="#f08080")
378        return self
379
380    def __contains__(self, obj):
381        """Allows to use `in` to check if an object is in the `Assembly`."""
382        return obj in self.objects
383
384    def __getitem__(self, i):
385        """Return i-th object."""
386        if isinstance(i, int):
387            return self.objects[i]
388        elif isinstance(i, str):
389            for m in self.objects:
390                if i in m.name:
391                    return m
392        return None
393
394    def __len__(self):
395        """Return nr. of objects in the assembly."""
396        return len(self.objects)
397
398    # TODO ####
399    # def propagate_transform(self):
400    #     """Propagate the transformation to all parts."""
401    #     # navigate the assembly and apply the transform to all parts
402    #     # and reset position, orientation and scale of the assembly
403    #     for i in range(self.GetNumberOfPaths()):
404    #         path = self.GetPath(i)
405    #         obj = path.GetLastNode().GetViewProp()
406    #         obj.SetUserTransform(self.transform.T)
407    #         obj.SetPosition(0, 0, 0)
408    #         obj.SetOrientation(0, 0, 0)
409    #         obj.SetScale(1, 1, 1)
410    #     raise NotImplementedError()
411
412    def unpack(self, i=None):
413        """Unpack the list of objects from a `Assembly`.
414
415        If `i` is given, get `i-th` object from a `Assembly`.
416        Input can be a string, in this case returns the first object
417        whose name contains the given string.
418
419        Examples:
420            - [custom_axes4.py](https://github.com/marcomusy/vedo/tree/master/examples/pyplot/custom_axes4.py)
421        """
422        if i is None:
423            return self.objects
424        elif isinstance(i, int):
425            return self.objects[i]
426        elif isinstance(i, str):
427            for m in self.objects:
428                if i in m.name:
429                    return m
430
431    def recursive_unpack(self):
432        """Flatten out an Assembly."""
433
434        def _genflatten(lst):
435            if not lst:
436                return []
437            ##
438            if isinstance(lst[0], Assembly):
439                lst = lst[0].unpack()
440            ##
441            for elem in lst:
442                if isinstance(elem, Assembly):
443                    apos = elem.GetPosition()
444                    asum = np.sum(apos)
445                    for x in elem.unpack():
446                        if asum:
447                            yield x.clone().shift(apos)
448                        else:
449                            yield x
450                else:
451                    yield elem
452
453        return list(_genflatten([self]))
454
455    def pickable(self, value=True):
456        """Set/get the pickability property of an assembly and its elements"""
457        self.SetPickable(value)
458        # set property to each element
459        for elem in self.recursive_unpack():
460            elem.pickable(value)
461        return self
462
463    def clone(self):
464        """Make a clone copy of the object. Same as `copy()`."""
465        newlist = []
466        for a in self.objects:
467            newlist.append(a.clone())
468        return Assembly(newlist)
469    
470    def copy(self):
471        """Return a copy of the object. Alias of `clone()`."""
472        return self.clone()
473
474    def write(self, filename="assembly.npy"):
475        """
476        Write the object to file in `numpy` format.
477        """
478        objs = []
479        for ob in self.unpack():
480            d = vedo.file_io._to_numpy(ob)
481            objs.append(d)
482        np.save(filename, objs)
483        return self