DUICUO

MONAIモデルに基づく医用画像の生成

翻訳者|朱賢宗

校正者 | Chonglou

皆さん、こんにちは!この記事では、MONAIの最新のオープンソース拡張機能であるMONAI Generative Modelsを使用して潜在拡散モデルを作成し、このモデルを使用して胸部X線画像を生成します。

導入

生成型人工知能は、トレーニングデータセットの基本的なパターンと構造を学習するモデルを作成できるため、ヘルスケア分野において計り知れない可能性を秘めています。このように、これらの生成モデルを用いることで、実データと同じ詳細度と特徴を持つ無数の合成データセットを、実データの特性に制限されることなく作成することができます。その重要性に鑑み、私たちはMONAIプラットフォームのオープンソース拡張機能であるMONAI生成モデルを開発しました。このモデルには、拡散モデル、自己回帰トランスフォーマー、生成的敵対ネットワークといった今日の最先端モデルに加え、これらの生成モデルのトレーニングと評価を容易にするコンポーネントが含まれています。

MONAI生成モデル

この記事では、完全なプロジェクトを通して、放射線診断レポートから胸部X線(CXR)画像を生成するポテンシャル拡散モデル(安定拡散モデルと同じタイプ)を作成します。この記事では、コードを理解しやすく、さまざまな開発環境に適応できるように努めました。そのため、最も効率的ではないかもしれませんが、楽しんでいただければ幸いです。

この記事で言及したオープンソースプロジェクトの完全版は、GitHubリポジトリでご覧いただけます。なお、この記事ではバージョン0.2について言及していることにご注意ください。

データセットをダウンロード

まず、データセットから始めましょう。このプロジェクトでは、 MIMICデータセットを使用します。このデータセットにアクセスするには、PhysioNetポータルでアカウントを作成する必要があります。MIMIC-CXR-JPG(JPGファイルを含む)とMIMIC-CXR(放射線レポートを含む)を使用します。どちらのデータセットも、PhysioNet Certified Health Data License 1.5.0に基づいて公開されています。無料トレーニングコースを修了すると、データセットページの下部にある指示に従ってデータセットを自由にダウンロードできます。CXR画像のサイズは約1000x1000ピクセルです。そのため、ダウンロードには時間がかかる場合があります。

胸部X線画像は、肺、心臓、血管など、胸腔内の構造や臓器に関する非常に貴重な情報を提供できる重要なツールです。ダウンロード後、35万枚以上の画像が揃いました!これらの画像は、前後(PA)、前後(AP)、外側(LAT)の3つの異なる投影角度のいずれかに対応しています。

本論文の実験プロジェクトでは、最も一般的な投影角度であるPA投影のみを対象とします。PA投影は、放射線診断報告書(96,162枚の画像を含む)に記載されている特徴のほとんどを視覚化できます。これらの報告書に関するファイルは合計85,882個あり、それぞれに複数のテキストセクションが含まれています。ここでは、所見(主に画像の内容の解釈)と印象(報告書の内容を要約し、結論を導き出すなど)のデータを使用します。

モデルとトレーニングプロセスをより管理しやすくするために、画像のサイズを最小軸で512ピクセルに変更します。これらの初期ステップを自動化するスクリプトのリストは、 https://github.com/Warvito/generative_chestxray#preprocessingにあります

開発モデル

潜在拡散モデルアーキテクチャ: オートエンコーダは入力画像 x を潜在表現 z に圧縮し、拡散モデルを使用して z の確率分布を推定します。

使用する潜在的拡散モデルは、次の部分で構成されています。

  1. オートエンコーダ: 入力画像をより小さな潜在表現に圧縮するために使用されます。
  2. 拡散モデル:CXRの潜在表現の確率データ分布を学習します。
  3. テキストエンコーダー:サンプリング処理を調整するための埋め込みベクトルを作成します。この例では、事前学習済みのサンプルを使用します。

MONAI生成モデルを使えば、これらのモデルを簡単に作成・学習できます。それでは、オートエンコーダから始めましょう!

KL正則化を用いたオートエンコーダの使用

KL正則化を用いたオートエンコーダ(AE-KL、または一部のプロジェクトではVAEと略される)の主な目的は、小さな潜在表現を作成し、高い忠実度(可能な限り多くの詳細を保持)で画像を再構成することです。このプロジェクトでは、64128128128個のチャネルで構成される4階層のオートエンコーダを作成し、各階層間でダウンサンプリングブロックを適用することで、最深層に進むにつれて特徴マップを小さくします。

