[REQ] Let wp.Volume.load_from_numpy accept tuples for voxel_size

[eigenvivek]

Description

Currently, volumes created with wp.Volume.load_from_numpy can only have a single voxel_size. It would be great if this method (and the underlying wp.Volume.allocate) could accept tuples as well.

Happy to make a PR for this issue if it's desirable.

warp/warp/_src/types.py

Lines 5501 to 5634 in 46f8fbf

	@classmethod
	def load_from_numpy(
	cls, ndarray: np.ndarray, min_world=(0.0, 0.0, 0.0), voxel_size=1.0, bg_value=0.0, device=None
	) -> Volume:
	"""Create a :class:`Volume` object from a dense 3D NumPy array.
	This function is only supported for CUDA devices.
	Args:
	min_world: The 3D coordinate of the lower corner of the volume.
	voxel_size: The size of each voxel in spatial coordinates.
	bg_value: Background value
	device: The CUDA device to create the volume on, e.g.: "cuda" or "cuda:0".
	Returns:
	A ``warp.Volume`` object.
	"""
	target_shape = (
	math.ceil(ndarray.shape[0] / 8) * 8,
	math.ceil(ndarray.shape[1] / 8) * 8,
	math.ceil(ndarray.shape[2] / 8) * 8,
	)
	if hasattr(bg_value, "__len__"):
	# vec3, assuming the numpy array is 4D
	padded_array = np.full(
	shape=(target_shape[0], target_shape[1], target_shape[2], 3), fill_value=bg_value, dtype=np.single
	)
	padded_array[0 : ndarray.shape[0], 0 : ndarray.shape[1], 0 : ndarray.shape[2], :] = ndarray
	else:
	padded_amount = (
	math.ceil(ndarray.shape[0] / 8) * 8 - ndarray.shape[0],
	math.ceil(ndarray.shape[1] / 8) * 8 - ndarray.shape[1],
	math.ceil(ndarray.shape[2] / 8) * 8 - ndarray.shape[2],
	)
	# Manual padding to avoid np.pad compatibility issues with code coverage tools (e.g., coverage.py)
	target_shape = (
	ndarray.shape[0] + padded_amount[0],
	ndarray.shape[1] + padded_amount[1],
	ndarray.shape[2] + padded_amount[2],
	)
	padded_array = np.full(target_shape, bg_value, dtype=ndarray.dtype)
	padded_array[: ndarray.shape[0], : ndarray.shape[1], : ndarray.shape[2]] = ndarray
	shape = padded_array.shape
	volume = warp.Volume.allocate(
	min_world,
	[
	min_world[0] + (shape[0] - 1) * voxel_size,
	min_world[1] + (shape[1] - 1) * voxel_size,
	min_world[2] + (shape[2] - 1) * voxel_size,
	],
	voxel_size,
	bg_value=bg_value,
	points_in_world_space=True,
	translation=min_world,
	device=device,
	)
	# Populate volume
	if hasattr(bg_value, "__len__"):
	warp.launch(
	warp._src.utils.copy_dense_volume_to_nano_vdb_v,
	dim=(shape[0], shape[1], shape[2]),
	inputs=[volume.id, warp.array(padded_array, dtype=warp.vec3, device=device)],
	device=device,
	)
	elif isinstance(bg_value, int):
	warp.launch(
	warp._src.utils.copy_dense_volume_to_nano_vdb_i,
	dim=shape,
	inputs=[volume.id, warp.array(padded_array, dtype=warp.int32, device=device)],
	device=device,
	)
	else:
	warp.launch(
	warp._src.utils.copy_dense_volume_to_nano_vdb_f,
	dim=shape,
	inputs=[volume.id, warp.array(padded_array, dtype=warp.float32, device=device)],
	device=device,
	)
	return volume
	@classmethod
	def allocate(
	cls,
	min: list[int],
	max: list[int],
	voxel_size: float,
	bg_value=0.0,
	translation=(0.0, 0.0, 0.0),
	points_in_world_space=False,
	device: warp.DeviceLike = None,
	) -> Volume:
	"""Allocate a new Volume based on the bounding box defined by min and max.
	This function is only supported for CUDA devices.
	Allocate a volume that is large enough to contain voxels [min[0], min[1], min[2]] - [max[0], max[1], max[2]], inclusive.
	If points_in_world_space is true, then min and max are first converted to index space with the given voxel size and
	translation, and the volume is allocated with those.
	The smallest unit of allocation is a dense tile of 8x8x8 voxels, the requested bounding box is rounded up to tiles, and
	the resulting tiles will be available in the new volume.
	Args:
	min (array-like): Lower 3D coordinates of the bounding box in index space or world space, inclusive.
	max (array-like): Upper 3D coordinates of the bounding box in index space or world space, inclusive.
	voxel_size: Voxel size of the new volume.
	bg_value (float or array-like): Value of unallocated voxels of the volume, also defines the volume's type,
	a :class:`warp.vec3` volume is created if this is `array-like`, otherwise a float volume is created
	translation (array-like): Translation between the index and world spaces.
	device: The CUDA device to create the volume on, e.g.: ``"cuda"`` or ``"cuda:0"``.
	"""
	if points_in_world_space:
	min = np.around((np.array(min, dtype=np.float32) - translation) / voxel_size)
	max = np.around((np.array(max, dtype=np.float32) - translation) / voxel_size)
	tile_min = np.array(min, dtype=np.int32) // 8
	tile_max = np.array(max, dtype=np.int32) // 8
	tiles = np.array(
	[
	[i, j, k]
	for i in range(tile_min[0], tile_max[0] + 1)
	for j in range(tile_min[1], tile_max[1] + 1)
	for k in range(tile_min[2], tile_max[2] + 1)
	],
	dtype=np.int32,
	)
	tile_points = array(tiles * 8, device=device)
	return cls.allocate_by_tiles(tile_points, voxel_size, bg_value, translation, device)

Context

Many volumetric data are anisotropic (e.g., in medical images, such CT scans, the spacing between 2D slices is often greater than the spacings within slices).

[gdaviet]

Hi @eigenvivek, the underlying sparse allocation functions (allocate_by_tiles / allocate_by_voxels) do accept anisotropic voxel sizes. But I don't see any reason to restrict the higher-level from_numpy to isotropic oxel size, so we will happily accept a PR if you want to make one!

[shi-eric]

@Glitchfix added support for anisotropic voxel sizes, which was merged in c64a6c3 Since we just released Warp 1.12.0 last week, it will be available starting with Warp v1.13.0 and the next nightly build

[eigenvivek]

awesome, thank you very much for adding this feature! looking forward to the new release.