Grand-challenge向けの環境準備 Part.3 (Google Colab上で推論プログラムを動かす)

[Part1](https://qiita.com/chihiro1209/items/6310a3e26e46115938f9)、[Part2](https://qiita.com/chihiro1209/items/5a46089d1217b2ef06ac)から引き続き、Grand challengeのチュートリアルになっている 網膜画像から血管の領域をセグメンテーションする例を動かしてみた際のメモです。

Grand Challengeのコンペに参加する際は、以下のように進めたいと考えています。
1. Google Colab上で学習プログラムを作って機械学習の重みファイルを作成する
2. Google Colab上で推論プログラムを作って推論できるかを確認する
3. Grand Challengeに提出できるようにローカルPCで推論プログラムを整える

[前記事(Part2)](https://qiita.com/chihiro1209/items/5a46089d1217b2ef06ac)では、上記のNo.3のやり方を確認しました。この記事では上記のNo.2のやり方を確認します。

# 5 Google Colabでの推論

提出用ファイルは前記事の4.3.1項でprocess.pyとして作成済みです。できるだけGoogle Colab上でもこのまま動くようにします。

## 5.1 必要なファイルのGoolge driveへの保存

### 5.1.1 学習＆テスト用のファイル
データは以下のページからダウンロードできます。Grand Challengeに登録し、「join」する必要がありました。
https://drive.grand-challenge.org/
データは28MB程度と軽いです。一旦、ローカルPCにダウンロードしてから、解凍して、Google Driveにアップロードしました。Google Driveの保存先は以下のようにしました。

```Text:Google Driveの保存先
マイドライブ/
└ Colab Notebooks/
	└ drive-vessels-net/
		└ datasets
			├ training
				├ images
				├ mask
				└ 1st_manual
			└ test
				├ images
				└ mask

```

### 5.1.2 requirement.txtと学習済み重みファイル
[前記事(Part2)](https://qiita.com/chihiro1209/items/5a46089d1217b2ef06ac)の 4.3.2項 で作ったrequirement.txtと、学習済みの重みファイル[best_metric_model_segmentation2d_dict.pth](https://github.com/DIAGNijmegen/drive-vessels-unet/blob/master/best_metric_model_segmentation2d_dict.pth)をGoogle Colabから読み込んで使えるようにGoogle Driveに置きます。
```Text:Google Driveの保存先
マイドライブ/
└ Colab Notebooks/
	└ drive-vessels-net/
		├ datasets/
			├ (省略)
		├ requirement.txt
		└ best_metric_model_segmentation2d_dict.pth
```

## 5.2 Google ColabからGoogle Driveのマウント
学習用のファイルや、学習結果の重みファイルを保存するために、Google ColabからGoogle Driveにアクセスできるようにマウントします。

```Python:Colabセル1
#Google driveをマウントする
from google.colab import drive
drive.mount('./gdrive')
drive_root_dir="/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/"
```

## 5.3 必要なモジュール等のpip/import
Colabノートブックを新規作成します。
必要なモジュール等のpip/importを行います。

### 5.3.1 pip
Google Colab環境ではmonaiがインストールされていないため、追加でインストールする必要があります。
`pip install monai`としてもよいですが、前頁で作ったrequirement.txtをGoogle Colabでも使うこととしました。
```Python:Colabセル2
!pip install -r "/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/requirements.txt"
```

### 5.3.2 import
必要なモジュールをimportします。4.3.1項で作成したprocess.pyの冒頭でimportしている部分をコピーしました。ただし、3行追加しています（理由は5.4項で記載します）。

```Python:Colabセル2
import SimpleITK
import numpy as np
import torch
import monai
from skimage import transform
from scipy.special import expit

from evalutils import SegmentationAlgorithm
from evalutils.validators import (
    UniquePathIndicesValidator,
    UniqueImagesValidator,
)

import sys                # 追加
from os import PathLike   # 追加
from pathlib import Path  # 追加
```

# 5.4 推論部

[前記事](https://qiita.com/chihiro1209/items/5a46089d1217b2ef06ac)4.3.1項で作成したprocess.pyの推論部分をコピーしました。ただし、Google Colab上で実行中かどうかで参照フォルダを変えるように修正しました。
* 提出時には、デフォルト設定である./input/および./output/フォルダから入出力画像を読み込む必要があります。
* Google Colab上ではGoogle drive上の特定の場所から入出力画像を読み込む必要があります。
* 重みファイルも同様です。

Google Colab上で実行中かどうかの判定は、`if 'google.colab' in sys.modules:`で行います。（参考:[Webサイト](https://it-ojisan.tokyo/colab-sys-modules/)）
* `import sys`が必要になります。

入出力画像のPathは、[evalutilsのソースコード](https://github.com/comic/evalutils/blob/master/evalutils/evalutils.py)から、`Algorism`クラスのコンストラクタ（`__init__`）で指定できます。pathの指定の際に、`Path`と`PathLike`を使っているためimportしました。

```Python:Colabセル3
class Vesselsegmentationcontainer(SegmentationAlgorithm):
    def __init__(
        self,
        *,                            # 追加
        input_path: Path,             # 追加
        output_path: Path,            # 追加
        output_file: PathLike,        # 追加
        weight_path                   # 追加
        ):
        super().__init__(
            input_path=input_path,    # 追加
            output_path=output_path,  # 追加
            output_file=output_file,  # 追加
            validators=dict(
                input_image=(
                    UniqueImagesValidator(),
                    UniquePathIndicesValidator(),
                )
            ),
        )
                
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        self.model = monai.networks.nets.UNet(
            dimensions=2,
            in_channels=3,
            out_channels=1,
            channels=(16, 32, 64, 128, 256),
            strides=(2, 2, 2, 2),
            num_res_units=2,
        ).to(self.device)
  
        self.model.load_state_dict(
            torch.load(
                weight_path,          # 修正
                map_location=torch.device("cuda" if torch.cuda.is_available() else "cpu")
            )
        )

    def predict(self, *, input_image: SimpleITK.Image) -> SimpleITK.Image:
        # Read the image
        image = SimpleITK.GetArrayFromImage(input_image)
        image = np.array(image)
        shape = image.shape

        # Pre-process the image
        image = transform.resize(image, (512, 512), order=3)
        image = image.astype(np.float32) / 255.
        image = image.transpose((2, 0, 1))
        image = torch.from_numpy(image).to(self.device).reshape(1, 3, 512, 512)
        
        # Do the forward pass
        out = self.model(image).squeeze().data.cpu().numpy()
        
        # Post-process the image
        out = transform.resize(out, shape[:-1], order=3)
        out = (expit(out) > 0.99)
        out = (out * 255).astype(np.uint8)
        
        out = SimpleITK.GetImageFromArray(out)
        return out


if __name__ == "__main__":

    # 追加:Google colab上で実行されているのか、ローカルPCで実行されているのかを判定して、各ファイルのパスを変える
    if 'google.colab' in sys.modules: 
        input_path = Path("/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/datasets/test/images/")
        output_path = Path("/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/output/images/")
        output_file = Path("/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/output/results.json")
        weight_path = "/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/best_metric_model_segmentation2d_dict.pth"
    else:
        input_path = Path("/input/")
        output_path = Path("/output/images/")
        output_file = Path("/output/results.json")
        weight_path = "./best_metric_model_segmentation2d_dict.pth"

    Vesselsegmentationcontainer(
        input_path=input_path,
        output_path=output_path,
        output_file=output_file,
        weight_path=weight_path
        ).process()
```

実行すると、Google Driveの"/content/gdrive/MyDrive/Colab Notebooks/drive-vessels-unet/output/images/"フォルダに推論結果の画像が保存されます。
これでGoogle Colab上で推論プログラムを動かすことが出来ました。
次記事ではGoogle Colab上で学習プログラムを動かします。