[PYTHON] Apprenez à coloriser les images monochromes avec Chainer
Aperçu
―― Implémentons facilement l'apprentissage de la colorisation d'images monochromes avec Chainer. ――Recent Chainer peut être écrit de manière abstraite, je vais donc l'implémenter en utilisant Trainer, etc.
- Comparons la différence de finition due à la différence de NN entièrement connecté, de convolution NN et de fonction d'activation.
base de données
Utilisez CIFAR-100 pour l'ensemble de données. Chainer dispose de plusieurs ensembles de données faciles à utiliser, tels que MNIST et CIFAR-100. Cette fois, j'ai choisi CIFAR-100, qui peut préparer rapidement de nombreuses images couleur, mais je pense qu'un jeu de données contenant plus de types d'images est à l'origine approprié.
Utiliser CIFAR-100 avec Chainer est implémenté comme suit.
train, test = chainer.datasets.get_cifar100(withlabel=False)
L'ensemble de données d'apprentissage supervisé de Chainer est implémenté en passant un tableau de (données de formation, données d'enseignant) à l'itérateur. Pour apprendre la colorisation des images monochromes, vous avez besoin d'un tableau de «(images monochromes, images couleur)». À partir du chainer.datasets.get_cifar100 () utilisé cette fois, vous pouvez obtenir un tableau de taples (comme) (image, label), mais comme aucune étiquette n'est requise, définissez withlabel sur False et seulement le tableau d'images. Obtenir
Ensuite, modifiez le jeu de données d'image couleur obtenu en (image monochrome, image couleur). Lorsque vous créez votre propre jeu de données avec Chainer, vous l'implémenterez souvent avec chainer.datasets.TupleDataset (tableau de données d'entraînement, tableau de données d'enseignant). Cependant, comme il est facile de générer une image monochrome à partir d'une image couleur, nous allons cette fois utiliser une classe qui hérite de chainer.dataset.DatasetMixin et l'implémenter en préparant une image monochrome juste avant de créer un mini-lot. Cet exemple sera utile. Il semble qu'il soit facile d'effectuer un traitement tel que le recadrage de l'image et l'ajout de bruit de cette manière.
class PreprocessedDataset(chainer.dataset.DatasetMixin):
def __init__(self, base_image_dataset):
self.base = base_image_dataset
def __len__(self):
return len(self.base)
def get_example(self, i):
color_image = self.base[i]
gray_image = np.ndarray((32, 32), dtype=np.float32)
for ch in range(3):
gray_image = (
0.298912*color_image[0]
+ 0.586611*color_image[1]
+ 0.114478*color_image[2]
)
return gray_image, color_image
Configuration du réseau
Cette fois, je vais essayer deux types de NN. L'un connecte simplement les couches entièrement connectées et l'autre connecte la couche Deconvolution derrière la couche Convolution.
class AIC_FC(chainer.Chain):
def __init__(self, n_units):
initializer = chainer.initializers.HeNormal()
super(AIC_FC, self).__init__(
fc_in = L.Linear(None, n_units),
bn1 = L.BatchNormalization(n_units),
fc2 = L.Linear(None, n_units),
bn2 = L.BatchNormalization(n_units),
fc_out = L.Linear(None, 32*32*3)
)
def __call__(self, x, t):
y = self.colorize(x)
loss = F.mean_squared_error(y, t)
chainer.reporter.report({
'loss': loss
})
return loss
def colorize(self, x, test=False):
h = F.elu(self.bn1(self.fc_in(x), test=test))
h = F.elu(self.bn2(self.fc2(h), test=test))
y = F.reshape(self.fc_out(h), (h.shape[0], 3, 32, 32))
return y
class AIC_DC(chainer.Chain):
def __init__(self, n_ch):
initializer = chainer.initializers.HeNormal()
super(AIC_DC, self).__init__(
cv_in = L.Convolution2D(1, n_ch//4, 4, 2, 1),
bn1 = L.BatchNormalization(n_ch//4),
cv1 = L.Convolution2D(n_ch//4, n_ch//2, 4, 2, 1),
bn2 = L.BatchNormalization(n_ch//2),
cv2 = L.Convolution2D(n_ch//2, n_ch, 4, 2, 1),
bn3 = L.BatchNormalization(n_ch),
cv3 = L.Convolution2D(n_ch, n_ch, 4, 2, 1),
bn4 = L.BatchNormalization(n_ch),
dc1 = L.Deconvolution2D(n_ch, n_ch, 4, 2, 1),
bn5 = L.BatchNormalization(n_ch),
dc2 = L.Deconvolution2D(n_ch, n_ch//2, 4, 2, 1),
bn6 = L.BatchNormalization(n_ch//2),
dc3 = L.Deconvolution2D(n_ch//2, n_ch//4, 4, 2, 1),
bn7 = L.BatchNormalization(n_ch//4),
dc_out = L.Deconvolution2D(n_ch//4, 3, 4, 2, 1, outsize=(32, 32))
)
def __call__(self, x, t):
y = self.colorize(x)
loss = F.mean_squared_error(y, t)
chainer.reporter.report({
'loss': loss
})
return loss
def colorize(self, x, test=False):
h = F.reshape(x, (x.shape[0], 1, 32, 32))
h = F.elu(self.bn1(self.cv_in(h), test=test))
h = F.elu(self.bn2(self.cv1(h), test=test))
h = F.elu(self.bn3(self.cv2(h), test=test))
h = F.elu(self.bn4(self.cv3(h), test=test))
h = F.elu(self.bn5(self.dc1(h), test=test))
h = F.elu(self.bn6(self.dc2(h), test=test))
h = F.elu(self.bn7(self.dc3(h), test=test))
y = self.dc_out(h)
return y
Lorsque j'ai essayé d'approfondir le NN entièrement connecté, seule une image floue a été générée (la convergence est lente?), Donc je l'ai rendue moins profonde.
À propos, le NN colorisé semble avoir une conception plus compliquée. (Référence)
Mise en œuvre de l'apprentissage
Le flux de mise en œuvre de l'apprentissage dans Chainer
- Créez un modèle
- Paramètres de l'optimiseur
- Préparation de l'ensemble de données
- Créer un itérateur à partir d'un jeu de données
- Définir le programme de mise à jour
- Configurer le formateur
Ce sera. En plus d'Adam, Optimizer a également des bases telles que SGD et Momentum SGD, et il semble que Standard Updater soit suffisant pour les problèmes où Updater n'effectue pas de calcul de perte compliqué comme GAN. Comparé à l'époque où j'écrivais tout par moi-même, je peux beaucoup l'apprécier car ceux qui se sentent bien par défaut sont préparés. Étant donné que seules les parties nécessaires doivent être modifiées, il est plus facile à comprendre lorsqu'il est vu par d'autres personnes, et je vous en suis reconnaissant.
Mettre en œuvre des tests de modèle
Des problèmes comme celui-ci sont difficiles à comprendre à partir de la seule valeur de la perte, je veux donc les visualiser. Mettez donc en œuvre le test du modèle avec l'extension Trainer. Les extensions sont créées en héritant de chainer.training.extention.Extension etc. ou en utilisant chainer.training.make_extension (). Implémentez le test de modèle et la sauvegarde d'image avec chainer.training.make_extension () comme suit. (Imsave de scipy.misc est importé.)
Une image de test est préparée séparément des données d'apprentissage.
@chainer.training.make_extension(trigger=(1, 'epoch'))
def test_model(trainer):
colorized_img = chainer.cuda.to_cpu(F.clipped_relu(model.colorize(test_img, test=True), z=1.0).data)
imsave(
'test_colorized{}.png'.format(trainer.updater.epoch),
colorized_img
.transpose(0, 2, 3, 1)
.reshape((8, 8, 32, 32, 3))
.transpose(1, 2, 0, 3, 4)
.reshape(8*32, 8*32, 3)
)
trainer.extend(test_model)
Résultat d'apprentissage
Je présenterai le résultat de l'apprentissage. Cette fois, j'ai essayé deux NN, un NN entièrement connecté et un NN alambiqué, donc je vais les comparer avec l'état d'apprentissage au moment de l'apprentissage 30 epoch. Nous avons appris que n_units de NN entièrement connecté est 2048 et n_ch de NN alambiqué est 512.
Image couleur
Image monochrome
Après 30 époque
NN entièrement connecté
Pliage NN
NN entièrement combiné semble avoir tendance à être une image granuleuse dans l'ensemble. J'ai été surpris que même un NN simple ait beaucoup de couleurs, mais je pense que le NN plié peut produire des images plus vives et plus belles.
En regardant chaque image, il semble que le ciel et la mer ont tendance à être bien reconnus et que les couleurs sont bien placées. Cependant, il semble difficile de faire la distinction entre le ciel bleu et le coucher du soleil. Dans l'image en bas à droite, qui semble montrer un petit animal, le sol est reconnu et les couleurs marron et herbe sont reproduites. Cependant, c'est un peu trop herbeux par rapport à la bonne réponse.
Modifications dues à des différences dans les fonctions d'activation
L'image ci-dessus provient de l'utilisation d'elu pour la fonction d'activation. J'ai également étudié comment la finition change lorsque relu ou leaky_relu est utilisé comme fonction d'activation. Seuls les résultats avec convolution NN, 30 epoch avec n_ch de 512 sont introduits.
relu
leaky_relu
elu
Le relu est parfaitement fini, mais il a une impression légèrement floue. elu et relu qui fuit semblent être supérieurs à relu en termes de clarté d'image. relu qui fuit donne une impression entre elu et relu.
Code entier utilisé
#! /usr/bin/env python
# coding : utf-8
import argparse
import numpy as np
from scipy.misc import imsave
import chainer
import chainer.functions as F
import chainer.links as L
from chainer.training import extensions
class PreprocessedDataset(chainer.dataset.DatasetMixin):
def __init__(self, base_image_dataset):
self.base = base_image_dataset
def __len__(self):
return len(self.base)
def get_example(self, i):
color_image = self.base[i]
gray_image = np.ndarray((32, 32), dtype=np.float32)
for ch in range(3):
#Calculez la luminosité et créez une image monochrome
gray_image = (
0.298912*color_image[0]
+ 0.586611*color_image[1]
+ 0.114478*color_image[2]
)
return gray_image, color_image
class AIC_FC(chainer.Chain):
def __init__(self, n_units):
initializer = chainer.initializers.HeNormal()
super(AIC_FC, self).__init__(
fc_in = L.Linear(None, n_units),
bn1 = L.BatchNormalization(n_units),
fc2 = L.Linear(None, n_units),
bn2 = L.BatchNormalization(n_units),
fc_out = L.Linear(None, 32*32*3)
)
def __call__(self, x, t):
y = self.colorize(x)
loss = F.mean_squared_error(y, t)
chainer.reporter.report({
'loss': loss
})
return loss
def colorize(self, x, test=False):
h = F.elu(self.bn1(self.fc_in(x), test=test))
h = F.elu(self.bn2(self.fc2(h), test=test))
y = F.reshape(self.fc_out(h), (h.shape[0], 3, 32, 32))
return y
class AIC_DC(chainer.Chain):
def __init__(self, n_ch):
initializer = chainer.initializers.HeNormal()
super(AIC_DC, self).__init__(
cv_in = L.Convolution2D(1, n_ch//4, 4, 2, 1),
bn1 = L.BatchNormalization(n_ch//4),
cv1 = L.Convolution2D(n_ch//4, n_ch//2, 4, 2, 1),
bn2 = L.BatchNormalization(n_ch//2),
cv2 = L.Convolution2D(n_ch//2, n_ch, 4, 2, 1),
bn3 = L.BatchNormalization(n_ch),
cv3 = L.Convolution2D(n_ch, n_ch, 4, 2, 1),
bn4 = L.BatchNormalization(n_ch),
dc1 = L.Deconvolution2D(n_ch, n_ch, 4, 2, 1),
bn5 = L.BatchNormalization(n_ch),
dc2 = L.Deconvolution2D(n_ch, n_ch//2, 4, 2, 1),
bn6 = L.BatchNormalization(n_ch//2),
dc3 = L.Deconvolution2D(n_ch//2, n_ch//4, 4, 2, 1),
bn7 = L.BatchNormalization(n_ch//4),
dc_out = L.Deconvolution2D(n_ch//4, 3, 4, 2, 1, outsize=(32, 32))
)
def __call__(self, x, t):
y = self.colorize(x)
loss = F.mean_squared_error(y, t)
chainer.reporter.report({
'loss': loss
})
return loss
def colorize(self, x, test=False):
#Remodeler pour que ndim soit 4 pour l'entrée dans la couche de convolution
h = F.reshape(x, (x.shape[0], 1, 32, 32))
h = F.elu(self.bn1(self.cv_in(h), test=test))
h = F.elu(self.bn2(self.cv1(h), test=test))
h = F.elu(self.bn3(self.cv2(h), test=test))
h = F.elu(self.bn4(self.cv3(h), test=test))
h = F.elu(self.bn5(self.dc1(h), test=test))
h = F.elu(self.bn6(self.dc2(h), test=test))
h = F.elu(self.bn7(self.dc3(h), test=test))
y = self.dc_out(h)
return y
def main():
parser = argparse.ArgumentParser(description='Automatic Image Colorization')
parser.add_argument('--batchsize', '-b', type=int, default=64,
help='Number of images in each mini-batch')
parser.add_argument('--epoch', '-e', type=int, default=30,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu', '-g', type=int, default=0,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--resume', '-r', default='',
help='Resume the training from snapshot')
parser.add_argument('--n_ch', '-nc', type=int, default=1024,
help='Number of channels')
parser.add_argument('--n_units', '-nu', type=int, default=0,
help='Number of units')
args = parser.parse_args()
print('# GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
if args.n_units > 0:
print('# n_units: {}\n'.format(args.n_units))
model = AIC_FC(args.n_units)
else:
print('# n_ch: {}\n'.format(args.n_ch))
model = AIC_DC(args.n_ch)
if args.gpu >= 0:
chainer.cuda.get_device().use()
model.to_gpu()
opt = chainer.optimizers.Adam()
opt.setup(model)
train, test = chainer.datasets.get_cifar100(withlabel=False)
test_img = (
0.298912*test[:64,0]
+ 0.586611*test[:64,1]
+ 0.114478*test[:64,2]
)
#Enregistrez 64 images en une seule image 8x8
imsave(
'test.png',
test[:64]
.transpose(0, 2, 3, 1)
.reshape((8, 8, 32, 32, 3))
.transpose(1, 2, 0, 3, 4)
.reshape(8*32, 8*32, 3)
)
imsave(
'test_gray.png',
test_img
.reshape((8, 8, 32, 32))
.transpose(1, 2, 0, 3)
.reshape(8*32, 8*32)
)
if args.gpu >= 0:
test_img = chainer.cuda.to_gpu(test_img)
dataset = PreprocessedDataset(train)
iterator = chainer.iterators.MultiprocessIterator(dataset, args.batchsize)
updater = chainer.training.StandardUpdater(iterator, opt, device=args.gpu)
trainer = chainer.training.Trainer(updater, (args.epoch, 'epoch'))
trainer.extend(extensions.LogReport())
trainer.extend(extensions.PrintReport([
'epoch', 'loss', 'elapsed_time'
]))
@chainer.training.make_extension(trigger=(10, 'epoch'))
def test_model(trainer):
#Plage de valeurs 0~Coupé pour être 1_Grâce à relu
colorized_img = chainer.cuda.to_cpu(F.clipped_relu(model.colorize(test_img, test=True), z=1.0).data)
imsave(
'test_colorized{}.png'.format(trainer.updater.epoch),
colorized_img
.transpose(0, 2, 3, 1)
.reshape((8, 8, 32, 32, 3))
.transpose(1, 2, 0, 3, 4)
.reshape(8*32, 8*32, 3)
)
trainer.extend(test_model)
trainer.extend(extensions.ProgressBar(update_interval=100))
trainer.run()
if __name__ == '__main__':
main()
Recommended Posts