Fast AI Assisted Annotation and Transfer Learning Powered by the Clara Train SDK

pasted image 0 28

The growing volume of clinical data in medical imaging slows down identification and analysis of specific features in an image. This reduces the annotation speed at which radiologists and imaging technicians capture, screen, and diagnose patient data.

The demand for artificial intelligence in medical image analysis has drastically grown in the last few years. AI-assisted solutions have emerged to speed up annotation workflows and increase productivity.

NVIDIA announced the AI-Assisted Annotation and Transfer Learning at the Radiological Society of North America (RSNA) 2018, a premier conference which showcases the latest research and technology advances making headlines in radiology. Clara Train SDK contains these tools to speed up AI-assisted workflows for medical imaging.

This post introduces to developers & data scientists how fast annotation works with transfer learning workflow for Medical Imaging. 

Annotation tools in medical imaging applications help radiologists annotate organs and abnormalities. These applications consist of software tools to draw boundaries around the organs of interest, isolating them from the rest of the dataset. These applications typically operate on 2D slices. A radiologist takes a single 2D image slice from a patient’s scan and marks annotation boundaries and corrections as needed. This manual process repeats for the next slice in which the organ of interest is annotated and corrected again. Some applications support 3D tools for annotation with region growing techniques that perform fairly well with certain organs but poorly on other organs. Moreover, both these 2D and 3D tools for annotation do not learn and adapt to the dataset being annotated which becomes crucial for medical imaging applications as the datasets grow.

Since these workflows are completely manual, the annotation process takes time. A normal Computed Tomography (CT) or Magnetic Resonance Imaging (MRI) scan of a patient has hundreds of 2D slices to cover a particular organ in 3D; new generation scanners producing increasingly higher resolution scans.

AI Assisted Annotation

NVIDIA brings its leadership position in modern artificial intelligence and deep learning to help automate the processing and understanding of images generated by medical scanners. The NVIDIA AI-assisted Annotation enables deep learning based applications by providing developers with tools that make it possible to speed up the annotation process, helping radiologists save time, and increase productivity, as figure 1 shows.

Annotation performance increase chart
Figure 1. Annotation speedup for different organs.

Clara Train SDK’s AI Assisted Annotation accelerates the annotation process by enabling application developers to integrate deep learning tools built into the SDK with existing medical imaging viewers , such as MITK, ITK-Snap, 3D Slicer or a custom-built application viewer. This uses a simple API and requires no prior deep learning knowledge.

As a result, radiologists can increase their productivity by analyzing more patient data while still using their existing workflows and familiar tools.

Generating Segmentation with Deep Learning Models

With AI-assisted annotation, radiologists don’t have to manually draw on each 2D slice to annotate an organ or an abnormality. Instead, they simply click a few extreme points on a particular organ of interest in a 3D MRI or CT scan and receive auto-annotated results for all the 2D slices of that particular organ.

When a user sends the extreme points to the SDK, a deep learning model receives this as input and returns the inference results of the segmented organ or abnormality Figure 2 highlights the steps in the process.

Auto-annotated results image
Figure 2. Auto-Annotated results showing liver as an example. Reference to algorithm: https://arxiv.org/abs/1711.09081

Model Registry for Organs

The SDK incorporates deep learning models to perform AI-assisted annotation on a particular organ or abnormality. NVIDIA provides 13 different organ models for the early access release which have been pre-trained on public datasets by NVIDIA’s data scientists. These models can be used for annotation and can kick start your development effort by enabling faster annotation of datasets for AI algorithms. Additional organ models are under development.

Polygon Editing 

Since deep learning models, by their nature, are sensitive to the data used to train them, annotation accuracy might be lower than originally achieved with the training data.

If in case there are 2D slices with incorrect or partial annotation, 2D smart polygon editing feature helps in correcting these slices. When a user moves a single polygon point on the 2D slice, all other points around a radius will automatically snap to the organ boundaries thereby making it easier and efficient when correcting, shown in figure 3.

Corrections with smart polygon editing image
Figure 3. Corrections made easy with smart polygon editing

The following section describes how to save these corrected slices as new ground-truth and used to retrain the annotation model, thereby consistently increasing annotation accuracy.

