Pythonのkeras, tensorflowによるグレースケール画像からカラー画像の生成

実現したいこと

今は次のような流れでプログラムしています。
最終的に、与えたグレースケール画像から学習データで与えたようなカラー画像を生成してほしいです。

①カラー画像を与えて、その画像を塗り絵をする前のような画像(被写体の色と色の境界線が残るような感じ、以降グレースケール画像と呼ぶ)を生成する。
②グレースケール画像を入力、カラー画像を出力としてデータ設定
③畳み込みニューラルネットワークの設定
④モデルのコンパイル及び学習
⑤グレースケール画像のテストデータを入力として、カラー画像を生成する

発生している問題・分からないこと

わからないことは2点あります。(初めて自分で機械学習プログラムを作っているのでわからないことばかりですが)

①まず、Sequential()モデルの定義において、畳み込み層はあるがプーリング層の実装をしていない。
普通だと畳み込み層→プーリング層を繰り返して全結合層という流れだと思うのですが、プーリング層をするとサイズが縮小されるため、入力サイズが256x256から変化してしまう(Upsampling2D等を用いて2倍にすればいいのでしょうか？？)。
プーリング層を実装する際、MaxPool2D→UpSampling2Dを用いて元のサイズを変化させないという流れでも大丈夫なのでしょうか？

②最終結果(カラー画像の生成結果)が一色しか塗られない(メイン)
最終的にはグレースケール画像を与えて学習データで与えたようなカラー画像を生成してほしいのですが、被写体に近い色一色で薄く塗りつぶしているような感じの画像が出てきます(元のグレースケール画像の境界線は見える)。
これを複数の色を使ってカラー画像を再現するようにしてほしいです。

epochs = 20, batch_size = 30となっています。(画像のサンプル数が少ないことも少し考えました)
①、②の問題についてわかる方、回答よろしくお願いします。

該当のソースコード

