メディアチーム白黒画像を 畳み込みニューラルネットワーク（CNN）を用いてカラー化する
メディアチームIT技術
畳み込みニューラルネットワーク（CNN）で白黒画像をカラー化しよう
カメラができた当初、撮影された画像は、今のようなカラー画像ではなく、明度のみを考慮した白黒画像でした。
しかし、当時でも、今と変わらず現実世界に色は付いていたはずです。
その色を白黒画像から復元する技術が、白黒画像のカラー化です。
医療における白黒画像のカラー化
普段、医療で使われている CT 画像やレントゲン写真は、もちろん白黒画像です。
それらは、どちらも放射線を照射し、その吸収率から画像を作成します。
そのため、どの部分がどの構成要素か、といったことが視認しずらくなっています。
例えば、骨と筋肉部分をカラー化できるようになれば、それだけで医師の負担も下がりますし、患者側も理解しやすくなるでしょう。
色付けを行った動画
以下は、脂肪、水、カルシウムなどの放射線吸収量を比較し、色付けを行う研究の動画です。
しかし、この動画の方法でも、ある程度は人為的な閾値が必要となります。
そのような数値なしに、与えられた画像のみから色を塗ることができれば、大幅に研究が進歩することでしょう。
以上の点から、白黒画像のカラー化の有用性がわかっていただけたのではないでしょうか。
今回の記事では、畳み込みニューラルネットワーク（Convolutional neural network: CNN）を用いて、顔画像のカラー復元を実験してみたいと思います。
実験の流れ
実験は、以下の流れに従って実施します。
- カラー画像を用意する。
- カラー画像を白黒画像に変換する。
- 白黒画像を入力として、カラー画像を生成するような CNN を構成する。
- CNN の出力と元のカラー画像との二乗誤差が最小になるようにモデルの学習を行う。
このようにして、白黒画像のカラー化を機械に学習させます。
こちらの記事もオススメ！
2020.07.28機械学習 特集知識編人工知能・機械学習でよく使われるワード徹底まとめ！機械学習の元祖「パーセプトロン」とは？【人工知能】ニューラルネ...
2020.07.17ライトコード的「やってみた！」シリーズ「やってみた！」を集めました！（株）ライトコードが今まで作ってきた「やってみた！」記事を集めてみました！※作成日が新し...
畳み込みニューラルネットワークのネットワーク構成
正直、ネットワークは何でもいいのですが、表現力を失わないようオートエンコーダのような形はやめました。
オートエンコーダは、次元圧縮して分類などをするときのものです。
そのため、画像そのものの変換を学習したいときは逆効果になります。（と私は理解しています）
今回提案するネットワーク
私が提案したのは、以下のネットワークです
U-net などを参考にしました。
この構造にしたのは、ResNet のように、恒等写像を追加し、浅い層の特徴を深層でも使用することで、学習を早めたり、勾配消失を抑えられると考えたからです。
もっと良いモデルが有るとは思います。
ですが今回は、どんなネットワークでも出来るということを伝えたいため、構造については深く考えていません。
例えば、敵対的生成ネットワーク（GAN） を用いれば、より綺麗なカラー画像が得られると思います。
敵対的生成ネットワーク（GAN）の関連記事はこちら
畳み込みニューラルネットワークで実験
早速、実験を始めましょう！
実験環境
顔画像を約「50,000枚」、画像サイズは「32×32」で学習しました。
バッチサイズは「100」、学習率は「1000分の1」、「左右反転画像」を使い、データの水増しを行っています。
Google Colab 上で実行
Google Colab 上で実行しました。
データセットは、以下にあります。
【Google Drive】
https://drive.google.com/drive/folders/1JBjwANEpvZJUPTcfOJzs29mz-rYSJsld?usp=sharing
Drive をマウントして、Drive 上のデータを Google Colab で読み取れるようにします。
1from google.colab import drive
2drive.mount('/content/drive', force_remount=True)次に、ランタイムを GPU に変えれば、環境設定は終わりです。
ソースコード全体
1import torch
2import torchvision
3
4import torchvision.datasets as dset
5from torch import nn
6from torch.autograd import Variable
7from torch.utils.data import DataLoader
8from torchvision import transforms
9from torchvision.utils import save_image
10from mpl_toolkits.mplot3d import axes3d
11from torchvision.datasets import MNIST
12import os
13import math
14import pylab
15import matplotlib.pyplot as plt
16
17
18
19num_epochs = 100 #エポック数
20batch_size = 100 #バッチサイズ
21learning_rate = 1e-3 #学習率
22
23train = True #学習を行うかどうかのフラグ
24pretrained =True #事前に学習したモデルがあるならそれを使う
25save_img = True #ネットワークによる生成画像を保存するかどうのフラグ
26
27def to_img(x):
28 x = 0.5 * (x + 1)
29 x = x.clamp(0, 1)
30 x = x.view(x.size(0), 3, 32, 32)
31 return x
32def to_img_mono(x):
33 x = 0.5 * (x + 1)
34 x = x.clamp(0, 1)
35 x = x.view(x.size(0), 1, 32,32)
36 return x
37
38#データセットを調整する関数
39transform = transforms.Compose(
40 [transforms.ToTensor(),
41 transforms.Normalize((0.5, ), (0.5, ))])
42
43#訓練用データセット
44dataset = dset.ImageFolder(root='./drive/My Drive/face/',
45 transform=transforms.Compose([
46 transforms.RandomResizedCrop(32, scale=(1.0, 1.0), ratio=(1., 1.)),
47 transforms.RandomHorizontalFlip(),
48 transforms.ColorJitter(brightness=0.05, contrast=0.05, saturation=0.05, hue=0.05),
49 transforms.ToTensor(),
50 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
51 ]))
52
53#データセットをdataoaderで読み込み
54dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
55
56#モノクロ画像をカラー化するニューラルネットワーク
57class ColorDecoder(nn.Module):
58 def __init__(self):
59 super(ColorDecoder, self).__init__()
60 nch_g = 64
61 self.layer1 = nn.ModuleDict({
62 'layer0': nn.Sequential(
63 nn.Conv2d(1, nch_g , 3, 1, 1),
64 nn.BatchNorm2d(nch_g),
65 nn.ReLU()
66 ), # (1, 32, 32) -> (64, 32, 32)
67
68
69 })
70
71 self.layer2 = nn.ModuleDict({
72
73 'layer0': nn.Sequential(
74 nn.Conv2d(nch_g , nch_g*2 , 3, 1, 1),
75 nn.BatchNorm2d(nch_g*2),
76 nn.ReLU()
77 ), # (64, 32, 32) -> (128, 32, 32)
78 })
79 self.layer3 = nn.ModuleDict({
80
81 'layer0': nn.Sequential(
82 nn.Conv2d(nch_g*2 , nch_g*4 , 3, 1, 1),
83 nn.BatchNorm2d(nch_g*4),
84 nn.ReLU()
85 ), # (128, 32, 32) -> (256, 32, 32)
86
87 })
88 self.layer4= nn.ModuleDict({
89 'layer0': nn.Sequential(
90 nn.Conv2d(nch_g*4 , nch_g*8 , 3, 1, 1),
91 nn.BatchNorm2d(nch_g*8),
92 nn.ReLU()
93 ), # (256, 32, 32) -> (512, 32, 32)
94
95 })
96
97
98 self.layer7 = nn.ModuleDict({
99 'layer0': nn.Sequential(
100 nn.Conv2d(nch_g*8 , nch_g*4 , 3, 1, 1),
101 nn.BatchNorm2d(nch_g*4),
102 nn.ReLU()
103 ), # (512, 32, 32) -> (256, 32, 32)
104 })
105 self.layer8 = nn.ModuleDict({
106 'layer0': nn.Sequential(
107 nn.Conv2d(nch_g*4 , nch_g*2 , 3, 1, 1),
108 nn.BatchNorm2d(nch_g*2),
109 nn.ReLU()
110 ), # (256, 32, 32) -> (128, 32, 32)
111 })
112 self.layer9= nn.ModuleDict({
113 'layer0': nn.Sequential(
114 nn.Conv2d(nch_g*2 , nch_g , 3, 1, 1),
115 nn.BatchNorm2d(nch_g),
116 nn.ReLU()
117 ), # (128, 32, 32) -> (64, 32, 32)
118 })
119 self.layer10 = nn.ModuleDict({
120 'layer0': nn.Sequential(
121 nn.Conv2d(nch_g , 3 , 3, 1, 1),
122 nn.BatchNorm2d(3),
123 nn.Tanh()
124 ), # (64, 32, 32) -> (3, 32, 32)
125 })
126
127 def forward(self, z):
128
129 for layer in self.layer1.values():
130 z = layer(z)
131 z1 =z
132
133 for layer in self.layer2.values():
134 z = layer(z)
135 z2 =z
136
137 for layer in self.layer3.values():
138 z = layer(z)
139 z3=z
140 for layer in self.layer4.values():
141 z = layer(z)
142
143
144 for layer in self.layer7.values():
145 z = layer(z)
146 z =z+z3
147 for layer in self.layer8.values():
148 z = layer(z)
149 z =z+z2
150 for layer in self.layer9.values():
151 z = layer(z)
152 z =z+z1
153 for layer in self.layer10.values():
154 z = layer(z)
155 return z
156
157
158def main():
159 #もしGPUがあるならGPUを使用してないならCPUを使用
160 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
161
162 #ネットワークを呼び出し
163 model = ColorDecoder().to(device)
164
165
166 #事前に学習しているモデルがあるならそれを読み込む
167 if pretrained:
168 param = torch.load('./Color_Decoder.pth')
169 model.load_state_dict(param)
170
171 #誤差関数には二乗誤差を使用
172 criterion = nn.MSELoss()
173
174 #更新式はAdamを適用
175 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate,
176
177 weight_decay=1e-5)
178
179
180 loss_train_list = []
181 loss_test_list= []
182 for epoch in range(num_epochs):
183
184 print(epoch)
185
186 for data in dataloader:
187
188 img, num = data
189 #img --> [batch_size,1,32,32]
190 #imgは元画像
191 #imgをGPUに載せる
192 img = Variable(img).to(device)
193
194 # ===================forward=====================
195
196 #_imgはカラー画像をモノクロに変換した画像
197 _img = torch.Tensor(batch_size,1,32,32)
198 _img[:,0,:,:] = (img[:,0,:,:]+img[:,1,:,:]+img[:,2,:,:])/3
199 _img =_img.to(device)
200
201 #ネットワークの出力結果
202 output = model(_img)
203
204 #もし学習するなら
205 if train:
206 #ネットワークの出力と元のカラー画像との誤差を損失として学習
207
208
209 # ===================backward====================
210 loss = criterion(output, img)
211 print(loss)
212 #勾配を初期化
213 optimizer.zero_grad()
214
215 #微分値を計算
216 loss.backward()
217
218 #パラメータを更新
219 optimizer.step()
220
221
222 else:#学習しないなら
223 break
224 # ===================log========================
225
226 if train == True:
227 #モデルを保存
228 torch.save(model.state_dict(), './Color_Decoder.pth')
229
230
231 #もし生成画像を保存するなら
232 if save_img:
233 value = int(math.sqrt(batch_size))
234
235 pic = to_img_mono(_img.cpu().data)
236 pic = torchvision.utils.make_grid(pic,nrow = value)
237 save_image(pic, './mono_image_{}.png'.format(epoch)) #モノクロ画像を保存
238
239 pic = to_img(img.cpu().data)
240 pic = torchvision.utils.make_grid(pic,nrow = value)
241 save_image(pic, './real_image_{}.png'.format(epoch)) #元画像を保存
242
243 pic = to_img(output.cpu().data)
244 pic = torchvision.utils.make_grid(pic,nrow = value)
245 save_image(pic, './image_{}.png'.format(epoch)) #生成画像
246if __name__ == '__main__':
247 main()実験結果
訓練用誤差
検証用誤差
「epoch144」あたりで検証用誤差が大きくなってきているため、この付近のモデルを学習後のモデルとします。
歴史上の人物の白黒画像をカラー復元
学習したモデルを使って、歴史上の人物に対してカラー復元を試してみました。
- 左：入力
- 右：今回学習したモデルを用いて出力した結果
あえてアインシュタインの画像を３枚使っていますが、顔の向きと色の塗り方に関係があるかみたかったからです。
しかし、この結果を見ると、顔の向きで得意不得意が決まることはなさそうです。
復元結果は、ヒゲの有無や帽子の有無などに影響を受けているように見えます。
ヒゲが少なく帽子無しの白黒画像をカラー復元
次は、これらの要素がない偉人で、カラー復元をしてみましょう。
やはり、ヒゲが長すぎず、帽子をかぶっていない場合のカラー化は、うまくできているように見えます。
使用したColab関連の記事はこちら
さいごに
今回の実験から、顔画像のカラー復元では、思ったよりも装飾品が邪魔をするという結果になりました。
そのため、もし顔の色付けを本気でしたいなら、顔画像かそうでないかで分類しておく以外の分類も必要だと思います。
例えば…「メガネはあるか」「ヒゲはあるか」「帽子はあるか」などの分類を、あらかじめ行っていたほうがいいでしょう。
その上で、入力となる白黒画像を分類器にかけ、それぞれの特徴に応じて色の塗り分けを行うべきです。
次回予告
今回は、「32×32」のサイズの画像のみでカラー化を行いました。
では、今回学習した解像度より大きい「64×64」でカラー化を試してみるとどうなるでしょうか？
ぱっと見、できているように見えますが、だいぶ色塗りが荒く「まだら」になっています。
ですが、実際の場面においては、「32×32」では圧倒的にサイズが足りません。
しかし、現状、どこの研究室でも計算資源に限界があり、学習する際には小さな画像サイズで実験せざるを得ません。
そんな限られたリソースの中で、大きな画像サイズでも、同じような結果を得られたら嬉しいですよね。
というわけで、次回は、画像の高解像度化と今回学習したカラー化を組み合わせた実験をしたいと思います。
ではお楽しみに！
次回の記事はこちら
こちらの記事もオススメ！
2020.07.28機械学習 特集知識編人工知能・機械学習でよく使われるワード徹底まとめ！機械学習の元祖「パーセプトロン」とは？【人工知能】ニューラルネ...
2020.07.17ライトコード的「やってみた！」シリーズ「やってみた！」を集めました！（株）ライトコードが今まで作ってきた「やってみた！」記事を集めてみました！※作成日が新し...
ライトコードでは、エンジニアを積極採用中！
ライトコードでは、エンジニアを積極採用しています！社長と一杯しながらお話しする機会もご用意しております。そのほかカジュアル面談等もございますので、くわしくは採用情報をご確認ください。
採用情報へ
「好きを仕事にするエンジニア集団」の（株）ライトコードです！ ライトコードは、福岡、東京、大阪、名古屋の４拠点で事業展開するIT企業です。 現在は、国内を代表する大手IT企業を取引先にもち、ITシステムの受託事業が中心。 いずれも直取引で、月間PV数1億を超えるWebサービスのシステム開発・運営、インフラの構築・運用に携わっています。 システム開発依頼・お見積もり大歓迎！ また、現在「WEBエンジニア」「モバイルエンジニア」「営業」「WEBデザイナー」を積極採用中です！ インターンや新卒採用も行っております。 以下よりご応募をお待ちしております！ https://rightcode.co.jp/recruit
