😇

MergeNet: 単眼デプス推定モデルの出力を高解像度化するpix2pixモデル

2021/06/12に公開約5,200字

MergeNet: 単眼デプス推定モデルの出力を高解像度化するpix2pixモデル

Boosting Monocular Depth Estimation Models to High-Resolution via Content-Adaptive Multi-Resolution Mergingという、単眼デプス推定を高解像度化するモデルが公開されました。単眼デプス推定が高解像度化されることで、できることは結構色々とあって、カメラしか搭載できないロボットが掴みたいものの形状を捉えたり、ものを避けたりとか、カメラのみから3D環境を認識したいときとかに使えます。私は、普通の動画や画像をLooking Glaasに立体表示するアプリを公開していて、そのアプリの性能アップに使う予定です。

本記事では、このモデルを他の言語に移植したり、アプリに組み込んだり、量子化したりするのに必要な知識を解説をします。評価とか分析に興味がある方は論文を読んでください。
ベースとして使われている単眼デプス推定手法は、MiDASと、SGRNetであり、新しいモデルと組み合わせても面白いかもしれません。

なぜMergeNetが必要か?

現段階では、単眼デプス推定モデルでは、詳細な形状を捉えることと、全体の幾何的な整合性を捉えることの双方をうまく満たすのは難しいです。
幾何的な整合性が取れるとは、例えば、大きな平面が平面であると表現できているということです。

例えば、上記画像は、解像度と対応したデプス画像で、解像度が低いほど、全体の整合性が取れていて、詳細が取れていない。解像度が高いほど、詳細は捉えられているが、壁が平面でなくなっていることがわかります。

この両方を満たすのは容易ではなかったが、MergeNetを使うことで実現できたとのことです。結果的に、単眼デプス推定の推定結果の高解像度化に成功したとのことです。

MergeNet: 手法においてコアとなるモデル

このモデルは、低解像度に適用した単眼デプス推定結果と高解像度に適用した単眼デプス推定結果をマージするためのネットワークです。

構造は10レイヤのunetで、pix2pixのような学習方法で学習させます。
このネットワークは、次の図のように、詳細な形状を捉えたデプス画像と、全体の幾何的な整合性を捉えたデプス画像を入力として、詳細な形状と全体の幾何学的な整合性の双方を捉えたデプス画像を出力します。(左上が全体を捉えたデプス画像、右上が詳細を捉えたデプス画像、下部の画像がいいとこ取りをした画像)

提案された手法では、このMergeNetを何度も使います。

手法の流れ

手法の全体像ですが、まずは、入力画像を低解像度にリサイズしたもの①と、高解像度にリサイズしたもの②の、デプス画像を推定します。そして、その画像をMergeNetに入力して、詳細と全体整合性のいいとこ取りをしたデプス画像③を出力します。
次に、その画像に対して、修正が必要となる箇所を修正します。修正するべきパッチに対して、同様に低解像度と高解像度のデプス画像を推定してマージします。その結果が④、⑤、⑥です。これらを更にMergeNetで③にマージすることで、最終的なデプス画像⑦ができるという流れです。

実装の話

次にコードレベルの話を説明していきます。

MergeNetの入出力

MergeNetの入力は、(H,W,2)の三次元テンソルです。出力は(H,W, 1)の三次元テンソルです。サンプルプログラムでは、(H,W) = (1024,1024)と割と高解像度の画像を受け付けるようになっていて、そこそこ、高解像度な画像に対応できるようになっていそうな感じはします。

