augmentation.py 文件源码

def augment(X, Y, n=5, elastic_alpha_range=20, rotation_range=30,
            shift_x_range=0.1, shift_y_range=0.1):
    """I admit this code is very ugly (much worse than Keras ImageDataGenerator API
       Note to myself : dont reinvent the wheel...
    """


    X_transformed = np.zeros([X.shape[0] * n * 3] + list(X.shape[1:]))
    Y_transformed = np.zeros([X.shape[0] * n * 3] + list(X.shape[1:]))

    z = 0
    for i in tnrange(n):

        for j, (XX, YY) in tqdm_notebook(enumerate(zip(X, Y)), total=len(X), disable=False, leave=False):
            X_transformed[z], Y_transformed[z] = elastic_transform([XX, YY], elastic_alpha_range, sigma=10)
            z += 1

        for j, (XX, YY) in tqdm_notebook(enumerate(zip(X, Y)), total=len(X), disable=False, leave=False):
            X_transformed[z], Y_transformed[z] = random_rotation([XX, YY], rg=rotation_range,
                                                                  row_index=1, col_index=2,
                                                                  channel_index=0,
                                                                  fill_mode='nearest', cval=0.)
            z += 1

        for j, (XX, YY) in tqdm_notebook(enumerate(zip(X, Y)), total=len(X), disable=False, leave=False):
            X_transformed[z], Y_transformed[z] = random_shift([XX, YY], shift_x_range, shift_y_range,
                                                               row_index=1, col_index=2,
                                                               channel_index=0,
                                                               fill_mode='nearest', cval=0.)
            z += 1

    return X_transformed, Y_transformed