Transfer Learning Workflow

Deep learning models are sensitive to the data used to train them, as described earlier. Factors such as varying scanner configurations, age differences of patients, and so on, need to be taken into account. This makes it hard to train the deep learning models on a specific dataset and deploy them on a different dataset.

Using transfer learning is well suited for medical image analysis. Since medical image analysis is a computer vision task, CNNs represent the best performing methods for this. Yet getting a large well-annotated dataset is considerably harder in the medical domain compared to general computer vision domain because of the domain expertise required to annotate the medical images. This makes transfer learning a natural fit for medical image analysis, use pre-trained CNN on larger database and then apply transfer learning to a target domain of medical images with limited availability.

As a solution, Transfer Learning for medical imaging helps users adapt pre-trained models into their datasets. Deep learning models used for annotation can be tuned and improved by re-training these pre-trained models based on new datasets. Figure 4 highlights this process.

TLT Workflow diagram
Figure 4. Workflow with Transfer Learning Toolkit to improve annotation accuracy

Clara Train SDK contains Annotation SDK and Transfer Learning Toolkit and is available in a docker container.

NVIDIA Transfer Learning Toolkit (TLT) provides domain specific pre-trained models for Intelligent Video Analytics. Same underlying technology allows users to adapt Medical Imaging Segmentation DNNs. Clara Train SDK for Medical Imaging contains several models from the public domain, optimized in-house to achieve high accuracy. Medical imaging developers can choose from one of the provided models as a starting point and use a high level python SDK to retrain and adapt these models with custom data.

We provide python wrappers for model retraining and exporting. These models are trained using Tensorflow framework on public datasets. Complexity of adaptation is considerably reduced because of easy to use shell commands, data conversion guidance, and an export wrapper for converting trained models to a TensorRT optimized graph.

A deep learning image segmentation approach is used for fine-grained predictions needed in medical imaging. Image segmentation algorithms partition input image into multiple segments. Images become divided down to the voxel level (volumetric pixel is the 3-D equivalent of a pixel) and each pixel gets assigned a label or is classified. For every pixel in the image, the network is trained to predict the pixel class. This allows the network to both identify several object classes in each image and determine the location of objects. Image segmentation generates a label image the same size as the input in which the pixels are color-coded according to their classes. 

Accurate segmentation is key in various medical applications.

Using Pre-trained Models

The medical segmentation decathlon challenge site provides a reliable dataset starting point for segmentation model development. All segmentation models in the SDK are trained from the identical segmentation training pipeline, with configurations for brain, heart, pancreas and spleen segmentation. Since several models have been trained with images 1x1x1mm resolution, data needs to be converted to 1x1x1mm NIfTI format. The site provides easy-to-use conversion utilities are provided

Models can be integrated into other applications. The toolkit supports applications built using NVIDIA’s Clara developer platform and deployed for inference with the TensorRT Inference Server.

Integrate AI-Assisted Annotation into your Existing Application in Four Steps

Annotation SDK follows a client-server approach to integrate into an application. Following these are three steps enables you to integrate the SDK into an existing medical imaging application.

  1. Downloading the models – NVIDIA’s model registry hosts the deep learning models used for annotation.
  2. Starting annotation server – Annotation server is made available as a container from NVIDIA GPU cloud registry.
  3. Client integration – Integrating the annotation SDK using either C++ or Python client source code.

Downloading the Models

AI-Assisted Annotation relies on deep learning models developed in-house at NVIDIA. These models are included in NVIDIA’s model registry and can be downloaded using simple shell commands.

All the commands below are part of the container. Start the container and access the terminal to use these commands.

docker run --runtime=nvidia -it -p 5000:5000 -v /path/to/save/model:/var/tmp/models -e NVIDIA_VISIBLE_DEVICES=0 nvcr.io/nvtltea/med/tlt-annotate:v0.1-py3 /bin/bash

A user can list all the models available on the model registry using the following command.

API_KEY=yourAPIkey
tlt-pull -lm -k $API_KEY -o nvtltea -t med

All the models starting with “annotation” can be used for annotation purposes. The following command shows how to pull a single model from the model registry.

