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

Question

我遇到了一个问题，使用 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)

标签： pythontensorflowkerasneural-networkdeep-learning

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。