以前から構造だけは考えていたので、重い腰を上げて実装してみました。まぁ基本は既存モデルの流用なので実装というほどでもないのですが、最新技術を簡単に利用できるのは素晴らしいですね。
JPEGの圧縮ノイズ除去を行う手法も色々ありますが、JPEGのQualityをラベルとして与えてあげれば少し改善されるのではないかというのがアイデアです。未知の画像はQualityが分かりませんので、それを推定する必要があります。全体構造としては下記の図のようになります。

1. JPEG画像をEfficientNet-b7[1]に入力してJPEGのQualityを推定(10クラス分類)
2. JPEG画像をパッチサイズ(128x128x3)に分割
3. JPEGのQualityをnn.Embeddingを使ってベクトル化→128x128x1にreshape
4. 画像とラベルをtorch.catで結合して128x128x4の入力テンソルに変形
5. 4ch入力に変更したS-Net[2]に入力してJPEGノイズを除去
6. SPSR[3]に入力して画像を4倍に拡大
7. 拡大したパッチを結合して画像として保存
といった流れになります。
ポイント1EfficientNet-b7への入力は600x600でしたので、入力画像から切り出す必要がありますが、torchivisionにtransforms.FiveCropという4隅+中央を指定したサイズで切り出すという処理がありましたので、画像から5か所切り出しでEfficientNet-b7に入力し、結果として得られた5個の結果の中央値を出力クラスとしました。かなり精度が高い印象でした。
ポイント2nn.Embeddingはラベルをベクトルに変換してくれますが、毎回同じベクトルにするために乱数シードの固定が必要です。また、ベクトルの次元が変わると生成される値も当然変わってきますので、推論時は画像を丸々S-Netに入力するのではなく、パッチサイズに分割してサイズを固定してあげる必要があります。
ポイント3分割してからノイズ除去→拡大を実行すると、かなり切れ目が目立ちました。下記のようにオフセットとして畳み込みで影響しそうな領域を大きめに切り取ってノイズ除去→拡大したのち、余分な部分をカットしました。
ポイント4パッチ分割する利点ですが、画像が大きすぎるとSPSRで拡大するときにメモリが足りなくなるので、その点でも有効だったと思います。
実験条件
EfficientNet-b7は出力層を10層に変更[4]して、DIV2Kを使ってファインチューニングしました。LossはCrossEntropy誤差を使用しました。
S-NetはEncoder部分の入力を4chに変更して、DIV2Kを使ってLossをSSIMLossで学習させました。元々原画像とJPEG圧縮したものをDatasetとして使うようになっていましたので、圧縮時のQuality÷10をラベル情報として追加で出力するように変更しました。
SPSRは全く変更しておらず、学習済みモデルをそのまま使用させていただきました。
実行結果
拡大しないと分かりにくいかもしれませんが結構きれいにできていると思います。

入力画像:Set14の画像をQuality=40で保存したJPEG画像

(左)Bicubic補間で入力画像を4倍したもの(右)1.~7.を実行したもの(一部切り出して拡大)
[1]
https://github.com/lukemelas/EfficientNet-PyTorch[2]
"S-Net: A Scalable Convolutional Neural Network for JPEG Compression Artifact Reduction"[3]
Structure-Preserving Super Resolution with Gradient Guidance (CVPR 2020)[4]
EfficientNetの事前学習モデルをPytorchで使う
posted by シンドラー at 15:37|
Comment(0)
|
Machine Learning
|