MODEL_NAME=annotation_spleen
VERSION=1
tlt-pull -m $MODEL_NAME -v $VERSION -k $API_KEY -d /var/tmp/models/$MODEL_NAME -o nvtltea -t med

$MODEL_NAME is the name of the organ model, $VERSION is the model version, and $API_KEY is the API key that a user obtains through NGC.

Convert the checkpoint to TensorRT optimized model

Once a model is downloaded, use the tlt-export command to export the model for use with the server.

CHECKPOINT=/var/tmp/models/annotation_spleen
tlt-export --model_name 'model' --model_file_path ${CHECKPOINT} --input_node_names 'cond/Merge' --output_node_names 'NV_OUTPUT_LABEL'

Starting the Annotation Server

The annotation server needs to know the location of the downloaded models. The configuration file engine_pool.json file must be configured based on the models that were downloaded by the user.

An example configuration file engine_pool.json is shown below, assuming the user downloaded two organ models, spleen and liver.

[
{
"name": "MyAnnotationSpleenModel",
"class": "Dextr3dCroppedEngine",
"module": "medical.aas.dextr3dcropped",
"config_file": "/var/tmp/models/spleen.json",
"labels": ["spleen"]
},
{
"name": "MyAnnotationLiverModel",
"class": "Dextr3dCroppedEngine",
"module": "medical.aas.dextr3dcropped",
"config_file": "/var/tmp/models/liver.json",
"labels": ["liver"]
}
]

Each model has been trained with certain configuration parameters defined the config_file. Create the config file /var/tmp/models/spleen.json with the contents shown below. Other examples can be found inside our docker container under /opt/nvidia/medical/aas/configs.

{
    "model_format": "trtmodel",
    "model_path": "/var/tmp/models/annotation_spleen/model.trt.pb",
    "thresh": 0.5,
    "result_dir": "/var/tmp/experiments/aiaa",

    "input_nodes":
    {
        "image": "cond/Merge"
    },

    "output_nodes":
    {
        "model": "NV_OUTPUT_LABEL"
    },

    "batch_size": 1,

    "pre_transforms":
    [    
        {
            "name": "transforms.VolumeTo4dArray",
            "fields": ["image", "label"]
        },           
        {
            "name": "transforms.ScaleIntensityRange",
            "field": "image",
            "a_min": -1024,
            "a_max": 1024,
            "b_min": -1.0,
            "b_max": 1.0,
            "clip": true
        },        
        {
            "name": "transforms.AddExtremePointsChannel",
            "image_field": "image",
            "label_field": "label",
            "sigma": 3,
            "pert": 0
        }                
    ],

    "post_transforms":
    [ 
        {
            "name": "transforms.SplitAcrossChannels",
            "applied_key": "model", 
            "channel_names": ["background", "prediction"]
        }        
    ],

    "writers":
    [
        {
            "applied_key": "prediction",
            "name": "writers.WriteNiftyResults",
            "dtype": "float32"
        }               
    ]
}

Use the script to start the annotation server.

/opt/nvidia/medical/aas/start_aas.sh --config /var/tmp/models/engine_pool.json

Once the annotation server is up and running, you can verify using the following curl command and check the server’s response.

curl http://$MACHINE_IP:5000/v1/models

Client Integration

Once the annotation server is setup and started with specific organ models, the client component of the SDK is delivered as an open source reference implementation to demonstrate the integration process.

Client code and libraries are provided for both C++ and Python languages on NVIDIA Github page.

Let’s use the C++ implementation as an example. The following section demonstrates how to integrate the AI-Assisted Annotation SDK into other medical imaging applications.

Public interfaces are defined in the open source example at https://github.com/NVIDIA/ai-assisted-annotation-client; the client.h interface defines the API calls that interface with the annotation server.

Begin by instantiating a “Client” object that interfaces with the server, giving server ip-address and API version as parameters. (Replace the example IP address shown with the correct one for your server).

nvidia::aiaa::Client client("http://110.120.15.26:5000/v1");

In order to know which organ models the annotation server supports, you can query the client and receive a ModelList with information about all the models.

nvidia::aiaa::ModelList modelList = client.models();

Taking “spleen” as one of the organ labels, retrieve that particular model.

