prepkit.py 文件源码

def elastictransform(sigma, alpha, randomstate=None, invert=False, padding=0, ignorechannels=True):

    # rng is going to end up in _elastictransform's closure, which should guarantee persistence over function calls
    if isinstance(randomstate, int):
        rng = np.random.RandomState(randomstate)
    elif isinstance(randomstate, np.random.RandomState):
        rng = randomstate
    else:
        rng = np.random.RandomState(None)

    # Define function on image
    def _elastictransform(image):
        assert image.ndim == 2, "Can only transform 2D images."

        # Pad image if required
        if not invert and padding > 0:
            # Pad
            image = np.pad(image, padding, mode='reflect')

        # Take measurements
        imshape = image.shape
        # Make random fields
        dx = rng.uniform(-1, 1, imshape) * alpha
        dy = rng.uniform(-1, 1, imshape) * alpha

        if __debug__ and False:
            print("RNG Debug on _elastictransform: ")
            print("Invert: {}, dx[0, 0]: {}".format(invert, dx[0, 0]))
            print("Invert: {}, dy[0, 0]: {}".format(invert, dy[0, 0]))

        # Smooth dx and dy
        sdx = gaussian_filter(dx, sigma=sigma, mode='reflect')
        sdy = gaussian_filter(dy, sigma=sigma, mode='reflect')
        # Make meshgrid
        x, y = np.meshgrid(np.arange(imshape[1]), np.arange(imshape[0]))
        # Distort meshgrid indices (invert if required)
        if not invert:
            distinds = (y + sdy).reshape(-1, 1), (x + sdx).reshape(-1, 1)
        else:
            distinds = (y - sdy).reshape(-1, 1), (x - sdx).reshape(-1, 1)
        # Map cooordinates from image to distorted index set
        transformedimage = map_coordinates(image, distinds, mode='reflect').reshape(imshape)

        # Crop image if required
        if invert and padding > 0:
            transformedimage= transformedimage[padding:-padding, padding:-padding]

        return transformedimage

    # Convert image function to batch function and return
    return image2batchfunc(_elastictransform, ignorechannels=ignorechannels)


#: Function for random rotations of the image