この入力は、値が0.0から1.0の範囲に正規化されている必要があるようで、そのように正規化しています。
MiDASの出力の正規化例は次の通りです。(オリジナルコード

# Normalization
depth_min = prediction.min()
depth_max = prediction.max()

if depth_max - depth_min > np.finfo("float").eps:
    prediction = (prediction - depth_min) / (depth_max - depth_min)
else:
    prediction = 0

また、MergeNetの出力も[0.0, 1.0]の範囲に正規化されているわけではないので、
必要に応じて正規化する必要があるようです。

MergeNetモデル自体の初期化や入力の仕方はこんな感じになっていて、
pytorchのモジュールではなく、独自のAbstractクラスを定義しています。

pix2pixmodel = Pix2Pix4DepthModel(opt)
pix2pixmodel.save_dir = './pix2pix/checkpoints/mergemodel'
pix2pixmodel.load_networks('latest')
pix2pixmodel.eval()

...

pix2pixmodel.set_input(estimate1, estimate2)
pix2pixmodel.test()
visuals = pix2pixmodel.get_current_visuals()
prediction_mapped = visuals['fake_B']

画像③を出力する処理

特にひねりが無い処理です。

画像⑦を作る処理

パッチのデプス推定を洗練する処理は画像③の生成とほぼ同じですが、
画像⑦を作る際は、パッチを元の画像に貼り付ける処理が必要です。

まず、値の範囲が貼り付け先の画像の値の範囲と異なるため、補正する必要があります。
その処理をnumpyのpolyfitを使って実装しています。

p_coef = np.polyfit(mapped.reshape(-1), patch_whole_estimate_base.reshape(-1), deg=1)

この処理は数式にすると\argmin_{p} \sum_{i} p_1 * (x_i - y_i) + p _0となるpを求める処理で、
単純にスケールとオフセットを求めているだけです。この処理一行で書けるのが面白い。

また、画像は単純に貼り付けるだけでなく、ガウスブレンディングを使っていて、
貼り付けた時のパッチの境界線が目立たないようになっています。

tobemergedto[h1:h2, w1:w2] = np.multiply(tobemergedto[h1:h2, w1:w2], 1 - mask) + np.multiply(merged, mask)

高解像度と低解像度の画像サイズを決定する処理

さて、今まで、低解像度と高解像度の画像を何度も生成したと思いますが、
この画像サイズを決定するための方法も必要です。

全体の整合性を保つためには、すべてのピクセルとcontextual cueの距離がreceptive fieldに大きさの半分より離れてはいけないらしく、その半分以上離れない解像度を求めるためのアルゴリズムも提案されています。
元画像のエッジ画像にカーネルサイズがreceptive fieldと同じで膨張を適用して、
すべてが1なるギリギリの解像度は、Receptive Fieldのサイズより小さくなるらしく、
低解像度の画像のサイズとしては、Receptive Fieldサイズが使われています。
また、実験的に同様の膨張を適用して、画像の20%のピクセルが0が解像度までなら、ContextCueの性能が落ちないことがわかっていて
、この解像度を高解像度の画像のサイズとして使っています。

また、次のパッチをRefineする処理において、一部にエッジが偏っている画像などではうまくいかないため、
上記の計算結果を使って、入力画像のサイズを補正しています。
下記のコードでベースサイズの補正をしています。
論文のSection 6に説明はありますが、読んでも腑に落ちませんでした...

global factor
factor = max(min(1, 4 * patch_scale * whole_image_optimal_size / whole_size_threshold), 0.2)
print('Adjust factor is:', 1/factor)

# Compute the target resolution.
if img.shape[0] > img.shape[1]:
    a = 2*whole_image_optimal_size
    b = round(2*whole_image_optimal_size*img.shape[1]/img.shape[0])
else:
    a = round(2*whole_image_optimal_size*img.shape[0]/img.shape[1])
    b = 2*whole_image_optimal_size
)
img = cv2.resize(img, (round(b/factor), round(a/factor)), interpolation=cv2.INTER_CUBIC

更にRefineするパッチの選択

また、処理するパッチ領域を選択するアルゴリズムも必要です。
このアルゴリズムは単純で、下記のような流れです。

  1. 枠を動かしていって、パッチの候補を作る
  2. そのパッチの候補のうち、GradientがFlatすぎるものとか、エッジの密度が元画像より低いものを除外
  3. エッジの密度が元画像より高ければ、エッジの密度が同じくらいになるまで画像のサイズを上げる

感想

以上が、アルゴリズムと実装でした。

全体的にデプス画像に対する分析が積み重なって作られていて、
どのアルゴリズムも正しく動くことを理解したり、調整するのが難しそうでした。
といっても、あまり汎用性が高くないことはやりたくないので、この辺くらいまでかなと思いました。

Questions

ちょっと、わからないことがいくつかあったので、筆者に質問を投げました
そのうち、回答されるかもしれません。

Discussion

ログインするとコメントできます