From 993e43d305afc875394f4bf33678edd01d0053bf Mon Sep 17 00:00:00 2001 From: builder Date: Tue, 28 Jul 2026 20:30:20 +0000 Subject: [PATCH] Samples: Add DLPack GPU interop sample and stitching ROI masking New sample: - capture_and_convert_to_dlpack_tensor_on_cuda: hand a Zivid DeviceArray to PyTorch on the GPU without a CPU round-trip. It shows two paths: torch.as_tensor through the CUDA array interface (the recommended one-liner), and building a DLPack capsule directly from the DeviceArray's device pointer, shape, strides and element type for consumers that expect the __dlpack__ protocol. Both hand-offs are verified zero-copy by comparing data pointers, and both are scoped to the CUDA backend. Requires PyTorch with CUDA support. - The README lists the new sample among the cuda samples. Stitching GUI region of interest: - The stitching verification tab now has an optional Region of Interest box. Enable it and set min/max X, Y and Z in mm to keep background clutter out of the stitched point cloud. - The extents are given in the frame that stays fixed relative to the scanned object: robot base frame for eye-in-hand, robot flange frame for eye-to-hand. The group box title states which frame is in use. - Masking is applied per capture before the hand-eye transform, both for new captures and for captures loaded from disk. - If the box would remove every point of a capture, the unmasked data is shown instead and a warning tells you to adjust the extents. --- README.md | 4 + .../capture_at_pose_selection_widget.py | 87 +++++- .../gui/verification/stitch_gui.py | 114 +++++++- ...re_and_convert_to_dlpack_tensor_on_cuda.py | 255 ++++++++++++++++++ 4 files changed, 453 insertions(+), 7 deletions(-) create mode 100644 source/applications/advanced/cuda/capture_and_convert_to_dlpack_tensor_on_cuda.py diff --git a/README.md b/README.md index 260f3891..81fd36ed 100644 --- a/README.md +++ b/README.md @@ -155,6 +155,9 @@ from the camera can be used. clouds, with color, from the Zivid camera, and visualize them in a loop. Press 'q' to exit. - **cuda** + - [capture_and_convert_to_dlpack_tensor_on_cuda] - Demonstrate + zero-copy GPU interop by handing a Zivid DeviceArray to PyTorch + on the GPU, showing two paths. - [capture_and_process_image_with_cupy_on_cuda] - Demonstrate GPU interop with CuPy: wrap a Zivid GPU image buffer as a CuPy array without copying it through CPU memory. @@ -391,6 +394,7 @@ Zivid Samples are distributed under the [BSD license]. [mask_point_cloud]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/mask_point_cloud.py [capture_vis_3d_in_loop]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/visualization/capture_vis_3d_in_loop.py [capture_vis_3d_in_loop_with_keypress_exit]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/visualization/capture_vis_3d_in_loop_with_keypress_exit.py + [capture_and_convert_to_dlpack_tensor_on_cuda]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/cuda/capture_and_convert_to_dlpack_tensor_on_cuda.py [capture_and_process_image_with_cupy_on_cuda]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/cuda/capture_and_process_image_with_cupy_on_cuda.py [capture_and_render_point_cloud_with_opengl_on_cuda]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/cuda/capture_and_render_point_cloud_with_opengl_on_cuda.py [capture_and_segment_image_with_pytorch_on_cuda]: https://github.com/zivid/zivid-python-samples/tree/master/source/applications/advanced/cuda/capture_and_segment_image_with_pytorch_on_cuda.py 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) diff --git a/source/applications/advanced/cuda/capture_and_convert_to_dlpack_tensor_on_cuda.py b/source/applications/advanced/cuda/capture_and_convert_to_dlpack_tensor_on_cuda.py new file mode 100644 index 00000000..1292daa8 --- /dev/null +++ b/source/applications/advanced/cuda/capture_and_convert_to_dlpack_tensor_on_cuda.py @@ -0,0 +1,255 @@ +""" +Demonstrate zero-copy GPU interop by handing a Zivid DeviceArray to PyTorch on the GPU, showing two paths. + +DISCLAIMER: Zivid does not provide or support the third-party library used in this sample. PyTorch is an external +dependency that the user is responsible for installing, configuring, and maintaining. This sample exists solely to +demonstrate how to expose Zivid GPU data to a GPU framework without a CPU round-trip. + +There are two ways to get a Zivid DeviceArray into PyTorch on the GPU: + +1. Primary, recommended path -- the CUDA Array Interface. zivid-python already implements ``__cuda_array_interface__`` + on DeviceArray (CUDA backend only), so ``torch.as_tensor(device_array, device="cuda")`` imports the Zivid GPU + buffer into PyTorch zero-copy, with no manual pointer, shape, or stride handling and no ctypes or PyCapsule code. + This is the simple path. + +2. Optional path -- build a DLPack capsule yourself. Only needed when you must produce a DLPack capsule directly from + the DeviceArray, for a consumer that expects the ``__dlpack__`` protocol directly (for example from C++). This + path constructs a ``DLManagedTensor`` from the raw building blocks the DeviceArray exposes (device + pointer, shape, strides in elements, element data type, CUDA device) using ctypes, wraps it in the standard + ``__dlpack__`` / ``__dlpack_device__`` protocol, and hands it to ``torch.from_dlpack``. The DeviceArray is kept + alive by the tensor's manager context and released only when the consumer is done with it. + +Only the hand-off is zero-copy in both paths: PyTorch views Zivid's GPU memory directly (verified below by comparing +data pointers). Any subsequent PyTorch operation allocates new GPU buffers as usual. See the related samples +``capture_and_process_image_with_cupy_on_cuda.py``, ``capture_and_segment_image_with_pytorch_on_cuda.py``, and +``capture_and_render_point_cloud_with_opengl_on_cuda.py`` for other ways to consume Zivid GPU data without a CPU +round-trip. + +Requirements: +- CUDA-capable GPU +- PyTorch with CUDA support: pip install torch + +""" + +from __future__ import annotations + +import ctypes +from typing import Optional, Tuple + +import zivid + +try: + import torch +except ImportError: + print("⚠️ Failed to import PyTorch. It is installed via `pip install torch`.") + raise + + +_KDLCUDA = 2 +_KDLINT = 0 +_KDLUINT = 1 +_KDLFLOAT = 2 + +_DLTENSOR_CAPSULE_NAME = b"dltensor" + + +class _DLDevice(ctypes.Structure): # pylint: disable=too-few-public-methods + _fields_ = [("device_type", ctypes.c_int), ("device_id", ctypes.c_int32)] + + +class _DLDataType(ctypes.Structure): # pylint: disable=too-few-public-methods + _fields_ = [("code", ctypes.c_uint8), ("bits", ctypes.c_uint8), ("lanes", ctypes.c_uint16)] + + +class _DLTensor(ctypes.Structure): # pylint: disable=too-few-public-methods + _fields_ = [ + ("data", ctypes.c_void_p), + ("device", _DLDevice), + ("ndim", ctypes.c_int32), + ("dtype", _DLDataType), + ("shape", ctypes.POINTER(ctypes.c_int64)), + ("strides", ctypes.POINTER(ctypes.c_int64)), + ("byte_offset", ctypes.c_uint64), + ] + + +class _DLManagedTensor(ctypes.Structure): # pylint: disable=too-few-public-methods + pass + + +_DLManagedTensorDeleter = ctypes.CFUNCTYPE(None, ctypes.POINTER(_DLManagedTensor)) + +_DLManagedTensor._fields_ = [ # pylint: disable=protected-access + ("dl_tensor", _DLTensor), + ("manager_ctx", ctypes.c_void_p), + ("deleter", _DLManagedTensorDeleter), +] + + +_pythonapi = ctypes.pythonapi +_PyCapsule_New = _pythonapi.PyCapsule_New +_PyCapsule_New.restype = ctypes.py_object +_PyCapsule_New.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + + +_MANAGER_CONTEXTS: dict = {} + + +@_DLManagedTensorDeleter +def _release_managed_tensor(managed_pointer: "ctypes._Pointer[_DLManagedTensor]") -> None: + address = ctypes.addressof(managed_pointer.contents) + _MANAGER_CONTEXTS.pop(address, None) + + +def _dl_data_type(device_array: zivid.DeviceArray) -> _DLDataType: + typestr = device_array.__cuda_array_interface__["typestr"] + kind = typestr[1] + bits = int(typestr[2:]) * 8 + code = {"f": _KDLFLOAT, "u": _KDLUINT, "i": _KDLINT}[kind] + return _DLDataType(code=code, bits=bits, lanes=1) # pylint: disable=no-value-for-parameter + + +def _build_dl_managed_tensor(device_array: zivid.DeviceArray, device_id: int) -> ctypes.c_void_p: + """Build a DLManagedTensor that references the DeviceArray's GPU buffer. + + Args: + device_array: The Zivid DeviceArray whose device memory is exposed. + device_id: The CUDA device ordinal the buffer lives on. + + Returns: + The address of the DLManagedTensor as a ctypes void pointer. + + """ + shape = tuple(int(dimension) for dimension in device_array.shape) + strides = tuple(int(stride) for stride in device_array.strides) + + shape_array = (ctypes.c_int64 * len(shape))(*shape) + strides_array = (ctypes.c_int64 * len(strides))(*strides) + + managed = _DLManagedTensor() + managed.dl_tensor.data = ctypes.c_void_p(device_array.device_pointer()) + managed.dl_tensor.device = _DLDevice(_KDLCUDA, device_id) # pylint: disable=no-value-for-parameter + managed.dl_tensor.ndim = len(shape) + managed.dl_tensor.dtype = _dl_data_type(device_array) + managed.dl_tensor.shape = shape_array + managed.dl_tensor.strides = strides_array + managed.dl_tensor.byte_offset = 0 + managed.deleter = _release_managed_tensor # pylint: disable=attribute-defined-outside-init + + _MANAGER_CONTEXTS[ctypes.addressof(managed)] = (managed, shape_array, strides_array, device_array) + return ctypes.cast(ctypes.byref(managed), ctypes.c_void_p) + + +class ZividDeviceArrayDLPack: # pylint: disable=too-few-public-methods + """Expose a Zivid DeviceArray to DLPack consumers via the ``__dlpack__`` / ``__dlpack_device__`` protocol. + + The DeviceArray is kept alive until the consumer releases the imported tensor. + """ + + def __init__(self, device_array: zivid.DeviceArray, device_id: int) -> None: + """Wrap a DeviceArray for DLPack consumption. + + Args: + device_array: The Zivid DeviceArray to expose. Must be on the CUDA backend. + device_id: The CUDA device ordinal the buffer lives on. + + Raises: + RuntimeError: If the DeviceArray is not on the CUDA backend. + + """ + if device_array.backend != zivid.ComputeBackend.cuda: + raise RuntimeError("DLPack construction in this sample is scoped to the CUDA backend") + self._device_array = device_array + self._device_id = device_id + + def __dlpack_device__(self) -> Tuple[int, int]: + """Return the DLPack device tuple for the buffer. + + Returns: + A ``(device_type, device_id)`` tuple with ``device_type`` set to kDLCUDA. + + """ + return (_KDLCUDA, self._device_id) + + def __dlpack__(self, stream: Optional[int] = None, **kwargs: object) -> object: # pylint: disable=unused-argument + """Return a PyCapsule named "dltensor" wrapping a DLManagedTensor for the DeviceArray's GPU buffer. + + Args: + stream: The consumer's CUDA stream, accepted for protocol compatibility. The DeviceArray was already + synchronized against the stream it was acquired on, so no extra synchronization is done here. + kwargs: Additional protocol keyword arguments (e.g. ``max_version``), accepted and ignored. + + Returns: + A PyCapsule named ``"dltensor"`` wrapping a DLManagedTensor. + + """ + managed = _build_dl_managed_tensor(self._device_array, self._device_id) + return _PyCapsule_New(managed, _DLTENSOR_CAPSULE_NAME, None) + + +def _main() -> None: + if not torch.cuda.is_available(): + raise RuntimeError("This sample requires PyTorch with CUDA support") + + torch_device = torch.device("cuda") + print(f"PyTorch CUDA device: {torch.cuda.get_device_name()}") + torch.zeros(1, device=torch_device) + + print("Initializing Zivid application (will reuse PyTorch's CUDA context)") + app = zivid.Application() + + print("Verifying that CUDA backend is available") + compute_device = app.compute_device() + if compute_device.backend != zivid.ComputeBackend.cuda: + raise RuntimeError("This sample requires CUDA backend") + print(f"Using GPU: {compute_device.model}") + + print("Connecting to camera") + camera = app.connect_camera() + + print("Configuring settings") + settings_2d = zivid.Settings2D(acquisitions=[zivid.Settings2D.Acquisition()]) + + print("Capturing 2D frame") + frame_2d = camera.capture_2d(settings_2d) + + print("Creating a PyTorch CUDA stream for GPU operations") + pytorch_stream = torch.cuda.Stream() + print(f"Using CUDA stream: {hex(pytorch_stream.cuda_stream)}") + + print("Getting GPU device buffer in float format (RGBAf32), synchronized into the PyTorch stream") + zivid_stream = zivid.CUDAStreamPtr(pytorch_stream.cuda_stream) + device_array = frame_2d.image_device_array(zivid_stream, zivid.PixelFormat.RGBAF) + print(f"DeviceArray: shape={tuple(device_array.shape)}, strides(elements)={tuple(device_array.strides)}") + print(f"DeviceArray device pointer: {hex(device_array.device_pointer())}") + + with torch.cuda.stream(pytorch_stream): + print("Primary path: CUDA Array Interface (no ctypes, no PyCapsule)") + torch_image = torch.as_tensor(device_array, device="cuda") # zero-copy via __cuda_array_interface__ + print( + f"PyTorch tensor: shape={tuple(torch_image.shape)}, dtype={torch_image.dtype}, device={torch_image.device}" + ) + if torch_image.data_ptr() != device_array.device_pointer(): + raise RuntimeError("Zero-copy check failed: PyTorch tensor does not share the Zivid device pointer") + print("Zero-copy verified: PyTorch tensor shares the Zivid device pointer") + + print("Optional path: build a DLPack capsule directly from the DeviceArray and import it") + dlpack_source = ZividDeviceArrayDLPack(device_array, torch.cuda.current_device()) + torch_image_via_capsule = torch.from_dlpack(dlpack_source) + if torch_image_via_capsule.data_ptr() != device_array.device_pointer(): + raise RuntimeError("Zero-copy check failed: capsule tensor does not share the Zivid device pointer") + print("Zero-copy verified: DLPack-capsule tensor shares the Zivid device pointer") + + print("Example: computing mean color on the GPU from the imported tensor") + mean_rgba = torch_image.float().mean(dim=(0, 1)) + + pytorch_stream.synchronize() + mean_values = mean_rgba.cpu().tolist() + print( + "Mean RGBA values: " + f"R={mean_values[0]:.3f}, G={mean_values[1]:.3f}, B={mean_values[2]:.3f}, A={mean_values[3]:.3f}" + ) + + +if __name__ == "__main__": + _main()