augmentation.py 文件源码

def elastic_transform(images, alpha_range=200, sigma=10, random_state=None):
    """Elastic deformation of images as described in [Simard2003]_.
    .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for
       Convolutional Neural Networks applied to Visual Document Analysis", in
       Proc. of the International Conference on Document Analysis and
       Recognition, 2003.
    """

    alpha = np.random.uniform(0, alpha_range)

    if random_state is None:
        random_state = np.random.RandomState(None)

    shape = images[0].shape
    if len(shape) == 3:
        shape = images[0].shape[1:]

    dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha
    dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma, mode="constant", cval=0) * alpha

    x, y = np.meshgrid(np.arange(shape[0]), np.arange(shape[1]), indexing='ij')
    indices = np.reshape(x+dx, (-1, 1)), np.reshape(y+dy, (-1, 1))

    results = []
    for image in images:

        if len(images[0].shape) == 3:
            im = np.zeros(image.shape)
            for i, c_image in enumerate(image):
                im[i] = map_coordinates(c_image, indices, order=1).reshape(shape)
        else:
            im = map_coordinates(image, indices, order=1).reshape(shape)

        results.append(im)

    return results