SelectionGizmo#

class SelectionGizmo:
    """Selection gizmo for the 3D viewer.

    This gizmo allows to select objects in the 3d viewer by drawing a rectangle.

    Parameters
    ----------
    viewer : Viewer
        The viewer to which the gizmo is attached.
    callback : callable, optional
        A function to call when the selection is done. Must accept the selection as a dictionary (see
        [`SelectionGizmo._find_selected_objects`][octarine.selection.SelectionGizmo._find_selected_objects]
        for details). Can also be added post-hoc using `SelectionGizmo.add_callback().
    ignore_invisible :  bool
        Whether to ignore invisible objects. Default is True.
    **kwargs
        Keyword arguments passed to BaseSelectionGizmo.

    """

    def __init__(self, viewer, callback=None, ignore_invisible=True, **kwargs):
        self.callbacks = []
        if not isinstance(callback, type(None)):
            self.add_callback(callback)

        if "callback_after" in kwargs:
            raise ValueError(
                "The `callback_after` keyword argument is reserved for internal use!"
            )
        kwargs["callback_after"] = self.handle_post_selection
        self.gizmo = BaseSelectionGizmo(viewer=viewer, **kwargs)
        self.ignore_invisible = ignore_invisible

    def add_callback(self, callback):
        if not callable(callback):
            raise ValueError("`callback` must be a single callable")
        self.callbacks.append(callback)

    def handle_post_selection(self, gizmo):
        """Handle things after selection."""
        # If no callback is set, we're done here
        if not self.callbacks and not self.gizmo._debug:
            return

        # Find the selected objects
        selection = self.find_selected_objects()
        if self.gizmo._debug:
            n_clipped = len([k for k in selection if selection[k]["clipped"]])
            n_contained = len([k for k in selection if selection[k]["contained"]])
            print(f"Selected objects: {n_clipped} clipped, {n_contained} contained:")
            print(selection)

        # Run the callbacks
        for func in self.callbacks:
            func(selection)

    def find_selected_objects(self):
        """Find objects that are selected by the selection box.

        Notes
        -----
        - currently only Meshes and Points are supported; other objects are returns as `None`

        Returns
        -------
        dict
                Returns a dictionary matching `Viewer.objects`, i.e. keyed by object groups:

                    {
                        "group_name": {
                            "clipped": bool,  # whether the object(s) is at least clipped
                            "contained": bool,  # whether the object(s) is fully contained
                            "objects": [  # an entry for each object in the object group
                                {
                                "clipped": bool,
                                "contained": bool,
                                "data": np.ndarray | None (if not clipped or fully contained)
                                }, ...
                            ]
                        }
                    }

        """
        # If no selection box, return right away
        if not self.gizmo._sel_info:
            return {}

        # Next, we need to iterate over all objects, (1) project their data to screen space
        # and (2) determine whether they are inside the selection box
        selection = {}
        for name in (
            self.gizmo._viewer.visible
            if self.ignore_invisible
            else self.gizmo._viewer.objects
        ):
            selection[name] = {"clipped": False, "contained": True, "objects": []}
            objects = self.gizmo._viewer.objects[name]
            for ob in objects:
                # Extract the data depending on the type of object
                if not isinstance(ob, (gfx.Mesh, gfx.Points, gfx.Line)):
                    # Note to self: we could use object boundaries where no data is available
                    if self.gizmo._debug:
                        print(f"Object {ob} not supported")
                    selection[name]["contained"] = None
                    selection[name]["objects"].append(
                        {"clipped": None, "contained": None, "data": None}
                    )
                    continue

                # Get the data representing the object
                data = ob.geometry.positions.data

                # Apply world transform
                data = la.vec_transform(data, ob.world.matrix)

                # Project the data to screen space
                # (N.B. were using the viewer's normal camera, not the overlay's NDC camera)
                data_screen = la.vec_transform(
                    data, self.gizmo._viewer.camera.camera_matrix
                )

                # Check whether the object is in the selection box
                mask = self.is_data_in_selection_box(data_screen)

                # Store the results
                is_clipped = bool(np.any(mask))  # avoid getting np.True_/np.False_

                if isinstance(ob, gfx.Line):
                    # Lines can have breaks where the data will be `nan` - these will
                    # always count as not inside the selection box. We need to ignore
                    # these when checking for full containment.
                    data_not_nan = np.all(~np.isnan(data), axis=1)
                    is_contained = bool(np.all(mask[data_not_nan]))
                else:
                    is_contained = bool(np.all(mask))

                # If the object is either fully contained or fully outside, we don't
                # really need to pass the mask
                if is_contained or not is_clipped:
                    mask = None

                selection[name]["objects"].append(
                    {"clipped": is_clipped, "contained": is_contained, "data": mask}
                )
                # Propagate results to parent group
                if not is_contained:
                    selection[name]["contained"] = False
                if is_clipped:
                    selection[name]["clipped"] = True

        return selection

    def is_data_in_selection_box(self, data_screen):
        """Check whether selected obejct is inside selection box.

        Parameters
        ----------
        data_screen : (N, 3) array
            The data in screen coordinates.

        Returns
        -------
        bool
            Whether the object is inside the selection box.

        """
        if not self.gizmo._sel_info:
            return np.zeros(data_screen.shape[0], dtype=bool)

        # Check whether the object is in the selection box
        return np.all(
            (data_screen[:, [0, 1]] >= self.gizmo.bounds[0, :-1])
            & (data_screen[:, [0, 1]] <= self.gizmo.bounds[1, :-1]),
            axis=1,
        )