biapy.data.generators.pair_base_data_generator

Base data generator for paired image and mask data in BiaPy.

This module provides the PairBaseDataGenerator class, which supports flexible data loading, augmentation, and normalization for deep learning workflows. It includes a wide range of augmentation options for both 2D and 3D data, and is designed to work with BiaPyDataset objects and normalization modules.

class biapy.data.generators.pair_base_data_generator.PairBaseDataGenerator(ndim: int, X: BiaPyDataset, Y: BiaPyDataset, norm_module: Dict, seed: int = 0, da: bool = True, da_prob: float = 0.5, rotation90: bool = False, rand_rot: bool = False, rnd_rot_range: Tuple[int, int] = (-180, 180), shear: bool = False, shear_range: Tuple[int, int] = (-20, 20), zoom: bool = False, zoom_range: Tuple[float, float] = (0.8, 1.2), zoom_in_z: bool = False, shift: bool = False, shift_range: Tuple[float, float] = (0.1, 0.2), affine_mode: Literal['constant', 'edge', 'symmetric', 'reflect', 'wrap'] = 'constant', vflip: bool = False, hflip: bool = False, elastic: bool = False, e_alpha: Tuple[int, int] = (240, 250), e_sigma: int = 25, e_mode: Literal['constant', 'edge', 'symmetric', 'reflect', 'wrap'] = 'constant', g_blur: bool = False, g_sigma: Tuple[float, float] = (1.0, 2.0), median_blur: bool = False, mb_kernel: Tuple[int, int] = (3, 7), motion_blur: bool = False, motb_k_range: Tuple[int, int] = (3, 8), gamma_contrast: bool = False, gc_gamma: Tuple[float, float] = (1.25, 1.75), brightness: bool = False, brightness_factor: Tuple[int, int] = (1, 3), brightness_mode: str = '2D', contrast: bool = False, contrast_factor: Tuple[int, int] = (1, 3), contrast_mode: str = '2D', dropout: bool = False, drop_range: Tuple[float, float] = (0.0, 0.2), cutout: bool = False, cout_nb_iterations: Tuple[int, int] = (1, 3), cout_size: Tuple[float, float] = (0.2, 0.4), cout_cval: int = 0, cout_apply_to_mask: bool = False, cutblur: bool = False, cblur_size: Tuple[float, float] = (0.1, 0.5), cblur_down_range: Tuple[int, int] = (2, 8), cblur_inside: bool = True, cutmix: bool = False, cmix_size: Tuple[float, float] = (0.2, 0.4), cutnoise: bool = False, cnoise_scale: Tuple[float, float] = (0.1, 0.2), cnoise_nb_iterations: Tuple[int, int] = (1, 3), cnoise_size: Tuple[float, float] = (0.2, 0.4), misalignment: bool = False, ms_displacement: int = 16, ms_rotate_ratio: float = 0.0, missing_sections: bool = False, missp_iterations: Tuple[int, int] = (30, 40), missp_channel_pb: float = 0.5, grayscale: bool = False, channel_shuffle: bool = False, gridmask: bool = False, grid_ratio: float = 0.6, grid_d_range: Tuple[float, float] = (0.4, 1.0), grid_rotate: int = 1, grid_invert: bool = False, gaussian_noise: bool = False, gaussian_noise_mean: int = 0, gaussian_noise_var: float = 0.01, gaussian_noise_use_input_img_mean_and_var: bool = False, poisson_noise: bool = False, salt: bool = False, salt_amount: float = 0.05, pepper: bool = False, pepper_amount: float = 0.05, salt_and_pepper: bool = False, salt_pep_amount: float = 0.05, salt_pep_proportion: float = 0.5, shape: Tuple[int, int, int] = (256, 256, 1), resolution: Tuple[int, ...] = (-1,), val: bool = False, n_classes: int = 1, ignore_index: int | None = None, n2v: bool = False, n2v_perc_pix: float = 0.198, n2v_manipulator='uniform_withCP', n2v_neighborhood_radius: int = 5, n2v_structMask=array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]]), n2v_load_gt: bool = False, instance_problem: bool = False, random_crop_scale: Tuple[int, ...] = (1, 1), convert_to_rgb: bool = False, preprocess_f=None, preprocess_cfg=None)[source]

Bases: Dataset

Custom BaseDataGenerator to transform paired image and mask data.

Parameters:

