LoginSignup
6
4

More than 1 year has passed since last update.

nnU-Net をColab で実装

Posted at

1 nnU-Netとは

医用画像処理のために、最適なセグメンテーションパイプラインを自動的に提供する方法。深層学習を適用する際のパラメータを「固定」、「ルールベース」、「経験的」の3種類に分類。固定パラメータの例としては、Learning rateがあって、この値は0.01と決まっている。ルールベースのパラメータは互いに関連しているものがまとまっていて、例えば、ネットワーク構造、パッチサイズ、バッチサイズがあり、これらは、GPUメモリの容量などから決定される。経験的パラメータの例としては利用するモデルを2D U-Netにするとか、3D U-Netにする等。
利用できるネットワーク構造は3種類で、2D U-Net, 3D U-Net と3D U-Net cascade(初めにダウンサンプルした画像に対して粗いセグメンテーションを行い、次に2つ目のU-Netで改善する)。

2 インストール

https://github.com/MIC-DKFZ/nnUNet
ここのInstalationに従い、Colab上にインストール。
テストのため、以下の2つのExamplesを実施。

  • run a training with the 3d full resolution U-Net on the Hippocampus dataset.
  • run inference with nnU-Net's pretrained models on the Prostate dataset.

1つ目はトレーニングプロセスの例で、デカスロンのデータをダウンロードしてきてフォーマット変換、トレーニングを走らせる。2つ目はトレーニング済みのネットワークを別のデカスロンデータに適用する。両方とも問題なく動いた。

3 ファイル構築

自分のデータに適用するには、最初に決まったファイル構造とファイル名、ファイル形式に変換する必要がある。
https://github.com/MIC-DKFZ/nnUNet#Examples
各ファイルは3次元のnifti形式で、case_identifier_XXXX.nii.gzというファイル名にする。xxxxは4桁の数字でチャネルを表す。さらにメタ情報を記載したjsonファイルも必要。

画像ファイルからこのファイル構造とjsonファイルを自動的に作成するプログラム(generate_dataset_json)も最近になって作られたようで、以下のリンクに使い方の例が書かれている。
https://github.com/MIC-DKFZ/nnUNet/blob/master/nnunet/dataset_conversion/Task120_Massachusetts_RoadSegm.py
ここでは、医療画像ではなく、地図関連の画像とそのセグメンテーション結果を利用している。
基本的には、convert_2d_image_to_niftiで画像をnifti形式に変換、generate_dataset_jsonでjsonファイルを作成している。
ちょうど良いので、リンク先のkaggleからデータをダウンロードしてきて一部をGoogle driveにアップロードして試してみる。
少しバグが有ったが修正して動作することを確認。

4 前処理

前処理では、彼らが「dataset fingerprint」と呼ぶ画像特徴を抽出する。この情報はネットワーク構造を決めるのに使われる。
基本構文は以下の通り。

!nnUNet_plan_and_preprocess -t XXX --verify_dataset_integrity

XXXはデータのIDで、前出の地図データの場合は120なので、以下のようになる。

!nnUNet_plan_and_preprocess -t 120 -pl3d None --verify_dataset_integrity

-pl3d None は「no need to run preprocessing for 3D U-Nets」を表す

5 トレーニング

nnU-Netでは5-fold cross-validationでトレーニングすることになっている。基本構文は以下の通り。

!nnUNet_train CONFIGURATION TRAINER_CLASS_NAME TASK_NAME_OR_ID FOLD  --npz (additional options)

CONFIGURATION は、2d, 3d_fullres, 3d_lowres, 3d_cascade_fullres 。3d_cascade_fullres は先に3d_lowres を計算しておく必要がある。
FOLDは5-foldのうち、何番目を使うか。地図データに適用する場合は次のようになる(fold=0)。

!nnUNet_train 2d nnUNetTrainerV2 120 0

ただし、時間がかかるので最後まで終わらなかった。

6 推論

学習済みモデルは以下のコマンドで一覧を出力できる。

!nnUNet_print_available_pretrained_models

リストの中から利用するモデルをダウンロードするには以下のコマンド。

!nnUNet_download_pretrained_model model_name

model_nameにnnUNet_print_available_pretrained_modelsで調べたモデルの名前を入れる。Task004_Hippocampusというモデルをダウンロードする場合は次の通り。

!nnUNet_download_pretrained_model Task004_Hippocampus

ダウンロード先は「~\nnUNet\nnunet\nnUNet_raw_data_base\nnUNet_raw_data」。

このモデルをデカスロンデータに適用してみる。
データは以下のウェブサイトから手動でダウンロードする。
https://drive.google.com/drive/folders/1HqEgzS8BV2c7xYNrZdEAnrHk7osJJ--2
ダウンロードしたファイルはそのままでは使えないので、以下のコマンドでフォーマッティングする。

!nnUNet_convert_decathlon_task -i data_folder

data_folderにダウンロードしたデータのフォルダを指定する。例えば、Task05_Prostateのデータをダウンロードしてフォーマッティングする場合は次の通り。

!nnUNet_convert_decathlon_task -i 'Download_folder/Task05_Prostate'

データとモデルが揃ったので、以下のコマンドで推論する。

!nnUNet_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -t TASK_NAME_OR_ID -m CONFIGURATION --save_npz

--save_npはensembleを使用する場合のみ。例えば、Task005_ProstateのデータをTask004_Hippocampusで推論する場合は次の通り。

!nnUNet_predict -i ~/nnUNet/nnunet/nnUNet_raw_data_base/nnUNet_raw_data/Task005_Prostate/imagesTs/ -o ./OUTPUT_DIRECTORY/Task05_Hippocampus_by_Task004 -t 4 -m 3d_fullres

ただし、Task005_Prostateのjsonファイルを編集する必要があった。具体的にどの部分が効いているかは今後確認の予定。

推論結果の例は次の通り。
Composite.jpg
学習画像(下図)と比べると見当違いの領域を拾っていることが分かる。
Composite_tr.jpg

7 最後に

ローカルでの実装も試してみたが、nnU-NetはLinuxでの動作しか保証していないため、Windows環境では色々と不具合が出てきた。特に、フォルダの取り扱いで問題が多い。今後時間があればLinuxでも試す予定。

6
4
1

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
6
4