423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-08-02 11:52:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Volumes data structure.

Browse Source 
Summary: Implemented a data structure for volumes.

Reviewed By: gkioxari

Differential Revision: D20342920

fbshipit-source-id: ccc23eaa183ed8a4e9cd7674b4dcf31e8a65c3c6
This commit is contained in:
David Novotny 2021-01-05 03:37:38 -08:00 committed by Facebook GitHub Bot
parent 1e4a2e8624
commit 03ee1dbf82
3 changed files with 1553 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
pytorch3d/structures/__init__.py
View File

@ -3,6 +3,7 @@
from .meshes import Meshes, join_meshes_as_batch, join_meshes_as_scene from .meshes import Meshes, join_meshes_as_batch, join_meshes_as_scene
from .pointclouds import Pointclouds from .pointclouds import Pointclouds
from .utils import list_to_packed, list_to_padded, packed_to_list, padded_to_list from .utils import list_to_packed, list_to_padded, packed_to_list, padded_to_list
from .volumes import Volumes
__all__ = [k for k in globals().keys() if not k.startswith("_")] __all__ = [k for k in globals().keys() if not k.startswith("_")]

703
pytorch3d/structures/volumes.py Normal file
View File

@ -0,0 +1,703 @@
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import copy
from typing import List, Optional, Tuple, Union
import torch
from ..transforms import Scale, Transform3d
from . import utils as struct_utils
class Volumes(object):
"""
This class provides functions for working with batches of volumetric grids
of possibly varying spatial sizes.
VOLUME DENSITIES
The Volumes class can be either constructed from a 5D tensor of
`densities` of size `batch x density_dim x depth x height x width` or
from a list of differently-sized 4D tensors `[D_1, ..., D_batch]`,
where each `D_i` is of size `[density_dim x depth_i x height_i x width_i]`.
In case the `Volumes` object is initialized from the list of `densities`,
the list of tensors is internally converted to a single 5D tensor by
zero-padding the relevant dimensions. Both list and padded representations can be
accessed with the `Volumes.densities()` or `Volumes.densities_list()` getters.
The sizes of the individual volumes in the structure can be retrieved
with the `Volumes.get_grid_sizes()` getter.
The `Volumes` class is immutable. I.e. after generating a `Volumes` object,
one cannot change its properties, such as `self._densities` or `self._features`
anymore.
VOLUME FEATURES
While the `densities` field is intended to represent various measures of the
"density" of the volume cells (opacity, signed/unsigned distances
from the nearest surface, ...), one can additionally initialize the
object with the `features` argument. `features` are either a 5D tensor
of shape `batch x feature_dim x depth x height x width` or a list of
of differently-sized 4D tensors `[F_1, ..., F_batch]`,
where each `F_i` is of size `[feature_dim x depth_i x height_i x width_i]`.
`features` are intended to describe other properties of volume cells,
such as per-voxel 3D vectors of RGB colors that can be later used
for rendering the volume.
VOLUME COORDINATES
Additionally, the `Volumes` class keeps track of the locations of the
centers of the volume cells in the local volume coordinates as well as in
the world coordinates.
Local coordinates:
- Represent the locations of the volume cells in the local coordinate
frame of the volume.
- The center of the voxel indexed with `[·, ·, 0, 0, 0]` in the volume
has its 3D local coordinate set to `[-1, -1, -1]`, while the voxel
at index `[·, ·, depth_i-1, height_i-1, width_i-1]` has its
3D local coordinate set to `[1, 1, 1]`.
- The first/second/third coordinate of each of the 3D per-voxel
XYZ vector denotes the horizontal/vertical/depth-wise position
respectively. I.e the order of the coordinate dimensions in the
volume is reversed w.r.t. the order of the 3D coordinate vectors.
- The intermediate coordinates between `[-1, -1, -1]` and `[1, 1, 1]`.
are linearly interpolated over the spatial dimensions of the volume.
- Note that the convention is the same as for the 5D version of the
`torch.nn.functional.grid_sample` function called with
`align_corners==True`.
- Note that the local coordinate convention of `Volumes`
(+X = left to right, +Y = top to bottom, +Z = away from the user)
is *different* from the world coordinate convention of the
renderer for `Meshes` or `Pointclouds`
(+X = right to left, +Y = bottom to top, +Z = away from the user).
World coordinates:
- These define the locations of the centers of the volume cells
in the world coordinates.
- They are specifiied with the following mapping that converts
points `x_local` in the local coordinates to points `x_world`
in the world coordinates:
```
x_world = (
x_local * (volume_size - 1) * 0.5 * voxel_size
) - volume_translation,
```
here `voxel_size` specifies the size of each voxel of the volume,
and `volume_translation` is the 3D offset of the central voxel of
the volume w.r.t. the origin of the world coordinate frame.
Both `voxel_size` and `volume_translation` are specified in
the world coordinate units. `volume_size` is the spatial size of
the volume in form of a 3D vector `[width, height, depth]`.
- Given the above definition of `x_world`, one can derive the
inverse mapping from `x_world` to `x_local` as follows:
```
x_local = (
(x_world + volume_translation) / (0.5 * voxel_size)
) / (volume_size - 1)
```
- For a trivial volume with `volume_translation==[0, 0, 0]`
with `voxel_size=-1`, `x_world` would range
from -(volume_size-1)/2` to `+(volume_size-1)/2`.
Coordinate tensors that denote the locations of each of the volume cells in
local / world coordinates (with shape `(depth x height x width x 3)`)
can be retrieved by calling the `Volumes.get_coord_grid()` getter with the
appropriate `world_coordinates` argument.
Internally, the mapping between `x_local` and `x_world` is represented
as a `Transform3D` object `Volumes._local_to_world_transform`.
Users can access the relevant transformations with the
`Volumes.get_world_to_local_coords_transform()` and
`Volumes.get_local_to_world_coords_transform()`
functions.
Example coordinate conversion:
- For a "trivial" volume with `voxel_size = 1.`,
`volume_translation=[0., 0., 0.]`, and the spatial size of
`DxHxW = 5x5x5`, the point `x_world = (-2, 0, 2)` gets mapped
to `x_local=(-1, 0, 1)`.
- For a "trivial" volume `v` with `voxel_size = 1.`,
`volume_translation=[0., 0., 0.]`, the following holds:
```
torch.nn.functional.grid_sample(
v.densities(),
v.get_coord_grid(world_coordinates=False),
align_corners=True,
) == v.densities(),
```
i.e. sampling the volume at trivial local coordinates
(no scaling with `voxel_size`` or shift with `volume_translation`)
results in the same volume.
"""
def __init__(
self,
densities: Union[List[torch.Tensor], torch.Tensor],
features: Optional[Union[List[torch.Tensor], torch.Tensor]] = None,
voxel_size: Union[float, torch.Tensor, Tuple, List] = 1.0,
volume_translation: Union[torch.Tensor, Tuple, List] = (0.0, 0.0, 0.0),
):
"""
Args:
**densities**: Batch of input feature volume occupancies of shape
`(minibatch, density_dim, depth, height, width)`, or a list
of 4D tensors `[D_1, ..., D_minibatch]` where each `D_i` has
shape `(density_dim, depth_i, height_i, width_i)`.
Typically, each voxel contains a non-negative number
corresponding to its opaqueness.
**features**: Batch of input feature volumes of shape:
`(minibatch, feature_dim, depth, height, width)` or a list
of 4D tensors `[F_1, ..., F_minibatch]` where each `F_i` has
shape `(feature_dim, depth_i, height_i, width_i)`.
The field is optional and can be set to `None` in case features are
not required.
**voxel_size**: Denotes the size of each volume voxel in world units.
Has to be one of:
a) A scalar (square voxels)
b) 3-tuple or a 3-list of scalars
c) a Tensor of shape (3,)
d) a Tensor of shape (minibatch, 3)
e) a Tensor of shape (minibatch, 1)
f) a Tensor of shape (1,) (square voxels)
**volume_translation**: Denotes the 3D translation of the center
of the volume in world units. Has to be one of:
a) 3-tuple or a 3-list of scalars
b) a Tensor of shape (3,)
c) a Tensor of shape (minibatch, 3)
d) a Tensor of shape (1,) (square voxels)
"""
# handle densities
densities, grid_sizes = self._convert_densities_features_to_tensor(
densities, "densities"
)
# take device from densities
self.device = densities.device
# assign to the internal buffers
self._densities = densities
self._grid_sizes = grid_sizes
# handle features
self._features = None
if features is not None:
self._set_features(features)
# set the local_to_world transform
self._set_local_to_world_transform(
voxel_size=voxel_size, volume_translation=volume_translation
)
def _convert_densities_features_to_tensor(
self, x: Union[List[torch.Tensor], torch.Tensor], var_name: str
):
"""
Handle the `densities` or `features` arguments to the constructor.
"""
if isinstance(x, (list, tuple)):
x_tensor = struct_utils.list_to_padded(x)
if any(x_.ndim != 4 for x_ in x):
raise ValueError(
f"`{var_name}` has to be a list of 4-dim tensors of shape: "
f"({var_name}_dim, height, width, depth)"
)
if any(x_.shape[0] != x[0].shape[0] for x_ in x):
raise ValueError(
f"Each entry in the list of `{var_name}` has to have the "
"same number of channels (first dimension in the tensor)."
)
x_shapes = torch.stack(
[
torch.tensor(
list(x_.shape[1:]), dtype=torch.long, device=x_tensor.device
)
for x_ in x
],
dim=0,
)
elif torch.is_tensor(x):
if x.ndim != 5:
raise ValueError(
f"`{var_name}` has to be a 5-dim tensor of shape: "
f"(minibatch, {var_name}_dim, height, width, depth)"
)
x_tensor = x
x_shapes = torch.tensor(
list(x.shape[2:]), dtype=torch.long, device=x.device
)[None].repeat(x.shape[0], 1)
else:
raise ValueError(
f"{var_name} must be either a list or a tensor with "
f"shape (batch_size, {var_name}_dim, H, W, D)."
)
return x_tensor, x_shapes
def _vsize_translation_to_transform(
self, voxel_size, volume_translation, batch_size
):
"""
Converts the `voxel_size` and `volume_translation` constructor arguments
to the internal `Transform3D` object `local_to_world_transform`.
"""
volume_size_zyx = self.get_grid_sizes().float()
volume_size_xyz = volume_size_zyx[:, [2, 1, 0]]
# x_local = (
# (x_world + volume_translation) / (0.5 * voxel_size)
# ) / (volume_size - 1)
# x_world = (
# x_local * (volume_size - 1) * 0.5 * voxel_size
# ) - volume_translation
local_to_world_transform = Scale(
(volume_size_xyz - 1) * voxel_size * 0.5, device=self.device
).translate(-volume_translation)
return local_to_world_transform
def _handle_voxel_size(self, voxel_size, batch_size):
"""
Handle the `voxel_size` argument to the `Volumes` constructor.
"""
err_msg = (
"voxel_size has to be either a 3-tuple of scalars, or a scalar, or"
" a torch.Tensor of shape (3,) or (1,) or (minibatch, 3) or (minibatch, 1)."
)
if isinstance(voxel_size, (float, int)):
# convert a scalar to a 3-element tensor
voxel_size = (voxel_size,) * 3
if torch.is_tensor(voxel_size):
if voxel_size.numel() == 1:
# convert a single-element tensor to a 3-element one
voxel_size = voxel_size.view(-1).repeat(3)
elif len(voxel_size.shape) == 2 and (
voxel_size.shape[0] == batch_size and voxel_size.shape[1] == 1
):
voxel_size = voxel_size.repeat(1, 3)
return self._convert_volume_property_to_tensor(voxel_size, batch_size, err_msg)
def _handle_volume_translation(self, translation, batch_size):
"""
Handle the `volume_translation` argument to the `Volumes` constructor.
"""
err_msg = (
"`volume_translation` has to be either a 3-tuple of scalars, or"
" a Tensor of shape (1,3) or (minibatch, 3) or (3,)`."
)
return self._convert_volume_property_to_tensor(translation, batch_size, err_msg)
def _convert_volume_property_to_tensor(
self, x, batch_size, err_msg
) -> torch.Tensor:
"""
Handle the `volume_translation` or `voxel_size` argument to
the Volumes constructor.
Return a tensor of shape (N, 3) where N is the batch_size or 1
if batch_size is None.
"""
if isinstance(x, (list, tuple)):
if len(x) != 3:
raise ValueError(err_msg)
x = torch.tensor(x, device=self.device, dtype=torch.float32)[None]
x = x.repeat((batch_size, 1))
elif torch.is_tensor(x):
ok = (
(x.shape[0] == 1 and x.shape[1] == 3)
or (x.shape[0] == 3 and len(x.shape) == 1)
or (x.shape[0] == batch_size and x.shape[1] == 3)
)
if not ok:
raise ValueError(err_msg)
if x.device != self.device:
x = x.to(self.device)
if x.shape[0] == 3 and len(x.shape) == 1:
x = x[None]
if x.shape[0] == 1:
x = x.repeat((batch_size, 1))
else:
raise ValueError(err_msg)
return x
def get_coord_grid(self, world_coordinates: bool = True) -> torch.Tensor:
"""
Return the 3D coordinate grid of the volumetric grid
in local (`world_coordinates=False`) or world coordinates
(`world_coordinates=True`).
The grid records location of each center of the corresponding volume voxel.
Local coordinates are scaled s.t. the values along one side of the
volume are in range [-1, 1].
Args:
**world_coordinates**: if `True`, the method
returns the grid in the world coordinates,
otherwise, in local coordinates.
Returns:
**coordinate_grid**: The grid of coordinates of shape
`(minibatch, depth, height, width, 3)`, where `minibatch`,
`height`, `width` and `depth` are the batch size, height, width
and depth of the volume `features` or `densities`.
"""
# TODO(dnovotny): Implement caching of the coordinate grid.
return self._calculate_coordinate_grid(world_coordinates=world_coordinates)
def _calculate_coordinate_grid(
self, world_coordinates: bool = True
) -> torch.Tensor:
"""
Calculate the 3D coordinate grid of the volumetric grid either in
in local (`world_coordinates=False`) or
world coordinates (`world_coordinates=True`) .
"""
densities = self.densities()
ba, _, de, he, wi = densities.shape
grid_sizes = self.get_grid_sizes()
# generate coordinate axes
vol_axes = [
torch.linspace(-1.0, 1.0, r, dtype=torch.float32, device=self.device)
for r in (de, he, wi)
]
# generate per-coord meshgrids
Z, Y, X = torch.meshgrid(vol_axes)
# stack the coord grids ... this order matches the coordinate convention
# of torch.nn.grid_sample
vol_coords_local = torch.stack((X, Y, Z), dim=3)[None].repeat(ba, 1, 1, 1, 1)
# get grid sizes relative to the maximal volume size
grid_sizes_relative = (
torch.tensor([[de, he, wi]], device=grid_sizes.device, dtype=torch.float32)
- 1
) / (grid_sizes - 1).float()
if (grid_sizes_relative != 1.0).any():
# if any of the relative sizes != 1.0, adjust the grid
grid_sizes_relative_reshape = grid_sizes_relative[:, [2, 1, 0]][
:, None, None, None
]
vol_coords_local *= grid_sizes_relative_reshape
vol_coords_local += grid_sizes_relative_reshape - 1
if world_coordinates:
vol_coords = self.local_to_world_coords(vol_coords_local)
else:
vol_coords = vol_coords_local
return vol_coords
def get_local_to_world_coords_transform(self) -> Transform3d:
"""
Return a Transform3d object that converts points in the
the local coordinate frame of the volume to world coordinates.
Local volume coordinates are scaled s.t. the coordinates along one
side of the volume are in range [-1, 1].
Returns:
**local_to_world_transform**: A Transform3d object converting
points from local coordinates to the world coordinates.
"""
return self._local_to_world_transform
def get_world_to_local_coords_transform(self) -> Transform3d:
"""
Return a Transform3d object that converts points in the
world coordinates to the local coordinate frame of the volume.
Local volume coordinates are scaled s.t. the coordinates along one
side of the volume are in range [-1, 1].
Returns:
**world_to_local_transform**: A Transform3d object converting
points from world coordinates to local coordinates.
"""
return self.get_local_to_world_coords_transform().inverse()
def world_to_local_coords(self, points_3d_world: torch.Tensor) -> torch.Tensor:
"""
Convert a batch of 3D point coordinates `points_3d_world` of shape
(minibatch, ..., dim) in the world coordinates to
the local coordinate frame of the volume. Local volume
coordinates are scaled s.t. the coordinates along one side of the volume
are in range [-1, 1].
Args:
**points_3d_world**: A tensor of shape `(minibatch, ..., 3)`
containing the 3D coordinates of a set of points that will
be converted from the local volume coordinates (ranging
within [-1, 1]) to the world coordinates
defined by the `self.center` and `self.voxel_size` parameters.
Returns:
**points_3d_local**: `points_3d_world` converted to the local
volume coordinates of shape `(minibatch, ..., 3)`.
"""
pts_shape = points_3d_world.shape
return (
self.get_world_to_local_coords_transform()
.transform_points(points_3d_world.view(pts_shape[0], -1, 3))
.view(pts_shape)
)
def local_to_world_coords(self, points_3d_local: torch.Tensor) -> torch.Tensor:
"""
Convert a batch of 3D point coordinates `points_3d_local` of shape
(minibatch, ..., dim) in the local coordinate frame of the volume
to the world coordinates.
Args:
**points_3d_local**: A tensor of shape `(minibatch, ..., 3)`
containing the 3D coordinates of a set of points that will
be converted from the local volume coordinates (ranging
within [-1, 1]) to the world coordinates
defined by the `self.center` and `self.voxel_size` parameters.
Returns:
**points_3d_world**: `points_3d_local` converted to the world
coordinates of the volume of shape `(minibatch, ..., 3)`.
"""
pts_shape = points_3d_local.shape
return (
self.get_local_to_world_coords_transform()
.transform_points(points_3d_local.view(pts_shape[0], -1, 3))
.view(pts_shape)
)
def __len__(self) -> int:
return self._densities.shape[0]
def __getitem__(
self, index: Union[int, List[int], Tuple[int], slice, torch.Tensor]
) -> "Volumes":
"""
Args:
index: Specifying the index of the volume to retrieve.
Can be an int, slice, list of ints or a boolean or a long tensor.
Returns:
Volumes object with selected volumes. The tensors are not cloned.
"""
if isinstance(index, int):
index = torch.LongTensor([index])
elif isinstance(index, (slice, list, tuple)):
pass
elif torch.is_tensor(index):
if index.dim() != 1 or index.dtype.is_floating_point:
raise IndexError(index)
else:
raise IndexError(index)
new = self.__class__(
features=self.features()[index] if self._features is not None else None,
densities=self.densities()[index],
)
# dont forget to update grid_sizes!
new._grid_sizes = self.get_grid_sizes()[index]
new._local_to_world_transform = self._local_to_world_transform[index]
return new
def features(self) -> Optional[torch.Tensor]:
"""
Returns the features of the volume.
Returns:
**features**: The tensor of volume features.
"""
return self._features
def densities(self) -> torch.Tensor:
"""
Returns the densities of the volume.
Returns:
**densities**: The tensor of volume densities.
"""
return self._densities
def densities_list(self) -> List[torch.Tensor]:
"""
Get the list representation of the densities.
Returns:
list of tensors of densities of shape (dim_i, D_i, H_i, W_i).
"""
return self._features_densities_list(self.densities())
def features_list(self) -> List[torch.Tensor]:
"""
Get the list representation of the features.
Returns:
list of tensors of features of shape (dim_i, D_i, H_i, W_i)
or `None` for feature-less volumes.
"""
if self._features is None:
# No features provided so return None
return None
return self._features_densities_list(self.features())
def _features_densities_list(self, x) -> List[torch.Tensor]:
"""
Retrieve the list representation of features/densities.
Args:
x: self.features() or self.densities()
Returns:
list of tensors of features/densities of shape (dim_i, D_i, H_i, W_i).
"""
x_dim = x.shape[1]
pad_sizes = torch.nn.functional.pad(
self.get_grid_sizes(), [1, 0], mode="constant", value=x_dim
)
x_list = struct_utils.padded_to_list(x, pad_sizes.tolist())
return x_list
def get_grid_sizes(self) -> torch.LongTensor:
"""
Returns the sizes of individual volumetric grids in the structure.
Returns:
**grid_sizes**: Tensor of spatial sizes of each of the volumes
of size (batchsize, 3), where i-th row holds (D_i, H_i, W_i).
"""
return self._grid_sizes
def update_padded(self, new_densities, new_features=None) -> "Volumes":
"""
Returns a Volumes structure with updated padded tensors and copies of
the auxiliary tensors `self._local_to_world_transform`,
`device` and `self._grid_sizes`. This function allows for an update of
densities (and features) without having to explicitly
convert it to the list representation for heterogeneous batches.
Args:
new_densities: FloatTensor of shape (N, dim_density, D, H, W)
new_features: (optional) FloatTensor of shape (N, dim_feature, D, H, W)
Returns:
Volumes with updated features and densities
"""
new = copy.copy(self)
new._set_densities(new_densities)
if new_features is None:
new._features = None
else:
new._set_features(new_features)
return new
def _set_features(self, features):
self._set_densities_features(features, "features")
def _set_densities(self, densities):
self._set_densities_features(densities, "densities")
def _set_densities_features(self, x, var_name):
x_tensor, grid_sizes = self._convert_densities_features_to_tensor(x, var_name)
if x_tensor.device != self.device:
raise ValueError(
f"`{var_name}` have to be on the same device as `self.densities`."
)
if len(x_tensor.shape) != 5:
raise ValueError(
f"{var_name} has to be a 5-dim tensor of shape: "
f"(minibatch, {var_name}_dim, height, width, depth)"
)
if not (
(self.get_grid_sizes().shape == grid_sizes.shape)
and torch.allclose(self.get_grid_sizes(), grid_sizes)
):
raise ValueError(
f"The size of every grid in `{var_name}` has to match the size of"
"the corresponding `densities` grid."
)
setattr(self, "_" + var_name, x_tensor)
def _set_local_to_world_transform(
self,
voxel_size: Union[float, torch.Tensor, Tuple, List] = 1.0,
volume_translation: Union[torch.Tensor, Tuple, List] = (0.0, 0.0, 0.0),
):
"""
Sets the internal representation of the transformation between the
world and local volume coordinates by specifying
`voxel_size` and `volume_translation`
Args:
**voxel_size**: Denotes the size of input voxels. Has to be one of:
a) A scalar (square voxels)
b) 3-tuple or a 3-list of scalars
c) a Tensor of shape (3,)
d) a Tensor of shape (minibatch, 3)
e) a Tensor of shape (1,) (square voxels)
**volume_translation**: Denotes the 3D translation of the center
of the volume in world units. Has to be one of:
a) 3-tuple or a 3-list of scalars
b) a Tensor of shape (3,)
c) a Tensor of shape (minibatch, 3)
d) a Tensor of shape (1,) (square voxels)
"""
# handle voxel size and center
# here we force the tensors to lie on self.device
voxel_size = self._handle_voxel_size(voxel_size, len(self))
volume_translation = self._handle_volume_translation(
volume_translation, len(self)
)
self._local_to_world_transform = self._vsize_translation_to_transform(
voxel_size, volume_translation, len(self)
)
def clone(self) -> "Volumes":
"""
Deep copy of Volumes object. All internal tensors are cloned
individually.
Returns:
new Volumes object.
"""
return copy.deepcopy(self)
def to(self, device, copy: bool = False) -> "Volumes":
"""
Match the functionality of torch.Tensor.to()
If copy = True or the self Tensor is on a different device, the
returned tensor is a copy of self with the desired torch.device.
If copy = False and the self Tensor already has the correct torch.device,
then self is returned.
Args:
**device**: Device id for the new tensor.
**copy**: Boolean indicator whether or not to clone self. Default False.
Returns:
Volumes object.
"""
if not copy and self.device == device:
return self
other = self.clone()
if self.device != device:
other.device = device
other._densities = self._densities.to(device)
if self._features is not None:
other._features = self.features().to(device)
other._local_to_world_transform = (
self.get_local_to_world_coords_transform().to(device)
)
other._grid_sizes = self._grid_sizes.to(device)
return other
def cpu(self) -> "Volumes":
return self.to(torch.device("cpu"))
def cuda(self) -> "Volumes":
return self.to(torch.device("cuda"))

