Depth Anythingの論文とPyTorch実装の確認

当記事では深度推定(Depth Estimation)に用いられるDepth Anythingの論文と`PyTorch`実装の確認を行いました。

## Depth Anythingの論文の確認
### Depth Anythingの論文の概要
![DepthAnything1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/e8d06828-f4e7-41b7-bc3a-d3f5294948e4.png)
[Depth Anything](https://arxiv.org/abs/2401.10891v2) Figure 1

Depth Anythingの論文にはバージョン1と2が存在し、それぞれ下記が対応します。

https://arxiv.org/abs/2401.10891v2

https://arxiv.org/abs/2406.09414v2

Depth Anythingの処理の概要は下図を元に確認できます。

![DepthAnything2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/eaf3fdf3-a200-45dd-b15c-b39420cc87ab.png)
[Depth Anything](https://arxiv.org/abs/2401.10891v2) Figure 2


### Depth Anythingの学習
Depth Anythingではよりクオリティの高い単眼深度推定(MDE; Monocular Depth Estimation)を行うにあたってラベル付き画像とラベルなし画像を組み合わせて学習を行います(半教師あり学習)。

```math
\begin{align}
\mathcal{D}^{l} &= \{ (x_m, d_m) \}_{m=1}^{M} \\
\mathcal{D}^{u} &= \{ u_n \}_{n=1}^{N}
\end{align}
```

ラベル付き画像のデータセット$\mathcal{D}^{l}$とラベルなしの画像のデータセット$\mathcal{D}^{u}$は上記のように定義されます。論文ではサンプルのインデックスが$i$で定義されていますが、論文のSecition3.1(Learning Labeled Images)で画像の位置も$i$で表されているので文字を変えました。

このデータセットからサンプリングした$(x, d)=(x_m, d_m)$についてaffine-invariant lossの$\mathcal{L}_{l}$は下記のように定義されます。

```math
\begin{align}
\mathcal{L}_{l} &= \frac{1}{HW} \sum_{i=1}^{HW} \rho(d_{i}^{*}, d_{i}) \\
\rho(d_{i}^{*}, d_{i}) &= |\hat{d}_{i}^{*} - \hat{d}_{i}| \\
\hat{d}_{i} &= \frac{d_{i}-t(d)}{s(d)} \\
t(d) &= \mathrm{median}(d) \\
s(d) &= \frac{1}{HW} \sum_{i=1}^{HW} |d_{i} - t(d)|
\end{align}
```

上記より、affine-invariant lossを用いることでground truthの$d \in \mathbb{R}^{HW}$に正規化に近い処理を行ったものをネットワークが予測するように学習が行われることが確認できます。

ラベルのないデータセットの処理は下記のように$T(u_{n})$を計算し学習を行います。

```math
\begin{align}
\hat{\mathcal{D}}^{u} = \{ (u_{n}, T(u_{n})) | u_{n} \in \mathcal{D}^{u} \}_{n=1}^{N}
\end{align}
```

上記における$T$はラベル付き画像(深度のラベル)を用いて学習させたTeacher modelのTです。このように構築した$\hat{\mathcal{D}}^{u}$とラベル付き画像の$\mathcal{D}^{l}$を元にStudent modelの学習を行います。Student modelの学習にあたっては一般的な蒸留(Distillation)の枠組みに基づく先行研究([St++](https://arxiv.org/abs/2106.05095))ではTeacher modelのfine-tuningに基づいてStudent modelの学習を行っているのに対し、Depth AnythingではStudent modelを初期化した上で学習が行われています。


## Depth AnythingのPyTorch実装の確認
### リポジトリのインストールとdemoの実行
Depth Anythingの`PyTorch`実装は下記から入手することができます。

https://github.com/LiheYoung/Depth-Anything

基本的には下記のコマンドを実行することでdemoを実行することができます。

・インストール
```
$ git clone https://github.com/LiheYoung/Depth-Anything
$ cd Depth-Anything
$ pip install -r requirements.txt
```

・デモの実行
```
$ python run.py --encoder vitl --img-path assets/examples --outdir depth_vis
```

上記を実行することで`Depth-Anything/assets/examples`以下の画像の深度推定を行い、`Depth-Anything/depth_vis`に出力することができます。

![demo2.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/31fa35b3-3af5-4ea4-a841-18d9d4c6bcc6.png)
`Depth-Anything/assets/examples/demo2.png`

![demo2_img_depth.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/6eea8d9a-ab77-4a22-9e72-462101c3405c.png)
`Depth-Anything/depth_vis/demo2_img_depth.png`


### demoスクリプトの処理の確認
前項で動かした`run.py`は少し分量が多いので、当項では`README.md`の記載を参考に下記スクリプトについて確認します。

```python
import numpy as np

from depth_anything.dpt import DepthAnything
from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet

import cv2
import torch
from torchvision.transforms import Compose


encoder = 'vitl' # can also be 'vitb' or 'vitl'
depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_{:}14'.format(encoder)).eval()


transform = Compose([
    Resize(
        width=518,
        height=518,
        resize_target=False,
        keep_aspect_ratio=True,
        ensure_multiple_of=14,
        resize_method='lower_bound',
        image_interpolation_method=cv2.INTER_CUBIC,
    ),
    NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    PrepareForNet(),
])


image = cv2.cvtColor(cv2.imread('assets/examples/demo1.png'), cv2.COLOR_BGR2RGB) / 255.0
image = transform({'image': image})['image']
image = torch.from_numpy(image).unsqueeze(0)

depth = depth_anything(image)

depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0        
depth = depth.permute(1,2,0).detach().numpy().astype(np.uint8)

depth = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
        
filename = "output.png"
cv2.imwrite(filename, depth)
```

![output.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/0d4c12f7-21a1-4276-913a-6dad290c1e04.png)
実行結果(`output.png`)

スクリプト全体の処理は下記を元に把握すると良いです。

1) 学習済みmodelのロードと深度推定の実行
→ 学習済みmodelをロードして`depth_anything`オブジェクトが作成されている。推論処理は`depth_anything(imagte)`のように実行する。

2) 画像の前処理
→ `Depth-Anything/depth_anything/util/transform.py`に実装されている`Resize`、`NormalizeImage`、`PrepareForNet`を用いて前処理を行う。

3) 実行結果の加工
→ 実行結果を0〜255の整数に変換し配列の順序の入れ替え(`permute`)や`cv2.applyColorMap`を用いてグレースケールのカラー化を行う。



### 学習済みmodelのロードと深度推定の実行
当項では以下、学習済みmodelのロードと深度推定の実行について確認します。

```python
from depth_anything.dpt import DepthAnything
from depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet

...

encoder = 'vitl' # can also be 'vitb' or 'vitl'
depth_anything = DepthAnything.from_pretrained('LiheYoung/depth_anything_{:}14'.format(encoder)).eval()

...

depth = depth_anything(image)
```

前項のdemo用のコードから深度推定に関連する実装を抽出すると上記のようになります。上記より基本的には`Depth-Anything/depth_anything/dpt.py`に実装されている`DepthAnything`クラスを確認すれば良いことが確認できます。

```python:Depth-Anything/depth_anything/dpt.py
...

from huggingface_hub import PyTorchModelHubMixin, hf_hub_download

from depth_anything.blocks import FeatureFusionBlock, _make_scratch

...

class DPT_DINOv2(nn.Module):
    def __init__(self, encoder='vitl', features=256, out_channels=[256, 512, 1024, 1024], use_bn=False, use_clstoken=False, localhub=True):
        super(DPT_DINOv2, self).__init__()
        
        assert encoder in ['vits', 'vitb', 'vitl']
        
        # in case the Internet connection is not stable, please load the DINOv2 locally
        if localhub:
            self.pretrained = torch.hub.load('torchhub/facebookresearch_dinov2_main', 'dinov2_{:}14'.format(encoder), source='local', pretrained=False)
        else:
            self.pretrained = torch.hub.load('facebookresearch/dinov2', 'dinov2_{:}14'.format(encoder))
        
        dim = self.pretrained.blocks[0].attn.qkv.in_features
        
        self.depth_head = DPTHead(1, dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
        
    def forward(self, x):
        h, w = x.shape[-2:]
        
        features = self.pretrained.get_intermediate_layers(x, 4, return_class_token=True)
        
        patch_h, patch_w = h // 14, w // 14

        depth = self.depth_head(features, patch_h, patch_w)
        depth = F.interpolate(depth, size=(h, w), mode="bilinear", align_corners=True)
        depth = F.relu(depth)

        return depth.squeeze(1)


class DepthAnything(DPT_DINOv2, PyTorchModelHubMixin):
    def __init__(self, config):
        super().__init__(**config)

...
```

上記より`DepthAnything`クラスは同じファイルに実装されている`DPT_DINOv2`と`huggingface_hub`の`PyTorchModelHubMixin`を継承して定義されていることが確認できます。`PyTorchModelHubMixin`を継承することで前項のように`DepthAnything.from_pretrained`を実行できることを抑えておくと良いです。

`DPT_DINOv2`クラスでは`forward`メソッドの`self.pretrained.get_intermediate_layers`でViT(Vision Transformer)の中間レイヤーを抽出し(encoder)、`self.depth_head`(`DPTHead`クラスに基づいて得られるオブジェクト)で深度推定を行います(decoder)。`DPTHead`クラスは下記のように実装されています。

```python:Depth-Anything/depth_anything/dpt.py
class DPTHead(nn.Module):
    def __init__(self, nclass, in_channels, features=256, use_bn=False, out_channels=[256, 512, 1024, 1024], use_clstoken=False):
        super(DPTHead, self).__init__()
        
        self.nclass = nclass
        self.use_clstoken = use_clstoken
        
        self.projects = nn.ModuleList([
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channel,
                kernel_size=1,
                stride=1,
                padding=0,
            ) for out_channel in out_channels
        ])
        
        self.resize_layers = nn.ModuleList([
            nn.ConvTranspose2d(
                in_channels=out_channels[0],
                out_channels=out_channels[0],
                kernel_size=4,
                stride=4,
                padding=0),
            nn.ConvTranspose2d(
                in_channels=out_channels[1],
                out_channels=out_channels[1],
                kernel_size=2,
                stride=2,
                padding=0),
            nn.Identity(),
            nn.Conv2d(
                in_channels=out_channels[3],
                out_channels=out_channels[3],
                kernel_size=3,
                stride=2,
                padding=1)
        ])
        
        if use_clstoken:
            self.readout_projects = nn.ModuleList()
            for _ in range(len(self.projects)):
                self.readout_projects.append(
                    nn.Sequential(
                        nn.Linear(2 * in_channels, in_channels),
                        nn.GELU()))
        
        self.scratch = _make_scratch(
            out_channels,
            features,
            groups=1,
            expand=False,
        )

        self.scratch.stem_transpose = None
        
        self.scratch.refinenet1 = _make_fusion_block(features, use_bn)
        self.scratch.refinenet2 = _make_fusion_block(features, use_bn)
        self.scratch.refinenet3 = _make_fusion_block(features, use_bn)
        self.scratch.refinenet4 = _make_fusion_block(features, use_bn)

        head_features_1 = features
        head_features_2 = 32
        
        if nclass > 1:
            self.scratch.output_conv = nn.Sequential(
                nn.Conv2d(head_features_1, head_features_1, kernel_size=3, stride=1, padding=1),
                nn.ReLU(True),
                nn.Conv2d(head_features_1, nclass, kernel_size=1, stride=1, padding=0),
            )
        else:
            self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
            
            self.scratch.output_conv2 = nn.Sequential(
                nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
                nn.ReLU(True),
                nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
                nn.ReLU(True),
                nn.Identity(),
            )
            
    def forward(self, out_features, patch_h, patch_w):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]
            
            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
            
            x = self.projects[i](x)
            x = self.resize_layers[i](x)
            
            out.append(x)
        
        layer_1, layer_2, layer_3, layer_4 = out
        
        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)
        
        path_4 = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, size=layer_2_rn.shape[2:])
        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, size=layer_1_rn.shape[2:])
        path_1 = self.scratch.refinenet1(path_2, layer_1_rn)
        
        out = self.scratch.output_conv1(path_1)
        out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
        out = self.scratch.output_conv2(out)
        
        return out
```

`self.scratch.layer1_rn`は`Depth-Anything/depth_anything/blocks.py`からロードする`_make_scratch`を元に用意される点は注意しておくと良いです。また、`self.scratch.refinenet1`は下記の`_make_fusion_block`を実行することで用意されます。

```python:Depth-Anything/depth_anything/dpt.py
from depth_anything.blocks import FeatureFusionBlock, _make_scratch

def _make_fusion_block(features, use_bn, size = None):
    return FeatureFusionBlock(
        features,
        nn.ReLU(False),
        deconv=False,
        bn=use_bn,
        expand=False,
        align_corners=True,
        size=size,
    )
```

上記より、`_make_fusion_block`関数は`Depth-Anything/depth_anything/blocks.py`から下記のように実装される`FeatureFusionBlock`クラスに基づくインスタンスを実行時に出力することが確認できます。

```python:Depth-Anything/depth_anything/blocks.py
class FeatureFusionBlock(nn.Module):
    """Feature fusion block.
    """

    def __init__(self, features, activation, deconv=False, bn=False, expand=False, align_corners=True, size=None):
        """Init.

        Args:
            features (int): number of features
        """
        super(FeatureFusionBlock, self).__init__()

        self.deconv = deconv
        self.align_corners = align_corners

        self.groups=1

        self.expand = expand
        out_features = features
        if self.expand==True:
            out_features = features//2
        
        self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1)

        self.resConfUnit1 = ResidualConvUnit(features, activation, bn)
        self.resConfUnit2 = ResidualConvUnit(features, activation, bn)
        
        self.skip_add = nn.quantized.FloatFunctional()

        self.size=size

    def forward(self, *xs, size=None):
        """Forward pass.

        Returns:
            tensor: output
        """
        output = xs[0]

        if len(xs) == 2:
            res = self.resConfUnit1(xs[1])
            output = self.skip_add.add(output, res)

        output = self.resConfUnit2(output)

        if (size is None) and (self.size is None):
            modifier = {"scale_factor": 2}
        elif size is None:
            modifier = {"size": self.size}
        else:
            modifier = {"size": size}

        output = nn.functional.interpolate(
            output, **modifier, mode="bilinear", align_corners=self.align_corners
        )

        output = self.out_conv(output)

        return output
```

上記の`forward`メソッドにおける`if len(xs) == 2:`はDecoderにおけるUNetライクな処理を実現させるにあたって用いられています。

![DepthAnything3.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3936087/17a3f6a0-8d6d-4d49-a67e-d2059fa1b8d8.png)

`if`文の内部では`nn.quantized.FloatFunctional()`に基づいて作成されたインスタンスの`add`メソッドを用いており、`add`メソッドは上記のような演算を行うことも合わせて確認しておくと良いです。