# Copyright 2022 Cerebras Systems.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Adapted from: https://github.com/MIC-DKFZ/batchgenerators (commit id: 01f225d)
#
# Copyright 2021 Division of Medical Image Computing, German Cancer Research Center (DKFZ)
# and Applied Computer Vision Lab, Helmholtz Imaging Platform
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from builtins import range
import numpy as np
from cerebras.modelzoo.data.internal.vision.segmentation.transforms.augmentation_utils import (
uniform,
)
from cerebras.modelzoo.data.vision.segmentation.transforms.utils import ( # original nnUnet used skimage.resize
nd_resize,
)
[docs]def augment_linear_downsampling_scipy(
data_sample,
zoom_range=(0.5, 1),
per_channel=True,
p_per_channel=1,
channels=None,
order_downsample=1,
order_upsample=0,
ignore_axes=None,
):
'''
Downsamples each sample (linearly) by a random factor and upsamples to original resolution again (nearest neighbor)
Info:
* Uses scipy zoom for resampling. A bit faster than nilearn.
* Resamples all dimensions (channels, x, y, z) with same downsampling factor (like isotropic=True from
linear_downsampling_generator_nilearn)
Args:
zoom_range: can be either tuple/list/np.ndarray or tuple of tuple. If tuple/list/np.ndarray, then the zoom
factor will be sampled from zoom_range[0], zoom_range[1] (zoom < 0 = downsampling!). If tuple of tuple then
each inner tuple will give a sampling interval for each axis (allows for different range of zoom values for
each axis
p_per_channel: probability for downsampling/upsampling a channel
per_channel (bool): whether to draw a new zoom_factor for each channel or keep one for all channels
channels (list, tuple): if None then all channels can be augmented. If list then only the channel indices can
be augmented (but may not always be depending on p_per_channel)
order_downsample:
order_upsample:
ignore_axes: tuple/list
'''
if not isinstance(zoom_range, (list, tuple, np.ndarray)):
zoom_range = [zoom_range]
shp = np.array(data_sample.shape[1:])
dim = len(shp)
if not per_channel:
if isinstance(zoom_range[0], (tuple, list, np.ndarray)):
assert len(zoom_range) == dim
zoom = np.array([uniform(i[0], i[1]) for i in zoom_range])
else:
zoom = uniform(zoom_range[0], zoom_range[1])
target_shape = np.round(shp * zoom).astype(int)
if ignore_axes is not None:
for i in ignore_axes:
target_shape[i] = shp[i]
if channels is None:
channels = list(range(data_sample.shape[0]))
for c in channels:
if np.random.uniform() < p_per_channel:
if per_channel:
if isinstance(zoom_range[0], (tuple, list, np.ndarray)):
assert len(zoom_range) == dim
zoom = np.array([uniform(i[0], i[1]) for i in zoom_range])
else:
zoom = uniform(zoom_range[0], zoom_range[1])
target_shape = np.round(shp * zoom).astype(int)
if ignore_axes is not None:
for i in ignore_axes:
target_shape[i] = shp[i]
downsampled = nd_resize(
data_sample[c].astype(float),
target_shape,
order=order_downsample,
mode='edge',
anti_aliasing=False,
)
data_sample[c] = nd_resize(
downsampled,
shp,
order=order_upsample,
mode='edge',
anti_aliasing=False,
)
return data_sample