Python
1#PS C:\Users\owner\.vscode\code\Python\ml\programs> 内で実行する
2
3import os
4os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
5
6import numpy as np
7import tensorflow as tf
8import cv2
9import glob
10import matplotlib.pyplot as plt
11from icrawler.builtin import BingImageCrawler
12from sklearn.model_selection import train_test_split
13from keras.models import Sequential
14from keras.layers import Conv2D, UpSampling2D, InputLayer, MaxPool2D, Dropout, BatchNormalization, Activation
15from tensorflow.python.keras.optimizers import adam_v2
16
17# カラー画像と白黒画像のペアを用意する関数
18def prepare_data(img_paths, img_size=(96, 96)):
19    color_imgs = []
20    gray_imgs = []
21
22    for img_path in img_paths:
23        # カラー画像
24        color_img = cv2.imread(img_path)
25        color_img = cv2.cvtColor(color_img, cv2.COLOR_BGR2RGB)
26        color_img = cv2.resize(color_img, img_size)
27
28        # 白黒画像
29        
30        gray_img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
31        gray_img = cv2.resize(gray_img, img_size)
32        gray_img = np.expand_dims(gray_img, axis=-1)  # チャンネル次元を追加
33        
34
35        '''
36        img = cv2.imread(img_path)
37        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
38
39        neiborhood = np.ones((5, 5), dtype=np.uint8)
40        dilated = cv2.dilate(img, neiborhood, iterations=5)
41
42        diff = cv2.absdiff(dilated, img)
43
44        #5. 白黒反転
45        gray_img = 255 - diff
46
47        gray_img = cv2.resize(gray_img, img_size)
48        gray_img = np.expand_dims(gray_img, axis=-1)  # チャンネル次元を追加
49        '''
50
51        color_imgs.append(color_img)
52        gray_imgs.append(gray_img)
53
54    gray_imgs = np.array(gray_imgs) / 255.0
55    color_imgs = np.array(color_imgs) / 255.0
56
57    #return np.array(color_imgs), np.array(gray_imgs)
58    return color_imgs, gray_imgs
59
60# データセットのパス
61dataset_paths = glob.glob('Python/ml/image3/slime2/*.png') 
62
63# カラー画像と白黒画像のペアを用意
64color_images, gray_images = prepare_data(dataset_paths)
65
66# データを訓練データとテストデータに分割
67X_train, X_test, y_train, y_test = train_test_split(gray_images, color_images, test_size=0.2, random_state=42)
68
69# モデルの定義
70model = Sequential()
71
72#モデルの初期化
73model.add(InputLayer(shape=(96, 96, 1)))
74
75model.add(Conv2D(3, (3, 3), padding='same', strides=1))
76model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
77model.add(Activation('relu'))
78model.add(Dropout(0.5))
79
80model.add(Conv2D(64, (3, 3), padding='same', strides=1))
81model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
82model.add(Activation('relu'))
83model.add(Dropout(0.5))
84
85model.add(Conv2D(128, (3, 3), padding='same', strides=1))
86model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
87model.add(Activation('relu'))
88model.add(Dropout(0.5))
89
90model.add(Conv2D(256, (3, 3), padding='same', strides=1))
91model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
92model.add(Activation('relu'))
93model.add(Dropout(0.5))
94
95model.add(Conv2D(128, (3, 3), padding='same', strides=1))
96model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
97model.add(Activation('relu'))
98model.add(Dropout(0.5))
99
100model.add(Conv2D(64, (3, 3), padding='same', strides=1))
101model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
102model.add(Activation('relu'))
103model.add(Dropout(0.5))
104
105model.add(Conv2D(3, (3, 3), padding='same', strides=1))  # Sigmoid から Linear に変更
106model.add(BatchNormalization(momentum=0.9, epsilon=1e-5))
107model.add(Activation('sigmoid'))
108
109# モデルのコンパイル
110#初期値：0.0002
111optimizer = tf.keras.optimizers.Adam(learning_rate=0.002)
112model.compile(optimizer=optimizer, loss='mae', metrics=['accuracy'])
113
114
115print(model.summary())
116
117# モデルの学習
118history = model.fit(X_train, y_train, epochs=20, batch_size=150, validation_data=(X_test, y_test))
119
120# 学習過程の損失と精度の情報を取得
121train_loss = history.history['loss']
122val_loss = history.history['val_loss']
123train_accuracy = history.history['accuracy']
124val_accuracy = history.history['val_accuracy']
125
126# Epochごとのグラフを描画
127epochs = range(1, len(train_loss) + 1)
128
129# 損失のグラフ
130plt.figure(figsize=(12, 4))
131plt.subplot(1, 2, 1)
132plt.plot(epochs, train_loss, 'bo-', label='Training Loss')
133plt.plot(epochs, val_loss, 'ro-', label='Validation Loss')
134plt.title('Training and Validation Loss')
135plt.xlabel('Epochs')
136plt.ylabel('Loss')
137plt.legend()
138
139# 精度のグラフ
140plt.subplot(1, 2, 2)
141plt.plot(epochs, train_accuracy, 'bo-', label='Training Accuracy')
142plt.plot(epochs, val_accuracy, 'ro-', label='Validation Accuracy')
143plt.title('Training and Validation Accuracy')
144plt.xlabel('Epochs')
145plt.ylabel('Accuracy')
146plt.legend()
147
148plt.show()
149
150# ある白黒データに対する予測
151input_gray_image = X_test[0].reshape(1, 96, 96, 1)
152predicted_color_image = model.predict(input_gray_image)
153
154# 結果の表示
155plt.subplot(1, 3, 1)
156plt.title('Input Grayscale Image')
157plt.imshow(X_test[0].reshape(96, 96), cmap='gray')  # グレースケール画像なので256x256に修正
158
159plt.subplot(1, 3, 2)
160plt.title('Input RGB Image')
161plt.imshow(y_test[0])
162
163plt.subplot(1, 3, 3)
164plt.title('Predicted Color Image')
165plt.imshow((predicted_color_image[0] * 255).astype(np.uint8))
166
167plt.show()
168