std::string label = "spleen";
nvidia::aiaa::Model model = modelList.getMatchingModel(label);

Input file and output results in this case are supplied in NIfTI file format. A user supplies an input dataset (inputImageFile) from the application which will be annotated for a particular organ.

In addition to the input dataset, the application should also supply the extreme points of an organ that the user clicks in 3D (pointSet). A padding region (PAD) around the organ will be automatically added to ensure that our model sees enough context around the organ of interest. The padded region is then cropped and resampled to ROI_SIZE which by default this is 128x128x128 in x, y, z directions.

std::string inputImageFile = "image.nii.gz";
  std::string outputDextr3dImageFile = "result.nii.gz";
  nvidia::aiaa::Point3DSet pointSet = nvidia::aiaa::Point3DSet::fromJson("[[1,2,3]"]);
  int PAD = 20; // Padding boundary 
  std::string ROI_SIZE = "128x128x128";
  int SIGMA = 3;

Once the required parameters are declared, the API call shown below can be used to do the actual inferencing, returning a 3D binary mask of the organ that has been segmented. This result is written into the file specified in the parameter outputdextr3dImageFile in NIfTI file format.

int ret = client.dextr3d(model, pointSet, inputImageFile, outputdextr3dImageFile, PAD, ROI_SIZE, SIGMA);

The 3D binary mask from the Inferencing result above can be converted into polygons using this API.

int pointRatio = 10;
nvidia::aiaa::PolygonsList result = client.mask2Polygon(pointRatio, outputdextr3dImageFile);

PolygonsList is a list of Polygon objects returned as a result from the mask2Polygon conversion API.

If a set of slices requires correction, the smart polygon editing feature can be used to automatically adjust neighboring points on a slice by moving a single point.

The API call fixPolygon takes a new polygon (p1) after the user change, old polygon (p2) (from before the user change), indices of the polygon, and the vertex that was changed. Parameter input2DImageFile is the input slice. Results are written back to output2DImageFile.

  int neighborhoodSize = 10;
  nvidia::aiaa::Polygons p1 = nvidia::aiaa::Polygons::fromJson("[[10,20],[20,30]");
  nvidia::aiaa::Polygons p2 = nvidia::aiaa::Polygons::fromJson("[[10,20],[21,31]]");
  int polygonIndex = 0;
  int vertexIndex = 1;
  std::string input2DImageFile = "image_slice_2D.png";
  std::string output2DImageFile = "updated_slice_2D.png";
  // fixPolygon():: Call Fix Polygon with new Polygon Points
  nvidia::aiaa::Polygons result = client.fixPolygon(p1, p2, neighborhoodSize, polygonIndex, vertexIndex, input2DImageFile, output2DImageFile);

Integration with MITK

The Medical Imaging Interaction Toolkit (MITK) is a free open-source software toolkit for development of interactive medical image processing software by the German Cancer Research Center (DKFZ).  NVIDIA and DKFZ closely collaborated to bring Clara Train SDK into MITK. Clara Train SDK annotation will be part of the MITK Workbench Segmentation plugin, providing 2D and 3D tools for polygon editing and segmentation respectively.  Starting with MITK 2018.04.2, users can take advantage of Clara Train SDK when using MITK in their workflows.  Users can download the latest MITK Workbench from the official website and apply for Early Access to use the CLARA Train SDK container from NVIDIA GPU Cloud.  MITK plugin integration source code has been open sourced by NVIDIA on Github.

Apply for early access

The need for improved annotation speed and accuracy is rising. NVIDIA research has made big strides in medical imaging, including the recent BraTS (Multimodal Brain Tumor Segmentation) Challenge model. Clara Train SDK for Medical Imaging contains several models from the public domain, optimized in-house to achieve high accuracy. Medical imaging developers can choose from one of the provided models as a starting point and use a high level python SDK to retrain and adapt these models with custom data.

With these solutions for NVIDIA’s Clara developer platform, we aim to accelerate the medical imaging workflow, enabling faster annotation, training, and deployment for medical imaging solutions. If you are a data scientist or software engineer developing medical imaging applications including annotation, we’d love to get your thoughts. Please apply for early access to get started today.

No Comments