オートエンコーダには自己注意のためのブロックが含まれている可能性がありますが、この例では、脳画像に関する以前の研究と同様の構造を採用し、メモリ節約のための注意メカニズムは使用していません。最後に、潜在表現には3つのチャネルがあります。

 from generative.networks.nets import AutoencoderKL ... model = AutoencoderKL( spatial_dims=2, in_channels=1, out_channels=1, num_channels=[64, 128, 128, 128], latent_channels=3, num_res_blocks=2, attention_levels=[False, False, False, False], with_encoder_nonlocal_attn=False, with_decoder_nonlocal_attn=False, )

注:スクリプトではモデルのハイパーパラメータを保存するためにOmegaConfパッケージを使用しています。以前の設定はこのファイル確認できます。まとめると、OmegaConfはPythonプロジェクト、特にディープラーニングやその他の複雑なソフトウェアシステムを含むプロジェクトの設定管理に強力なツールです。OmegaConfを使用すると、ハイパーパラメータを.yamlファイルに簡単に整理し、スクリプトで読み込むことができます。

トレーニング AE-KL

次に、学習プロセスのいくつかの要素を定義します。まず、 KL正則化 を使用します。この部分は、拡散モデルの潜在空間分布とガウス分布との間の距離を評価する役割を果たします。Rombachら によって提案されたようにこれは潜在空間の分散を制約するために使用され、拡散モデルを学習する際に非常に役立ちます(詳細は後述)。このモデルの順方向メソッドは再構成と、KLダイバージェンスの計算に使用する潜在空間表現のμベクトルとσベクトルを返します。

 #位于训练循环内reconstruction, z_mu, z_sigma = model(x=images) … kl_loss = 0.5 * torch.sum(z_mu.pow(2) + z_sigma.pow(2) - torch.log(z_sigma.pow(2)) - 1, dim=[1, 2, 3]) kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]

第二に、私たちのモデルはピクセルレベルの損失アルゴリズムを使用しています。このプロジェクトでは、L1距離を用いてAE - kl再構成と元の画像との差を評価します

 l1_loss = F.l1_loss(reconstruction.float(), images.float())

