😖

RSNA 2025 3D脳動脈セグメンテーションの作り方

に公開
AI
ディープラーニング
Kaggle
segmentation
rsna
tech

はじめに

kaggle RSNA2025 では、脳内動脈瘤（intracranial aneurysm）を検出・局所化することを目的とした AI コンペティションが開催されている。
配布されているデータには、以下が含まれている：

  • 画像（CT や MRI 等、4000件くらい）
  • どの部位に動脈瘤があるかのラベル（１３クラス＋有無、予測対象）
  • 多クラスのセグメンテーションラベル（13クラス、補助データ、200件ほど）
  • 動脈瘤の座標情報（13クラスxy座標、補助データ）

この中で、3D で多クラスセグメンテーションを学習させる アプローチに苦戦したため、本記事ではその過程と工夫を記録しておきたい。

f値で0.85くらいなのでそこそこ認識できてると思います。

前提

もしかしたら普通に1段の学習できてる人もいるかもしれないので、私の対処法としてみてください。その場合やり方教えてください・・・

前処理

配布データの性質上、撮像範囲（スライス数・空間解像度など）がケースごとにばらつきがある。
このため、後続の学習を安定させるべく、動脈領域を抽出してクロップ する前処理を入れた。
こうすることで、画素強度・構造位置のばらつきが抑えられ、セグメンテーション精度が向上しやすくなる。
動脈瘤セグメンテーションは、このクロップ後領域を対象に行うことを前提とする。

図：クロップする場所のイメージ。配布された動脈のセグマスクをBBOXで囲ったもの

最終的なアプローチ

まず全体像です。

最終的に採用した構成は、3D U-Net を 2 段のカスケード構造でつなぎ、
その中間段階で「血管（動脈）としてのクラス統合」出力を導入して Loss を取るという手法です。
こうすることで、最終的な多クラス動脈瘤セグメンテーションの精度を安定化できました。

構成：2段カスケード構造

  1. 1段目：大きいモデル。monaiで公開されている脳データで事前学習済みモデル。
    • 少ないデータで性能を出すために決め打ちでこれにした
    • これ以外だとぜんぜん学習できなかった
  2. 2段目：ちっちゃいモデル。事前学習はされてない。

学習手順は次の通り：

  1. 第1段階モデル（2 クラス：前景／背景） を学習
  2. 第2段階モデル を K クラス（各動脈＋背景）を予測
  3. 学習時には、第1段階・第2段階で異なる重み（Loss weight）を設定

こうした多段階学習（または段階的ファインチューニング）により、最終的に安定した性能が得られた。

試行錯誤の流れ

以下のようなステップを経て、最終構成に至った：

  1. 最初から多クラスで直接学習
    　→ 結果：精度が不安定（収束しづらい、False 部分が散発）

  2. クラスを統合して 1 クラスで学習（背景 vs すべての動脈）
    　→ 結果：構造を捉えるモデルは完成。

  3. 段階分割：血管を検出するモデルとクラス分割モデルを別段階に（カスケード構造）
    　→ 結果：安定して動かせるようになった

    • インスピレーション元として、OpenPose などが段階的に前値を入力して修正していく構造を持つ点を参考にした
    • 一度に両方を学習させると収束が難しいため、第1段階 → 第2段階 に分けて学習させる方式を採用

結果として、1 回のエンドツーエンド学習ではうまく動かず、2 段階フェーズでの学習という方式に落ち着いた。

コード例（モデル定義部分）

以下は、Monai の事前学習モデルと自作 U-Net モデルを“無理やり”結合してカスケード構造を組んだ部分の例だ。

from monai.networks.nets import UNet
from monai.bundle import load

def create_model(config):
    device = config.device
    K = getattr(config, "num_artery_classes", 14)  # 最終クラス数

    # ----- Stage1: 2クラス（背景 / 前景） -----
    stage1 = load(
        name="brats_mri_segmentation",
        source="monaihosting",
        device=device,
        net_override={
            "in_channels": config.in_channels,
            "out_channels": 2,  # 背景／前景の 2 クラス
        },
    )

    # ----- Stage2: 入力 = 元画像 + Stage1 の前景確率（1チャネル付加） -----
    stage2_in_ch = config.in_channels + 1
    stage2 = UNet(
        spatial_dims=3,
        in_channels=stage2_in_ch,
        out_channels=K,
        channels=(16, 32, 64, 128),
        strides=(2, 2, 2),
        num_res_units=1,
        norm=Norm.BATCH,
    ).to(device)

    class CascadeSeg(nn.Module):
        def __init__(self, s1, s2):
            super().__init__()
            self.stage1 = s1
            self.stage2 = s2

        def _take_tensor(self, out):
            # bundle の出力が dict のケースにも対応
            if torch.is_tensor(out):
                return out
            if isinstance(out, dict):
                for v in out.values():
                    if torch.is_tensor(v):
                        return v
            raise RuntimeError("Stage1 の出力テンソルが見つかりません")

        def forward(self, x):
            # ---- Stage1 ----
            logit1 = self._take_tensor(self.stage1(x))  # [B,2,D,H,W]
            prob1  = F.softmax(logit1, dim=1)           # softmax 確率
            fg1    = prob1[:, 1:2, ...]                 # 2ch 目（前景）を抽出 → 形状 [B,1,D,H,W]

            # ---- Stage2 ----
            x2 = torch.cat([x, fg1], dim=1)             # 入力チャネルを拡張
            logit2 = self.stage2(x2)                    # [B,K,D,H,W]

            return {
                "coarse": logit1,   # Stage1 出力ロジット（BCE/Dice に使える）
                "fine": logit2,     # 最終 K クラス出力ロジット
            }

    model = CascadeSeg(stage1.to(device), stage2).to(device)
    return model

補足・注意点

  • stage1 の出力にはモジュール固有の構造（辞書型返却など）があるため、 _take_tensor 関数で最初に見つかるテンソルを取り出すようにしている
  • softmax をかけて前景確率を得たうえで、そのチャネルを入力に結合
  • stage2 側はオリジナル U-Net 構造で構成
  • coarse（2 クラス出力）と fine（K クラス出力）双方で Loss を取る設計にすることが可能

Discussion