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