vedo.assembly

  1#!/usr/bin/env python3
  2# -*- coding: utf-8 -*-
  3from weakref import ref as weak_ref_to
  4from typing import List, Union, Any
  5import numpy as np
  6
  7import vedo.vtkclasses as vtki  # a wrapper for lazy imports
  8
  9import vedo
 10from vedo.transformations import LinearTransform
 11from vedo.visual import CommonVisual, Actor3DHelper
 12
 13__docformat__ = "google"
 14
 15__doc__ = """
 16Submodule for managing groups of vedo objects
 17
 18![](https://vedo.embl.es/images/basic/align4.png)
 19"""
 20
 21__all__ = ["Group", "Assembly", "procrustes_alignment"]
 22
 23
 24#################################################
 25def procrustes_alignment(sources: List["vedo.Mesh"], rigid=False) -> "Assembly":
 26    """
 27    Return an `Assembly` of aligned source meshes with the `Procrustes` algorithm.
 28    The output `Assembly` is normalized in size.
 29
 30    The `Procrustes` algorithm takes N set of points and aligns them in a least-squares sense
 31    to their mutual mean. The algorithm is iterated until convergence,
 32    as the mean must be recomputed after each alignment.
 33
 34    The set of average points generated by the algorithm can be accessed with
 35    `algoutput.info['mean']` as a numpy array.
 36
 37    Arguments:
 38        rigid : bool
 39            if `True` scaling is disabled.
 40
 41    Examples:
 42        - [align4.py](https://github.com/marcomusy/vedo/tree/master/examples/basic/align4.py)
 43
 44        ![](https://vedo.embl.es/images/basic/align4.png)
 45    """
 46
 47    group = vtki.new("MultiBlockDataGroupFilter")
 48    for source in sources:
 49        if sources[0].npoints != source.npoints:
 50            vedo.logger.error("sources have different nr of points")
 51            raise RuntimeError()
 52        group.AddInputData(source.dataset)
 53    procrustes = vtki.new("ProcrustesAlignmentFilter")
 54    procrustes.StartFromCentroidOn()
 55    procrustes.SetInputConnection(group.GetOutputPort())
 56    if rigid:
 57        procrustes.GetLandmarkTransform().SetModeToRigidBody()
 58    procrustes.Update()
 59
 60    acts = []
 61    for i, s in enumerate(sources):
 62        poly = procrustes.GetOutput().GetBlock(i)
 63        mesh = vedo.mesh.Mesh(poly)
 64        mesh.actor.SetProperty(s.actor.GetProperty())
 65        mesh.properties = s.actor.GetProperty()
 66        if hasattr(s, "name"):
 67            mesh.name = s.name
 68        acts.append(mesh)
 69    assem = Assembly(acts)
 70    assem.transform = procrustes.GetLandmarkTransform()
 71    assem.info["mean"] = vedo.utils.vtk2numpy(procrustes.GetMeanPoints().GetData())
 72    return assem
 73
 74
 75#################################################
 76class Group(vtki.vtkPropAssembly):
 77    """Form groups of generic objects (not necessarily meshes)."""
 78
 79    def __init__(self, objects=()):
 80        """Form groups of generic objects (not necessarily meshes)."""
 81        super().__init__()
 82
 83        if isinstance(objects, dict):
 84            for name in objects:
 85                objects[name].name = name
 86            objects = list(objects.values())
 87
 88        self.actor = self
 89
 90        self.name = "Group"
 91        self.filename = ""
 92        self.trail = None
 93        self.trail_points = []
 94        self.trail_segment_size = 0
 95        self.trail_offset = None
 96        self.shadows = []
 97        self.info = {}
 98        self.rendered_at = set()
 99        self.scalarbar = None
