ClaraでCOVID-19病変箇所をセグメンテーションする方法

記事の概要

• Clara Train SDKと呼ばれるアノテーション支援ツールがある
• 胸部CTの病変箇所のセグメンテーションもできるけど，どうもファイルに欠損があるっぽい
• 設定ファイルを追加し，AIAAサーバーに登録することで解決できたので共有

検証環境

docker pull nvcr.io/nvidia/clara-train-sdk:v3.1.01

cf. https://ngc.nvidia.com/catalog/containers/nvidia:clara-train-sdk

学習済みモデルのダウンロード

ngc registry model download-version "nvidia/clara_train_covid19_ct_lesion_seg:1"

cf. https://ngc.nvidia.com/catalog/models/nvidia:clara_train_covid19_ct_lesion_seg

設定ファイルの追加

AIAAサーバーで動かすための設定ファイルconfig_aiaa.jsonがないので，下記のファイルを作成する.

• clara_train_covid19_ct_lesion_seg_v1/config/config_aiaa.json

{
  "version": "3",
  "type": "segmentation",
  "labels": [
    "label_class0",
    "label_class1"
  ],
  "description": "The model described in this card is used to segment the COVID-19 affected region from the 3D chest CT images",
  "pre_transforms": [
    {
      "name": "LoadNifti",
      "args": {
        "fields": "image",
        "as_closest_canonical": true
      }
    },
    {
      "name": "ConvertToChannelsFirst",
      "args": {
        "fields": "image"
      }
    },
    {
      "name": "ScaleByResolution",
      "args": {
        "fields": "image",
        "target_resolution": [0.8, 0.8, 5.0]
      }
    },
    {
      "name": "ScaleIntensityRange",
      "args": {
        "fields": "image",
        "a_min": -1500,
        "a_max": 500,
        "b_min": 0.0,
        "b_max": 1.0,
        "clip": true
      }
    }
  ],
  "inference": {
    "image": "image",
    "name": "TRTISInference",
    "args": {
      "batch_size": 1,
      "roi": [
        384,
        384,
        32
      ],
      "scanning_window": true
    },
    "trtis": {
      "platform": "tensorflow_graphdef",
      "max_batch_size": 1,
      "input": [
        {
          "name": "NV_MODEL_INPUT",
          "data_type": "TYPE_FP32",
          "dims": [
            1,
            384,
            384,
            32
          ]
        }
      ],
      "output": [
        {
          "name": "NV_MODEL_OUTPUT",
          "data_type": "TYPE_FP32",
          "dims": [
            2,
            384,
            384,
            32
          ]
        }
      ],
      "instance_group": [
        {
          "count": 1,
          "kind": "KIND_AUTO"
        }
      ]
    },
    "tf": {
      "input_nodes": {
        "image": "NV_MODEL_INPUT"
      },
      "output_nodes": {
        "model": "NV_MODEL_OUTPUT"
      }
    }
  },
  "post_transforms": [
    {
      "name": "ArgmaxAcrossChannels",
      "args": {
        "fields": "model"
      }
    },
    {
      "name": "FetchExtremePoints",
      "args": {
        "image_field": "image",
        "label_field": "model",
        "points": "points"
      }
    },
    {
      "name": "CopyProperties",
      "args": {
        "fields": [
          "model"
        ],
        "from_field": "image",
        "properties": [
          "affine",
          "original_affine",
          "as_canonical"
        ]
      }
    },
    {
      "name": "RestoreOriginalShape",
      "args": {
        "field": "model",
        "src_field": "image",
        "is_label": true
      }
    }
  ],
  "writer": {
    "name": "WriteNifti",
    "args": {
      "field": "model",
      "dtype": "uint8",
      "revert_canonical": true
    }
  }
}

AIAAサーバーへ登録

$LOCAL_PORT=80
curl -X PUT "http://127.0.0.1:$LOCAL_PORT/admin/model/clara_train_covid19_ct_lesion_seg" \
     -F "config=@clara_train_covid19_ct_lesion_seg_v1/config/config_aiaa.json;type=application/json" \
     -F "data=@clara_train_covid19_ct_lesion_seg_v1/models/model.trt.pb"

cf. https://docs.nvidia.com/clara/tlt-mi/aiaa/server_apis.html

これで，3D Slicerで検証できるようになります．