def randomize(self, img_size):
self._size = [self.roi_size] * len(img_size) if not isinstance(self.roi_size, (list, tuple)) else self.roi_size
if self.random_size:
self._size = [self.R.randint(low=self._size[i], high=img_size[i] + 1) for i in range(len(img_size))]
if self.random_center:
valid_size = get_valid_patch_size(img_size, self._size)
self._slices = ensure_tuple(slice(None)) + get_random_patch(img_size, valid_size, self.R)
def randomize(self, img_size) -> None: # type: ignore # see issue #495
self._size = fall_back_tuple(self.roi_size, img_size)
if self.random_size:
self._size = [self.R.randint(low=self._size[i], high=img_size[i] + 1) for i in range(len(img_size))]
if self.random_center:
valid_size = get_valid_patch_size(img_size, self._size)
self._slices = (slice(None),) + get_random_patch(img_size, valid_size, self.R)
def randomize(self, img_size: Sequence[int]) -> None:
self._size = fall_back_tuple(self.roi_size, img_size)
if self.random_size:
self._size = tuple((self.R.randint(low=self._size[i], high=img_size[i] + 1) for i in range(len(img_size))))
if self.random_center:
valid_size = get_valid_patch_size(img_size, self._size)
self._slices = (slice(None),) + get_random_patch(img_size, valid_size, self.R)
def randomize(self, img_size):
self._size = ensure_tuple_rep(self.roi_size, len(img_size))
if self.random_size:
self._size = [self.R.randint(low=self._size[i], high=img_size[i] + 1) for i in range(len(img_size))]
if self.random_center:
valid_size = get_valid_patch_size(img_size, self._size)
self._slices = ensure_tuple(slice(None)) + get_random_patch(img_size, valid_size, self.R)
def __call__(self, data):
d = dict(data)
image_shape = d[self.keys[0]].shape # image shape from the first data key
patch_spatial_size = get_valid_patch_size(image_shape, self.patch_spatial_size)
self.randomize(image_shape, patch_spatial_size)
for key in self.keys:
d[key] = d[key][self._slices]
return d
def randomize(self, img_size: Sequence[int]) -> None:
self._size = fall_back_tuple(self.roi_size, img_size)
if self.random_size:
max_size = img_size if self.max_roi_size is None else fall_back_tuple(self.max_roi_size, img_size)
if any([i > j for i, j in zip(self._size, max_size)]):
raise ValueError(f"min ROI size: {self._size} is bigger than max ROI size: {max_size}.")
self._size = [self.R.randint(low=self._size[i], high=max_size[i] + 1) for i in range(len(img_size))]
if self.random_center:
valid_size = get_valid_patch_size(img_size, self._size)
self._slices = (slice(None),) + get_random_patch(img_size, valid_size, self.R)
def sliding_window_inference(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[[torch.Tensor], torch.Tensor],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Sliding window inference on `inputs` with `predictor`.
When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
To maintain the same spatial sizes, the output image will be cropped to the original input size.
Args:
inputs: input image to be processed (assuming NCHW[D])
roi_size: the spatial window size for inferences.
When its components have None or non-positives, the corresponding inputs dimension will be used.
if the components of the `roi_size` are non-positive values, the transform will use the
corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
to `(32, 64)` if the second spatial dimension size of img is `64`.
sw_batch_size: the batch size to run window slices.
predictor: given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)`
should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D];
where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`.
overlap: Amount of overlap between scans.
mode: {``"constant"``, ``"gaussian"``}
How to blend output of overlapping windows. Defaults to ``"constant"``.
- ``"constant``": gives equal weight to all predictions.
- ``"gaussian``": gives less weight to predictions on edges of windows.
padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}
Padding mode when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
See also: https://pytorch.org/docs/stable/nn.functional.html#pad
cval: fill value for 'constant' padding mode. Default: 0
device: device running the concatenation of the windows.
By default the device and accordingly the memory of the input device is used. If for example
set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
input and roi_size parameter. Output is on the device set or if not set the inputs device.
Raises:
NotImplementedError: When ``inputs`` does not have batch size = 1.
Note:
- input must be channel-first and have a batch dim, support both spatial 2D and 3D.
- currently only supports `inputs` with batch_size=1.
"""
num_spatial_dims = len(inputs.shape) - 2
assert 0 <= overlap < 1, "overlap must be >= 0 and < 1."
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
# TODO: Enable batch sizes > 1 in future
if batch_size > 1:
raise NotImplementedError("Currently only inputs with batch size = 1 are supported.")
if device is None:
device = inputs.device
roi_size = fall_back_tuple(roi_size, image_size_)
# in case that image size is smaller than roi size
image_size = tuple(max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs, pad=pad_size, mode=PytorchPadMode(padding_mode).value, value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims, overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
slice_batches = []
for slice_index in range(0, len(slices), sw_batch_size):
slice_index_range = range(slice_index, min(slice_index + sw_batch_size, len(slices)))
input_slices = []
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
if len(curr_slice) == 3:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1], curr_slice[2]])
else:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1]])
slice_batches.append(torch.stack(input_slices))
# Perform predictions
output_rois = list()
for data in slice_batches:
seg_prob = predictor(data) # batched patch segmentation
output_rois.append(seg_prob.to(device))
# stitching output image
output_classes = output_rois[0].shape[1]
output_shape = [batch_size, output_classes] + list(image_size)
# Create importance map
importance_map = compute_importance_map(get_valid_patch_size(image_size, roi_size), mode=mode, device=device)
# allocate memory to store the full output and the count for overlapping parts
output_image = torch.zeros(output_shape, dtype=torch.float32, device=device)
count_map = torch.zeros(output_shape, dtype=torch.float32, device=device)
for window_id, slice_index in enumerate(range(0, len(slices), sw_batch_size)):
slice_index_range = range(slice_index, min(slice_index + sw_batch_size, len(slices)))
# store the result in the proper location of the full output. Apply weights from importance map.
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
if len(curr_slice) == 3:
output_image[0, :, curr_slice[0], curr_slice[1], curr_slice[2]] += (
importance_map * output_rois[window_id][curr_index - slice_index, :]
)
count_map[0, :, curr_slice[0], curr_slice[1], curr_slice[2]] += importance_map
else:
output_image[0, :, curr_slice[0], curr_slice[1]] += (
importance_map * output_rois[window_id][curr_index - slice_index, :]
)
count_map[0, :, curr_slice[0], curr_slice[1]] += importance_map
# account for any overlapping sections
output_image = output_image / count_map
if num_spatial_dims == 3:
return output_image[
...,
pad_size[4] : image_size_[0] + pad_size[4],
pad_size[2] : image_size_[1] + pad_size[2],
pad_size[0] : image_size_[2] + pad_size[0],
]
return output_image[
..., pad_size[2] : image_size_[0] + pad_size[2], pad_size[0] : image_size_[1] + pad_size[0]
] # 2D
def sliding_window_inference(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor],
Dict[Any, torch.Tensor]]],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
sigma_scale: Union[Sequence[float], float] = 0.125,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
sw_device: Union[torch.device, str, None] = None,
device: Union[torch.device, str, None] = None,
progress: bool = False,
roi_weight_map: Union[torch.Tensor, None] = None,
*args: Any,
**kwargs: Any,
) -> Union[torch.Tensor, Tuple[torch.Tensor, ...], Dict[Any, torch.Tensor]]:
"""
Sliding window inference on `inputs` with `predictor`.
The outputs of `predictor` could be a tensor, a tuple, or a dictionary of tensors.
Each output in the tuple or dict value is allowed to have different resolutions with respect to the input.
e.g., the input patch spatial size is [128,128,128], the output (a tuple of two patches) patch sizes
could be ([128,64,256], [64,32,128]).
In this case, the parameter `overlap` and `roi_size` need to be carefully chosen to ensure the output ROI is still
an integer. If the predictor's input and output spatial sizes are not equal, we recommend choosing the parameters
so that `overlap*roi_size*output_size/input_size` is an integer (for each spatial dimension).
When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
To maintain the same spatial sizes, the output image will be cropped to the original input size.
Args:
inputs: input image to be processed (assuming NCHW[D])
roi_size: the spatial window size for inferences.
When its components have None or non-positives, the corresponding inputs dimension will be used.
if the components of the `roi_size` are non-positive values, the transform will use the
corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
to `(32, 64)` if the second spatial dimension size of img is `64`.
sw_batch_size: the batch size to run window slices.
predictor: given input tensor ``patch_data`` in shape NCHW[D],
The outputs of the function call ``predictor(patch_data)`` should be a tensor, a tuple, or a dictionary
with Tensor values. Each output in the tuple or dict value should have the same batch_size, i.e. NM'H'W'[D'];
where H'W'[D'] represents the output patch's spatial size, M is the number of output channels,
N is `sw_batch_size`, e.g., the input shape is (7, 1, 128,128,128),
the output could be a tuple of two tensors, with shapes: ((7, 5, 128, 64, 256), (7, 4, 64, 32, 128)).
In this case, the parameter `overlap` and `roi_size` need to be carefully chosen
to ensure the scaled output ROI sizes are still integers.
If the `predictor`'s input and output spatial sizes are different,
we recommend choosing the parameters so that ``overlap*roi_size*zoom_scale`` is an integer for each dimension.
overlap: Amount of overlap between scans.
mode: {``"constant"``, ``"gaussian"``}
How to blend output of overlapping windows. Defaults to ``"constant"``.
- ``"constant``": gives equal weight to all predictions.
- ``"gaussian``": gives less weight to predictions on edges of windows.
sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``.
Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``.
When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding
spatial dimensions.
padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}
Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
cval: fill value for 'constant' padding mode. Default: 0
sw_device: device for the window data.
By default the device (and accordingly the memory) of the `inputs` is used.
Normally `sw_device` should be consistent with the device where `predictor` is defined.
device: device for the stitched output prediction.
By default the device (and accordingly the memory) of the `inputs` is used. If for example
set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
`inputs` and `roi_size`. Output is on the `device`.
progress: whether to print a `tqdm` progress bar.
roi_weight_map: pre-computed (non-negative) weight map for each ROI.
If not given, and ``mode`` is not `constant`, this map will be computed on the fly.
args: optional args to be passed to ``predictor``.
kwargs: optional keyword args to be passed to ``predictor``.
Note:
- input must be channel-first and have a batch dim, supports N-D sliding window.
"""
compute_dtype = inputs.dtype
num_spatial_dims = len(inputs.shape) - 2
if overlap < 0 or overlap >= 1:
raise ValueError("overlap must be >= 0 and < 1.")
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
batch_size, _, *image_size_ = inputs.shape
if device is None:
device = inputs.device
if sw_device is None:
sw_device = inputs.device
roi_size = fall_back_tuple(roi_size, image_size_)
# in case that image size is smaller than roi size
image_size = tuple(
max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs,
pad=pad_size,
mode=look_up_option(padding_mode, PytorchPadMode),
value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims,
overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
num_win = len(slices) # number of windows per image
total_slices = num_win * batch_size # total number of windows
# Create window-level importance map
valid_patch_size = get_valid_patch_size(image_size, roi_size)
if valid_patch_size == roi_size and (roi_weight_map is not None):
importance_map = roi_weight_map
else:
try:
importance_map = compute_importance_map(valid_patch_size,
mode=mode,
sigma_scale=sigma_scale,
device=device)
except BaseException as e:
raise RuntimeError(
"Seems to be OOM. Please try smaller patch size or mode='constant' instead of mode='gaussian'."
) from e
importance_map = convert_data_type(importance_map, torch.Tensor, device,
compute_dtype)[0] # type: ignore
# handle non-positive weights
min_non_zero = max(importance_map[importance_map != 0].min().item(), 1e-3)
importance_map = torch.clamp(importance_map.to(torch.float32),
min=min_non_zero).to(compute_dtype)
# Perform predictions
dict_key, output_image_list, count_map_list = None, [], []
_initialized_ss = -1
is_tensor_output = True # whether the predictor's output is a tensor (instead of dict/tuple)
# for each patch
for slice_g in tqdm(range(0, total_slices,
sw_batch_size)) if progress else range(
0, total_slices, sw_batch_size):
slice_range = range(slice_g, min(slice_g + sw_batch_size,
total_slices))
unravel_slice = [
[slice(int(idx / num_win),
int(idx / num_win) + 1),
slice(None)] + list(slices[idx % num_win]) for idx in slice_range
]
window_data = torch.cat([
convert_data_type(inputs[win_slice], torch.Tensor)[0]
for win_slice in unravel_slice
]).to(sw_device)
seg_prob_out = predictor(window_data, *args,
**kwargs) # batched patch segmentation
# convert seg_prob_out to tuple seg_prob_tuple, this does not allocate new memory.
seg_prob_tuple: Tuple[torch.Tensor, ...]
if isinstance(seg_prob_out, torch.Tensor):
seg_prob_tuple = (seg_prob_out, )
elif isinstance(seg_prob_out, Mapping):
if dict_key is None:
dict_key = sorted(
seg_prob_out.keys()) # track predictor's output keys
seg_prob_tuple = tuple(seg_prob_out[k] for k in dict_key)
is_tensor_output = False
else:
seg_prob_tuple = ensure_tuple(seg_prob_out)
is_tensor_output = False
# for each output in multi-output list
for ss, seg_prob in enumerate(seg_prob_tuple):
seg_prob = seg_prob.to(device) # BxCxMxNxP or BxCxMxN
# compute zoom scale: out_roi_size/in_roi_size
zoom_scale = []
for axis, (img_s_i, out_w_i, in_w_i) in enumerate(
zip(image_size, seg_prob.shape[2:],
window_data.shape[2:])):
_scale = out_w_i / float(in_w_i)
if not (img_s_i * _scale).is_integer():
warnings.warn(
f"For spatial axis: {axis}, output[{ss}] will have non-integer shape. Spatial "
f"zoom_scale between output[{ss}] and input is {_scale}. Please pad inputs."
)
zoom_scale.append(_scale)
if _initialized_ss < ss: # init. the ss-th buffer at the first iteration
# construct multi-resolution outputs
output_classes = seg_prob.shape[1]
output_shape = [batch_size, output_classes] + [
int(image_size_d * zoom_scale_d)
for image_size_d, zoom_scale_d in zip(
image_size, zoom_scale)
]
# allocate memory to store the full output and the count for overlapping parts
output_image_list.append(
torch.zeros(output_shape,
dtype=compute_dtype,
device=device))
count_map_list.append(
torch.zeros([1, 1] + output_shape[2:],
dtype=compute_dtype,
device=device))
_initialized_ss += 1
# resizing the importance_map
resizer = Resize(spatial_size=seg_prob.shape[2:],
mode="nearest",
anti_aliasing=False)
# store the result in the proper location of the full output. Apply weights from importance map.
for idx, original_idx in zip(slice_range, unravel_slice):
# zoom roi
original_idx_zoom = list(
original_idx) # 4D for 2D image, 5D for 3D image
for axis in range(2, len(original_idx_zoom)):
zoomed_start = original_idx[axis].start * zoom_scale[axis -
2]
zoomed_end = original_idx[axis].stop * zoom_scale[axis - 2]
if not zoomed_start.is_integer() or (
not zoomed_end.is_integer()):
warnings.warn(
f"For axis-{axis-2} of output[{ss}], the output roi range is not int. "
f"Input roi range is ({original_idx[axis].start}, {original_idx[axis].stop}). "
f"Spatial zoom_scale between output[{ss}] and input is {zoom_scale[axis - 2]}. "
f"Corresponding output roi range is ({zoomed_start}, {zoomed_end}).\n"
f"Please change overlap ({overlap}) or roi_size ({roi_size[axis-2]}) for axis-{axis-2}. "
"Tips: if overlap*roi_size*zoom_scale is an integer, it usually works."
)
original_idx_zoom[axis] = slice(int(zoomed_start),
int(zoomed_end), None)
importance_map_zoom = resizer(
importance_map.unsqueeze(0))[0].to(compute_dtype)
# store results and weights
output_image_list[ss][
original_idx_zoom] += importance_map_zoom * seg_prob[
idx - slice_g]
count_map_list[ss][original_idx_zoom] += (
importance_map_zoom.unsqueeze(0).unsqueeze(0).expand(
count_map_list[ss][original_idx_zoom].shape))
# account for any overlapping sections
for ss in range(len(output_image_list)):
output_image_list[ss] = (output_image_list[ss] /
count_map_list.pop(0)).to(compute_dtype)
# remove padding if image_size smaller than roi_size
for ss, output_i in enumerate(output_image_list):
if torch.isnan(output_i).any() or torch.isinf(output_i).any():
warnings.warn(
"Sliding window inference results contain NaN or Inf.")
zoom_scale = [
seg_prob_map_shape_d / roi_size_d
for seg_prob_map_shape_d, roi_size_d in zip(
output_i.shape[2:], roi_size)
]
final_slicing: List[slice] = []
for sp in range(num_spatial_dims):
slice_dim = slice(
pad_size[sp * 2],
image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
slice_dim = slice(
int(
round(slice_dim.start *
zoom_scale[num_spatial_dims - sp - 1])),
int(
round(slice_dim.stop *
zoom_scale[num_spatial_dims - sp - 1])),
)
final_slicing.insert(0, slice_dim)
while len(final_slicing) < len(output_i.shape):
final_slicing.insert(0, slice(None))
output_image_list[ss] = output_i[final_slicing]
if dict_key is not None: # if output of predictor is a dict
final_output = dict(zip(dict_key, output_image_list))
else:
final_output = tuple(output_image_list) # type: ignore
final_output = final_output[
0] if is_tensor_output else final_output # type: ignore
if isinstance(inputs, MetaTensor):
final_output = convert_to_dst_type(final_output,
inputs)[0] # type: ignore
return final_output
def __call__(self, img):
patch_size = get_valid_patch_size(img.shape, self.patch_size)
slices = get_random_patch(img.shape, patch_size)
return img[slices]
def sliding_window_inference(
inputs: torch.Tensor,
roi_size,
sw_batch_size: int,
predictor: Callable,
overlap: float = 0.25,
blend_mode: str = "constant",
padding_mode: str = "constant",
cval=0,
):
"""
Sliding window inference on `inputs` with `predictor`.
When roi_size is 0, the corresponding inputs dimension will be used as the window dimension.
When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
To maintain the same spatial sizes, the output image will be cropped to the original input size.
Args:
inputs: input image to be processed (assuming NCHW[D])
roi_size (list, tuple): the spatial window size for inferences.
sw_batch_size: the batch size to run window slices.
predictor: given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)`
should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D];
where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`.
overlap: Amount of overlap between scans.
blend_mode: How to blend output of overlapping windows. Options are 'constant', 'gaussian'. 'constant'
gives equal weight to all predictions while gaussian gives less weight to predictions on edges of windows.
padding_mode: padding mode when roi_size is larger than inputs.
options are 'constant', 'reflect', 'replicate' or 'circular'. Default: 'constant'
cval: fill value for 'constant' padding mode. Default: 0
Note:
- input must be channel-first and have a batch dim, support both spatial 2D and 3D.
- currently only supports `inputs` with batch_size=1.
"""
num_spatial_dims = len(inputs.shape) - 2
assert len(
roi_size
) == num_spatial_dims, f"roi_size {roi_size} does not match input dims."
assert 0 <= overlap < 1, "overlap must be >= 0 and < 1."
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
# TODO: Enable batch sizes > 1 in future
if batch_size > 1:
raise NotImplementedError
original_image_size = [image_size_[i] for i in range(num_spatial_dims)]
# in case that image size is smaller than roi size
image_size = tuple(
max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs, pad=pad_size, mode=padding_mode, value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims,
overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
slice_batches = []
for slice_index in range(0, len(slices), sw_batch_size):
slice_index_range = range(
slice_index, min(slice_index + sw_batch_size, len(slices)))
input_slices = []
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
if len(curr_slice) == 3:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]])
else:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1]])
slice_batches.append(torch.stack(input_slices))
# Perform predictions
output_rois = list()
for data in slice_batches:
seg_prob = predictor(data) # batched patch segmentation
output_rois.append(seg_prob)
# stitching output image
output_classes = output_rois[0].shape[1]
output_shape = [batch_size, output_classes] + list(image_size)
# Create importance map
importance_map = compute_importance_map(get_valid_patch_size(
image_size, roi_size),
mode=blend_mode,
device=inputs.device)
# allocate memory to store the full output and the count for overlapping parts
output_image = torch.zeros(output_shape,
dtype=torch.float32,
device=inputs.device)
count_map = torch.zeros(output_shape,
dtype=torch.float32,
device=inputs.device)
for window_id, slice_index in enumerate(
range(0, len(slices), sw_batch_size)):
slice_index_range = range(
slice_index, min(slice_index + sw_batch_size, len(slices)))
# store the result in the proper location of the full output. Apply weights from importance map.
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
if len(curr_slice) == 3:
output_image[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += (
importance_map *
output_rois[window_id][curr_index -
slice_index, :])
count_map[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += importance_map
else:
output_image[0, :, curr_slice[0], curr_slice[1]] += (
importance_map *
output_rois[window_id][curr_index - slice_index, :])
count_map[0, :, curr_slice[0], curr_slice[1]] += importance_map
# account for any overlapping sections
output_image /= count_map
if num_spatial_dims == 3:
return output_image[...,
pad_size[4]:original_image_size[0] + pad_size[4],
pad_size[2]:original_image_size[1] + pad_size[2],
pad_size[0]:original_image_size[2] + pad_size[0], ]
return output_image[..., pad_size[2]:original_image_size[0] + pad_size[2],
pad_size[0]:original_image_size[1] + pad_size[0]] # 2D
def sliding_window_classification(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[[torch.Tensor], torch.Tensor],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
device: Optional[torch.device] = None,
) -> torch.Tensor:
num_spatial_dims = len(inputs.shape) - 2
assert 0 <= overlap < 1, "overlap must be >= 0 and < 1."
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
# TODO: Enable batch sizes > 1 in future
if batch_size > 1:
raise NotImplementedError(
"Currently only inputs with batch size = 1 are supported.")
if device is None:
device = inputs.device
roi_size = fall_back_tuple(roi_size, image_size_)
# in case that image size is smaller than roi size
image_size = tuple(
max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs,
pad=pad_size,
mode=PytorchPadMode(padding_mode).value,
value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims,
overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
slice_batches = []
for slice_index in range(0, len(slices), sw_batch_size):
slice_index_range = range(
slice_index, min(slice_index + sw_batch_size, len(slices)))
input_slices = []
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
if len(curr_slice) == 3:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]])
else:
input_slices.append(inputs[0, :, curr_slice[0], curr_slice[1]])
slice_batches.append(torch.stack(input_slices))
# Perform predictions
output_rois = list()
for data in slice_batches:
cls_prob = predictor(data) # batched patch classification
output_rois.append(cls_prob.to(device))
# stitching output image
output_classes = output_rois[0].shape[1]
output_shape = [batch_size, output_classes] + list(image_size)
# Create importance map
importance_map = compute_importance_map(get_valid_patch_size(
image_size, roi_size),
mode=mode,
device=device)
# allocate memory to store the full output and the count for overlapping parts
output_image = torch.zeros(output_shape,
dtype=torch.float32,
device=device)
count_map = torch.zeros(output_shape, dtype=torch.float32, device=device)
for window_id, slice_index in enumerate(
range(0, len(slices), sw_batch_size)):
slice_index_range = range(
slice_index, min(slice_index + sw_batch_size, len(slices)))
# store the result in the proper location of the full output. Apply weights from importance map.
for curr_index in slice_index_range:
curr_slice = slices[curr_index]
curr_roi = output_rois[window_id][curr_index - slice_index, :]
if len(curr_slice) == 3:
output_image[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += (importance_map *
curr_roi.view(-1, 1, 1, 1))
count_map[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += importance_map
else:
output_image[0, :, curr_slice[0],
curr_slice[1]] += (importance_map *
curr_roi.view(-1, 1, 1))
count_map[0, :, curr_slice[0], curr_slice[1]] += importance_map
# account for any overlapping sections
output_image = output_image / count_map
if num_spatial_dims == 3:
return output_image[..., pad_size[4]:image_size_[0] + pad_size[4],
pad_size[2]:image_size_[1] + pad_size[2],
pad_size[0]:image_size_[2] + pad_size[0], ]
return output_image[..., pad_size[2]:image_size_[0] + pad_size[2],
pad_size[0]:image_size_[1] + pad_size[0]] # 2D
def sliding_window_inference(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[..., torch.Tensor],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
sigma_scale: Union[Sequence[float], float] = 0.125,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
sw_device: Union[torch.device, str, None] = None,
device: Union[torch.device, str, None] = None,
*args: Any,
**kwargs: Any,
) -> torch.Tensor:
"""
Sliding window inference on `inputs` with `predictor`.
When roi_size is larger than the inputs' spatial size, the input image are padded during inference.
To maintain the same spatial sizes, the output image will be cropped to the original input size.
Args:
inputs: input image to be processed (assuming NCHW[D])
roi_size: the spatial window size for inferences.
When its components have None or non-positives, the corresponding inputs dimension will be used.
if the components of the `roi_size` are non-positive values, the transform will use the
corresponding components of img size. For example, `roi_size=(32, -1)` will be adapted
to `(32, 64)` if the second spatial dimension size of img is `64`.
sw_batch_size: the batch size to run window slices.
predictor: given input tensor `patch_data` in shape NCHW[D], `predictor(patch_data)`
should return a prediction with the same spatial shape and batch_size, i.e. NMHW[D];
where HW[D] represents the patch spatial size, M is the number of output channels, N is `sw_batch_size`.
overlap: Amount of overlap between scans.
mode: {``"constant"``, ``"gaussian"``}
How to blend output of overlapping windows. Defaults to ``"constant"``.
- ``"constant``": gives equal weight to all predictions.
- ``"gaussian``": gives less weight to predictions on edges of windows.
sigma_scale: the standard deviation coefficient of the Gaussian window when `mode` is ``"gaussian"``.
Default: 0.125. Actual window sigma is ``sigma_scale`` * ``dim_size``.
When sigma_scale is a sequence of floats, the values denote sigma_scale at the corresponding
spatial dimensions.
padding_mode: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}
Padding mode for ``inputs``, when ``roi_size`` is larger than inputs. Defaults to ``"constant"``
See also: https://pytorch.org/docs/stable/nn.functional.html#pad
cval: fill value for 'constant' padding mode. Default: 0
sw_device: device for the window data.
By default the device (and accordingly the memory) of the `inputs` is used.
Normally `sw_device` should be consistent with the device where `predictor` is defined.
device: device for the stitched output prediction.
By default the device (and accordingly the memory) of the `inputs` is used. If for example
set to device=torch.device('cpu') the gpu memory consumption is less and independent of the
`inputs` and `roi_size`. Output is on the `device`.
args: optional args to be passed to ``predictor``.
kwargs: optional keyword args to be passed to ``predictor``.
Note:
- input must be channel-first and have a batch dim, supports N-D sliding window.
"""
num_spatial_dims = len(inputs.shape) - 2
assert 0 <= overlap < 1, "overlap must be >= 0 and < 1."
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
if device is None:
device = inputs.device
if sw_device is None:
sw_device = inputs.device
roi_size = fall_back_tuple(roi_size, image_size_)
# in case that image size is smaller than roi size
image_size = tuple(
max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs,
pad=pad_size,
mode=PytorchPadMode(padding_mode).value,
value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims,
overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
num_win = len(slices) # number of windows per image
total_slices = num_win * batch_size # total number of windows
# Create window-level importance map
importance_map = compute_importance_map(get_valid_patch_size(
image_size, roi_size),
mode=mode,
sigma_scale=sigma_scale,
device=device)
importance_map = importance_map.to('cpu')
# Perform predictions
output_image = torch.tensor(0.0) #, device=device)
count_map = torch.tensor(0.0) #, device=device)
_initialized = False
for slice_g in range(0, total_slices, sw_batch_size):
slice_range = range(slice_g, min(slice_g + sw_batch_size,
total_slices))
unravel_slice = [
[slice(int(idx / num_win),
int(idx / num_win) + 1),
slice(None)] + list(slices[idx % num_win]) for idx in slice_range
]
window_data = torch.cat(
[inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
seg_prob = predictor(window_data, *args,
**kwargs).to(device) # batched patch segmentation
seg_prob = seg_prob.to('cpu')
if not _initialized: # init. buffer at the first iteration
output_classes = seg_prob.shape[1]
output_shape = [batch_size, output_classes] + list(image_size)
# allocate memory to store the full output and the count for overlapping parts
output_image = torch.zeros(output_shape,
dtype=torch.float32) #, device=device)
count_map = torch.zeros(output_shape,
dtype=torch.float32) #, device=device)
_initialized = True
# store the result in the proper location of the full output. Apply weights from importance map.
for idx, original_idx in zip(slice_range, unravel_slice):
output_image[original_idx] += importance_map * seg_prob[idx -
slice_g]
count_map[original_idx] += importance_map
# account for any overlapping sections
output_image = output_image / count_map
final_slicing: List[slice] = []
for sp in range(num_spatial_dims):
slice_dim = slice(
pad_size[sp * 2],
image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
final_slicing.insert(0, slice_dim)
while len(final_slicing) < len(output_image.shape):
final_slicing.insert(0, slice(None))
if str(sw_device)[0:3] == 'hpu':
return output_image[final_slicing].to(sw_device).contiguous(
memory_format=torch.channels_last_3d)
return output_image[final_slicing]
def sliding_window_2d_inference_3d(
inputs: torch.Tensor,
roi_size: Union[Sequence[int], int],
sw_batch_size: int,
predictor: Callable[..., torch.Tensor],
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
sigma_scale: Union[Sequence[float], float] = 0.125,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval: float = 0.0,
sw_device: Union[torch.device, str, None] = None,
device: Union[torch.device, str, None] = None,
z_axis: int = 2,
*args: Any,
**kwargs: Any,
) -> torch.Tensor:
num_spatial_dims = len(inputs.shape) - 2
if overlap < 0 or overlap >= 1:
raise AssertionError("overlap must be >= 0 and < 1.")
# determine image spatial size and batch size
# Note: all input images must have the same image size and batch size
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
if device is None:
device = inputs.device
if sw_device is None:
sw_device = inputs.device
roi_size = ensure_list(roi_size)
if len(roi_size) == 2:
roi_size.insert(z_axis, 1)
elif len(roi_size) == 3 and 1 not in roi_size:
raise ValueError(
f'Need one dimension = 1, e.g. (16,16,1), but got {roi_size}')
roi_size = fall_back_tuple(roi_size, image_size_)
# in case that image size is smaller than roi size
image_size = tuple(
max(image_size_[i], roi_size[i]) for i in range(num_spatial_dims))
pad_size = []
for k in range(len(inputs.shape) - 1, 1, -1):
diff = max(roi_size[k - 2] - inputs.shape[k], 0)
half = diff // 2
pad_size.extend([half, diff - half])
inputs = F.pad(inputs,
pad=pad_size,
mode=PytorchPadMode(padding_mode).value,
value=cval)
scan_interval = _get_scan_interval(image_size, roi_size, num_spatial_dims,
overlap)
# Store all slices in list
slices = dense_patch_slices(image_size, roi_size, scan_interval)
num_win = len(slices) # number of windows per image
total_slices = num_win * batch_size # total number of windows
# Create window-level importance map
importance_map = compute_importance_map(get_valid_patch_size(
image_size, roi_size),
mode=mode,
sigma_scale=sigma_scale,
device=device)
# Perform predictions
output_image, count_map = torch.tensor(0.0, device=device), torch.tensor(
0.0, device=device)
_initialized = False
for slice_g in range(0, total_slices, sw_batch_size):
slice_range = range(slice_g, min(slice_g + sw_batch_size,
total_slices))
unravel_slice = [
[slice(int(idx / num_win),
int(idx / num_win) + 1),
slice(None)] + list(slices[idx % num_win]) for idx in slice_range
]
window_data = torch.cat(
[inputs[win_slice] for win_slice in unravel_slice]).to(sw_device)
window_data = window_data.squeeze(z_axis + 2)
seg_prob = predictor(window_data, *args,
**kwargs).to(device) # batched patch segmentation
seg_prob = seg_prob.unsqueeze(z_axis + 2)
if not _initialized: # init. buffer at the first iteration
output_classes = seg_prob.shape[1]
output_shape = [batch_size, output_classes] + list(image_size)
# allocate memory to store the full output and the count for overlapping parts
output_image = torch.zeros(output_shape,
dtype=torch.float32,
device=device)
count_map = torch.zeros(output_shape,
dtype=torch.float32,
device=device)
_initialized = True
# store the result in the proper location of the full output. Apply weights from importance map.
for idx, original_idx in zip(slice_range, unravel_slice):
output_image[original_idx] += importance_map * seg_prob[idx -
slice_g]
count_map[original_idx] += importance_map
# account for any overlapping sections
output_image = output_image / count_map
final_slicing: List[slice] = []
for sp in range(num_spatial_dims):
slice_dim = slice(
pad_size[sp * 2],
image_size_[num_spatial_dims - sp - 1] + pad_size[sp * 2])
final_slicing.insert(0, slice_dim)
while len(final_slicing) < len(output_image.shape):
final_slicing.insert(0, slice(None))
return output_image[final_slicing]
