Source code for spateo.tdr.widgets.clip

from typing import Literal, Optional, Union

import numpy as np
import vtk
from pyvista import MultiBlock, PolyData, UnstructuredGrid

from ..models import collect_models, multiblock2model
from .pick import _interactive_pick
from .utils import _interactive_plotter


[docs]def _interactive_rectangle_clip( plotter, model, picking_list, picking_r_list, ): """Add a 2D rectangle widget to the scene.""" legend = [] def _split_model(original_model): """Adds a new model to the plotter each time cells are picked, and removes them from the original model""" try: original_model = original_model[0] except: pass # If nothing selected. if not original_model.n_cells: return # Remove the picked cells from main grid. ghost_cells = np.zeros(model.n_cells, np.uint8) ghost_cells[original_model["orig_extract_id"]] = 1 model.cell_data[vtk.vtkDataSetAttributes.GhostArrayName()] = ghost_cells model.RemoveGhostCells() # Add the selected model this to the main plotter. color = np.random.random(3) legend.append(["picked model %d" % len(picking_list), color]) plotter.add_mesh(original_model, color=color) plotter.add_legend(legend, bcolor=None, face="circle", loc="lower right") # Track the picked models and label them. original_model["picked_index"] = np.ones(original_model.n_points) * len(picking_list) picking_list.append(original_model) picking_r_list.append(model) plotter.enable_cell_picking(mesh=model, callback=_split_model, show=False, color="black") plotter.add_text( "Please double-click the camera orientation widget in the upper right corner first." "\nPress `r` to enable retangle based selection. Press `r` again to turn it off." "\nPress `q` to exit the interactive window. ", font_size=12, color="black", font="arial", )
[docs]def interactive_rectangle_clip( model: Union[PolyData, UnstructuredGrid, MultiBlock], key: str = "groups", model_style: Union[Literal["points", "surface", "wireframe"], list] = "points", model_size: Union[float, list] = 8.0, colormap: str = "Spectral", invert: bool = False, bg_model=None, ) -> MultiBlock: """ Pick the interested part of a model using a 2D rectangle widget. Multiple models can be generated at the same time. Args: model: Reconstructed 3D model. key: The key under which are the labels. invert: Flag on whether to flip/invert the pick. bg_model: A visualization-only background model to help clip our target model. Returns: picked_model: A MultiBlock that contains all models you picked. """ # Check input model. model = ( multiblock2model(model=model, message="rectangle clipping") if isinstance(model, MultiBlock) else model.copy() ) # Create an interactive window for using widgets. p = _interactive_plotter() # Add a visualization-only background model with checkbox button widgets label_models = [] if bg_model is not None: bg_model = multiblock2model(model=bg_model, message=None) if isinstance(bg_model, MultiBlock) else bg_model for label in np.unique(bg_model[key]): if label not in np.unique(model[key]): label_models.append(bg_model.remove_cells(bg_model[key] != label)) checkbox_size = 27 checkbox_start_pos = 5.0 for label_model in label_models: _interactive_pick( plotter=p, model=label_model, key=key, checkbox_color=label_model[f"{key}_rgba"][0][:3], checkbox_position=(5.0, checkbox_start_pos), ) checkbox_start_pos = checkbox_start_pos + checkbox_size + (checkbox_size // 10) # Clip model via a 2D rectangle widget. picked_models, picking_r_list = [], [] if model_style == "points": render_spheres, render_tubes, smooth_shading = True, False, True elif model_style == "wireframe": render_spheres, render_tubes, smooth_shading = False, True, False else: render_spheres, render_tubes, smooth_shading = False, False, True if f"{key}_rgba" in model.array_names: p.add_mesh( model, scalars=f"{key}_rgba", rgba=True, style=model_style, render_points_as_spheres=render_spheres, render_lines_as_tubes=render_tubes, point_size=model_size, line_width=model_size, smooth_shading=smooth_shading, ) else: p.add_mesh( model, scalars=key, style=model_style, render_points_as_spheres=render_spheres, render_lines_as_tubes=render_tubes, point_size=model_size, line_width=model_size, smooth_shading=smooth_shading, cmap=colormap, ) _interactive_rectangle_clip( plotter=p, model=model, picking_list=picked_models, picking_r_list=picking_r_list, ) p.show(cpos="iso") # Obtain final picked model picked_model = collect_models([picking_r_list[0]]) if invert else collect_models(picked_models) return picked_model
[docs]def interactive_box_clip( model: Union[PolyData, UnstructuredGrid, MultiBlock], key: str = "groups", invert: bool = False, ) -> MultiBlock: """ Pick the interested part of a model using a 3D box widget. Only one model can be generated. Args: model: Reconstructed 3D model. key: The key under which are the labels. invert: Flag on whether to flip/invert the pick. Returns: picked_model: A MultiBlock that contains all models you picked. """ # Check input model. model = multiblock2model(model=model, message="box clipping") if isinstance(model, MultiBlock) else model.copy() # Create an interactive window for using widgets. p = _interactive_plotter() # Clip a model using a 3D box widget. p.add_mesh_clip_box( model, invert=invert, scalars=f"{key}_rgba", rgba=True, render_points_as_spheres=True, point_size=10, widget_color="black", ) p.show(cpos="iso") # obtain final picked models picked_model = collect_models(models=p.box_clipped_meshes) return picked_model