*Interp移植録 - 単眼視深度推定 / HR-Depth (OnnxInterp)

# 0.Prologue
暇つぶしに、興味を引いた DNNアプリを *Interpに移植して遊んでいる。
本稿はその雑記&記録。

先日、単眼視深度推定モデル SC-Depthを *Interpに移植してみた。その続きとして、同じく単眼視深度推定の HR-Depthモデルで遊んでみることにする。モデルが提案された時期は、HR-Depthの方が SC-Depthよりも少し前になるようだ。

# 1.Original Work
HR-Depthの HRは "High Resolution"(高解像度)を意味する。HR-Depthが提案された当時、"self-supervised depth estimation"の高解像度化は難しいという課題があった。その要因の一つは物体境界の推定精度とのことだが、これは彼らが明らかにした事実のようだ。

物体境界は主に semantic情報と spatial情報で既定される。前者のsemantic情報は深度推定モデル(U-Net)のダウンサンプリング部においてより深い層に行くほど鮮明になるという傾向がある。そこでHR-Depthでは、下図の様に U-Netのスキップ接続に層間を跨ぐ情報の伝搬を付け加えることにより、深度推定の高解像度化を試みている。詳細は論文を参照のこと。

- HR-Depth: High Resolution Self-Supervised Monocular Depth Estimation
https://arxiv.org/abs/2012.07356

- GitHub: HR-Depth: High Resolution Self-Supervised Monocular Depth Estimation
https://github.com/shawLyu/HR-Depth

<img width="100%" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/2035fd5f-9111-f063-8c28-71eea205ef67.png">


# 2.準備
HR-Depthの学習済みパラメタ(Pytorchのcheckpoint)は、本家プロジェクトREADME.mdの "Pretrained Model"セクションに張り付けらたリンクから入手できる。これを ONNXモデルに変換する。

https://github.com/shawLyu/HR-Depth#pretrained-model
<img width="100%" style="border:solid 1px #000000"
src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/22d2b1f2-c7f6-ea40-5ffa-2416f39e4408.png">

Pytorchの checkpointにはモデルのグラフ情報は含まれていないため、本家プロジェクトに同梱されているモデル定義のコードが必要になる。プロジェクトをまるっと `git clone`しよう。

```shell
$ git clone https://github.com/shawLyu/HR-Depth
$ cd HR-Depth
```

残念ながら、今回も ONNXモデルをエクスポートするスクリプトは添付されていない。自前でなんとする。

ダウンロードした学習済みパラメタ "HR_Depth_K_M_1280x384"を./modelsディレクトリに置いたとすると、PytorchモデルからONNXモデルに変換するスクリプトは下記の様になる。HR-Depthは、面白いことに encoder(ダウンストリーム)部とdecoder(アップストリーム)部の二つにモデルが分割されている。変換スクリプトを実行すると、それぞれに対応する `depth_encoder.onnx`と`depth_decoder.onnx`がカレント・ディレクトに出来上がる。

```python: export_onnx.py
import torch
import torch.onnx
import networks

# pretrained model parameter files:
depth_encoder_path = "./models/HR_Depth_K_M_1280x384/encoder.pth"
depth_decoder_path = "./models/HR_Depth_K_M_1280x384/depth.pth"

# convert the encoder to ONNX
depth_encoder = networks.ResnetEncoder(18, False)

encoder_dict = torch.load(depth_encoder_path, map_location=torch.device('cpu'))
img_height = encoder_dict["height"]
img_width = encoder_dict["width"]
print("Test image height is:", img_height)
print("Test image width is:", img_width)

load_dict = {k: v for k, v in encoder_dict.items() if k in depth_encoder.state_dict()}
depth_encoder.load_state_dict(load_dict)
depth_encoder.eval()

dummy_input_a = torch.randn(1, 3, img_height, img_width)

torch.onnx.export(depth_encoder,
    dummy_input_a,
    "./depth_encoder.onnx",
    export_params=True,
    #opset_version=10,
    do_constant_folding=True,
    input_names=["input.0"],
    output_names=["output.0", "output.1", "output.2", "output.3", "output.4"],
    #dynamic_axes={}
    )

# convert the decoder to ONNX
depth_decoder = networks.HRDepthDecoder(depth_encoder.num_ch_enc)

decoder_dict = torch.load(depth_decoder_path, map_location=torch.device('cpu'))
depth_decoder.load_state_dict(decoder_dict)
depth_decoder.eval()

dummy_input0 = torch.randn(1,  64, 192, 640)
dummy_input1 = torch.randn(1,  64,  96, 320)
dummy_input2 = torch.randn(1, 128,  48, 160)
dummy_input3 = torch.randn(1, 256,  24,  80)
dummy_input4 = torch.randn(1, 512,  12,  40)
dummy_input_b  = [dummy_input0,dummy_input1,dummy_input2,dummy_input3,dummy_input4]

torch.onnx.export(depth_decoder,
    dummy_input_b,
    "./depth_decoder.onnx",
    export_params=True,
    #opset_version=10,
    do_constant_folding=True,
    input_names=["input.0", "input.1", "input.2", "input.3", "input.4"],
    output_names=["scale0", "scale1", "scale2", "sacle3"],
    #dynamic_axes={}
    )
```

