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

Part1、Part2から引き続き、Grand challengeのチュートリアルになっている 網膜画像から血管の領域をセグメンテーションする例を動かしてみた際のメモです。

Grand Challengeのコンペに参加する際は、以下のように進めたいと考えています。

1. Google Colab上で学習プログラムを作って機械学習の重みファイルを作成する
2. Google Colab上で推論プログラムを作って推論できるかを確認する
3. Grand Challengeに提出できるようにローカルPCで推論プログラムを整える

前記事(Part2)では、上記の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の保存先は以下のようにしました。

Google Driveの保存先

マイドライブ/
└ Colab Notebooks/
	└ drive-vessels-net/
		└ datasets
			├ training
				├ images
				├ mask
				└ 1st_manual
			└ test
				├ images
				└ mask

5.1.2 requirement.txtと学習済み重みファイル

前記事(Part2)の 4.3.2項 で作ったrequirement.txtと、学習済みの重みファイルbest_metric_model_segmentation2d_dict.pthをGoogle Colabから読み込んで使えるようにGoogle Driveに置きます。

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にアクセスできるようにマウントします。

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でも使うこととしました。

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項で記載します）。

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 推論部

前記事4.3.1項で作成したprocess.pyの推論部分をコピーしました。ただし、Google Colab上で実行中かどうかで参照フォルダを変えるように修正しました。

• 提出時には、デフォルト設定である./input/および./output/フォルダから入出力画像を読み込む必要があります。
• Google Colab上ではGoogle drive上の特定の場所から入出力画像を読み込む必要があります。
• 重みファイルも同様です。

Google Colab上で実行中かどうかの判定は、if 'google.colab' in sys.modules:で行います。（参考:Webサイト）

• import sysが必要になります。

入出力画像のPathは、evalutilsのソースコードから、Algorismクラスのコンストラクタ（__init__）で指定できます。pathの指定の際に、PathとPathLikeを使っているためimportしました。

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上で学習プログラムを動かします。