ndim (int) – Dimensions of the data (2 for 2D and 3 for 3D).
X (BiaPyDataset) – X dataset.
Y (BiaPyDataset) – Y dataset.
norm_module (Dict) – Normalization module that defines the normalization steps to apply.
seed (int, optional) – Seed for random functions.
da (bool, optional) – To activate the data augmentation.
da_prob (float, optional) – Probability of doing each transformation.
rotation90 (bool, optional) – To make square (90, 180,270) degree rotations.
rand_rot (bool, optional) – To make random degree range rotations.
rnd_rot_range (tuple of float, optional) – Range of random rotations. E. g. (-180, 180).
shear (bool, optional) – To make shear transformations.
shear_range (tuple of int, optional) – Degree range to make shear. E. g. (-20, 20).
zoom (bool, optional) – To make zoom on images.
zoom_range (tuple of floats, optional) – Zoom range to apply. E. g. (0.8, 1.2).
zoom_in_z (bool, optional) – Whether to apply or not zoom in Z axis.
shift (float, optional) – To make shifts.
shift_range (tuple of float, optional) – Range to make a shift. E. g. (0.1, 0.2).
affine_mode (str, optional) – Method to use when filling in newly created pixels. Same meaning as in skimage (and numpy.pad()). E.g. constant, reflect etc.
vflip (bool, optional) – To activate vertical flips.
hflip (bool, optional) – To activate horizontal flips.
elastic (bool, optional) – To make elastic deformations.
e_alpha (tuple of ints, optional) – Strength of the distortion field. E. g. (240, 250).
e_sigma (int, optional) – Standard deviation of the gaussian kernel used to smooth the distortion fields.
e_mode (str, optional) – Parameter that defines the handling of newly created pixels with the elastic transformation.
g_blur (bool, optional) – To insert gaussian blur on the images.
g_sigma (tuple of floats, optional) – Standard deviation of the gaussian kernel. E. g. (1.0, 2.0).
median_blur (bool, optional) – To blur an image by computing median values over neighbourhoods.
mb_kernel (tuple of ints, optional) – Median blur kernel size. E. g. (3, 7).
motion_blur (bool, optional) – Blur images in a way that fakes camera or object movements.
motb_k_range (int, optional) – Kernel size to use in motion blur.
gamma_contrast (bool, optional) – To insert gamma constrast changes on images.
gc_gamma (tuple of floats, optional) – Exponent for the contrast adjustment. Higher values darken the image. E. g. (1.25, 1.75).
brightness (bool, optional) – To aply brightness to the images as PyTorch Connectomics.
brightness_factor (tuple of 2 floats, optional) – Strength of the brightness range, with valid values being 0 <= brightness_factor <= 1. E.g. (0.1, 0.3).
contrast (boolen, optional) – To apply contrast changes to the images as PyTorch Connectomics.
contrast_factor (tuple of 2 floats, optional) – Strength of the contrast change range, with valid values being 0 <= contrast_factor <= 1. E.g. (0.1, 0.3).
dropout (bool, optional) – To set a certain fraction of pixels in images to zero.
drop_range (tuple of floats, optional) – Range to take a probability p to drop pixels. E.g. (0, 0.2) will take a p folowing 0<=p<=0.2 and then drop p percent of all pixels in the image (i.e. convert them to black pixels).
cutout (bool, optional) – To fill one or more rectangular areas in an image using a fill mode.
cout_nb_iterations (tuple of ints, optional) – Range of number of areas to fill the image with. E. g. (1, 3).
cout_size (tuple of floats, optional) – Range to select the size of the areas in % of the corresponding image size. Values between 0 and 1. E. g. (0.2, 0.4).
cout_cval (int, optional) – Value to fill the area of cutout with.
cout_apply_to_mask (boolen, optional) – Whether to apply cutout to the mask.
cutblur (boolean, optional) – Blur a rectangular area of the image by downsampling and upsampling it again.
cblur_size (tuple of floats, optional) – Range to select the size of the area to apply cutblur on. E. g. (0.2, 0.4).
cblur_inside (boolean, optional) – If True only the region inside will be modified (cut LR into HR image). If False the 50% of the times the region inside will be modified (cut LR into HR image) and the other 50% the inverse will be done (cut HR into LR image). See Figure 1 of the official paper.
cutmix (boolean, optional) – Combine two images pasting a region of one image to another.
cmix_size (tuple of floats, optional) – Range to select the size of the area to paste one image into another. E. g. (0.2, 0.4).
cnoise (boolean, optional) – Randomly add noise to a cuboid region in the image.
cnoise_scale (tuple of floats, optional) – Range to choose a value that will represent the % of the maximum value of the image that will be used as the std of the Gaussian Noise distribution. E.g. (0.1, 0.2).
cnoise_nb_iterations (tuple of ints, optional) – Number of areas with noise to create. E.g. (1, 3).
cnoise_size (tuple of floats, optional) – Range to choose the size of the areas to transform. E.g. (0.2, 0.4).
misalignment (boolean, optional) – To add miss-aligment augmentation.
ms_displacement (int, optional) – Maximum pixel displacement in xy-plane for misalignment.
ms_rotate_ratio (float, optional) – Ratio of rotation-based mis-alignment
missing_sections (boolean, optional) – Augment the image by creating a black line in a random position.
missp_iterations (tuple of 2 ints, optional) – Iterations to dilate the missing line with. E.g. (30, 40).
missp_channel_pb (float, optional) – Probability of applying missing section to each channel. E.g. 0.5.
grayscale (bool, optional) – Whether to augment images converting partially in grayscale.
gridmask (bool, optional) – Whether to apply gridmask to the image. See the official paper for more information about it and its parameters.
grid_ratio (float, optional) – Determines the keep ratio of an input image (r in the original paper).
grid_d_range (tuple of floats, optional) – Range to choose a d value. It represents the % of the image size. E.g. (0.4,1).
grid_rotate (float, optional) – Rotation of the mask in GridMask. Needs to be between [0,1] where 1 is 360 degrees.
grid_invert (bool, optional) – Whether to invert the mask of GridMask.
channel_shuffle (bool, optional) – Whether to shuflle the channels of the images.
gaussian_noise (bool, optional) – To apply Gaussian noise to the images.
gaussian_noise_mean (tuple of ints, optional) – Mean of the Gaussian noise.
gaussian_noise_var (tuple of ints, optional) – Variance of the Gaussian noise.
gaussian_noise_use_input_img_mean_and_var (bool, optional) – Whether to use the mean and variance of the input image instead of gaussian_noise_mean and gaussian_noise_var.
poisson_noise (bool, optional) – To apply Poisson noise to the images.
salt (tuple of ints, optional) – Mean of the gaussian noise.
salt_amount (tuple of ints, optional) – Variance of the gaussian noise.
pepper (bool, optional) – To apply poisson noise to the images.
pepper_amount (tuple of ints, optional) – Mean of the gaussian noise.
salt_and_pepper (bool, optional) – To apply poisson noise to the images.
salt_pep_amount (tuple of ints, optional) – Mean of the gaussian noise.
salt_pep_proportion (bool, optional) – To apply poisson noise to the images.
resolution (2D tuple of floats, optional) – Resolution of the given data (y,x). E.g. (8,8).
val (bool, optional) – Advise the generator that the images will be to validate the model to not make random crops (as the val. data must be the same on each epoch).
n_classes (int, optional) – Number of classes.
ignore_index (int, optional) – Value to ignore in the loss/metrics.
n2v (bool, optional) – Whether to create Noise2Void mask. Used in DENOISING problem type.
n2v_perc_pix (float, optional) – Input image pixels to be manipulated.
n2v_manipulator (str, optional) – How to manipulate the input pixels. Most pixel manipulators will compute the replacement value based on a neighborhood. Possible options: normal_withoutCP: samples the neighborhood according to a normal gaussian distribution, but without the center pixel; normal_additive: adds a random number to the original pixel value. The random number is sampled from a gaussian distribution with zero-mean and sigma = n2v_neighborhood_radius ; normal_fitted: uses a random value from a gaussian normal distribution with mean equal to the mean of the neighborhood and standard deviation equal to the standard deviation of the neighborhood ; identity: performs no pixel manipulation.
n2v_neighborhood_radius (int, optional) – Neighborhood size to use when manipulating the values.
n2v_structMask (Array of ints, optional) – Masking kernel for StructN2V to hide pixels adjacent to main blind spot. Value 1 = ‘hidden’, Value 0 = ‘non hidden’. Nested lists equivalent to ndarray. Must have odd length in each dimension (center pixel is blind spot). None implies normal N2V masking.
instance_problem (bool, optional) – Advice the class that the workflow is of instance segmentation to divide the labels by channels.
random_crop_scale (tuple of ints, optional) – Scale factor the mask used in super-resolution workflow. E.g. (2,2).
convert_to_rgb (bool, optional) – In case RGB images are expected, e.g. if crop_shape channel is 3, those images that are grayscale are converted into RGB.
preprocess_f (function, optional) – The preprocessing function, is necessary in case you want to apply any preprocessing.
preprocess_cfg (dict, optional) – Configuration parameters for preprocessing, is necessary in case you want to apply any preprocessing.