次に、私たちのモデルは知覚レベルの損失計算も利用します。知覚損失の背後にある考え方は、入力画像と再構成画像の違いをピクセルレベルで評価する代わりに、これら 2 つの画像を事前トレーニング済みのモデルに渡すというものです。次に、内部アクティベーションと特徴マップ間の距離を測定します。MONAI 生成モデルでは、医療画像の事前トレーニング済みネットワーク (https://github.com/Project-MONAI/GenerativeModels/blob/main/generative/losses/perceptual.py で入手可能) に基づく知覚ネットワークを簡単に使用できます。130万枚を超える医療画像でトレーニングされた RadImageNet 研究 (Mei らによる) の2D ネットワークにアクセスできます。私たちは 2D 事前トレーニング済みネットワークを使用してスライスを評価することで 3D 画像を評価する 2.5D メソッドを実装しました。最後に、MedicalNet にアクセスして、純粋な 3D メソッドを使用して 3D 画像を評価できます。このプロジェクトでは、 Pinayaと同様のアプローチを使用し、学習知覚画像パッチ類似度 (LPIPS) メトリック (MONAI 生成モデルでも利用可能) を採用しました。

 # 实例化感知损失类perceptual_loss = PerceptualLoss( spatial_dims=2, network_type="squeeze", ) ... #在训练循环内部... p_loss = perceptual_loss(reconstruction.float(), images.float())

最後に、敵対的損失計算を用いて再構成の詳細を処理します。敵対的ネットワークとは、画像全体が本物か偽物かを単一の予測ではなく、画像内の複数のパッチについて予測を行うパッチ識別器(Pix2Pixの研究で最初に提案されたもの)です。

オリジナルの潜在拡散モデルと安定拡散とは異なり、最小二乗GANから導出された識別器損失を使用しました。より高度な敵対的損失ではありませんが、3D医用画像で学習させた際に有効性と安定性を示しました(ただし、まだ改善の余地はあります)。敵対的損失は非常に不安定になる可能性がありますが、知覚的損失と組み合わせることで、識別器損失と生成器損失の安定化にも役立ちます。

トレーニングループと評価手順は、 https://github.com/Warvito/generative_chestxray/blob/83f6c0892c63a1cdbf308ff654d40afc0af51bab/src/python/training/training_functions.py#L129 と https://github.com/Warvito/generative_chestxray/blob/83f6c0892c63a1cdbf308ff654d40afc0af51bab/src/python/training/training_functions.py#L236 でご覧いただけます。合計75回のトレーニング後、 MLflow パッケージを使用してモデルを保存しました。 MLflowパッケージを使用することで、実験をより適切に監視できるようになりました。MLflowは、gitハッシュや対応するパラメータなどの情報を整理し、実験グループごとに異なる実行結果を一意のIDで保存し、異なる結果の比較を容易にします(重みやバイアスなどの他のツールと同様です)。AE-KLログファイルは、 https://drive.google.com/drive/folders/1Ots9ujg4dSaTABKkyaUnFLpCKQ7kCgdK?usp=sharing のリンクからご覧いただけます

拡散モデルの使用

次に、拡散モデルをトレーニングする必要があります。

拡散モデルはU-Netのようなネットワークです。従来、このようなモデルはノイズ画像(または潜在表現)を入力として取り、そのノイズ成分を予測します。これらのモデルは反復的なノイズ除去メカニズムを用いることで、マルコフ連鎖上のノイズからわずか数ステップで画像を生成することができます。そのため、このモデルは、サンプリングプロセスのどの段階にあるかを定義する時間ステップにも依存します。

DiffusionModelUNet クラスを用いることで、拡散モデル用のU-Netのようなネットワークを作成できます。このプロジェクトでは設定ファイルで定義された設定パラメータを使用します。この設定パラメータは、3チャンネルの入力と出力(AE - kl が3チャンネルの潜在空間を持つため)と、256,512,768チャンネルの3つの異なるレベルを定義します。各レベルには2つの残差ブロックがあります。

前述の通り、使用するモデルのタイムステップ値を渡すことは非常に重要です。これにより、モデルはこれらの残差ブロックの挙動を容易に調整できます。最後に、ネットワーク内のアテンションメカニズムを定義します。今回のケースでは、第2レベルと第3レベルにアテンションブロック(「attention_levels」パラメータで表されます)があり、アテンションヘッドごとにそれぞれ512チャネルと768チャネルあります(つまり、各レベルに1つのアテンションヘッドがあります)。これらのアテンションメカニズムは、クロスアテンション法を介してネットワークに外部条件(放射性レポート)を適用できるため重要です。

外部条件(または「コンテキスト」)はU-Netの注意ブロックに適用される。

このプロジェクトでは、事前学習済みのテキストエンコーダを使用しました。簡略化のため、Stable Diffusion v2.1モデル(「stabilityai/Stable-Diffusion-2–1-base」)と同じモデルを使用して、テキストトークンをテキスト埋め込みに変換しました。テキスト埋め込みは、DiffusionModel UNetクロスアテンション層でキーと値のベクトルとして使用されます。テキスト埋め込みの各トークンは1024次元で、これは「with_conditing」および「cross_attention_dim」パラメータで定義しました。

 from generative.networks.nets import DiffusionModelUNet ... diffusion = DiffusionModelUNet( spatial_dims=2, in_channels=3, out_channels=3, num_res_blocks=2, num_channels=[256, 512, 768], attention_levels=[False, True, True], with_cnotallow=True, cross_attention_dim=1024, num_head_channels=[0, 512, 768], )

モデルの定義に加えて、拡散モデルからのノイズを学習中に入力画像にどのように追加し、サンプリング中にどのように除去するかを定義することも重要です。この目的のために、MONAI生成モデルにSchedulersクラスを実装し、ノイズスケジューラを定義しました。この例では、1000ステップのDDPMSchedulerと以下のハイパーパラメータを使用します。

 from generative.networks.schedulers import DDPMScheduler ... scheduler = DDPMScheduler( beta_schedule="scaled_linear", num_train_timesteps=1000, beta_start=0.0015, beta_end=0.0205, prediction_type="v_prediction", )

ここでは「v-予測」アプローチを選択しました。このアプローチでは、U-Netは追加ノイズだけでなく、速度成分(元の画像と追加ノイズの組み合わせ)を予測しようとします。このアプローチは、より安定した学習とより速い収束を示すことが示されています(リンクhttps://arxiv.org/abs/2210.02303の対応するアルゴリズムモデルでも使用されています)。

拡散モデルのトレーニング

拡散モデルをトレーニングする前に、適切なスケーリング係数を見つける必要があります。Rombach らが述べているように、潜在的な空間分布の標準偏差が高すぎると、信号対雑音比が LDM によって取得される結果に影響を与える可能性があります。潜在的な表現の値が高すぎると、それに追加する最大ガウスノイズの量がすべての情報を破壊するのに十分でない可能性があります。そのため、トレーニング中に、元の潜在的な表現からの情報が、必要のないときに表示され、後で純粋なノイズから画像をサンプリングすることが不可能になる可能性があります。KL 正則化はこの点で役立ちますが、スケーリング係数を使用して潜在的な表現の値を調整する方が適切です。このスクリプトでは、一連のトレーニング セットで潜在的な空間コンポーネントの標準偏差の大きさを検証しました。スケーリング係数は少なくとも 0.8221 である必要があることがわかりました。この場合、より保守的な値 0.3 (安定した拡散の値に近い) を使用しました。

スケーリング係数を定義することで、モデルをトレーニングできます。以下のコードで、対応するトレーニングループを見てみましょう。

 #位于训练循环内部... with torch.no_grad(): e = stage1(images) * scale_factor prompt_embeds = text_encoder(reports.squeeze(1))[0] timesteps = torch.randint(0, scheduler.num_train_timesteps, (images.shape[0],), device=device).long() noise = torch.randn_like(e).to(device) noisy_e = scheduler.add_noise(original_samples=e, noise=noise, timesteps=timesteps) noise_pred = model(x=noisy_e, timesteps=timesteps, cnotallow=prompt_embeds) if scheduler.prediction_type == "v_prediction": # 使用v-预测参数法target = scheduler.get_velocity(e, noise, timesteps) elif scheduler.prediction_type == "epsilon": target = noise loss = F.mse_loss(noise_pred.float(), target.float())

ご覧のとおり、まずデータローダーから画像とレポートを取得します。画像を処理するために、MONAIの変換機能を使用し、さらにカスタム変換を追加して、放射線レポートからランダムな文章を抽出し、入力テキストにタグを付けました。約10%のケースでは、空の文字列("" - 文章開始タグ(値 = 49406)とパディングタグ(値 = 49407)を含むベクトル)を使用して、サンプリング中に分類器を使用しないガイダンスを可能にしました。

次に、潜在表現と手がかりの埋め込みを取得します。追加するノイズ、反復処理で使用するランダムな時間ステップ、そして目標値(速度成分)を作成します。最後に、平均二乗誤差を用いて損失を計算します。

このトレーニングは500回の反復で完了しました関連するトレーニングログデータは、こちらのリンクからご覧いただけます: https://drive.google.com/drive/folders/1Ots9ujg4dSaTABKkyaUnFLpCKQ7kCgdK?usp=sharing

サンプル画像

2つのモデルをトレーニングした後、合成画像をサンプリングできます。この実験では、 https://github.com/Warvito/generative_chestxray/blob/main/src/python/testing/sample_images.pyで提供されているスクリプトを使用しました。

このスクリプトは、Hoらが提案した分類器を使用しないガイダンス法を用いて、画像生成においてテキストプロンプトの使用を強制します。このアプローチでは、生成されるデータの多様性を犠牲にして、テキストプロンプトへの忠実度の高いサンプルを取得するためのガイダンススケールを使用します。デフォルト値は7.0です。

下の図では、トレーニングされたモデルが臨床的特徴とその場所および重症度をどのように理解しているかがわかります。

評価する

このセクションでは、MONAI のメトリックを使用して、生成モデルのパフォーマンスをいくつかの側面から評価する方法を説明します。

MS-SSIM法に基づくオートエンコーダ再構成品質評価

まず、入力画像を再構成する際のオートエンコーダ kl の有効性を検証します。圧縮および再構成されたデータの品質によってサンプル品質の上限が決まることから、これはモデル開発において重要な考慮事項です。モデルが潜在表現から画像をデコードする方法をうまく学習しない場合、または潜在空間を適切にモデル化しない場合は、合成表現を現実的にデコードすることは不可能になります。 このスクリプトでは、テスト セットで合計 5000 枚の画像を使用してモデルを評価します。再構成のパフォーマンスを確認するには、マルチスケール構造類似性指標測定 (MS-SSIM) を使用できます。MS-SSIM は、2 つの画像間の類似性を測定するために広く使用されている画質評価方法です。PSNR や SSIM などの従来の画質評価方法とは異なり、MS-SSIM はさまざまなスケールで画像の構造情報を取得できます。

私たちのモデルでは、値が高いほどモデル性能が優れていることを示しています。現在のリリース(バージョン0.2)では、モデルのMS-SSIM再構築指標の平均は0.9789でした​​。

MS-SSIMサンプルの多様性

まず、モデルによって生成されるサンプルの多様性を評価します。そのために、異なる生成画像間のマルチスケール構造類似性指標を計算します。この実験プロジェクトでは、生成モデルが異なる画像を生成できる場合、合成画像ペアを比較した際に平均MS-SSIM値が低くなると仮定しています。例えば、モード崩壊などの問題が発生した場合、生成される画像は類似したものになり、MS-SSIM値は実際のデータセットで観測される値よりもはるかに低くなります。

私たちのプロジェクトでは、無条件サンプル(テキストプロンプトとして ""(空文字列)を使用して生成されたサンプル)を使用して、元のデータセットの自然な比率を維持しました。 このスクリプトに示すように、モデルから 1000 個の合成サンプルを選択し、MONAI のデータローダーを使用してすべての可能な画像ペアをロードしました。ネストされたループを使用してすべての可能な画像ペアを反復処理し、両方のデータローダーによって選択された画像が同じ場合は無視しました。ここでは、MS-SSIM 値が 0.4083 であることがわかります。テストセットの実際の画像を参照として、同じ評価を実行できます。 このスクリプトを使用すると、テストセットで MS-SSIM 値 0.4046 が得られました。これは、モデルによって生成された画像が、実際のデータで観測されるものと同様の多様性を持っていることを示しています。

しかし、多様性があるからといって、必ずしも画像が美しく、リアルに見えるとは限りません。そこで次のステップで画質を確認しましょう。

FID指数を使用して合成画像の品質を評価する

最後に、プロジェクトによって生成されたサンプルのFID(Fixed Identification Value:固定識別値)を測定します。FIDは、2つのグループ間の分布を評価し、それらの類似性を示す指標です。そのためには、距離(知覚損失に類似)を計算するための特徴を抽出できる、事前学習済みのニューラルネットワークが必要です。この例では、 torchxrayvisionパッケージで提供されているニューラルネットワークを使用することにします。

CXR画像の合成に文献で使用されているネットワークを近似するために、Dense121ネットワーク(「densenet121-res224-all」)を使用しました。このネットワークから、1024次元の特徴ベクトルを取得しました。元のFIDドキュメントで示唆されているように、特徴の数と同数のサンプルを使用することが重要です。そのため、無条件画像1000枚を使用し、テストセットの1000枚の画像と比較しました。FIDメトリックは値が低いほど良好であり、ここでは9.0237という妥当なFIDが得られました。

結論は

この記事では、MONAI生成モデルを用いたプロジェクト開発手法を、データのダウンロードから生成モデルの評価、そして最終的なデータ合成まで網羅的に紹介します。このバージョンのプロジェクトはより効率的で、ハイパーパラメータも改善されていますが、開発した拡張機能が提供する様々な機能を効果的に説明できていることを願っています。CXRモデルの改善に関するアイデアや、パッケージへの貢献をご希望の場合は、https://github.com/Warvito/generative_chestxray/issues または https://github.com/Project-MONAI/GenerativeModels/issues の課題セクションにコメントをお寄せください。

最後に、この記事で学習したモデルは、 3D Brain Generatorやその他のモデルとともに、 MONAI Model Zooウェブサイトで公開されています。Model Zooを利用すると、モデルの重みと推論コードのダウンロードがはるかに簡単になります。

最後に、チュートリアルやリソースについて詳しく知りたい場合は、https://github.com/Project-MONAI/GenerativeModels/tree/main/tutorials のチュートリアル ページにアクセスし、最新の更新情報や同様のコンテンツを入手するために私をフォローしてください。

[注記] 特に記載がない限り、本文中の画像はすべて著者提供です。

翻訳者紹介

Zhu Xianzong 氏は、51CTO コミュニティ エディター、51CTO 専門ブロガー兼講師、維坊の大学のコンピューター教師、そしてフリー プログラミング コミュニティのベテランです。

原題: Generating Medical Images with MONAI 、著者: Walter Hugo Lopez Pinaya