849
tests/test_volumes.py Normal file
View File

@ -0,0 +1,849 @@
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
import copy
import itertools
import random
import unittest
import numpy as np
import torch
from common_testing import TestCaseMixin
from pytorch3d.structures.volumes import Volumes
from pytorch3d.transforms import Scale
class TestVolumes(TestCaseMixin, unittest.TestCase):
def setUp(self) -> None:
np.random.seed(42)
torch.manual_seed(42)
random.seed(42)
@staticmethod
def _random_volume_list(
num_volumes, min_size, max_size, num_channels, device, rand_sizes=None
):
"""
Init a list of `num_volumes` random tensors of size [num_channels, *rand_size].
If `rand_sizes` is None, rand_size is a 3D long vector sampled
from [min_size, max_size]. Otherwise, rand_size should be a list
[rand_size_1, rand_size_2, ..., rand_size_num_volumes] where each
`rand_size_i` denotes the size of the corresponding `i`-th tensor.
"""
if rand_sizes is None:
rand_sizes = [
[random.randint(min_size, vs) for vs in max_size]
for _ in range(num_volumes)
]
volume_list = [
torch.randn(
size=[num_channels, *rand_size], device=device, dtype=torch.float32
)
for rand_size in rand_sizes
]
return volume_list, rand_sizes
def _check_indexed_volumes(self, v, selected, indices):
for selectedIdx, index in indices:
self.assertClose(selected.densities()[selectedIdx], v.densities()[index])
self.assertClose(
v._local_to_world_transform.get_matrix()[index],
selected._local_to_world_transform.get_matrix()[selectedIdx],
)
if selected.features() is not None:
self.assertClose(selected.features()[selectedIdx], v.features()[index])
def test_get_item(
self,
num_volumes=5,
num_channels=4,
volume_size=(10, 13, 8),
dtype=torch.float32,
):
device = torch.device("cuda:0")
# make sure we have at least 3 volumes to prevent indexing crash
num_volumes = max(num_volumes, 3)
features = torch.randn(
size=[num_volumes, num_channels, *volume_size],
device=device,
dtype=torch.float32,
)
densities = torch.randn(
size=[num_volumes, 1, *volume_size], device=device, dtype=torch.float32
)
features_list, rand_sizes = TestVolumes._random_volume_list(
num_volumes, 3, volume_size, num_channels, device
)
densities_list, _ = TestVolumes._random_volume_list(
num_volumes, 3, volume_size, 1, device, rand_sizes=rand_sizes
)
volume_translation = -torch.randn(num_volumes, 3).type_as(features)
voxel_size = torch.rand(num_volumes, 1).type_as(features) + 0.5
for features_, densities_ in zip(
(None, features, features_list), (densities, densities, densities_list)
):
# init the volume structure
v = Volumes(
features=features_,
densities=densities_,
volume_translation=volume_translation,
voxel_size=voxel_size,
)
# int index
index = 1
v_selected = v[index]
self.assertEqual(len(v_selected), 1)
self._check_indexed_volumes(v, v_selected, [(0, 1)])
# list index
index = [1, 2]
v_selected = v[index]
self.assertEqual(len(v_selected), len(index))
self._check_indexed_volumes(v, v_selected, enumerate(index))
# slice index
index = slice(0, 2, 1)
v_selected = v[0:2]
self.assertEqual(len(v_selected), 2)
self._check_indexed_volumes(v, v_selected, [(0, 0), (1, 1)])
# bool tensor
index = (torch.rand(num_volumes) > 0.5).to(device)
index[:2] = True # make sure smth is selected
v_selected = v[index]
self.assertEqual(len(v_selected), index.sum())
self._check_indexed_volumes(
v,
v_selected,
zip(
torch.arange(index.sum()),
torch.nonzero(index, as_tuple=False).squeeze(),
),
)
# int tensor
index = torch.tensor([1, 2], dtype=torch.int64, device=device)
v_selected = v[index]
self.assertEqual(len(v_selected), index.numel())
self._check_indexed_volumes(v, v_selected, enumerate(index.tolist()))
# invalid index
index = torch.tensor([1, 0, 1], dtype=torch.float32, device=device)
with self.assertRaises(IndexError):
v_selected = v[index]
index = 1.2 # floating point index
with self.assertRaises(IndexError):
v_selected = v[index]
def test_coord_transforms(self, num_volumes=3, num_channels=4, dtype=torch.float32):
"""
Test the correctness of the conversion between the internal
Transform3D Volumes._local_to_world_transform and the initialization
from the translation and voxel_size.
"""
device = torch.device("cuda:0")
# try for 10 sets of different random sizes/centers/voxel_sizes
for _ in range(10):
size = torch.randint(high=10, size=(3,), low=3).tolist()
densities = torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
)
# init the transformation params
volume_translation = torch.randn(num_volumes, 3)
voxel_size = torch.rand(num_volumes, 3) * 3.0 + 0.5
# get the corresponding Transform3d object
local_offset = torch.tensor(list(size), dtype=torch.float32, device=device)[
[2, 1, 0]
][None].repeat(num_volumes, 1)
local_to_world_transform = (
Scale(0.5 * local_offset - 0.5, device=device)
.scale(voxel_size)
.translate(-volume_translation)
)
# init the volume structures with the scale and translation,
# then get the coord grid in world coords
v_trans_vs = Volumes(
densities=densities,
voxel_size=voxel_size,
volume_translation=volume_translation,
)
grid_rot_trans_vs = v_trans_vs.get_coord_grid(world_coordinates=True)
# map the default local coords to the world coords
# with local_to_world_transform
v_default = Volumes(densities=densities)
grid_default_local = v_default.get_coord_grid(world_coordinates=False)
grid_default_world = local_to_world_transform.transform_points(
grid_default_local.view(num_volumes, -1, 3)
).view(num_volumes, *size, 3)
# check that both grids are the same
self.assertClose(grid_rot_trans_vs, grid_default_world, atol=1e-5)
# check that the transformations are the same
self.assertClose(
v_trans_vs.get_local_to_world_coords_transform().get_matrix(),
local_to_world_transform.get_matrix(),
atol=1e-5,
)
def test_coord_grid_convention(
self, num_volumes=3, num_channels=4, dtype=torch.float32
):
"""
Check that for a trivial volume with spatial size DxHxW=5x7x5:
1) xyz_world=(0, 0, 0) lands right in the middle of the volume
with xyz_local=(0, 0, 0).
2) xyz_world=(-2, 3, 2) results in xyz_local=(-1, 1, -1).
3) The centeral voxel of the volume coordinate grid
has coords x_world=(0, 0, 0) and x_local=(0, 0, 0)
4) grid_sampler(world_coordinate_grid, local_coordinate_grid)
is the same as world_coordinate_grid itself. I.e. the local coordinate
grid matches the grid_sampler coordinate convention.
"""
device = torch.device("cuda:0")
densities = torch.randn(
size=[num_volumes, num_channels, 5, 7, 5],
device=device,
dtype=torch.float32,
)
v_trivial = Volumes(densities=densities)
# check the case with x_world=(0,0,0)
pts_world = torch.zeros(num_volumes, 1, 3, device=device, dtype=torch.float32)
pts_local = v_trivial.world_to_local_coords(pts_world)
pts_local_expected = torch.zeros_like(pts_local)
self.assertClose(pts_local, pts_local_expected)
# check the case with x_world=(-2, 3, -2)
pts_world = torch.tensor([-2, 3, -2], device=device, dtype=torch.float32)[
None, None
].repeat(num_volumes, 1, 1)
pts_local = v_trivial.world_to_local_coords(pts_world)
pts_local_expected = torch.tensor(
[-1, 1, -1], device=device, dtype=torch.float32
)[None, None].repeat(num_volumes, 1, 1)
self.assertClose(pts_local, pts_local_expected)
# check that the central voxel has coords x_world=(0, 0, 0) and x_local(0, 0, 0)
grid_world = v_trivial.get_coord_grid(world_coordinates=True)
grid_local = v_trivial.get_coord_grid(world_coordinates=False)
for grid in (grid_world, grid_local):
x0 = grid[0, :, :, 2, 0]
y0 = grid[0, :, 3, :, 1]
z0 = grid[0, 2, :, :, 2]
for coord_line in (x0, y0, z0):
self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
# resample grid_world using grid_sampler with local coords
# -> make sure the resampled version is the same as original
grid_world_resampled = torch.nn.functional.grid_sample(
grid_world.permute(0, 4, 1, 2, 3), grid_local, align_corners=True
).permute(0, 2, 3, 4, 1)
self.assertClose(grid_world_resampled, grid_world, atol=1e-7)
def test_coord_grid_convention_heterogeneous(
self, num_channels=4, dtype=torch.float32
):
"""
Check that for a list of 2 trivial volumes with
spatial sizes DxHxW=(5x7x5, 3x5x5):
1) xyz_world=(0, 0, 0) lands right in the middle of the volume
with xyz_local=(0, 0, 0).
2) xyz_world=((-2, 3, -2), (-2, -2, 1)) results
in xyz_local=((-1, 1, -1), (-1, -1, 1)).
3) The centeral voxel of the volume coordinate grid
has coords x_world=(0, 0, 0) and x_local=(0, 0, 0)
4) grid_sampler(world_coordinate_grid, local_coordinate_grid)
is the same as world_coordinate_grid itself. I.e. the local coordinate
grid matches the grid_sampler coordinate convention.
"""
device = torch.device("cuda:0")
sizes = [(5, 7, 5), (3, 5, 5)]
densities_list = [
torch.randn(size=[num_channels, *size], device=device, dtype=torch.float32)
for size in sizes
]
# init the volume
v_trivial = Volumes(densities=densities_list)
# check the border point locations
pts_world = torch.tensor(
[[-2.0, 3.0, -2.0], [-2.0, -2.0, 1.0]], device=device, dtype=torch.float32
)[:, None]
pts_local = v_trivial.world_to_local_coords(pts_world)
pts_local_expected = torch.tensor(
[[-1.0, 1.0, -1.0], [-1.0, -1.0, 1.0]], device=device, dtype=torch.float32
)[:, None]
self.assertClose(pts_local, pts_local_expected)
# check that the central voxel has coords x_world=(0, 0, 0) and x_local(0, 0, 0)
grid_world = v_trivial.get_coord_grid(world_coordinates=True)
grid_local = v_trivial.get_coord_grid(world_coordinates=False)
for grid in (grid_world, grid_local):
x0 = grid[0, :, :, 2, 0]
y0 = grid[0, :, 3, :, 1]
z0 = grid[0, 2, :, :, 2]
for coord_line in (x0, y0, z0):
self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
x0 = grid[1, :, :, 2, 0]
y0 = grid[1, :, 2, :, 1]
z0 = grid[1, 1, :, :, 2]
for coord_line in (x0, y0, z0):
self.assertClose(coord_line, torch.zeros_like(coord_line), atol=1e-7)
# resample grid_world using grid_sampler with local coords
# -> make sure the resampled version is the same as original
for grid_world_, grid_local_, size in zip(grid_world, grid_local, sizes):
grid_world_crop = grid_world_[: size[0], : size[1], : size[2], :][None]
grid_local_crop = grid_local_[: size[0], : size[1], : size[2], :][None]
grid_world_crop_resampled = torch.nn.functional.grid_sample(
grid_world_crop.permute(0, 4, 1, 2, 3),
grid_local_crop,
align_corners=True,
).permute(0, 2, 3, 4, 1)
self.assertClose(grid_world_crop_resampled, grid_world_crop, atol=1e-7)
def test_coord_grid_transforms(
self, num_volumes=3, num_channels=4, dtype=torch.float32
):
"""
Test whether conversion between local-world coordinates of the
volume returns correct results.
"""
device = torch.device("cuda:0")
# try for 10 sets of different random sizes/centers/voxel_sizes
for _ in range(10):
size = torch.randint(high=10, size=(3,), low=3).tolist()
center = torch.randn(num_volumes, 3, dtype=torch.float32, device=device)
voxel_size = torch.rand(1, dtype=torch.float32, device=device) * 5.0 + 0.5
for densities in (
torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
),
TestVolumes._random_volume_list(
num_volumes, 3, size, num_channels, device, rand_sizes=None
)[0],
):
# init the volume structure
v = Volumes(
densities=densities,
voxel_size=voxel_size,
volume_translation=-center,
)
# get local coord grid
grid_local = v.get_coord_grid(world_coordinates=False)
# convert from world to local to world
grid_world = v.get_coord_grid(world_coordinates=True)
grid_local_2 = v.world_to_local_coords(grid_world)
grid_world_2 = v.local_to_world_coords(grid_local_2)
# assertions on shape and values of grid_world and grid_local
self.assertClose(grid_world, grid_world_2, atol=1e-5)
self.assertClose(grid_local, grid_local_2, atol=1e-5)
# check that the individual slices of the location grid have
# constant values along expected dimensions
for plane_dim in (1, 2, 3):
for grid_plane in grid_world.split(1, dim=plane_dim):
grid_coord_dim = {1: 2, 2: 1, 3: 0}[plane_dim]
grid_coord_plane = grid_plane.squeeze()[..., grid_coord_dim]
# check that all elements of grid_coord_plane are
# the same for each batch element
self.assertClose(
grid_coord_plane.reshape(num_volumes, -1).max(dim=1).values,
grid_coord_plane.reshape(num_volumes, -1).min(dim=1).values,
)
def test_clone(
self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
):
"""
Test cloning of a `Volumes` object
"""
device = torch.device("cuda:0")
features = torch.randn(
size=[num_volumes, num_channels, *size], device=device, dtype=torch.float32
)
densities = torch.rand(
size=[num_volumes, 1, *size], device=device, dtype=torch.float32
)
for has_features in (True, False):
v = Volumes(
densities=densities, features=features if has_features else None
)
vnew = v.clone()
vnew._densities.data[0, 0, 0, 0, 0] += 1.0
self.assertNotAlmostEqual(
float(
(vnew.densities()[0, 0, 0, 0, 0] - v.densities()[0, 0, 0, 0, 0])
.abs()
.max()
),
0.0,
)
if has_features:
vnew._features.data[0, 0, 0, 0, 0] += 1.0
self.assertNotAlmostEqual(
float(
(vnew.features()[0, 0, 0, 0, 0] - v.features()[0, 0, 0, 0, 0])
.abs()
.max()
),
0.0,
)
def _check_vars_on_device(self, v, desired_device):
for var_name, var in vars(v).items():
if var_name != "device":
if var is not None:
self.assertTrue(var.device.type == desired_device.type)
else:
self.assertTrue(var.type == desired_device.type)
def test_to(
self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
):
"""
Test the moving of the volumes from/to gpu and cpu
"""
device = torch.device("cuda:0")
device_cpu = torch.device("cpu")
features = torch.randn(
size=[num_volumes, num_channels, *size], device=device, dtype=torch.float32
)
densities = torch.rand(size=[num_volumes, 1, *size], device=device, dtype=dtype)
for features_ in (features, None):
v = Volumes(densities=densities, features=features_)
v_cpu = v.cpu()
v_cuda = v_cpu.cuda()
v_cuda_2 = v_cuda.cuda()
v_cpu_2 = v_cuda_2.cpu()
for v1, v2 in itertools.combinations(
(v, v_cpu, v_cpu_2, v_cuda, v_cuda_2), 2
):
if v1 is v_cuda and v2 is v_cuda_2:
# checks that we do not copy if the devices stay the same
assert_fun = self.assertIs
else:
assert_fun = self.assertSeparate
assert_fun(v1._densities, v2._densities)
if features_ is not None:
assert_fun(v1._features, v2._features)
for v_ in (v1, v2):
if v_ in (v_cpu, v_cpu_2):
self._check_vars_on_device(v_, device_cpu)
else:
self._check_vars_on_device(v_, device)
def _check_padded(self, x_pad, x_list, grid_sizes):
"""
Check that padded tensors x_pad are the same as x_list tensors.
"""
num_volumes = len(x_list)
for i in range(num_volumes):
self.assertClose(
x_pad[i][:, : grid_sizes[i][0], : grid_sizes[i][1], : grid_sizes[i][2]],
x_list[i],
)
def test_feature_density_setters(self):
"""
Tests getters and setters for padded/list representations.
"""
device = torch.device("cuda:0")
diff_device = torch.device("cpu")
num_volumes = 30
num_channels = 4
K = 20
densities = []
features = []
grid_sizes = []
diff_grid_sizes = []
for _ in range(num_volumes):
grid_size = torch.randint(K - 1, size=(3,)).long() + 1
densities.append(
torch.rand((1, *grid_size), device=device, dtype=torch.float32)
)
features.append(
torch.rand(
(num_channels, *grid_size), device=device, dtype=torch.float32
)
)
grid_sizes.append(grid_size)
diff_grid_size = (
copy.deepcopy(grid_size) + torch.randint(2, size=(3,)).long() + 1
)
diff_grid_sizes.append(diff_grid_size)
grid_sizes = torch.stack(grid_sizes).to(device)
diff_grid_sizes = torch.stack(diff_grid_sizes).to(device)
volumes = Volumes(densities=densities, features=features)
self.assertClose(volumes.get_grid_sizes(), grid_sizes)
# test the getters
features_padded = volumes.features()
densities_padded = volumes.densities()
features_list = volumes.features_list()
densities_list = volumes.densities_list()
for x_pad, x_list in zip(
(densities_padded, features_padded, densities_padded, features_padded),
(densities_list, features_list, densities, features),
):
self._check_padded(x_pad, x_list, grid_sizes)
# test feature setters
features_new = [
torch.rand((num_channels, *grid_size), device=device, dtype=torch.float32)
for grid_size in grid_sizes
]
volumes._set_features(features_new)
features_new_list = volumes.features_list()
features_new_padded = volumes.features()
for x_pad, x_list in zip(
(features_new_padded, features_new_padded),
(features_new, features_new_list),
):
self._check_padded(x_pad, x_list, grid_sizes)
# wrong features to update
bad_features_new = [
[
torch.rand(
(num_channels, *grid_size), device=diff_device, dtype=torch.float32
)
for grid_size in diff_grid_sizes
],
torch.rand(
(num_volumes, num_channels, K + 1, K, K),
device=device,
dtype=torch.float32,
),
None,
]
for bad_features_new_ in bad_features_new:
with self.assertRaises(ValueError):
volumes._set_densities(bad_features_new_)
# test density setters
densities_new = [
torch.rand((1, *grid_size), device=device, dtype=torch.float32)
for grid_size in grid_sizes
]
volumes._set_densities(densities_new)
densities_new_list = volumes.densities_list()
densities_new_padded = volumes.densities()
for x_pad, x_list in zip(
(densities_new_padded, densities_new_padded),
(densities_new, densities_new_list),
):
self._check_padded(x_pad, x_list, grid_sizes)
# wrong densities to update
bad_densities_new = [
[
torch.rand((1, *grid_size), device=diff_device, dtype=torch.float32)
for grid_size in diff_grid_sizes
],
torch.rand(
(num_volumes, 1, K + 1, K, K), device=device, dtype=torch.float32
),
None,
]
for bad_densities_new_ in bad_densities_new:
with self.assertRaises(ValueError):
volumes._set_densities(bad_densities_new_)
# test update_padded
volumes = Volumes(densities=densities, features=features)
volumes_updated = volumes.update_padded(
densities_new, new_features=features_new
)
densities_new_list = volumes_updated.densities_list()
densities_new_padded = volumes_updated.densities()
features_new_list = volumes_updated.features_list()
features_new_padded = volumes_updated.features()
for x_pad, x_list in zip(
(
densities_new_padded,
densities_new_padded,
features_new_padded,
features_new_padded,
),
(densities_new, densities_new_list, features_new, features_new_list),
):
self._check_padded(x_pad, x_list, grid_sizes)
self.assertIs(volumes.get_grid_sizes(), volumes_updated.get_grid_sizes())
self.assertIs(
volumes.get_local_to_world_coords_transform(),
volumes_updated.get_local_to_world_coords_transform(),
)
self.assertIs(volumes.device, volumes_updated.device)
def test_constructor_for_padded_lists(self):
"""
Tests constructor for padded/list representations.
"""
device = torch.device("cuda:0")
diff_device = torch.device("cpu")
num_volumes = 3
num_channels = 4
size = (6, 8, 10)
diff_size = (6, 8, 11)
# good ways to define densities
ok_densities = [
torch.randn(
size=[num_volumes, 1, *size], device=device, dtype=torch.float32
).unbind(0),
torch.randn(
size=[num_volumes, 1, *size], device=device, dtype=torch.float32
),
]
# bad ways to define features
bad_features = [
torch.randn(
size=[num_volumes + 1, num_channels, *size],
device=device,
dtype=torch.float32,
).unbind(
0
), # list with diff batch size
torch.randn(
size=[num_volumes + 1, num_channels, *size],
device=device,
dtype=torch.float32,
), # diff batch size
torch.randn(
size=[num_volumes, num_channels, *diff_size],
device=device,
dtype=torch.float32,
).unbind(
0
), # list with different size
torch.randn(
size=[num_volumes, num_channels, *diff_size],
device=device,
dtype=torch.float32,
), # different size
torch.randn(
size=[num_volumes, num_channels, *size],
device=diff_device,
dtype=torch.float32,
), # different device
torch.randn(
size=[num_volumes, num_channels, *size],
device=diff_device,
dtype=torch.float32,
).unbind(
0
), # list with different device
]
# good ways to define features
ok_features = [
torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
).unbind(
0
), # list of features of correct size
torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
),
]
for densities in ok_densities:
for features in bad_features:
self.assertRaises(
ValueError, Volumes, densities=densities, features=features
)
for features in ok_features:
Volumes(densities=densities, features=features)
def test_constructor(
self, num_volumes=3, num_channels=4, size=(6, 8, 10), dtype=torch.float32
):
"""
Test different ways of calling the `Volumes` constructor
"""
device = torch.device("cuda:0")
# all ways to define features
features = [
torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
), # padded tensor
torch.randn(
size=[num_volumes, num_channels, *size],
device=device,
dtype=torch.float32,
).unbind(
0
), # list of features
None, # no features
]
# bad ways to define features
bad_features = [
torch.randn(
size=[num_volumes, num_channels, 2, *size],
device=device,
dtype=torch.float32,
), # 6 dims
torch.randn(
size=[num_volumes, *size], device=device, dtype=torch.float32
), # 4 dims
torch.randn(
size=[num_volumes, *size], device=device, dtype=torch.float32
).unbind(
0
), # list of 4 dim tensors
]
# all ways to define densities
densities = [
torch.randn(
size=[num_volumes, 1, *size], device=device, dtype=torch.float32
), # padded tensor
torch.randn(
size=[num_volumes, 1, *size], device=device, dtype=torch.float32
).unbind(
0
), # list of densities
]
# bad ways to define densities
bad_densities = [
None, # omitted
torch.randn(
size=[num_volumes, 1, 1, *size], device=device, dtype=torch.float32
), # 6-dim tensor
torch.randn(
size=[num_volumes, 1, 1, *size], device=device, dtype=torch.float32
).unbind(
0
), # list of 5-dim densities
]
# all possible ways to define the voxels sizes
vox_sizes = [
torch.Tensor([1.0, 1.0, 1.0]),
[1.0, 1.0, 1.0],
torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes, 1),
torch.Tensor([1.0])[None].repeat(num_volumes, 1),
1.0,
torch.Tensor([1.0]),
]
# all possible ways to define the volume translations
vol_translations = [
torch.Tensor([1.0, 1.0, 1.0]),
[1.0, 1.0, 1.0],
torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes, 1),
]
# wrong ways to define voxel sizes
bad_vox_sizes = [
torch.Tensor([1.0, 1.0, 1.0, 1.0]),
[1.0, 1.0, 1.0, 1.0],
torch.Tensor([]),
None,
]
# wrong ways to define the volume translations
bad_vol_translations = [
torch.Tensor([1.0, 1.0]),
[1.0, 1.0],
1.0,
torch.Tensor([1.0, 1.0, 1.0])[None].repeat(num_volumes + 1, 1),
]
def zip_with_ok_indicator(good, bad):
return zip([*good, *bad], [*([True] * len(good)), *([False] * len(bad))])
for features_, features_ok in zip_with_ok_indicator(features, bad_features):
for densities_, densities_ok in zip_with_ok_indicator(
densities, bad_densities
):
for vox_size, size_ok in zip_with_ok_indicator(
vox_sizes, bad_vox_sizes
):
for vol_translation, trans_ok in zip_with_ok_indicator(
vol_translations, bad_vol_translations
):
if (
size_ok and trans_ok and features_ok and densities_ok
): # if all entries are good we check that this doesnt throw
Volumes(
features=features_,
densities=densities_,
voxel_size=vox_size,
volume_translation=vol_translation,
)
else: # otherwise we check for ValueError
self.assertRaises(
ValueError,
Volumes,
features=features_,
densities=densities_,
voxel_size=vox_size,
volume_translation=vol_translation,
)