100
101        for a in vedo.utils.flatten(objects):
102            if a:
103                self.AddPart(a.actor)
104
105        self.PickableOff()
106
107
108    def __str__(self):
109        """Print info about Group object."""
110        module = self.__class__.__module__
111        name = self.__class__.__name__
112        out = vedo.printc(
113            f"{module}.{name} at ({hex(id(self))})".ljust(75),
114            bold=True, invert=True, return_string=True,
115        )
116        out += "\x1b[0m"
117        if self.name:
118            out += "name".ljust(14) + ": " + self.name
119            if "legend" in self.info.keys() and self.info["legend"]:
120                out+= f", legend='{self.info['legend']}'"
121            out += "\n"
122
123        n = len(self.unpack())
124        out += "n. of objects".ljust(14) + ": " + str(n) + " "
125        names = [a.name for a in self.unpack() if a.name]
126        if names:
127            out += str(names).replace("'","")[:56]
128        return out.rstrip() + "\x1b[0m"
129
130    def __iadd__(self, obj):
131        """
132        Add an object to the group
133        """
134        if not vedo.utils.is_sequence(obj):
135            obj = [obj]
136        for a in obj:
137            if a:
138                self.AddPart(a)
139        return self
140
141    def _unpack(self):
142        """Unpack the group into its elements"""
143        elements = []
144        self.InitPathTraversal()
145        parts = self.GetParts()
146        parts.InitTraversal()
147        for i in range(parts.GetNumberOfItems()):
148            ele = parts.GetItemAsObject(i)
149            elements.append(ele)
150
151        # gr.InitPathTraversal()
152        # for _ in range(gr.GetNumberOfPaths()):
153        #     path  = gr.GetNextPath()
154        #     print([path])
155        #     path.InitTraversal()
156        #     for i in range(path.GetNumberOfItems()):
157        #         a = path.GetItemAsObject(i).GetViewProp()
158        #         print([a])
159
160        return elements
161
162    def clear(self) -> "Group":
163        """Remove all parts"""
164        for a in self._unpack():
165            self.RemovePart(a)
166        return self
167
168    def on(self) -> "Group":
169        """Switch on visibility"""
170        self.VisibilityOn()
171        return self
172
173    def off(self) -> "Group":
174        """Switch off visibility"""
175        self.VisibilityOff()
176        return self
177
178    def pickable(self, value=True) -> "Group":
179        """The pickability property of the Group."""
180        self.SetPickable(value)
181        return self
182    
183    def use_bounds(self, value=True) -> "Group":
184        """Set the use bounds property of the Group."""
185        self.SetUseBounds(value)
186        return self
187    
188    def print(self) -> "Group":
189        """Print info about the object."""
190        print(self)
191        return self
192
193
194#################################################
195class Assembly(CommonVisual, Actor3DHelper, vtki.vtkAssembly):
196    """
197    Group many objects and treat them as a single new object.
198    """
199
200    def __init__(self, *meshs):
201        """
202        Group many objects and treat them as a single new object,
203        keeping track of internal transformations.
204
205        Examples:
206            - [gyroscope1.py](https://github.com/marcomusy/vedo/tree/master/examples/simulations/gyroscope1.py)
207
208            ![](https://vedo.embl.es/images/simulations/39766016-85c1c1d6-52e3-11e8-8575-d167b7ce5217.gif)
209        """
210        super().__init__()
211
212        # Init by filename
213        if len(meshs) == 1 and isinstance(meshs[0], str):
214            filename = vedo.file_io.download(meshs[0], verbose=False)
215            data = np.load(filename, allow_pickle=True)
216            meshs = [vedo.file_io._from_numpy(dd) for dd in data]
217        # Name and load from dictionary
218        if len(meshs) == 1 and isinstance(meshs[0], dict):
219            meshs = meshs[0]
220            for name in meshs:
221                meshs[name].name = name
222            meshs = list(meshs.values())
223        else:
224            if len(meshs) == 1:
225                meshs = meshs[0]
226            else:
227                meshs = vedo.utils.flatten(meshs)
228
229        self.actor = self
230        self.actor.retrieve_object = weak_ref_to(self)
231
232        self.name = "Assembly"
233        self.filename = ""
234        self.rendered_at = set()
235        self.scalarbar = None
236        self.info = {}
237        self.time = 0
238
239        self.transform = LinearTransform()
240
241        self.objects = [m for m in meshs if m]
242        self.actors  = [m.actor for m in self.objects]
243
244        scalarbars = []
245        for a in self.actors:
246            if isinstance(a, vtki.get_class("Prop3D")): # and a.GetNumberOfPoints():
247                self.AddPart(a)
248            if hasattr(a, "scalarbar") and a.scalarbar is not None:
249                scalarbars.append(a.scalarbar)
250
251        if len(scalarbars) > 1:
252            self.scalarbar = Group(scalarbars)
253        elif len(scalarbars) == 1:
254            self.scalarbar = scalarbars[0]
255
256        self.pipeline = vedo.utils.OperationNode(
257            "Assembly",
258            parents=self.objects,
259            comment=f"#meshes {len(self.objects)}",
260            c="#f08080",
261        )
262        ##########################################
263
264    def __str__(self):
265        """Print info about Assembly object."""
266        module = self.__class__.__module__
267        name = self.__class__.__name__
268        out = vedo.printc(
269            f"{module}.{name} at ({hex(id(self))})".ljust(75),
270            bold=True, invert=True, return_string=True,
271        )
272        out += "\x1b[0m"
273
274        if self.name:
275            out += "name".ljust(14) + ": " + self.name
276            if "legend" in self.info.keys() and self.info["legend"]:
277                out+= f", legend='{self.info['legend']}'"
278            out += "\n"
279
280        n = len(self.unpack())
281        out += "n. of objects".ljust(14) + ": " + str(n) + " "
282        names = [a.name for a in self.unpack() if a.name]
283        if names:
284            out += str(names).replace("'","")[:56]
285        out += "\n"
286
287        pos = self.GetPosition()
288        out += "position".ljust(14) + ": " + str(pos) + "\n"
289
290        bnds = self.GetBounds()
291        bx1, bx2 = vedo.utils.precision(bnds[0], 3), vedo.utils.precision(bnds[1], 3)
292        by1, by2 = vedo.utils.precision(bnds[2], 3), vedo.utils.precision(bnds[3], 3)
293        bz1, bz2 = vedo.utils.precision(bnds[4], 3), vedo.utils.precision(bnds[5], 3)
294        out += "bounds".ljust(14) + ":"
295        out += " x=(" + bx1 + ", " + bx2 + "),"
296        out += " y=(" + by1 + ", " + by2 + "),"
297        out += " z=(" + bz1 + ", " + bz2 + ")\n"
298        return out.rstrip() + "\x1b[0m"
299
300    def _repr_html_(self):
301        """
302        HTML representation of the Assembly object for Jupyter Notebooks.
303
304        Returns:
305            HTML text with the image and some properties.
306        """
307        import io
308        import base64
309        from PIL import Image
310
311        library_name = "vedo.assembly.Assembly"
312        help_url = "https://vedo.embl.es/docs/vedo/assembly.html"
313
314        arr = self.thumbnail(zoom=1.1, elevation=-60)
315
316        im = Image.fromarray(arr)
317        buffered = io.BytesIO()
318        im.save(buffered, format="PNG", quality=100)
319        encoded = base64.b64encode(buffered.getvalue()).decode("utf-8")
320        url = "data:image/png;base64," + encoded
321        image = f"<img src='{url}'></img>"
322
323        # statisitics
324        bounds = "<br/>".join(
325            [
326                vedo.utils.precision(min_x, 4) + " ... " + vedo.utils.precision(max_x, 4)
327                for min_x, max_x in zip(self.bounds()[::2], self.bounds()[1::2])
328            ]
329        )
330
331        help_text = ""
332        if self.name:
333            help_text += f"<b> {self.name}: &nbsp&nbsp</b>"
334        help_text += '<b><a href="' + help_url + '" target="_blank">' + library_name + "</a></b>"
335        if self.filename:
336            dots = ""
337            if len(self.filename) > 30:
338                dots = "..."
339            help_text += f"<br/><code><i>({dots}{self.filename[-30:]})</i></code>"
340
341        allt = [
342            "<table>",
343            "<tr>",
344            "<td>",
345            image,
346            "</td>",
347            "<td style='text-align: center; vertical-align: center;'><br/>",
348            help_text,
349            "<table>",
350            "<tr><td><b> nr. of objects </b></td><td>"
351            + str(self.GetNumberOfPaths())
352            + "</td></tr>",
353            "<tr><td><b> position </b></td><td>" + str(self.GetPosition()) + "</td></tr>",
354            "<tr><td><b> diagonal size </b></td><td>"
355            + vedo.utils.precision(self.diagonal_size(), 5)
356            + "</td></tr>",
357            "<tr><td><b> bounds </b> <br/> (x/y/z) </td><td>" + str(bounds) + "</td></tr>",
358            "</table>",
359            "</table>",
360        ]
361        return "\n".join(allt)
362
363    def __add__(self, obj):
364        """
365        Add an object to the assembly
366        """
367        if isinstance(getattr(obj, "actor", None), vtki.get_class("Prop3D")):
368
369            self.objects.append(obj)
370            self.actors.append(obj.actor)
371            self.AddPart(obj.actor)
372
373            if hasattr(obj, "scalarbar") and obj.scalarbar is not None:
374                if self.scalarbar is None:
375                    self.scalarbar = obj.scalarbar
376                    return self
377
378                def unpack_group(scalarbar):
379                    if isinstance(scalarbar, Group):
380                        return scalarbar.unpack()
381                    else:
382                        return scalarbar
383
384                if isinstance(self.scalarbar, Group):
385                    self.scalarbar += unpack_group(obj.scalarbar)
386                else:
387                    self.scalarbar = Group([unpack_group(self.scalarbar), unpack_group(obj.scalarbar)])
388            self.pipeline = vedo.utils.OperationNode("add mesh", parents=[self, obj], c="#f08080")
389        return self
390
391    def __contains__(self, obj):
392        """Allows to use `in` to check if an object is in the `Assembly`."""
393        return obj in self.objects
394
395    def __getitem__(self, i):
396        """Return i-th object."""
397        if isinstance(i, int):
398            return self.objects[i]
399        elif isinstance(i, str):
400            for m in self.objects:
401                if i == m.name:
402                    return m
403        return None
404
405    def __len__(self):
406        """Return nr. of objects in the assembly."""
407        return len(self.objects)
408
409    # TODO ####
410    # def propagate_transform(self):
411    #     """Propagate the transformation to all parts."""
412    #     # navigate the assembly and apply the transform to all parts
413    #     # and reset position, orientation and scale of the assembly
414    #     for i in range(self.GetNumberOfPaths()):
415    #         path = self.GetPath(i)
416    #         obj = path.GetLastNode().GetViewProp()
417    #         obj.SetUserTransform(self.transform.T)
418    #         obj.SetPosition(0, 0, 0)
419    #         obj.SetOrientation(0, 0, 0)
420    #         obj.SetScale(1, 1, 1)
421    #     raise NotImplementedError()
422
423    def unpack(self, i=None) -> Union[List["vedo.Mesh"], "vedo.Mesh"]:
424        """Unpack the list of objects from a `Assembly`.
425
426        If `i` is given, get `i-th` object from a `Assembly`.
427        Input can be a string, in this case returns the first object
428        whose name contains the given string.
429
430        Examples:
431            - [custom_axes4.py](https://github.com/marcomusy/vedo/tree/master/examples/pyplot/custom_axes4.py)
432        """
433        if i is None:
434            return self.objects
435        elif isinstance(i, int):
436            return self.objects[i]
437        elif isinstance(i, str):
438            for m in self.objects:
439                if i == m.name:
440                    return m
441        return []
442
443    def recursive_unpack(self) -> List["vedo.Mesh"]:
444        """Flatten out an Assembly."""
445
446        def _genflatten(lst):
447            if not lst:
448                return []
449            ##
450            if isinstance(lst[0], Assembly):
451                lst = lst[0].unpack()
452            ##
453            for elem in lst:
454                if isinstance(elem, Assembly):
455                    apos = elem.GetPosition()
456                    asum = np.sum(apos)
457                    for x in elem.unpack():
458                        if asum:
459                            yield x.clone().shift(apos)
460                        else:
461                            yield x
462                else:
463                    yield elem
464
465        return list(_genflatten([self]))
466
467    def pickable(self, value=True) -> "Assembly":
468        """Set/get the pickability property of an assembly and its elements"""
469        self.SetPickable(value)
470        # set property to each element
471        for elem in self.recursive_unpack():
472            elem.pickable(value)
473        return self
474
475    def clone(self) -> "Assembly":
476        """Make a clone copy of the object. Same as `copy()`."""
477        newlist = []
478        for a in self.objects:
479            newlist.append(a.clone())
480        return Assembly(newlist)
481
482    def clone2d(self, pos="bottom-left", size=1, rotation=0, ontop=False, scale=None) -> Group:
483        """
484        Convert the `Assembly` into a `Group` of 2D objects.
485
486        Arguments:
487            pos : (list, str)
488                Position in 2D, as a string or list (x,y).
489                The center of the renderer is [0,0] while top-right is [1,1].
490                Any combination of "center", "top", "bottom", "left" and "right" will work.
491            size : (float)
492                global scaling factor for the 2D object.
493                The scaling is normalized to the x-range of the original object.
494            rotation : (float)
495                rotation angle in degrees.
496            ontop : (bool)
497                if `True` the now 2D object is rendered on top of the 3D scene.
498            scale : (float)
499                deprecated, use `size` instead.
500
501        Returns:
502            `Group` object.
503        """
504        if scale is not None:
505            vedo.logger.warning("clone2d(scale=...) is deprecated, use clone2d(size=...) instead")
506            size = scale
507
508        padding = 0.05
509        x0, x1 = self.xbounds()
510        y0, y1 = self.ybounds()
511        pp = self.pos()
512        x0 -= pp[0]
513        x1 -= pp[0]
514        y0 -= pp[1]
515        y1 -= pp[1]
516
517        offset = [x0, y0]
518        if "cent" in pos:
519            offset = [(x0 + x1) / 2, (y0 + y1) / 2]
520            position = [0., 0.]
521            if "right" in pos:
522                offset[0] = x1
523                position = [1 - padding, 0]
524            if "left" in pos:
525                offset[0] = x0
526                position = [-1 + padding, 0]
527            if "top" in pos:
528                offset[1] = y1
529                position = [0, 1 - padding]
530            if "bottom" in pos:
531                offset[1] = y0
532                position = [0, -1 + padding]
533        elif "top" in pos:
534            if "right" in pos:
535                offset = [x1, y1]
536                position = [1 - padding, 1 - padding]
537            elif "left" in pos:
538                offset = [x0, y1]
539                position = [-1 + padding, 1 - padding]
540            else:
541                raise ValueError(f"incomplete position pos='{pos}'")
542        elif "bottom" in pos:
543            if "right" in pos:
544                offset = [x1, y0]
545                position = [1 - padding, -1 + padding]
546            elif "left" in pos:
547                offset = [x0, y0]
548                position = [-1 + padding, -1 + padding]
549            else:
550                raise ValueError(f"incomplete position pos='{pos}'")
551        else:
552            position = pos
553
554        scanned : List[Any] = []
555        group = Group()
556        for a in self.recursive_unpack():
557            if a in scanned:
558                continue
559            if not isinstance(a, vedo.Points):
560                continue
561            if a.npoints == 0:
562                continue
563
564            s = size * 500 / (x1 - x0)
565            if a.properties.GetRepresentation() == 1:
566                # wireframe is not rendered correctly in 2d
567                b = a.boundaries().lw(1).c(a.color(), a.alpha())
568                if rotation:
569                    b.rotate_z(rotation, around=self.origin())
570                a2d = b.clone2d(size=s, offset=offset)
571            else:
572                if rotation:
573                    # around=self.actor.GetCenter()
574                    a.rotate_z(rotation, around=self.origin())
575                a2d = a.clone2d(size=s, offset=offset)
576            a2d.pos(position).ontop(ontop)
577            group += a2d
578
579        try: # copy info from Histogram1D
580            group.entries = self.entries
581            group.frequencies = self.frequencies
582            group.errors = self.errors
583            group.edges = self.edges
584            group.centers = self.centers
585            group.mean = self.mean
586            group.mode = self.mode
587            group.std = self.std
588        except AttributeError:
589            pass
590
591        group.name = self.name
592        return group
593
594    def copy(self) -> "Assembly":
595        """Return a copy of the object. Alias of `clone()`."""
596        return self.clone()
597
598    # def write(self, filename="assembly.npy") -> "Assembly":
599    #     """
600    #     Write the object to file in `numpy` format.
601    #     """
602    #     objs = []
603    #     for ob in self.unpack():
604    #         d = vedo.file_io._to_numpy(ob)
605    #         objs.append(d)
606    #     np.save(filename, objs)
607    #     return self