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