Мне действительно интересна эта статья, в которой используются VAE для окрашивания черно-белых изображений https://arxiv.org/abs/1612.01958. Однако я хотел бы воссоздать модель, так как я все еще новичок в разработке моделей глубокого обучения, и мне было довольно сложно понять связанный с ней код (https://github.com/aditya12agd5/divcolor).
Я начал с нуля с основ, и мне удалось создать VAE с одним кодером и одним декодером, который может кодировать и декодировать цветные изображения. Вот код:
# Define Encoder
x_in = Input(shape=(train_uv.shape[1], train_uv.shape[2], train_uv.shape[3]))
layer = x_in
layer = Conv2D(filters=8, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=16, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2D(filters=64, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
shape = K.int_shape(layer) # Shape before flattening
# Latent space
layer = Flatten()(layer)
layer = Dense(16, activation='relu')(layer)
z_mean = Dense(latent_dim)(layer)
z_log_var = Dense(latent_dim)(layer)
z = Lambda(sample, output_shape=(latent_dim,))([z_mean, z_log_var]) # Data passable to the decoder
# Define Decoder
latent_inputs = Input(shape=(latent_dim,))
layer = Dense(shape[1] * shape[2] * shape[3], activation='relu')(latent_inputs)
layer = Reshape((shape[1], shape[2], shape[3]))(layer)
layer = Conv2DTranspose(filters=64, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=32, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=16, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
layer = Conv2DTranspose(filters=8, kernel_size=kernel_size, activation='relu', strides=2, padding='same')(layer)
outputs = Conv2DTranspose(filters=2, kernel_size=kernel_size, activation='sigmoid', padding='same')(layer)
# VAE
vae = Model(x_in, x_out)
# Define loss
xent_loss = K.sum(K.binary_crossentropy(x_in, x_out), axis=[1, 2, 3])
kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
vae_loss = K.mean(xent_loss + kl_loss)
vae.add_loss(vae_loss)
vae.compile(optimizer='rmsprop')
Чтобы сделать его цветным, я понимаю, что мне нужно добавить второй кодировщик, который может взять изображение в градациях серого и использовать его для выборки из скрытого пространства, чтобы извлечь соответствующее цветное изображение от декодера. Я заблудился относительно того, как внедрить этот отдельный кодер в модель. Может ли кто-нибудь дать мне объяснение о том, как реализовать VAE, который имеет два кодера?