使用ImageDataGenerator的预处理功能转换色彩空间

时间:2018-11-26 12:12:33

标签: python tensorflow keras neural-network deep-learning

我遇到了一个问题,即使用ImageDataGenerator的预处理功能来转换图像的色彩空间。我试图将RGB转换为CMYK,并将其输入到CNN中。但是预处理功能只能返回与输入相同的形状。例如,RGB模式下的图像具有形状(224,224,3),在转换后,该形状应转换为(224,224,4)。这对于使用预处理功能似乎是一个问题。 有人可以帮忙吗?

import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib.image as img
from keras.models import Sequential
from keras.layers import Flatten, Dense, Input, Lambda
from keras.layers.convolutional import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.utils import to_categorical, multi_gpu_model
from keras.preprocessing.image import ImageDataGenerator, array_to_img, img_to_array, load_img
from keras import optimizers
from keras.applications.vgg16 import VGG16, preprocess_input
from keras.models import Model
import os
import keras.backend.tensorflow_backend as KTF
import numpy as np
import cv2
from PIL import Image
import pickle

os.environ["CUDA_VISIBLE_DEVICES"]="0"



#cmyk convetion
def rgb_to_cmyk(img_arr):
    aaa = []

    img_arr_min = (255 - img_arr)
    K = np.min(img_arr_min, axis=2)

    C_init = (img_arr_min[:, :, 0])
    M_init = (img_arr_min[:, :, 1])
    Y_init = (img_arr_min[:, :, 2])

    C = C_init - K
    M = M_init - K
    Y = Y_init - K

    aaa.append(C)
    aaa.append(M)
    aaa.append(Y)
    aaa.append(K)
    new_arr = np.array(aaa).transpose(1,2,0)
    return new_arr


#customize function used for color convetion
def myFunc(image):
    image = np.array(image)
    converted_img = rgb_to_cmyk(image)
    converted_img /= 255
    return converted_img




#parameters
batch_size = 32
n_epoch = 150
lr_rate = 0.001



#data augementation
train_datagen = ImageDataGenerator(rescale = None,
                                   rotation_range = 30,
                                   width_shift_range = 0.2,
                                   height_shift_range = 0.2,
                                   shear_range = 0.2,
                                   zoom_range = 0.2,
                                   horizontal_flip = True,
                                   preprocessing_function=myFunc)

test_datagen = ImageDataGenerator(rescale = None,
                                  preprocessing_function=myFunc)



with tf.device('/cpu:0'):
    # load data
    train_generator = train_datagen.flow_from_directory('train_dir',
                                                        target_size=(224,224),
                                                        color_mode='rgb',
                                                        batch_size=batch_size,
                                                        class_mode='categorical')

    test_generator = test_datagen.flow_from_directory('test_dir',
                                                      target_size=(224,224),
                                                      color_mode='rgb',
                                                      batch_size=batch_size,
                                                      class_mode='categorical')








#initialize the model
model = Sequential()

# Conv Block 1
model.add(Conv2D(64, (3, 3), input_shape=(224,224,1), activation='relu', padding='same'))
model.add(Conv2D(64, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

# Conv Block 2
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(128, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

# Conv Block 3
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(256, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

# Conv Block 4
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

# Conv Block 5
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(Conv2D(512, (3, 3), activation='relu', padding='same'))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))

# FC layers
model.add(Flatten())
model.add(Dense(4096, activation='relu'))
model.add(Dense(4096, activation='relu'))
model.add(Dense(101, activation='softmax'))






network = model

sgd = optimizers.SGD(lr=lr_rate, decay=1e-6, momentum=0.9, nesterov=True)
network.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])
network.summary()



# train
history_callback = network.fit_generator(train_generator,
                      steps_per_epoch=2000,
                      epochs=n_epoch,
                      validation_data=test_generator,
                      validation_steps=800, verbose=2)
                      #use_multiprocessing=False)




all_history = history_callback.history
pickle.dump(all_history, open('save_dir/results', 'wb'), 0)

1 个答案:

答案 0 :(得分:1)

Keras仅适用于rgb或灰度图像。 例如array_to_img()方法中的一段代码:

if x.shape[2] == 4:
    # RGBA
    return pil_image.fromarray(x.astype('uint8'), 'RGBA')
elif x.shape[2] == 3:
    # RGB
    return pil_image.fromarray(x.astype('uint8'), 'RGB')
elif x.shape[2] == 1:
    # grayscale

因此,您必须先将CMYK图像转换为RGB。