diff --git a/modules/zividsamples/gui/verification/capture_at_pose_selection_widget.py b/modules/zividsamples/gui/verification/capture_at_pose_selection_widget.py index 19eacac7..870c493b 100644 --- a/modules/zividsamples/gui/verification/capture_at_pose_selection_widget.py +++ b/modules/zividsamples/gui/verification/capture_at_pose_selection_widget.py @@ -1,7 +1,7 @@ import threading from collections import OrderedDict from pathlib import Path -from typing import List, Optional +from typing import List, NamedTuple, Optional import numpy as np import zivid @@ -47,6 +47,46 @@ ] +class RoiConfig(NamedTuple): + min_x: float + max_x: float + min_y: float + max_y: float + min_z: float + max_z: float + + +def _build_roi_box_in_camera_frame(roi_config, robot_pose, hand_eye_transform, eye_in_hand): + if eye_in_hand: + camera_to_reference = robot_pose * hand_eye_transform + else: + camera_to_reference = robot_pose.inv() * hand_eye_transform + reference_to_camera = camera_to_reference.inv() + transform = reference_to_camera.as_matrix() + rotation = transform[:3, :3] + translation = transform[:3, 3] + + point_o = np.array([roi_config.min_x, roi_config.min_y, 0.0]) + point_a = np.array([roi_config.max_x, roi_config.min_y, 0.0]) + point_b = np.array([roi_config.min_x, roi_config.max_y, 0.0]) + + return zivid.Settings.RegionOfInterest.Box( + enabled=True, + point_o=rotation @ point_o + translation, + point_a=rotation @ point_a + translation, + point_b=rotation @ point_b + translation, + extents=(roi_config.min_z, roi_config.max_z), + ) + + +def _try_apply_roi_mask(point_cloud, roi_config, robot_pose, hand_eye_transform, eye_in_hand): + roi_box = _build_roi_box_in_camera_frame(roi_config, robot_pose, hand_eye_transform, eye_in_hand) + if point_cloud.masked_by_region_of_interest(roi_box).to_unorganized_point_cloud().size > 0: + point_cloud.mask_by_region_of_interest(roi_box) + return True + return False + + class CaptureAtPose: def _translation_to_string(self, translation: NDArray[Shape["3"], Float32]) -> str: # type: ignore @@ -63,6 +103,7 @@ def __init__( eye_in_hand: bool, optimize_for_speed: bool = True, from_disk: bool = False, + roi_config: Optional[RoiConfig] = None, ): self.poseID = poseID @@ -72,6 +113,7 @@ def __init__( self.camera_frame_path: Path = self.directory / f"capture_{self.poseID}.zdf" self.robot_pose = robot_pose self.camera_frame = camera_frame + self.roi_all_points_masked = False if not from_disk: zivid.Matrix4x4(self.robot_pose.as_matrix()).save(self.robot_pose_yaml_path) @@ -80,6 +122,12 @@ def __init__( if optimize_for_speed: self.camera_frame.point_cloud().downsample(zivid.PointCloud.Downsampling.by2x2) + if roi_config is not None: + if not _try_apply_roi_mask( + self.camera_frame.point_cloud(), roi_config, robot_pose, hand_eye_transform, eye_in_hand + ): + self.roi_all_points_masked = True + if eye_in_hand: transform_robot_base_to_camera = self.robot_pose * hand_eye_transform self.camera_frame.point_cloud().transform(zivid.Matrix4x4(transform_robot_base_to_camera.as_matrix())) @@ -119,7 +167,7 @@ def robot_pose_yaml_text(self) -> str: class _CaptureAtPoseLoadWorker(QObject): """Loads capture-at-pose data from disk in a background thread.""" - item_loaded = pyqtSignal(int, int, object, object) + item_loaded = pyqtSignal(int, int, object, object, bool) finished = pyqtSignal(int) # pylint: disable=too-many-positional-arguments @@ -130,6 +178,7 @@ def __init__( pose_ids: List[int], hand_eye_transform: TransformationMatrix, eye_in_hand: bool, + roi_config: Optional[RoiConfig] = None, ) -> None: super().__init__() self._generation = generation @@ -137,6 +186,7 @@ def __init__( self._pose_ids = pose_ids self._hand_eye_transform = hand_eye_transform self._eye_in_hand = eye_in_hand + self._roi_config = roi_config self._cancel_event = threading.Event() def cancel(self) -> None: @@ -157,13 +207,24 @@ def run(self) -> None: camera_frame.point_cloud().downsample(zivid.PointCloud.Downsampling.by2x2) + roi_warning = False + if self._roi_config is not None: + if not _try_apply_roi_mask( + camera_frame.point_cloud(), + self._roi_config, + robot_pose, + self._hand_eye_transform, + self._eye_in_hand, + ): + roi_warning = True + if self._eye_in_hand: transform = robot_pose * self._hand_eye_transform else: transform = robot_pose.inv() * self._hand_eye_transform camera_frame.point_cloud().transform(zivid.Matrix4x4(transform.as_matrix())) - self.item_loaded.emit(self._generation, poseID, robot_pose, camera_frame) + self.item_loaded.emit(self._generation, poseID, robot_pose, camera_frame, roi_warning) except FileNotFoundError: continue self.finished.emit(self._generation) @@ -210,7 +271,12 @@ def __init__(self, directory: Path, parent: Optional[QWidget] = None) -> None: def set_directory(self, directory: Path) -> None: self.directory = directory - def load_capture_at_poses(self, hand_eye_transform: TransformationMatrix, eye_in_hand: bool) -> None: + def load_capture_at_poses( + self, + hand_eye_transform: TransformationMatrix, + eye_in_hand: bool, + roi_config: Optional[RoiConfig] = None, + ) -> None: if self.number_of_active_captures() > 0: reply = QMessageBox.question( self, @@ -245,6 +311,7 @@ def load_capture_at_poses(self, hand_eye_transform: TransformationMatrix, eye_in pose_ids=pose_ids, hand_eye_transform=hand_eye_transform, eye_in_hand=eye_in_hand, + roi_config=roi_config, ) self._loader_worker.moveToThread(self._loader_thread) assert self._loader_thread is not None @@ -263,8 +330,14 @@ def cancel_loading(self) -> None: self._loader_worker = None self._loader_thread = None + # pylint: disable=too-many-positional-arguments def _on_capture_loaded( - self, generation: int, poseID: int, robot_pose: TransformationMatrix, camera_frame: zivid.Frame + self, + generation: int, + poseID: int, + robot_pose: TransformationMatrix, + camera_frame: zivid.Frame, + roi_warning: bool, ) -> None: if generation != self._load_generation: return @@ -281,6 +354,7 @@ def _on_capture_loaded( eye_in_hand=True, from_disk=True, ) + capture_at_pose.roi_all_points_masked = roi_warning capture_at_pose_layout = QHBoxLayout() capture_at_pose.capture_pose_button.clicked.connect(lambda: self.on_capture_at_pose_clicked(capture_at_pose)) capture_at_pose.remove_capture_at_pose_button.clicked.connect( @@ -323,12 +397,14 @@ def remove_capture_at_pose(self, capture_at_pose: CaptureAtPose) -> None: def is_loading(self) -> bool: return self._loader_thread is not None and self._loader_thread.isRunning() + # pylint: disable=too-many-positional-arguments def add_capture_at_pose( self, robot_pose: TransformationMatrix, camera_frame: zivid.Frame, hand_eye_transform: TransformationMatrix, eye_in_hand: bool, + roi_config: Optional[RoiConfig] = None, ) -> None: if self.is_loading(): return @@ -352,6 +428,7 @@ def add_capture_at_pose( camera_frame=camera_frame, hand_eye_transform=hand_eye_transform, eye_in_hand=eye_in_hand, + roi_config=roi_config, ) capture_at_pose_layout = QHBoxLayout() capture_at_pose.capture_pose_button.clicked.connect(lambda: self.on_capture_at_pose_clicked(capture_at_pose)) diff --git a/modules/zividsamples/gui/verification/stitch_gui.py b/modules/zividsamples/gui/verification/stitch_gui.py index ab00a025..9eeb589c 100644 --- a/modules/zividsamples/gui/verification/stitch_gui.py +++ b/modules/zividsamples/gui/verification/stitch_gui.py @@ -15,11 +15,27 @@ from nptyping import NDArray, Shape, UInt8 from PyQt5.QtCore import pyqtSignal from PyQt5.QtGui import QCloseEvent, QImage -from PyQt5.QtWidgets import QCheckBox, QFileDialog, QHBoxLayout, QMessageBox, QPushButton, QVBoxLayout, QWidget +from PyQt5.QtWidgets import ( + QCheckBox, + QDoubleSpinBox, + QFileDialog, + QGridLayout, + QGroupBox, + QHBoxLayout, + QLabel, + QMessageBox, + QPushButton, + QVBoxLayout, + QWidget, +) from zivid.experimental.point_cloud_export import export_unorganized_point_cloud from zivid.experimental.point_cloud_export.file_format import PLY from zividsamples.gui.robot.robot_control import RobotTarget -from zividsamples.gui.verification.capture_at_pose_selection_widget import CaptureAtPose, CaptureAtPoseSelectionWidget +from zividsamples.gui.verification.capture_at_pose_selection_widget import ( + CaptureAtPose, + CaptureAtPoseSelectionWidget, + RoiConfig, +) from zividsamples.gui.widgets.pointcloud_visualizer import VisualizerWidget from zividsamples.gui.widgets.pose_widget import PoseWidget, PoseWidgetDisplayMode from zividsamples.gui.widgets.tab_with_robot_support import TabWidgetWithRobotSupport @@ -96,6 +112,45 @@ def create_widgets(self, initial_rotation_information: RotationInformation) -> N self.save_point_cloud_button.setEnabled(False) self.stitched_point_cloud: Optional[zivid.UnorganizedPointCloud] = None + self.roi_enabled_checkbox = QCheckBox("Enable ROI masking") + self.roi_enabled_checkbox.setChecked(False) + + self.roi_min_x_spinbox = self._create_roi_spinbox(-500) + self.roi_max_x_spinbox = self._create_roi_spinbox(500) + self.roi_min_y_spinbox = self._create_roi_spinbox(-500) + self.roi_max_y_spinbox = self._create_roi_spinbox(500) + self.roi_min_z_spinbox = self._create_roi_spinbox(-500) + self.roi_max_z_spinbox = self._create_roi_spinbox(500) + + self.roi_group_box = QGroupBox(self._roi_frame_text()) + roi_layout = QVBoxLayout() + roi_layout.addWidget(self.roi_enabled_checkbox) + + self.roi_extents_widget = QWidget() + roi_grid = QGridLayout() + roi_grid.setContentsMargins(0, 0, 0, 0) + roi_grid.addWidget(QLabel(""), 0, 0) + roi_grid.addWidget(QLabel("Min (mm)"), 0, 1) + roi_grid.addWidget(QLabel("Max (mm)"), 0, 2) + roi_grid.addWidget(QLabel("X"), 1, 0) + roi_grid.addWidget(self.roi_min_x_spinbox, 1, 1) + roi_grid.addWidget(self.roi_max_x_spinbox, 1, 2) + roi_grid.addWidget(QLabel("Y"), 2, 0) + roi_grid.addWidget(self.roi_min_y_spinbox, 2, 1) + roi_grid.addWidget(self.roi_max_y_spinbox, 2, 2) + roi_grid.addWidget(QLabel("Z"), 3, 0) + roi_grid.addWidget(self.roi_min_z_spinbox, 3, 1) + roi_grid.addWidget(self.roi_max_z_spinbox, 3, 2) + self.roi_extents_widget.setLayout(roi_grid) + self.roi_extents_widget.setVisible(False) + roi_layout.addWidget(self.roi_extents_widget) + self.roi_group_box.setLayout(roi_layout) + + self.roi_warning_label = QLabel() + self.roi_warning_label.setStyleSheet("color: orange;") + self.roi_warning_label.setWordWrap(True) + self.roi_warning_label.setVisible(False) + def setup_layout(self) -> None: layout = QVBoxLayout() left_panel = QVBoxLayout() @@ -110,6 +165,8 @@ def setup_layout(self) -> None: left_panel.addLayout(confirm_robot_pose_layout) left_panel.addWidget(self.hand_eye_pose_widget) right_panel.addWidget(self.capture_at_pose_selection_widget) + right_panel.addWidget(self.roi_group_box) + right_panel.addWidget(self.roi_warning_label) right_panel.addWidget(self.uniform_color_check_box) right_panel.addWidget(self.save_point_cloud_button) center_layout.addLayout(left_panel) @@ -126,6 +183,7 @@ def connect_signals(self) -> None: self.capture_at_pose_selection_widget.loading_finished.connect(self.loading_finished) self.uniform_color_check_box.stateChanged.connect(self.update_stitched_view) self.save_point_cloud_button.clicked.connect(self.on_save_point_cloud_clicked) + self.roi_enabled_checkbox.toggled.connect(self._on_roi_enabled_toggled) def update_instructions(self, captured: bool, robot_pose_confirmed: bool) -> None: self.has_confirmed_robot_pose = robot_pose_confirmed @@ -151,6 +209,7 @@ def on_pending_changes(self) -> None: self.capture_at_pose_selection_widget.load_capture_at_poses( self.hand_eye_pose_widget.get_transformation_matrix(), self.hand_eye_configuration.eye_in_hand, + roi_config=self._get_roi_config(), ) else: self.capture_at_pose_selection_widget.set_directory(self.data_directory) @@ -168,6 +227,7 @@ def hand_eye_configuration_update(self, hand_eye_configuration: HandEyeConfigura self.hand_eye_configuration = hand_eye_configuration self.hand_eye_pose_widget.on_eye_in_hand_toggled(self.hand_eye_configuration.eye_in_hand) self.robot_pose_widget.on_eye_in_hand_toggled(self.hand_eye_configuration.eye_in_hand) + self.roi_group_box.setTitle(self._roi_frame_text()) def rotation_format_update(self, rotation_information: RotationInformation) -> None: self.hand_eye_pose_widget.set_rotation_format(rotation_information) @@ -193,6 +253,16 @@ def update_stitched_view(self) -> None: if self.uniform_color_check_box.isChecked(): point_cloud_at_pose.paint_uniform_color(capture_at_pose.color + [128]) unorganized_point_cloud.extend(point_cloud_at_pose) + + has_roi_warnings = any(cap.roi_all_points_masked for cap in capture_at_poses) + if has_roi_warnings: + self.roi_warning_label.setText( + "ROI masking removed all points for one or more captures. Showing unmasked data — adjust ROI settings." + ) + self.roi_warning_label.setVisible(True) + else: + self.roi_warning_label.setVisible(False) + if unorganized_point_cloud.size > 0: unorganized_point_cloud = unorganized_point_cloud.voxel_downsampled(voxel_size=1, min_points_per_voxel=1) self.point_cloud_widget.set_point_cloud(unorganized_point_cloud) @@ -223,6 +293,7 @@ def process_capture(self, frame: zivid.Frame, _: NDArray[Shape["N, M, 4"], UInt8 camera_frame=frame, hand_eye_transform=self.hand_eye_pose_widget.get_transformation_matrix(), eye_in_hand=self.hand_eye_configuration.eye_in_hand, + roi_config=self._get_roi_config(), ) self.update_stitched_view() self.update_instructions(captured=True, robot_pose_confirmed=False) @@ -238,6 +309,45 @@ def get_tab_widgets_in_order(self) -> List[QWidget]: widgets.extend(self.hand_eye_pose_widget.get_tab_widgets_in_order()) return widgets + @staticmethod + def _create_roi_spinbox(default_value: float) -> QDoubleSpinBox: + spinbox = QDoubleSpinBox() + spinbox.setRange(-10000, 10000) + spinbox.setDecimals(1) + spinbox.setSingleStep(10) + spinbox.setValue(default_value) + return spinbox + + def _roi_spinboxes(self) -> List[QDoubleSpinBox]: + return [ + self.roi_min_x_spinbox, + self.roi_max_x_spinbox, + self.roi_min_y_spinbox, + self.roi_max_y_spinbox, + self.roi_min_z_spinbox, + self.roi_max_z_spinbox, + ] + + def _roi_frame_text(self) -> str: + if self.hand_eye_configuration.eye_in_hand: + return "Region of Interest (Robot Base Frame)" + return "Region of Interest (Robot Flange Frame)" + + def _on_roi_enabled_toggled(self, enabled: bool): + self.roi_extents_widget.setVisible(enabled) + + def _get_roi_config(self): + if not self.roi_enabled_checkbox.isChecked(): + return None + return RoiConfig( + min_x=self.roi_min_x_spinbox.value(), + max_x=self.roi_max_x_spinbox.value(), + min_y=self.roi_min_y_spinbox.value(), + max_y=self.roi_max_y_spinbox.value(), + min_z=self.roi_min_z_spinbox.value(), + max_z=self.roi_max_z_spinbox.value(), + ) + def closeEvent(self, event: QCloseEvent) -> None: # pylint: disable=C0103 self.point_cloud_widget.close() super().closeEvent(event)