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