abstract save_aug_samples(img: ndarray[tuple[int, ...], dtype[_ScalarType_co]], mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]], orig_images: Dict, i: int, pos: int, out_dir: str)[source]

Save transformed samples in order to check the generator.

Parameters:

img (3D/4D Numpy array) – Image to use as sample. E.g. (y, x, channels) for 2D and (z, y, x, channels) for 3D.
mask (3D/4D Numpy array) – Mask to use as sample. E.g. (y, x, channels) for 2D and (z, y, x, channels) for 3D.
orig_images (dict) – Dict where the original image and mask are saved in “o_x” and “o_y”, respectively.
i (int) – Number of the sample within the transformed ones.
pos (int) – Number of the sample within the dataset.
out_dir (str) – Directory to save the images.

load_sample(_idx: int, first_load: bool = False) → Tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]]][source]

Load one data sample given its corresponding index.

Parameters:

_idx (int) – Sample index counter.
first_load (bool, optional) – Whether its the first time a sample is loaded to prevent normalizing it.

Returns:

img (3D/4D Numpy array) – X element. E.g. (y, x, channels) in 2D and (z, y, x, channels) in 3D.
mask (3D/4D Numpy array) – Y element. E.g. (y, x, channels) in 2D and (z, y, x, channels) in 3D.

