nnU-Net をColab で実装

#1 nnU-Netとは

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

#2 インストール
https://github.com/MIC-DKFZ/nnUNet
ここのInstalationに従い、Colab上にインストール。
テストのため、以下の２つの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.

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

#3 ファイル構築
自分のデータに適用するには、最初に決まったファイル構造とファイル名、ファイル形式に変換する必要がある。
https://github.com/MIC-DKFZ/nnUNet#Examples
各ファイルは３次元のnifti形式で、case_identifier_XXXX.nii.gzというファイル名にする。xxxxは４桁の数字でチャネルを表す。さらにメタ情報を記載した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ファイルを編集する必要があった。具体的にどの部分が効いているかは今後確認の予定。

[](

"tensorImageSize": "3D",
"modality": { 
   "0": "MRI"
 }, 
 "labels": { 
   "0": "background", 
   "1": "Anterior", 
   "2": "Posterior"
 }, 

)

推論結果の例は次の通り。

学習画像（下図）と比べると見当違いの領域を拾っていることが分かる。

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