グレースケール画像からカラー画像への変換

私たちのプロジェクトでは、機械学習による画像変換技術を利用しています。ここでは、プロジェクトで開発した技術の1つとして、グレースケール画像からカラー画像を生成する方法について紹介します。

データセット

ここでは、公開されているデータセットであるCOCOを利用しました。約118,000枚の画像をトレーニングに利用しています。なお、強力なGPUを搭載したコンピュータを利用できる場合には、より多くの画像でトレーニングさせることが望ましいです。

DCGANについて

私たちの機械学習モデルは、DCGANと呼ばれるアーキテクチャを採用しています。Generatorがカラー画像を生成し、Discriminatorがカラー画像が本物か、偽物かを判別します。

GeneratorはDiscriminatorを騙せるように、より本物に近いカラー画像を生成するようにトレーニングします。一方で、DiscriminatorはGeneratorに騙されないように、本物と偽物を区別できるようにトレーニングします。両者を交互にトレーニングすることで、最終的に高い精度でカラー画像を生成するGeneratorを獲得できます。

モデル構造

カラー画像は通常、RGB形式で扱われます。しかし、私たちは後述の理由により、LAB形式を利用しています。ここで、Lはグレースケールの画像情報、ABはカラーの情報です。

Generatorは、グレースケール画像であるLから、カラー情報のABを生成します。その後、後処理をしてLAB形式の画像を作成し、Discriminatorに渡します。DiscriminatorはLAB形式の画像から、その画像が本物か、それともGeneratorにより生成された偽物かを判別します。Discriminatorをトレーニングするためには本物の画像も与える必要があるので、本物の画像についてもRGB形式からLAB形式に変換してDiscriminatorへと渡します。

トレーニング

トレーニングに使う画像のサイズは、大きければ大きいほど消費する計算リソースが多くなり、これがコンピュータの保有するリソースを超えてしまうと、計算が強制的に終了されてしまいます。そこで私たちは、計算リソースを抑えるために画像サイズを256 x 256ピクセルに変換してからトレーニングをしています。しかし、実際にはこれよりももっと大きな画像に対してもカラー画像に変換できなければ、使い物にならないでしょう。この問題は、どのように対処すればよいでしょうか。画像サイズを小さく変換してからカラー画像に変換し、その後もとのサイズに引き伸ばしてしまうと、多くの画素情報が失われて画像の品質が下がってしまいます。

この問題に対処するために、LAB形式を採用しました。Generatorに渡すL画像については、256 x 256ピクセルに変換してからGeneratorにて処理をして、カラー情報を含むAB画像を生成します。このAB画像を、サイズを変換する前のオリジナルL画像と結合させることで、鮮明な輪郭を維持したままカラー画像を得ることができます。本記事の冒頭に載せた画像は256 x 256ピクセルよりも大きなものですが、この方法により一定程度の画質を維持して変換できていることがわかります。

Generatorの損失関数にはMean Square Error (MSE)、Discriminatorの損失関数にはBinary Cross Entropy (BCE) を利用しました。私たちの利用したコードは、GitHubに公開しています。Dockerコンテナを起動すると、トレーニングやWebUI (Flask) での画像変換テストをすることができます。

https://github.com/vannponlork/colorize_image

結論

機械学習技術は用途が幅広くて、人々の生活に大変役立っています。今回は私たちのプロジェクトで開発した機械学習技術として、グレースケール画像をカラー画像に変換する例を紹介しました。他の技術についても、今後ご紹介していきたいと思います。