python - 使用ImageDataGenerator的预处理函数转换颜色空间

Question

我在使用 ImageDataGenerator 的预处理功能转换图像的颜色空间时遇到了问题。我尝试将 RGB 转换为 CMYK 并将其输入 CNN。但是预处理函数只能返回与输入相同的形状。例如RGB模式下的图像有一个shape(224,224,3)，转换后的shape应该是(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)

Answer 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。