def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
"""
Apply the transform to `img`.
"""
if isinstance(img, torch.Tensor):
return img.permute(self.indices or tuple(range(img.ndim)[::-1]))
return img.transpose(self.indices) # type: ignore
def __call__(self, prob_map: NdarrayOrTensor):
"""
prob_map: the input probabilities map, it must have shape (H[, W, ...]).
"""
if self.sigma != 0:
if not isinstance(prob_map, torch.Tensor):
prob_map = torch.as_tensor(prob_map, dtype=torch.float)
self.filter.to(prob_map.device)
prob_map = self.filter(prob_map)
prob_map_shape = prob_map.shape
outputs = []
while prob_map.max() > self.prob_threshold:
max_idx = unravel_index(prob_map.argmax(), prob_map_shape)
prob_max = prob_map[tuple(max_idx)]
max_idx = max_idx.cpu().numpy() if isinstance(
max_idx, torch.Tensor) else max_idx
prob_max = prob_max.item() if isinstance(
prob_max, torch.Tensor) else prob_max
outputs.append([prob_max] + list(max_idx))
idx_min_range = (max_idx - self.box_lower_bd).clip(0, None)
idx_max_range = (max_idx + self.box_upper_bd).clip(
None, prob_map_shape)
# for each dimension, set values during index ranges to 0
slices = tuple(
slice(idx_min_range[i], idx_max_range[i])
for i in range(self.spatial_dims))
prob_map[slices] = 0
return outputs
def __call__(self,
img: NdarrayOrTensor,
randomize: bool = True) -> NdarrayOrTensor:
"""
Apply the transform to `img`, if `randomize` randomizing the smooth field otherwise reusing the previous.
"""
img = convert_to_tensor(img, track_meta=get_track_meta())
if randomize:
self.randomize()
if not self._do_transform:
return img
img_min = img.min()
img_max = img.max()
img_rng = img_max - img_min
field = self.sfield()
rfield, *_ = convert_to_dst_type(field, img)
# everything below here is to be computed using the destination type (numpy, tensor, etc.)
img = (img - img_min) / (img_rng + 1e-10) # rescale to unit values
img = img**rfield # contrast is changed by raising image data to a power, in this case the field
out = (img * img_rng
) + img_min # rescale back to the original image value range
return out
def searchsorted(a: NdarrayOrTensor, v: NdarrayOrTensor, right=False, sorter=None):
side = "right" if right else "left"
if isinstance(a, np.ndarray):
return np.searchsorted(a, v, side, sorter) # type: ignore
if hasattr(torch, "searchsorted"):
return torch.searchsorted(a, v, right=right) # type: ignore
# if using old PyTorch, will convert to numpy array then compute
ret = np.searchsorted(a.cpu().numpy(), v.cpu().numpy(), side, sorter) # type: ignore
ret, *_ = convert_to_dst_type(ret, a)
return ret
def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
"""
Args:
img: numpy arrays with required dimension `dim` removed
"""
if self.dim is None:
return img.squeeze()
# for pytorch/numpy unification
if img.shape[self.dim] != 1:
raise ValueError("Can only squeeze singleton dimension")
return img.squeeze(self.dim)
def ravel(x: NdarrayOrTensor) -> NdarrayOrTensor:
"""`np.ravel` with equivalent implementation for torch.
Args:
x: array/tensor to ravel
Returns:
Return a contiguous flattened array/tensor.
"""
if isinstance(x, torch.Tensor):
if hasattr(torch, "ravel"): # `ravel` is new in torch 1.8.0
return x.ravel()
return x.flatten().contiguous()
return np.ravel(x)
def assert_allclose(
actual: NdarrayOrTensor,
desired: NdarrayOrTensor,
type_test: Union[bool, str] = True,
device_test: bool = False,
*args,
**kwargs,
):
"""
Assert that types and all values of two data objects are close.
Args:
actual: Pytorch Tensor or numpy array for comparison.
desired: Pytorch Tensor or numpy array to compare against.
type_test: whether to test that `actual` and `desired` are both numpy arrays or torch tensors.
if type_test == "tensor", it checks whether the `actual` is a torch.tensor or metatensor according to
`get_track_meta`.
device_test: whether to test the device property.
args: extra arguments to pass on to `np.testing.assert_allclose`.
kwargs: extra arguments to pass on to `np.testing.assert_allclose`.
"""
if isinstance(type_test, str) and type_test == "tensor":
if get_track_meta():
np.testing.assert_equal(isinstance(actual, MetaTensor), True,
"must be a MetaTensor")
else:
np.testing.assert_equal(
isinstance(actual, torch.Tensor)
and not isinstance(actual, MetaTensor), True,
"must be a torch.Tensor")
elif type_test:
# check both actual and desired are of the same type
np.testing.assert_equal(isinstance(actual, np.ndarray),
isinstance(desired, np.ndarray), "numpy type")
np.testing.assert_equal(isinstance(actual, torch.Tensor),
isinstance(desired, torch.Tensor),
"torch type")
if isinstance(desired, torch.Tensor) or isinstance(actual, torch.Tensor):
if device_test:
np.testing.assert_equal(str(actual.device), str(desired.device),
"torch device check") # type: ignore
actual = actual.detach().cpu().numpy() if isinstance(
actual, torch.Tensor) else actual
desired = desired.detach().cpu().numpy() if isinstance(
desired, torch.Tensor) else desired
np.testing.assert_allclose(actual, desired, *args, **kwargs)
def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
"""
Filter the image on the `applied_labels`.
Args:
img: Pytorch tensor or numpy array of any shape.
Raises:
NotImplementedError: The provided image was not a Pytorch Tensor or numpy array.
Returns:
Pytorch tensor or numpy array of the same shape as the input.
"""
if not isinstance(img, (np.ndarray, torch.Tensor)):
raise NotImplementedError(
f"{self.__class__} can not handle data of type {type(img)}.")
if isinstance(img, torch.Tensor):
if hasattr(torch, "isin"): # `isin` is new in torch 1.10.0
appl_lbls = torch.as_tensor(self.applied_labels,
device=img.device)
return torch.where(torch.isin(img, appl_lbls), img,
torch.tensor(0.0).to(img))
else:
out = self(img.detach().cpu().numpy())
out, *_ = convert_to_dst_type(out, img)
return out
return np.asarray(np.where(np.isin(img, self.applied_labels), img, 0))
def any_np_pt(x: NdarrayOrTensor,
axis: Union[int, Sequence[int]]) -> NdarrayOrTensor:
"""`np.any` with equivalent implementation for torch.
For pytorch, convert to boolean for compatibility with older versions.
Args:
x: input array/tensor
axis: axis to perform `any` over
Returns:
Return a contiguous flattened array/tensor.
"""
if isinstance(x, np.ndarray):
return np.any(x, axis) # type: ignore
# pytorch can't handle multiple dimensions to `any` so loop across them
axis = [axis] if not isinstance(axis, Sequence) else axis
for ax in axis:
try:
x = torch.any(x, ax)
except RuntimeError:
# older versions of pytorch require the input to be cast to boolean
x = torch.any(x.bool(), ax)
return x
def swapaxes_boxes(boxes: NdarrayOrTensor, axis1: int, axis2: int):
"""
Interchange two axes of boxes.
Args:
boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
axis1: First axis.
axis2: Second axis.
Returns:
boxes with two axes interchanged.
"""
spatial_dims: int = get_spatial_dims(boxes=boxes)
if isinstance(boxes, torch.Tensor):
boxes_swap = boxes.clone()
else:
boxes_swap = deepcopy(boxes) # type: ignore
boxes_swap[:, [axis1, axis2]] = boxes_swap[:, [axis2,
axis1]] # type: ignore
boxes_swap[:, [spatial_dims + axis1, spatial_dims +
axis2]] = boxes_swap[ # type: ignore
:, [spatial_dims + axis2, spatial_dims + axis1]]
return boxes_swap
def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
# if img has channel dim, squeeze it
if img.ndim == 4 and img.shape[0] == 1:
img = img.squeeze(0)
result = [(img == 1) | (img == 4),
(img == 1) | (img == 4) | (img == 2), img == 4]
# merge labels 1 (tumor non-enh) and 4 (tumor enh) and 2 (large edema) to WT
# label 4 is ET
return torch.stack(result, dim=0) if isinstance(
img, torch.Tensor) else np.stack(result, axis=0)
def assert_allclose(
actual: NdarrayOrTensor,
desired: NdarrayOrTensor,
type_test: bool = True,
device_test: bool = False,
*args,
**kwargs,
):
"""
Assert that types and all values of two data objects are close.
Args:
actual: Pytorch Tensor or numpy array for comparison.
desired: Pytorch Tensor or numpy array to compare against.
type_test: whether to test that `actual` and `desired` are both numpy arrays or torch tensors.
device_test: whether to test the device property.
args: extra arguments to pass on to `np.testing.assert_allclose`.
kwargs: extra arguments to pass on to `np.testing.assert_allclose`.
"""
if type_test:
# check both actual and desired are of the same type
np.testing.assert_equal(isinstance(actual, np.ndarray),
isinstance(desired, np.ndarray), "numpy type")
np.testing.assert_equal(isinstance(actual, torch.Tensor),
isinstance(desired, torch.Tensor),
"torch type")
if isinstance(desired, torch.Tensor) or isinstance(actual, torch.Tensor):
if device_test:
np.testing.assert_equal(str(actual.device), str(desired.device),
"torch device check") # type: ignore
actual = actual.detach().cpu().numpy() if isinstance(
actual, torch.Tensor) else actual
desired = desired.detach().cpu().numpy() if isinstance(
desired, torch.Tensor) else desired
np.testing.assert_allclose(actual, desired, *args, **kwargs)
def percentile(x: NdarrayOrTensor, q) -> Union[NdarrayOrTensor, float, int]:
"""`np.percentile` with equivalent implementation for torch.
Pytorch uses `quantile`, but this functionality is only available from v1.7.
For earlier methods, we calculate it ourselves. This doesn't do interpolation,
so is the equivalent of ``numpy.percentile(..., interpolation="nearest")``.
Args:
x: input data
q: percentile to compute (should in range 0 <= q <= 100)
Returns:
Resulting value (scalar)
"""
if np.isscalar(q):
if not 0 <= q <= 100:
raise ValueError
elif any(q < 0) or any(q > 100):
raise ValueError
result: Union[NdarrayOrTensor, float, int]
if isinstance(x, np.ndarray):
result = np.percentile(x, q)
else:
q = torch.tensor(q, device=x.device)
if hasattr(torch, "quantile"):
result = torch.quantile(x, q / 100.0)
else:
# Note that ``kthvalue()`` works one-based, i.e., the first sorted value
# corresponds to k=1, not k=0. Thus, we need the `1 +`.
k = 1 + (0.01 * q * (x.numel() - 1)).round().int()
if k.numel() > 1:
r = [x.view(-1).kthvalue(int(_k)).values.item() for _k in k]
result = torch.tensor(r, device=x.device)
else:
result = x.view(-1).kthvalue(int(k)).values.item()
return result
def flip_boxes(
boxes: NdarrayOrTensor,
spatial_size: Union[Sequence[int], int],
flip_axes: Optional[Union[Sequence[int], int]] = None,
):
"""
Flip boxes when the corresponding image is flipped
Args:
boxes: bounding boxes, Nx4 or Nx6 torch tensor or ndarray. The box mode is assumed to be ``StandardMode``
spatial_size: image spatial size.
flip_axes: spatial axes along which to flip over. Default is None.
The default `axis=None` will flip over all of the axes of the input array.
If axis is negative it counts from the last to the first axis.
If axis is a tuple of ints, flipping is performed on all of the axes
specified in the tuple.
Returns:
flipped boxes, with same data type as ``boxes``, does not share memory with ``boxes``
"""
spatial_dims: int = get_spatial_dims(boxes=boxes)
spatial_size = ensure_tuple_rep(spatial_size, spatial_dims)
if flip_axes is None:
flip_axes = tuple(range(0, spatial_dims))
flip_axes = ensure_tuple(flip_axes)
# flip box
if isinstance(boxes, torch.Tensor):
_flip_boxes = boxes.clone()
else:
_flip_boxes = deepcopy(boxes) # type: ignore
for axis in flip_axes:
_flip_boxes[:, axis + spatial_dims] = spatial_size[
axis] - boxes[:, axis] - TO_REMOVE
_flip_boxes[:, axis] = spatial_size[
axis] - boxes[:, axis + spatial_dims] - TO_REMOVE
return _flip_boxes
def __call__(self, image: NdarrayOrTensor) -> NdarrayOrTensor:
if self.grid_size == (1, 1) and self.patch_size is None:
if isinstance(image, torch.Tensor):
return torch.stack([image])
elif isinstance(image, np.ndarray):
return np.stack([image]) # type: ignore
else:
raise ValueError(
f"Input type [{type(image)}] is not supported.")
patch_size, steps = self.get_params(image.shape[1:])
patches: NdarrayOrTensor
if isinstance(image, torch.Tensor):
patches = (image.unfold(1, patch_size[0], steps[0]).unfold(
2, patch_size[1],
steps[1]).flatten(1, 2).transpose(0, 1).contiguous())
elif isinstance(image, np.ndarray):
x_step, y_step = steps
c_stride, x_stride, y_stride = image.strides
n_channels = image.shape[0]
patches = as_strided(
image,
shape=(*self.grid_size, n_channels, patch_size[0],
patch_size[1]),
strides=(x_stride * x_step, y_stride * y_step, c_stride,
x_stride, y_stride),
writeable=False,
)
# flatten the first two dimensions
patches = patches.reshape(np.prod(patches.shape[:2]),
*patches.shape[2:])
# make it a contiguous array
patches = np.ascontiguousarray(patches)
else:
raise ValueError(f"Input type [{type(image)}] is not supported.")
return patches
def write_nifti(
data: NdarrayOrTensor,
file_name: str,
affine: Optional[NdarrayOrTensor] = None,
target_affine: Optional[np.ndarray] = None,
resample: bool = True,
output_spatial_shape: Union[Sequence[int], np.ndarray, None] = None,
mode: Union[GridSampleMode, str] = GridSampleMode.BILINEAR,
padding_mode: Union[GridSamplePadMode, str] = GridSamplePadMode.BORDER,
align_corners: bool = False,
dtype: DtypeLike = np.float64,
output_dtype: DtypeLike = np.float32,
) -> None:
"""
Write numpy data into NIfTI files to disk. This function converts data
into the coordinate system defined by `target_affine` when `target_affine`
is specified.
If the coordinate transform between `affine` and `target_affine` could be
achieved by simply transposing and flipping `data`, no resampling will
happen. otherwise this function will resample `data` using the coordinate
transform computed from `affine` and `target_affine`. Note that the shape
of the resampled `data` may subject to some rounding errors. For example,
resampling a 20x20 pixel image from pixel size (1.5, 1.5)-mm to (3.0,
3.0)-mm space will return a 10x10-pixel image. However, resampling a
20x20-pixel image from pixel size (2.0, 2.0)-mm to (3.0, 3.0)-mma space
will output a 14x14-pixel image, where the image shape is rounded from
13.333x13.333 pixels. In this case `output_spatial_shape` could be specified so
that this function writes image data to a designated shape.
The saved `affine` matrix follows:
- If `affine` equals to `target_affine`, save the data with `target_affine`.
- If `resample=False`, transform `affine` to `new_affine` based on the orientation
of `target_affine` and save the data with `new_affine`.
- If `resample=True`, save the data with `target_affine`, if explicitly specify
the `output_spatial_shape`, the shape of saved data is not computed by `target_affine`.
- If `target_affine` is None, set `target_affine=affine` and save.
- If `affine` and `target_affine` are None, the data will be saved with an identity
matrix as the image affine.
This function assumes the NIfTI dimension notations.
Spatially it supports up to three dimensions, that is, H, HW, HWD for
1D, 2D, 3D respectively.
When saving multiple time steps or multiple channels `data`, time and/or
modality axes should be appended after the first three dimensions. For
example, shape of 2D eight-class segmentation probabilities to be saved
could be `(64, 64, 1, 8)`. Also, data in shape (64, 64, 8), (64, 64, 8, 1)
will be considered as a single-channel 3D image.
Args:
data: input data to write to file.
file_name: expected file name that saved on disk.
affine: the current affine of `data`. Defaults to `np.eye(4)`
target_affine: before saving
the (`data`, `affine`) as a Nifti1Image,
transform the data into the coordinates defined by `target_affine`.
resample: whether to run resampling when the target affine
could not be achieved by swapping/flipping data axes.
output_spatial_shape: spatial shape of the output image.
This option is used when resample = True.
mode: {``"bilinear"``, ``"nearest"``}
This option is used when ``resample = True``.
Interpolation mode to calculate output values. Defaults to ``"bilinear"``.
See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
padding_mode: {``"zeros"``, ``"border"``, ``"reflection"``}
This option is used when ``resample = True``.
Padding mode for outside grid values. Defaults to ``"border"``.
See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
align_corners: Geometrically, we consider the pixels of the input as squares rather than points.
See also: https://pytorch.org/docs/stable/nn.functional.html#grid-sample
dtype: data type for resampling computation. Defaults to ``np.float64`` for best precision.
If None, use the data type of input data.
output_dtype: data type for saving data. Defaults to ``np.float32``.
"""
if isinstance(data, torch.Tensor):
data, *_ = convert_data_type(data, np.ndarray)
if isinstance(affine, torch.Tensor):
affine, *_ = convert_data_type(affine, np.ndarray)
if not isinstance(data, np.ndarray):
raise AssertionError("input data must be numpy array or torch tensor.")
dtype = dtype or data.dtype
sr = min(data.ndim, 3)
if affine is None:
affine = np.eye(4, dtype=np.float64)
affine = to_affine_nd(sr, affine) # type: ignore
if target_affine is None:
target_affine = affine
target_affine = to_affine_nd(sr, target_affine)
if np.allclose(affine, target_affine, atol=1e-3):
# no affine changes, save (data, affine)
results_img = nib.Nifti1Image(data.astype(output_dtype),
to_affine_nd(3, target_affine))
nib.save(results_img, file_name)
return
# resolve orientation
start_ornt = nib.orientations.io_orientation(affine)
target_ornt = nib.orientations.io_orientation(target_affine)
ornt_transform = nib.orientations.ornt_transform(start_ornt, target_ornt)
data_shape = data.shape
data = nib.orientations.apply_orientation(data, ornt_transform)
_affine = affine @ nib.orientations.inv_ornt_aff(ornt_transform,
data_shape)
if np.allclose(_affine, target_affine, atol=1e-3) or not resample:
results_img = nib.Nifti1Image(data.astype(output_dtype),
to_affine_nd(3, _affine)) # type: ignore
nib.save(results_img, file_name)
return
# need resampling
affine_xform = AffineTransform(normalized=False,
mode=mode,
padding_mode=padding_mode,
align_corners=align_corners,
reverse_indexing=True)
transform = np.linalg.inv(_affine) @ target_affine
if output_spatial_shape is None:
output_spatial_shape, _ = compute_shape_offset(data.shape, _affine,
target_affine)
output_spatial_shape_ = list(
output_spatial_shape) if output_spatial_shape is not None else []
if data.ndim > 3: # multi channel, resampling each channel
while len(output_spatial_shape_) < 3:
output_spatial_shape_ = output_spatial_shape_ + [1]
spatial_shape, channel_shape = data.shape[:3], data.shape[3:]
data_np: np.ndarray = data.reshape(list(spatial_shape) +
[-1]) # type: ignore
data_np = np.moveaxis(data_np, -1, 0) # channel first for pytorch
data_torch = affine_xform(
torch.as_tensor(
np.ascontiguousarray(data_np).astype(dtype)).unsqueeze(0),
torch.as_tensor(np.ascontiguousarray(transform).astype(dtype)),
spatial_size=output_spatial_shape_[:3],
)
data_np = data_torch.squeeze(0).detach().cpu().numpy()
data_np = np.moveaxis(data_np, 0, -1) # channel last for nifti
data_np = data_np.reshape(
list(data_np.shape[:3]) + list(channel_shape))
else: # single channel image, need to expand to have batch and channel
while len(output_spatial_shape_) < len(data.shape):
output_spatial_shape_ = output_spatial_shape_ + [1]
data_torch = affine_xform(
torch.as_tensor(
np.ascontiguousarray(data).astype(dtype)[None, None]),
torch.as_tensor(np.ascontiguousarray(transform).astype(dtype)),
spatial_size=output_spatial_shape_[:len(data.shape)],
)
data_np = data_torch.squeeze(0).squeeze(0).detach().cpu().numpy()
results_img = nib.Nifti1Image(data_np.astype(output_dtype),
to_affine_nd(3, target_affine))
nib.save(results_img, file_name)
return