from typing import Optional, Union
import numpy as np
from pyvista import MultiBlock, PolyData, UnstructuredGrid
from ..models import collect_models, multiblock2model
from .utils import _interactive_plotter
##################################
# Picking by entering parameters #
##################################
[docs]def three_d_pick(
model: Union[PolyData, UnstructuredGrid, MultiBlock],
key: str = "groups",
picked_groups: Union[str, list] = None,
) -> MultiBlock:
"""Pick the desired groups."""
# Check input model.
model = multiblock2model(model=model, message=None) if isinstance(model, MultiBlock) else model.copy()
# Pick groups.
picked_groups = [picked_groups] if isinstance(picked_groups, str) else picked_groups
picked_models = []
for group in picked_groups:
if key in model.point_data.keys():
picked_model = model.extract_points(model.point_data[key] == group)
elif key in model.cell_data.keys():
picked_model = model.extract_cells(model.cell_data[key] == group)
else:
raise ValueError(
"`key` value is wrong." "\nThe `key` value must be contained in model.point_data or model.cell_data."
)
if isinstance(model, PolyData):
picked_model = picked_model.extract_surface()
picked_models.append(picked_model)
return collect_models(picked_models)
#################################
# Picking by interactive window #
#################################
[docs]def _interactive_pick(
plotter,
model,
picking_list: Optional[list] = None,
key: str = "groups",
label_size: int = 12,
checkbox_size: int = 27,
checkbox_color: Union[str, tuple, list] = "blue",
checkbox_position: tuple = (5.0, 5.0),
):
"""Add a checkbox button widget to the scene."""
def toggle_vis(flag):
actor.SetVisibility(flag)
if picking_list is not None:
if flag is True:
picking_list.append(model)
elif flag is False and model in picking_list:
picking_list.remove(model)
# Make a separate callback for each widget
actor = plotter.add_mesh(
model,
scalars=f"{key}_rgba",
rgba=True,
render_points_as_spheres=True,
point_size=10,
)
plotter.add_checkbox_button_widget(
toggle_vis,
value=True,
position=checkbox_position,
size=checkbox_size,
border_size=3,
color_on=checkbox_color,
color_off="white",
background_color=checkbox_color,
)
plotter.add_text(
f"\n {model[key][0]}",
position=checkbox_position,
font_size=label_size,
color="black",
font="arial",
)
[docs]def interactive_pick(
model: Union[PolyData, UnstructuredGrid, MultiBlock],
key: str = "groups",
checkbox_size: int = 27,
label_size: int = 12,
) -> MultiBlock:
"""
Add a checkbox button widget to pick the desired groups through the interactive window and output the picked groups.
Args:
model: Reconstructed 3D model.
key: The key under which are the groups.
checkbox_size: The size of the button in number of pixels.
label_size: The font size of the checkbox labels.
Returns:
A MultiBlock that contains all models you picked.
"""
# Check input model.
model = multiblock2model(model=model, message=None) if isinstance(model, MultiBlock) else model.copy()
label_models = []
if key in model.point_data.keys():
for label in np.unique(model.point_data[key]):
# label_model = model.remove_points(model.point_data[key] != label)
# print(label_model[0])
# label_models.append(label_model[0])
label_models.append(model.remove_cells(model[key] != label))
elif key in model.cell_data.keys():
for label in np.unique(model.cell_data[key]):
label_models.append(model.remove_cells(model.cell_data[key] != label))
else:
raise ValueError(
"`key` value is wrong." "\nThe `key` value must be contained in model.point_data or model.cell_data."
)
picked_models = label_models.copy()
# Create an interactive window for using widgets.
p = _interactive_plotter()
checkbox_start_pos = 5.0
for label_model in label_models:
_interactive_pick(
plotter=p,
model=label_model,
picking_list=picked_models,
key=key,
label_size=label_size,
checkbox_size=checkbox_size,
checkbox_color=tuple(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)
p.show(cpos="iso")
return collect_models(picked_models)
#############################
# Picking overlapping parts #
#############################
[docs]def overlap_pc_pick(
pc: PolyData,
mesh: PolyData,
) -> [PolyData, PolyData]:
"""
Pick the point cloud inside the mesh model and point cloud outside the mesh model.
Args:
pc: Reconstructed 3D point cloud model corresponding to mesh.
mesh: Reconstructed 3D mesh model.
Returns:
inside_pc: Point cloud inside the mesh model.
outside_pc: Point cloud outside the mesh model.
"""
select_pc = pc.select_enclosed_points(surface=mesh, check_surface=False)
inside_pc = select_pc.threshold(0.5, scalars="SelectedPoints")
outside_pc = select_pc.threshold(0.5, invert=True, scalars="SelectedPoints")
return inside_pc, outside_pc
[docs]def overlap_mesh_pick(
mesh1: PolyData,
mesh2: PolyData,
) -> PolyData:
"""
Pick the intersection between two mesh models.
Args:
mesh1: Reconstructed 3D mesh model.
mesh2: Reconstructed 3D mesh model.
Returns:
select_mesh: The intersection mesh model.
"""
select_mesh = mesh1.boolean_intersection(mesh2)
return select_mesh
[docs]def overlap_pick(
main_mesh: PolyData,
other_mesh: PolyData,
main_pc: Optional[PolyData] = None,
other_pc: Optional[PolyData] = None,
):
"""
Add a checkbox button widget to pick the desired groups through the interactive window and output the picked groups.
Args:
main_mesh: Reconstructed 3D mesh model.
other_mesh: Reconstructed 3D mesh model.
main_pc: Reconstructed 3D point cloud model corresponding to main_mesh.
other_pc: Reconstructed 3D point cloud model corresponding to other_mesh.
Returns:
A MultiBlock that contains all models you picked.
"""
select_mesh = overlap_mesh_pick(mesh1=main_mesh, mesh2=other_mesh)
if not (main_pc is None and other_pc is None):
inside_pc1, _ = overlap_pc_pick(pc=main_pc, mesh=select_mesh)
inside_pc2, _ = overlap_pc_pick(pc=other_pc, mesh=select_mesh)
select_pc = collect_models(models=[inside_pc1, inside_pc2])
else:
select_pc = None
return select_mesh, select_pc