deep_colorization.py 文件源码

def image_a_b_gen(batch_size):
    for batch in datagen.flow(Xtrain, batch_size=batch_size):
        if batch == None:
            break       
        lab_batch = rgb2lab(batch)
        X_batch = lab_batch[:,:,:,0]
        Y_batch = lab_batch[:,:,:,1:]
        yield (X_batch.reshape(X_batch.shape+(1,)), Y_batch)

# Train model