def __call__(
self,
img: np.ndarray,
msk: Optional[np.ndarray] = None,
center: Optional[tuple] = None,
z_axis: Optional[int] = None,
):
"""
Apply the transform to `img`, assuming `img` is channel-first and
slicing doesn't apply to the channel dim.
"""
if self.mask_data is None and msk is None:
raise ValueError("Unknown mask_data.")
mask_data_ = np.array([[1]])
if self.mask_data is not None and msk is None:
mask_data_ = self.mask_data > 0
if msk is not None:
mask_data_ = msk > 0
mask_data_ = np.asarray(mask_data_)
if mask_data_.shape[0] != 1 and mask_data_.shape[0] != img.shape[0]:
raise ValueError(
"When mask_data is not single channel, mask_data channels must match img, "
f"got img={img.shape[0]} mask_data={mask_data_.shape[0]}.")
z_axis_ = z_axis if z_axis is not None else self.z_axis
if center is None:
center = self.get_center_pos(mask_data_, z_axis_)
if self.crop_mode in ["single", "parallel"]:
size_ = self.get_new_spatial_size(z_axis_)
size_ = list(map(int, size_))
slice_ = SpatialCrop(roi_center=center, roi_size=size_)(img)
if np.any(slice_.shape[1:] != size_):
slice_ = ResizeWithPadOrCrop(spatial_size=size_)(slice_)
return np.moveaxis(slice_.squeeze(0), z_axis_, 0)
else:
cross_slices = np.zeros(shape=(3, ) + self.roi_size)
for k in range(3):
size_ = np.insert(self.roi_size, k, 1)
slice_ = SpatialCrop(roi_center=center, roi_size=size_)(img)
if np.any(slice_.shape[1:] != size_):
slice_ = ResizeWithPadOrCrop(spatial_size=size_)(slice_)
cross_slices[k] = slice_.squeeze()
return cross_slices
def __call__(
self, data: Mapping[Hashable,
np.ndarray]) -> List[Dict[Hashable, np.ndarray]]:
d = dict(data)
label = d[self.label_key]
image = d[self.image_key] if self.image_key else None
fg_indices = d.get(self.fg_indices_key,
None) if self.fg_indices_key is not None else None
bg_indices = d.get(self.bg_indices_key,
None) if self.bg_indices_key is not None else None
self.randomize(label, fg_indices, bg_indices, image)
assert isinstance(self.spatial_size, tuple)
assert self.centers is not None
results: List[Dict[Hashable, np.ndarray]] = [
dict() for _ in range(self.num_samples)
]
for key in data.keys():
if key in self.keys:
img = d[key]
for i, center in enumerate(self.centers):
if self.crop_mode in ["single", "parallel"]:
size_ = self.get_new_spatial_size()
slice_ = SpatialCrop(roi_center=tuple(center),
roi_size=size_)(img)
results[i][key] = np.moveaxis(slice_.squeeze(0),
self.z_axis, 0)
else:
cross_slices = np.zeros(shape=(3, ) +
self.spatial_size)
for k in range(3):
size_ = np.insert(self.spatial_size, k, 1)
slice_ = SpatialCrop(roi_center=tuple(center),
roi_size=size_)(img)
cross_slices[k] = slice_.squeeze()
results[i][key] = cross_slices
else:
for i in range(self.num_samples):
results[i][key] = data[key]
return results
