[solafune] 衛星画像の5倍超解像度化 (for OSS) の解法

こんにちは、manuriya（@lTlanual）と申します。
元々はコンピュータビジョン関連の仕事をしており、最近は航空写真や衛星画像を入力とした深層学習モデルの開発をしています。

Solafuneで開催された衛星画像の超解像コンペに参加しました。
諸事情によりコンペに参加できる日数が平日5日（学習時間を除くと10時間程度）しかなかったため、なるべく軽量なモデル（≒学習時間が短い）で効果的にモデルを作る必要がありました。
本記事では超解像モデルを作るにあたってどのような考え方で作成したかを順を追って記載したいと思います。

cf. @solafune(https://solafune.com) コンテストの参加以外を目的とした利用及び商用利用は禁止されています。商用利用・その他当コンテスト以外で利用したい場合はお問い合わせください。(https://solafune.com)

目標

上位20位以内に入ること

コンペをする上での問題点

本コンペでは主に拡大倍率と使用可能なライセンスの2つで問題があります。

まず、拡大倍率ですが一般的な超解像では2の冪乗がよく使われています。
公開されている超解像モデルの重みも2倍や4倍超解像のものが殆どです。
しかし、本コンペでは5倍超解像が求められているため、公開されている重みをそのまま使うことは難しいと考えられます。
そのためUpsample部分を5倍に対応させる必要があります。

次にライセンスについてですが、コンペページにあるルールに書いてありますが、使用可能なライセンスに制限があります。
かなり乱暴に言ってしまうと「商用利用可能でライセンスを変更しても問題ない」という条件です。
一般的に公開されている衛星画像データセットはResearch onlyやCC-BY-SAなど商用利用不可やライセンス変更不可のものが殆どです。
（例：Google Earth Engine、SpaceNet）
そのため、本コンペは外部データによる水増しが難しいコンペであると言えます。

上記2つを念頭に置きつつモデルの作成を行なっていきました。

使用モデル

ESRGANをベースに改良を加えたものを採用しています。
候補モデルとしては他にStable DiffusionやSR3といった拡散モデルやSwinIR、RCANなどがありますが、学習時間や作成難易度から不採用としました。

ESRGANの特徴はRRDBというResidualなブロックを何層も重ねているところです。層数を深くすることで質の良い画像が生成されると言う仕組みです。

実際のモデルでは2点ほど改良を加えています。
1点目はRRDBにあるDenseBlockのActivationをLReLUからSiLU(Swish)にしています。これはこちらの論文にあるように連続で非単調増加なActivationの方が生成品質が良くなるためです。
2点目はUpsampleを5倍に対応させている点です。対応のさせ方としては一気に5倍にせず2倍→2倍→1.25倍と言うように段階的に拡大するようにしています。2倍に大きくするところはESPCNで提案されたPixelShuffleを、最後の1.25倍するところはbilinearで拡大しています。

データ前処理

学習データを見てみると市街地だけではなく森林地帯や海上など様々なデータが含まれていることが確認できます。そのため、train-validを分ける時にランダムで適当に分けてしまうとvalidに森林や海上がなく評価されてしまう可能性があります。本当は土地被覆分類（対象の衛星画像が何に覆われているかを分類する問題）を実施したいところでしたが、時間がないので簡単に実行できるk-meansを用いて15クラス程度にクラスタリングし、15クラスが均等にtrain-validデータに含まれるようデータセットを分けました。k-meansには画像の最大/最小/平均/分散を特徴量として与えています。ランダムに分けた場合との比較を行なってはいないので精度にどの程度寄与したのかは今後調べてみたいです。

モデル学習

学習の方針としては先程提示したこちらを参考にしています。要約すると「大きなパッチサイズで長く学習すると良くなる」と言ったものです。そのため、入力パッチサイズを低解像度画像の最大サイズである130に、学習Epoch数は色々実験しながらすることを想定していたので少し短めの100Epochぐらいにしています。

AugmentationはFlipやRotate、明るさ変換など一般的なものを使用しています。正規化は単純に255で割って0〜1になるようにしました。

また、今回はGAN手法ではなく通常の学習手法を用いてモデルを作成しています。これはコンペの評価値がSSIMであり、知覚品質を重視したものではないためGANの採用を見送っています。後はGANの学習は不安定なため時間が限られている今回は採用できなかったという事情もあります。

Loss関数ですが、1-SSIMをLossにしたSSIM Lossを使用しています。これはSSIMは値が高いほど品質が良いこと、またSSIMは0〜1の範囲を値域としているため、Lossとしては1-SSIMにしてSSIMが高い（＝生成品質が良い）とLossが低くなるようにしています。元々は超解像タスクでよく使用されているL1 Lossを使用していましたが、ValidationデータのSSIMを見てみるとSSIM Lossの方が結果が良かったためそちらを採用しています。

結果

Public、Privateともに8位でした。
作業時間がなかった割には良い結果になったと思います。

不採用手法

評価結果が良くなかった手法を紹介します。

１つ目はTTA(Test Time Augmentations)です。分類タスクなどでは精度向上のためよく使われており、こちらの論文では上手くいっていたため期待していたのですが、本コンペでは今一つな結果でした。ただ、これは今回使用したモデルとの相性や私の実装ミスの可能性があるため、必ず悪くなる訳ではないかと思います。

２つ目は学習時に高解像度画像に対してUnsharpMaskをかけて学習すると言ったものです。UnsharpMaskとはエッジ強調の手法の1つです。こちらの論文で採用された方法で、高解像度画像にのみUnsharpMaskをかけることで生成画像に上手くエッジを復元するよう学習させ、生成品質を上げる手法です。しかし、パラメータ調整の影響か通常学習の方が結果がよくなってしまったため不採用となりました。

時間があったらやりたかったこと

今回、コンペが開催されていることを知ったのが終了10日前、また自分が作業できる時間が10時間程度しかなかったため色々妥協したものとなっています。コンペ期間中に思いついた試したかったことは以下になります。

• SwinIRの使用
• Wavelet変換を用いたモデルの作成
• 外部データの使用（ライセンスフリーな航空写真）
• 土地被覆モデルを用いたクラスタリング/分類クラスごとにモデルを作成

参加した感想

無事に20位以内に入れて良かったです。古い手法を組み合わせて入賞狙いでやってましたが、思ったより高い順位でビックリしました。もう少し時間があったら実験的な手法を試してみたかったです。次は8倍超解像や変化検出系のコンペがあったら参加してみたいと思います。