このままのONNXモデルでも利用可能だが、[ONNX Simplifier](https://github.com/daquexian/onnx-simplifier)でシェイプアップしておこう。

```shell
$ onnxsim "depth_encoder.onnx" "HR_Depth_K_M_1280x384_encoder.onnx"
$ onnxsim "depth_decoder.onnx" "HR_Depth_K_M_1280x384_decoder.onnx"
```

# 3.OnnxInterp用のLivebookノート

Mix.installの依存リストに記述するモジュールは下記の通り。

```elixir
File.cd!(__DIR__)
# for windows JP
#System.shell("chcp 65001")

Mix.install([
  {:onnx_interp, path: ".."},
  {:cimg, "~> 0.1.18"},
  {:nx, "~> 0.4.2"},
  {:kino, "~> 0.8.0"}
])
```

先にも触れた通り、HR-Depthはencoderとdecoderの二つのモデルから構成されている。つまり、OnnxInterpで HR-Depthの深度推定を実行するには、(1)encoderモデルのモジュールHRDepthEncoder、(2)decoderモデルのモジュールHRDepthDecoder、(3)そしてそれら二つを繋ぐメイン・モジュールHRDepthを用意することになる。

HRDepthEncoderは入力画像をエンコードし、後続のHRDepthDecoderへの入力となるtensorリストを出力する。

> [モデル・カード]
> - inputs:
> [0] f32:{1,3,384,1280} - RGB画像,NCHWレイアウト,画素はR,G,Bの各値を{0.0～1.0}に正規化
> - outputs:
> [0] f32:{1, 64,192,640} - decoderモデルの入力となる tensor
> [1] f32:{1, 64, 96,320} - 同上
> [2] f32:{1,128, 48,160} - 同上
> [3] f32:{1,256, 24, 80} - 同上
> [4] f32:{1,512, 12, 40} - 同上

```elixir
defmodule HRDepthEncoder do
  @width  1280
  @height 384

  alias OnnxInterp, as: NNInterp
  use NNInterp,
    model: "model/HR_Depth_K_M_1280x384_encoder.onnx",
    url: "https://github.com/shoz-f/onnx_interp/releases/download/models/HR_Depth_K_M_1280x384_encoder.onnx",
    inputs: [f32: {1,3,@height,@width}],
    outputs: [f32: {1,64,192,640}, f32: {1,64,96,320}, f32: {1,128,48,160}, f32: {1,256,24,80}, f32: {1,512,12,40}]

  def apply(img) do
    # preprocess
    input0 = img
      |> CImg.resize({@width, @height})
      |> CImg.to_binary([{:range, {0.0, 1.0}}, :nchw])

    # prediction
    session()
    |> NNInterp.set_input_tensor(0, input0)
    |> NNInterp.invoke()
    |> get_output_tensors(0..4)
  end

  def get_output_tensors(session, range) do
    for i <- range, do: NNInterp.get_output_tensor(session, i)
  end
end
```

HRDepthDecoderは、HRDepthEncoderから受け取ったtensorリストをモデルに入力し、その推定結果をグレイ画像(深度マップ)に加工して出力する。

> [モデル・カード]
> - inputs:
> [0] f32:{1, 64,192,640} - encoderモデルの出力tensor
> [1] f32:{1, 64, 96,320} - 同上
> [2] f32:{1,128, 48,160} - 同上
> [3] f32:{1,256, 24, 80} - 同上
> [4] f32:{1,512, 12, 40} - 同上
> - outputs:
> [0] f32:{1,1,384,1280} - 深度マップ。各要素の値は、その位置にある物体までの深度を表す。
> [1] f32:{1,1,192, 640} - 同上(未使用)
> [2] f32:{1,1, 96, 320} - 同上(未使用)
> [3] f32:{1,1, 48, 160} - 同上(未使用)

```elixir
defmodule HRDepthDecoder do
  @width  1280
  @height 384

  alias OnnxInterp, as: NNInterp
  use NNInterp,
    model: "model/HR_Depth_K_M_1280x384_decoder.onnx",
    url: "https://github.com/shoz-f/onnx_interp/releases/download/models/HR_Depth_K_M_1280x384_decoder.onnx",
    inputs: [f32: {1,64,192,640}, f32: {1,64,96,320}, f32: {1,128,48,160}, f32: {1,256,24,80}, f32: {1,512,12,40}],
    outputs: [f32: {1,1,@height,@width}, f32: {1,1,192,640}, f32: {1,1,96,320}, f32: {1,1,48,160}]

  def apply(inputs) do
    # prediction
    output0 = session()
      |> set_input_tensors(0, inputs)
      |> NNInterp.invoke()
      |> NNInterp.get_output_tensor(0)

    # postprocess
    CImg.from_binary(output0, @width, @height, 1, 1, range: min_max(output0), dtype: "<f4") # convert to the gray image
  end

  def set_input_tensors(session, from, items) when is_list(items) do
    Enum.with_index(items, from)
    |> Enum.reduce(session, fn {item, i}, session -> NNInterp.set_input_tensor(session, i, item) end)
  end

  defp min_max(bin) do
    t = Nx.from_binary(bin, :f32)
    {
      Nx.reduce_min(t) |> Nx.to_number(),
      Nx.reduce_max(t) |> Nx.to_number()
    }
  end
end
```

外部との窓口となるモジュール HRDepthの中身は HRDepthEncoderとHRDepthDecoderを順に呼び出しているだけ。

```elixir:hr_depth.ex
defmodule HRDepth do
  def apply(img) do
    img
    |> HRDepthEncoder.apply()
    |> HRDepthDecoder.apply()
  end
end
```

デモ・モジュール LiveHRDepthは、HRDepthから受け取った深度マップ(グレイ画像)に色付けを行い描画する。

```elixir
defmodule LiveHRDepth do
  def run(path) do
    img = CImg.load(path)
    {w,h,_,_} = CImg.shape(img)

    depth = HRDepth.apply(img)
      |> CImg.resize({w,h})
      |> CImg.color_mapping(:jet)

    Kino.Layout.grid(
      Enum.map([img, depth], &CImg.display_kino(&1, :jpeg)),
      columns: 2
    )
  end
end
```

# 4.デモンストレーション
`HRDepthEncoder`と`HRDepthDecoder`を起動する。

```elixir
HRDepthEncoder.start_link([])
HRDepthDecoder.start_link([])
```

画像を与え、HR-Depthを実行する。

```elixir
LiveHRDepth.run("sample.jpg")
```

<img width="100%" src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/14158/e70d7e13-7ee9-5591-7c01-c1f5db627c06.png">

# 5.Epilogue
単眼視深度推定の HR-Depthを移植して遊んでみた。

元来 U-Netは動作が重たいモデルなのだが、HR-Depthではそのスキップ接続部にさらに計算が追加されている。SC-Depthと比べると HR-Depthの動きはとても鈍重だなぁとの印象だ。実アプリで単眼視深度推定が必要な場合は、より新しいモデルを採用した方が良さそうに思う。リサーチャー達のたゆまぬ挑戦に感謝:kissing_heart:

# Appendix
- OnnxInterpのノート
https://github.com/shoz-f/onnx_interp/blob/main/demo_HR_depth/HRDepth.livemd