Conduct elastic transform to images.

  Elastic transform is performed on small patches of the images.

  Args:
    x: numpy array, uncorrupted image, assumed to have uint8 pixel in [0,255].
    severity: integer, severity of corruption.

  Returns:
    numpy array, image with uint8 pixels in [0,255]. Applied elastic transform.

def elastic(x, severity=1):
  """Conduct elastic transform to images.

  Elastic transform is performed on small patches of the images.

  Args:
    x: numpy array, uncorrupted image, assumed to have uint8 pixel in [0,255].
    severity: integer, severity of corruption.

  Returns:
    numpy array, image with uint8 pixels in [0,255]. Applied elastic transform.
  """
  c = [(244 * 2, 244 * 0.7, 244 * 0.1), (244 * 2, 244 * 0.08, 244 * 0.2),
       (244 * 0.05, 244 * 0.01, 244 * 0.02), (244 * 0.07, 244 * 0.01,
                                              244 * 0.02),
       (244 * 0.12, 244 * 0.01, 244 * 0.02)][severity - 1]

  image = np.array(x, dtype=np.float32) / 255.
  shape = image.shape
  shape_size = shape[:2]

  # random affine
  center_square = np.float32(shape_size) // 2
  square_size = min(shape_size) // 3
  pts1 = np.float32([
      center_square + square_size,
      [center_square[0] + square_size, center_square[1] - square_size],
      center_square - square_size
  ])
  pts2 = pts1 + np.random.uniform(
      -c[2], c[2], size=pts1.shape).astype(np.float32)
  affine_trans = tfds.core.lazy_imports.cv2.getAffineTransform(pts1, pts2)
  image = tfds.core.lazy_imports.cv2.warpAffine(
      image,
      affine_trans,
      shape_size[::-1],
      borderMode=tfds.core.lazy_imports.cv2.BORDER_REFLECT_101)

  dx = (tfds.core.lazy_imports.skimage.filters.gaussian(
      np.random.uniform(-1, 1, size=shape[:2]),
      c[1],
      mode='reflect',
      truncate=3) * c[0]).astype(np.float32)
  dy = (tfds.core.lazy_imports.skimage.filters.gaussian(
      np.random.uniform(-1, 1, size=shape[:2]),
      c[1],
      mode='reflect',
      truncate=3) * c[0]).astype(np.float32)
  dx, dy = dx[..., np.newaxis], dy[..., np.newaxis]

  x, y, z = np.meshgrid(
      np.arange(shape[1]), np.arange(shape[0]), np.arange(shape[2]))
  indices = np.reshape(y + dy, (-1, 1)), np.reshape(x + dx,
                                                    (-1, 1)), np.reshape(
                                                        z, (-1, 1))
  x_clip = np.clip(
      tfds.core.lazy_imports.scipy.ndimage.interpolation.map_coordinates(
          image, indices, order=1, mode='reflect').reshape(shape), 0, 1) * 255
  return around_and_astype(x_clip)