HMS-Netの実装

はじめに

こちらのブログは、2022年度映像メディア学のレポート課題として投稿しました。 HMS-Netというモデルを実装しました。

HMS-Netについて

HMS-Net(HMS-Net: Hierarchical Multi-Scale Sparsity-Invariant Network for Sparse Depth Completion)は、Depth Completionタスクに対するネットワークで、2020年にTIPに掲載された論文です。

Dpeth Completionとは、LiDAR等のセンサから得られる疎なデプスマップから、密なデプスマップを推定するタスクです。

HMS-Netでは、Sparsity Invariant CNNsという論文で提案された疎な画像に対応した畳み込み操作に加えて、新しくUpsampling,Average,Concatenationなどの操作を疎な画像に対応させ、マルチスケールな特徴を活用できるエンコーダ-デコーダネットワークを提案しました。新しく提案された操作と、ネットワークモデルについて順に説明します。

Sparsity Invariant Operations

LiDARから得られる深度画像は下のように疎な画像になります。

(KITTIデータセットより)

深度情報を含まない画素が多く含まれるため、そのまま通常の畳み込み演算をしてしまうのは不適切でしょう。そのため、深度情報の有無を反映した新しい演算方法が提案されました。

  1. Sparsity-Invariant Convolution

まず、Sparsity Invariant CNNsの方で提案された新しい畳み込み演算の式です。

 \displaystyle
z(u,v) = \frac{\Sigma_{i,j=-k}^{k} m_x(u+i, v+j) w(i, j) x(u+i, v+j)}{\Sigma_{i,j=-k}^{k} m_x(u+i, v+j) + \epsilon} + b

ここで、xが入力画像、m_xが入力画像の各画素に深度情報が含まれるかどうかを表したマスク画像です。深度情報の含まれない画素はm_xが0なので畳み込み演算に使われません。分母では画像にいくつ有効な深度情報があるかを計算しており、正規化しています。

  1. Sparsity-Invariant Bilinear Upsampling

さて、ここからが今回の論文で提案された操作です。

(論文から引用)

まずはアップサンプリングです。上図のように、画像とマスクを入力とし、サイズの大きくなった画像とマスクを出力します。新しいマスクへと、どの画素に深度情報が含まれているかが伝達されています。

Max-poolingやdownsamplingも同様の操作で実現できます。Fの部分を対応する操作に変更するだけだそうです。

  1. Sparsity-Invariant Average

(論文から引用)

続いて、2枚の画像・マスクの組を入力とする操作です。2枚の画素値を足し合わせた後、2枚の画像に深度情報が含まれる場合は2で割り、そうでない場合は情報がある方の画素値そのまま使うようになっており、深度情報の有無を反映した平均操作になっています。

論文内に記述が見当たらず自信ないですが、後述のネットワークで使われているSparsity-Invariant Summationもこれと同じ感じで実装しました。

  1. Joint Sparsity-Invariant Concatenation and Convolution

(論文から引用)

テンソルの結合と畳み込みを行う操作です。少し複雑です。図のように、2つのテンソルを結合した後、特殊なカーネルで1x1の畳み込みを行います。

この時のカーネルですが、マスク画像の情報で異なるカーネルを使うようにしています。3つのカーネルを用意しておき、両方の画素に深度情報がある場合、片方のみ深度情報がある場合(2パターン)の3通りでカーネルを使い分けます。

以上のように、疎な画像に対応した新しい操作を提案することで、畳み込み操作以外でもスパース性を加味したモデルを作ることができるようになります。

HMS-Net

この論文は上記の操作を利用したエンコーダ-デコーダモデルを提案したものです。デプスマップのみを入力とするものと、RGB画像も活用するものの二つが提案されています。Low, Middle, Highの3つのフローがあり、局所的な特徴と大域的な特徴を活用できるようになっています。また、Two-scale Block, Three-scale Blockの部分で、異なるフローの情報を融合しています。

(論文から引用)

スパースな入力に対応したDownsamplingやUpsamplingが可能になることで、エンコーダ-デコーダモデルを構成できるようになりました。画素ごとの深度予測により大域的な情報を利用できるようになり、特に深度情報が少ない疎な領域でも、予測ができるようになったそうです。また、マルチスケールの情報によって、境界線がはっきりし、ノイズに強いという特徴があります。

実装した部分

私はこのようなしっかりしたモデルを書いたことはないので、各種実装を参考にしました(もちろん論文のコアな部分は自分で)。Sparsity Invariant CNNsで既に提案されている部分の参考としてこちら、エンコーダ-デコーダモデルの参考としてU-Netを実装しているこちらなどを参考にしました。

自分でやった部分としては、上で説明したSparsity-Invariant Operationとネットワークモデルの部分です。

GitHubにあげています。

結果

学習についてですが、PCスペックとかGPUの容量とかデータセットとして使ったKITTIのサイズとかの問題で論文のようなしっかりした学習は出来ていません。データセットの一部分しか使っていない他、少し小さめにクリップしたりミニバッチのサイズを小さめにしたりしています。

一応、学習させてみた結果として以下にいくつか例を貼っておきます。 RGB画像、入力する深度情報、出力された深度画像、正解の深度情報の順です。

一つ目の例では、右の車の部分やその隣の自転車の周囲は、深度情報が少ないにも関わらずそれらしい出力になっているように思います。

二つ目の例は、論文にも貼られている例です。車や右の方にある丸い岩みたいなものの境界が綺麗になるはずだそうですが、今回は微妙かもしれません。

ちなみに、精度は次のようになりました。使用したデータ数が大幅に少ないので特に考察できることはありません。強いて言えば、これでもSparsity Invariant CNNsの結果には勝っており、提案された演算の有効性が確認できたと言えるかもしれません。

[mm] 今回 論文中の結果
RMSE 1429.35 937.48
MAE 449.66 258.48

まとめ

今回は、2020年にTIPに掲載されたHMS-Netを自分で実装してみました。 コードはこちらにあります。 データセット全体を使った学習は諸々の容量の限界でできませんでしたので、できるタイミングがあればやりたいと思っています。

参考文献