def randomize(self, img_size):
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 = ensure_tuple(slice(None)) + get_random_patch(img_size, valid_size, self.R)
def __call__(self, img):
"""
Apply the transform to `img`, assuming `img` is channel-first and
slicing doesn't apply to the channel dim.
"""
self.roi_size = fall_back_tuple(self.roi_size, img.shape[1:])
center = [i // 2 for i in img.shape[1:]]
cropper = SpatialCrop(roi_center=center, roi_size=self.roi_size)
return cropper(img)
def _determine_data_pad_width(self, data_shape):
self.spatial_size = fall_back_tuple(self.spatial_size, data_shape)
if self.method == Method.SYMMETRIC:
pad_width = list()
for i in range(len(self.spatial_size)):
width = max(self.spatial_size[i] - data_shape[i], 0)
pad_width.append((width // 2, width - (width // 2)))
return pad_width
else:
return [(0, max(self.spatial_size[i] - data_shape[i], 0)) for i in range(len(self.spatial_size))]
def __call__(self, data):
d = dict(data)
self.randomize()
sp_size = fall_back_tuple(self.rand_affine.spatial_size, data[self.keys[0]].shape[1:])
if self.rand_affine.do_transform:
grid = self.rand_affine.rand_affine_grid(spatial_size=sp_size)
else:
grid = create_grid(spatial_size=sp_size)
for idx, key in enumerate(self.keys):
d[key] = self.rand_affine.resampler(d[key], grid, mode=self.mode[idx], padding_mode=self.padding_mode[idx])
return d
def __call__(self, data):
d = dict(data)
sp_size = fall_back_tuple(self.rand_3d_elastic.spatial_size, data[self.keys[0]].shape[1:])
self.randomize(grid_size=sp_size)
grid = create_grid(spatial_size=sp_size)
if self.rand_3d_elastic.do_transform:
device = self.rand_3d_elastic.device
grid = torch.tensor(grid).to(device)
gaussian = GaussianFilter(spatial_dims=3, sigma=self.rand_3d_elastic.sigma, truncated=3.0).to(device)
offset = torch.tensor(self.rand_3d_elastic.rand_offset, device=device).unsqueeze(0)
grid[:3] += gaussian(offset)[0] * self.rand_3d_elastic.magnitude
grid = self.rand_3d_elastic.rand_affine_grid(grid=grid)
for idx, key in enumerate(self.keys):
d[key] = self.rand_3d_elastic.resampler(
d[key], grid, mode=self.mode[idx], padding_mode=self.padding_mode[idx]
)
return d
def __call__(self, img, mode: Optional[Union[NumpyPadMode, str]] = None):
"""
Args:
img: data to be transformed, assuming `img` is channel-first
and padding doesn't apply to the channel dim.
mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,
``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
One of the listed string values or a user supplied function. Defaults to ``self.mode``.
See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
"""
spatial_shape = img.shape[1:]
k = fall_back_tuple(self.k, (1,) * len(spatial_shape))
new_size = []
for k_d, dim in zip(k, spatial_shape):
new_dim = int(np.ceil(dim / k_d) * k_d) if k_d > 0 else dim
new_size.append(new_dim)
spatial_pad = SpatialPad(spatial_size=new_size, method=Method.SYMMETRIC, mode=mode or self.mode)
return spatial_pad(img)
def __call__(self, data):
d = dict(data)
sp_size = fall_back_tuple(self.rand_2d_elastic.spatial_size, data[self.keys[0]].shape[1:])
self.randomize(spatial_size=sp_size)
if self.rand_2d_elastic.do_transform:
grid = self.rand_2d_elastic.deform_grid(spatial_size=sp_size)
grid = self.rand_2d_elastic.rand_affine_grid(grid=grid)
grid = _torch_interp(
input=grid.unsqueeze(0),
scale_factor=list(self.rand_2d_elastic.deform_grid.spacing),
mode=InterpolateMode.BICUBIC.value,
align_corners=False,
)
grid = CenterSpatialCrop(roi_size=sp_size)(grid[0])
else:
grid = create_grid(spatial_size=sp_size)
for idx, key in enumerate(self.keys):
d[key] = self.rand_2d_elastic.resampler(
d[key], grid, mode=self.mode[idx], padding_mode=self.padding_mode[idx]
)
return d
def generate_pos_neg_label_crop_centers(
label: np.ndarray,
spatial_size,
num_samples: int,
pos_ratio: float,
image: Optional[np.ndarray] = None,
image_threshold: float = 0.0,
rand_state: np.random.RandomState = np.random,
):
"""Generate valid sample locations based on image with option for specifying foreground ratio
Valid: samples sitting entirely within image, expected input shape: [C, H, W, D] or [C, H, W]
Args:
label (numpy.ndarray): use the label data to get the foreground/background information.
spatial_size (sequence of int): spatial size of the ROIs to be sampled.
num_samples: total sample centers to be generated.
pos_ratio: ratio of total locations generated that have center being foreground.
image (numpy.ndarray): if image is not None, use ``label = 0 & image > image_threshold``
to select background. so the crop center will only exist on valid image area.
image_threshold: if enabled image_key, use ``image > image_threshold`` to
determine the valid image content area.
rand_state (random.RandomState): numpy randomState object to align with other modules.
Raises:
ValueError: no sampling location available.
"""
max_size = label.shape[1:]
spatial_size = fall_back_tuple(spatial_size, default=max_size)
if not (np.subtract(max_size, spatial_size) >= 0).all():
raise ValueError("proposed roi is larger than image itself.")
# Select subregion to assure valid roi
valid_start = np.floor_divide(spatial_size, 2)
valid_end = np.subtract(max_size + np.array(1), spatial_size / np.array(2)).astype(np.uint16) # add 1 for random
# int generation to have full range on upper side, but subtract unfloored size/2 to prevent rounded range
# from being too high
for i in range(len(valid_start)): # need this because np.random.randint does not work with same start and end
if valid_start[i] == valid_end[i]:
valid_end[i] += 1
def _correct_centers(center_ori, valid_start, valid_end):
for i, c in enumerate(center_ori):
center_i = c
if c < valid_start[i]:
center_i = valid_start[i]
if c >= valid_end[i]:
center_i = valid_end[i] - 1
center_ori[i] = center_i
return center_ori
centers = []
# Prepare fg/bg indices
if label.shape[0] > 1:
label = label[1:] # for One-Hot format data, remove the background channel
label_flat = np.any(label, axis=0).ravel() # in case label has multiple dimensions
fg_indices = np.nonzero(label_flat)[0]
if image is not None:
img_flat = np.any(image > image_threshold, axis=0).ravel()
bg_indices = np.nonzero(np.logical_and(img_flat, ~label_flat))[0]
else:
bg_indices = np.nonzero(~label_flat)[0]
if not len(fg_indices) or not len(bg_indices):
if not len(fg_indices) and not len(bg_indices):
raise ValueError("no sampling location available.")
warnings.warn(
f"N foreground {len(fg_indices)}, N background {len(bg_indices)},"
"unable to generate class balanced samples."
)
pos_ratio = 0 if not len(fg_indices) else 1
for _ in range(num_samples):
indices_to_use = fg_indices if rand_state.rand() < pos_ratio else bg_indices
random_int = rand_state.randint(len(indices_to_use))
center = np.unravel_index(indices_to_use[random_int], label.shape)
center = center[1:]
# shift center to range of valid centers
center_ori = [c for c in center]
centers.append(_correct_centers(center_ori, valid_start, valid_end))
return centers
def sliding_window_inference(
inputs: Union[torch.Tensor, tuple],
roi_size,
sw_batch_size: int,
predictor: Callable,
overlap: float = 0.25,
mode: Union[BlendMode, str] = BlendMode.CONSTANT,
padding_mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
cval=0,
uncertainty_flag=False):
"""
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 (list, tuple): 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
Raises:
NotImplementedError: inputs must 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.
"""
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
inputs_type = type(inputs)
if inputs_type == tuple:
phys_inputs = inputs[1]
inputs = inputs[0]
num_spatial_dims = len(inputs.shape) - 2
image_size_ = list(inputs.shape[2:])
batch_size = inputs.shape[0]
# TODO: Enable batch sizes > 1 in future
if batch_size > 1:
raise NotImplementedError("inputs must have batch_size=1.")
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)
# print(f'The slices are {slices}')
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
if not uncertainty_flag:
# No uncertainty, so only one prediction, so proceed normally
output_rois = list()
for data in slice_batches:
if not uncertainty_flag and inputs_type == tuple:
seg_prob, _ = predictor(
data, phys_inputs) # batched patch segmentation
output_rois.append(seg_prob)
elif inputs_type != tuple:
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=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:
# print(output_image.shape, curr_slice, importance_map.shape, output_rois[window_id].shape)
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]: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
else:
# Decide on number of histogram samples
num_hist_samples = 20
overall_stochastic_logits_hist = torch.empty(
(1, 2, 181, 217, 181, num_hist_samples))
overall_true_seg_net_out_hist = torch.empty(
(1, 2, 181, 217, 181, num_hist_samples))
output_rois = list()
unc_output_rois = list()
# Have uncertainty, therefore have MANY outputs, but only have ONE pass through network
for data in slice_batches:
if inputs_type == tuple:
seg_prob, unc_prob, _ = predictor(
data, phys_inputs) # batched patch segmentation
output_rois.append(seg_prob)
unc_output_rois.append(unc_prob)
elif inputs_type != tuple:
seg_prob, unc_prob, _ = predictor(
data) # batched patch segmentation
output_rois.append(seg_prob)
unc_output_rois.append(unc_prob)
# Get shape of logits
logits_shape = list(seg_prob.shape)
# Now want an array of randomly normally distributed samples size of logits x num samples
# logits_shape.append(num_hist_samples)
inf_ax = torch.distributions.Normal(
torch.tensor(0.0).to(device=torch.device("cuda:0")),
torch.tensor(1.0).to(device=torch.device("cuda:0")))
# inf_noise_array = torch.empty(logits_shape).normal_(mean=0, std=1)
# Loop through samples
for infpass in range(num_hist_samples):
true_output_rois = list()
true_unc_output_rois = list()
# print(f'The lengths of rois are {len(output_rois)}, {len(unc_output_rois)}')
for roi, unc_roi in zip(output_rois, unc_output_rois):
# output_rois = list()
# unc_output_rois = list()m
# Repeat steps above to get more samples
# noise_sample = inf_noise_array[..., infpass]
stochastic_logits = roi + unc_roi * inf_ax.sample(
logits_shape) # noise_sample
# print(f'The sigma mean is {torch.mean(unc_roi)}, logits mean is {torch.mean(roi)}')
# print(
# f'The logits, sigma, ax sizes are: {roi.shape}, {unc_roi.shape}, {inf_ax.sample(logits_shape).shape}')
# print(
# f'A little ax check: {inf_ax.sample(logits_shape)[0, 0, 0, 0, 0]}, {inf_ax.sample(logits_shape)[0, 1, 0, 0, 0]}')
true_seg_net_out = torch.softmax(stochastic_logits, dim=1)
# print(f'The stochastic logits shapes are {stochastic_logits.shape}, {true_seg_net_out.shape}')
true_output_rois.append(true_seg_net_out)
true_unc_output_rois.append(stochastic_logits)
# stitching output image
# print(f'The true output rois tensor shapes are {true_output_rois[0].shape}')
output_classes = true_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=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)
# slic_index, zero to len(slices) in increments of sw_batch_size
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:
# print(output_image.shape, curr_slice, importance_map.shape, true_output_rois[window_id].shape)
output_image[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += (
importance_map *
true_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 *
true_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:
output_image = 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], ]
overall_true_seg_net_out_hist[..., infpass] = output_image
else:
output_image = output_image[..., pad_size[2]:image_size_[0] +
pad_size[2],
pad_size[0]:image_size_[1] +
pad_size[0]] # 2D
overall_true_seg_net_out_hist[..., infpass] = output_image
# Uncertainty part
# stitching output image
output_classes = true_unc_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=inputs.device)
# allocate memory to store the full output and the count for overlapping parts
unc_output = 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:
unc_output[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += (
importance_map *
true_unc_output_rois[window_id][
curr_index - slice_index, :])
count_map[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += importance_map
else:
unc_output[0, :, curr_slice[0], curr_slice[1]] += (
importance_map *
true_unc_output_rois[window_id][curr_index -
slice_index, :])
count_map[0, :, curr_slice[0],
curr_slice[1]] += importance_map
# account for any overlapping sections
unc_output /= count_map
if num_spatial_dims == 3:
unc_output = unc_output[..., 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], ]
overall_stochastic_logits_hist[..., infpass] = unc_output
else:
unc_output = unc_output[..., pad_size[2]:image_size_[0] +
pad_size[2],
pad_size[0]:image_size_[1] +
pad_size[0]] # 2D
overall_stochastic_logits_hist[..., infpass] = unc_output
# Sigma part
# stitching output image
output_classes = unc_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=inputs.device)
# allocate memory to store the full output and the count for overlapping parts
sigma_output = 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:
sigma_output[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += (
importance_map *
unc_output_rois[window_id][curr_index -
slice_index, :]
)
count_map[0, :, curr_slice[0], curr_slice[1],
curr_slice[2]] += importance_map
else:
sigma_output[0, :, curr_slice[0], curr_slice[1]] += (
importance_map *
unc_output_rois[window_id][curr_index -
slice_index, :])
count_map[0, :, curr_slice[0],
curr_slice[1]] += importance_map
# account for any overlapping sections
sigma_output /= count_map
if num_spatial_dims == 3:
sigma_output = sigma_output[..., 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], ]
else:
sigma_output = sigma_output[..., pad_size[2]:image_size_[0] +
pad_size[2],
pad_size[0]:image_size_[1] +
pad_size[0]] # 2D
return overall_true_seg_net_out_hist, overall_stochastic_logits_hist, sigma_output
def randomize(self, label, image):
self.spatial_size = fall_back_tuple(self.spatial_size,
default=label.shape[1:])
self.centers = generate_pos_neg_label_crop_centers(
label, self.spatial_size, self.num_samples, self.pos_ratio, image,
self.image_threshold, self.R)