getitem(index: int) → Tuple[Tensor, Tensor][source]

Generate one pair of data.

Parameters:: index (int) – Index counter.
Returns:: item – X and Y (if avail) elements. Each one shape is (z, y, x, channels) if 2D or (y, x, channels) if 3D.
Return type:: 3D/4D Torch tensors

apply_transform(image: ndarray[tuple[int, ...], dtype[_ScalarType_co]], mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]], e_im: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | None, e_mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]] | None) → Tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]]][source]

Transform the input image and its mask at the same time with one of the selected choices based on a probability.

Parameters:

image (3D/4D Numpy array) – Image to transform. E.g. (y, x, channels) in 2D and (y, x, z, channels) in 3D.
mask (3D/4D Numpy array) – Mask to transform. E.g. (y, x, channels) in 2D and (y, x, z, channels) in 3D.
e_im (3D/4D Numpy array) – Extra image to help transforming image. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.
e_mask (3D/4D Numpy array) – Extra mask to help transforming mask. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.

Returns:

image (3D/4D Numpy array) – Transformed image. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.
mask (3D/4D Numpy array) – Transformed image mask. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.

get_transformed_samples(num_examples: int, random_images: bool = True, save_to_dir: bool = True, out_dir: str = 'aug', train: bool = False, draw_grid: bool = True)[source]

Apply selected transformations to a defined number of images from the dataset.

Parameters:

num_examples (int) – Number of examples to generate.
random_images (bool, optional) – Randomly select images from the dataset. If False the examples will be generated from the start of the dataset.
save_to_dir (bool, optional) – Save the images generated. The purpose of this variable is to check the images generated by data augmentation.
out_dir (str, optional) – Name of the folder where the examples will be stored.
train (bool, optional) – To avoid drawing a grid on the generated images. This should be set when the samples will be used for training.
draw_grid (bool, optional) – Draw a grid in the generated samples. Useful to see some types of deformations.

Returns:

sample_x (List of 3D/4D Numpy array) – Transformed images. E.g. list of (y, x, channels) in 2D and (z, y, x, channels) in 3D.
sample_y (List of 3D/4D Numpy array) – Transformed image mask. E.g. list of (y, x, channels) in 2D and (z, y, x, channels) in 3D.

Examples

Setting elastic=True an example output should be similar to the following:

Original crop	Original crop mask
Elastic transformation applied	Elastic transformation applied

The grid is only painted if train=False which should be used just to display transformations made. Selecting random rotations between 0 and 180 degrees should generate the following:

Original crop	Original crop mask
Random rotation [0, 180] applied	Random rotation [0, 180] applied

draw_grid(im: ndarray[tuple[int, ...], dtype[_ScalarType_co]], grid_width: int | None = None) → ndarray[tuple[int, ...], dtype[_ScalarType_co]][source]

Draw grid of the specified size on an image.

Parameters:

im (3D/4D Numpy array) – Image to draw the grid into. E.g. (y, x, channels) in 2D or (z, y, x, channels) in 3D.
grid_width (int, optional) – Grid’s width.

prepare_n2v(_img: ndarray[tuple[int, ...], dtype[_ScalarType_co]], _mask: ndarray[tuple[int, ...], dtype[_ScalarType_co]]) → Tuple[ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]]][source]

Create Noise2Void mask.

Parameters:

_img (3D/4D Numpy array) – Image to wipe some pixels from. E.g. (y, x, channels) in 2D or (z, y, x, channels) in 3D.
_mask (3D/4D Numpy array) – Mask to use values from. Only used when the ground truth is loaded with self.n2v_load_gt. E.g. (y, x, channels) in 2D or (z, y, x, channels) in 3D.

Returns:

img (3D/4D Numpy array) – Input image modified removing some pixels. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.
mask (3D/4D Numpy array) – Noise2Void mask created. E.g. (y, x, channels) in 2D or (y, x, z, channels) in 3D.