ImageAugmentation

class lib.training.data.augmentation.ImageAugmentation(batch_size: int, processing_size: int)

Bases: object

Performs augmentation on batches of training images.

Parameters:
  • batch_size (int) – The number of images that will be fed through the augmentation functions at once.

  • processing_size (int) – The largest input or output size of the model. This is the size that images are processed at.

Methods Summary

color_adjust(batch)

Perform color augmentation on the passed in batch.

random_flip(batch, points)

Perform random horizontal flipping on the passed in batch.

transform(batch, points)

Perform random transformation on the passed in batch and optional (x, y) points.

warp(batch[, to_landmarks, ...])

Perform random warping on the passed in batch by one of two methods.

Methods Documentation

color_adjust(batch: ndarray) ndarray

Perform color augmentation on the passed in batch.

The color adjustment parameters are set in config.train.ini

Parameters:

batch (ndarray) – The batch should be a 4-dimensional array of shape (batchsize, height, width, 3) and in BGR format of uint8 dtype.

Return type:

A 4-dimensional array of the same shape as batch with color augmentation applied.

random_flip(batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None) None

Perform random horizontal flipping on the passed in batch.

The probability of flipping an image is set in config.train.ini

Parameters:
  • batch (npt.NDArray[np.uint8]) – The batch should be a 4-dimensional array of shape (batchsize, height, width, channels) and in BGR format.

  • points (npt.NDArray[np.float32] | None) – Any (x, y) points to transform. Can be in any shape but the final dimension should be shape 2. None if there are no points to transform

Return type:

None

transform(batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None) None

Perform random transformation on the passed in batch and optional (x, y) points.

The transformation parameters are set in config.train.ini

Parameters:
  • batch (npt.NDArray[np.uint8]) – The batch should be a 4-dimensional array of shape (batchsize, height, width, channels) and in BGR format.

  • points (npt.NDArray[np.float32] | None) – Any (x, y) points to transform. in shape (batch_size, num_sides, 68, 2). None if there are no points to transform

Return type:

None

warp(batch: ndarray, to_landmarks: bool = False, batch_src_points: ndarray | None = None, batch_dst_points: ndarray | None = None) ndarray

Perform random warping on the passed in batch by one of two methods.

Parameters:
  • batch (ndarray) – The batch should be a 4-dimensional array of shape (batchsize, height, width, 3) and in BGR format.

  • to_landmarks (bool) – If False perform standard random warping of the input image. If True perform warping to semi-random similar corresponding landmarks from the other side. Default: False

  • batch_src_points (ndarray | None) – Only used when to_landmarks is True. A batch of 68 point landmarks for the source faces. This is a 3-dimensional array in the shape (batchsize, 68, 2). Default: None

  • batch_dst_points (ndarray | None) – Only used when to_landmarks is True. A batch of randomly chosen closest match destination faces landmarks. This is a 3-dimensional array in the shape (batchsize, 68, 2). Default None

Return type:

A 4-dimensional array of the same shape as batch with warping applied.

color_adjust(batch: ndarray) ndarray

Perform color augmentation on the passed in batch.

The color adjustment parameters are set in config.train.ini

Parameters:

batch (ndarray) – The batch should be a 4-dimensional array of shape (batchsize, height, width, 3) and in BGR format of uint8 dtype.

Return type:

A 4-dimensional array of the same shape as batch with color augmentation applied.

random_flip(batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None) None

Perform random horizontal flipping on the passed in batch.

The probability of flipping an image is set in config.train.ini

Parameters:
  • batch (npt.NDArray[np.uint8]) – The batch should be a 4-dimensional array of shape (batchsize, height, width, channels) and in BGR format.

  • points (npt.NDArray[np.float32] | None) – Any (x, y) points to transform. Can be in any shape but the final dimension should be shape 2. None if there are no points to transform

Return type:

None

transform(batch: npt.NDArray[np.uint8], points: npt.NDArray[np.float32] | None) None

Perform random transformation on the passed in batch and optional (x, y) points.

The transformation parameters are set in config.train.ini

Parameters:
  • batch (npt.NDArray[np.uint8]) – The batch should be a 4-dimensional array of shape (batchsize, height, width, channels) and in BGR format.

  • points (npt.NDArray[np.float32] | None) – Any (x, y) points to transform. in shape (batch_size, num_sides, 68, 2). None if there are no points to transform

Return type:

None

warp(batch: ndarray, to_landmarks: bool = False, batch_src_points: ndarray | None = None, batch_dst_points: ndarray | None = None) ndarray

Perform random warping on the passed in batch by one of two methods.

Parameters:
  • batch (ndarray) – The batch should be a 4-dimensional array of shape (batchsize, height, width, 3) and in BGR format.

  • to_landmarks (bool) – If False perform standard random warping of the input image. If True perform warping to semi-random similar corresponding landmarks from the other side. Default: False

  • batch_src_points (ndarray | None) – Only used when to_landmarks is True. A batch of 68 point landmarks for the source faces. This is a 3-dimensional array in the shape (batchsize, 68, 2). Default: None

  • batch_dst_points (ndarray | None) – Only used when to_landmarks is True. A batch of randomly chosen closest match destination faces landmarks. This is a 3-dimensional array in the shape (batchsize, 68, 2). Default None

Return type:

A 4-dimensional array of the same shape as batch with warping applied.