Fast AI Assisted Annotation and Transfer Learning with Clara Train

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. The Clara Train SDK contains these tools to speed up AI-assisted workflows for medical imaging.

This post introduces you to how fast annotation works with transfer learning workflow for Medical Imaging. Developers of medical image analysis applications for healthcare providers will learn about integrating AI-Assisted Annotation into existing applications. The Clara Train SDK part of the Clara Platform.

Demand for Better Annotation Workflow

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 that is 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 process of processing and understanding images generated by medical scanners. The NVIDIA AI-assisted Annotation helps bring deep learning based workflows and speed up the annotation process, helping radiologists save time, and increase productivity, as figure 1 shows.

AI-Assisted annotation performance 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 the deep learning tools built into the SDK with their existing medical imaging applications, such as MITK, ITK-Snap, 3D Slicer or a custom built application. This is accomplished using 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 they are interested in a 3D MRI or CT scan and receive auto-annotated results for all the 2D slices of that particular organ.

Deep learning techniques enable AI-assisted annotation. 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-annotation results image
Figure 2. Auto-Annotated results showing liver as an example

Model Registry for Organs

The SDK incorporates deep learning models to perform AI-assisted annotation on a particular organ or abnormality. NVIDIA develops and provides these deep learning models which have been pre-trained by data scientists who are experts in medical imaging. NVIDIA provides 13 different organ models for the early access release, which can be used for annotation. Additional organ models are under development.

Polygon Editing 

Since deep learning models, by their nature, are sensitive to the data that they are trained on, annotation accuracy might be slightly lower than what was 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, which is shown in figure 3.

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

The following section describes how these corrected slices can be saved 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 they are trained on, 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 to be used on a different dataset for annotation reproducing same accuracy.

Using transfer learning fits very well in medical image analysis. Since medical image analysis is seen as 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 transfer learn 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 retraining the pre-trained models based on new datasets. Figure 4 highlights this process.

TLT workflow improves annotation accuracy image
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.

Python wrappers are provided 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 is 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 Pretrained 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 SDK provides easy-to-use conversion utilities.

Pretrained Models provided with the Clara Train SDK can be integrated into other applications. These include organ and tumor segmentation models as well as several annotation models to be used for AI assisted annotation. These annotation models can be loaded into an annotation server included in SDK so that it listens to the client annotations performed as described above.

Models can be integrated into other applications. The toolkit supports applications built using NVIDIA CLARA Platform and deployed for inference with the TensorRT Inference Server.

Integrate AI-Assisted Annotation into your Existing Application in Three 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 – The NGC model registry hosts the deep learning models used for annotation.
  2. Starting annotation server – Annotation server is made available in the Clara Train SDK container from the NGC registry.
  3. Client integration – Integrating the annotation APIs 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 -t -p 5000:5000 -v /var/tmp/models:/var/tmp/models 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.

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.

tlt-pull -m $MODEL_NAME -v $VERSION -k $API_KEY -d ./path/to/save/model -o nvtltea -t dlmed

$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.

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 two organ models, spleen and liver were downloaded by the user.

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

Each model has been trained with certain configuration parameters defined in the config_file. Create the config file /var/tmp/models/spleen.json with the contents shown below.

{
   {
      "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.NormalizeNonzeroIntensities",
         "fields": "image"
      },
   {
   "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 start_aas.sh 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 loaded with specific organ models, the client component of the SDK is delivered as an open source reference implementation which can be used. Please note the implemetaion is presented 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([[70,172,86],[105,161,180],[125,147,164],[56,174,124],[91,119,143],[77,219,120]]);
  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 and pointRatio controls the density of points in each polygon.

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 p_new after the user change, old polygon p_old (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 = nvidia::aiaa::Polygons::fromJson([[[69,167],[73,156],[78,146],[87,137],[98,131],[108,130],[118,132],[123,141],[139,155],[119,161],[109,166],[98,170],[89,176],[80,183],[71,182],[69,172]]]);
  nvidia::aiaa::Polygons = nvidia::aiaa::Polygons::fromJson([[[69,167],[74,157],[78,146],[87,137],[98,131],[108,130],[118,132],[123,141],[139,155],[119,161],[109,166],[98,170],[89,176],[80,183],[71,182],[69,172]]]);
  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(p_new, p_old, 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 here and apply 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, which won the 2018 challenge. 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. You can also learn more about deploying Clara platform solutions using the Clara Deploy SDK.

With these solutions for NVIDIA Clara platform we aim to accelerate the medical imaging workflow, enabling faster annotation, training, and deployment for medical imaging solutions. If you are a 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