Tutorial: Brain Segmentation PyTorch

We are demonstrating from importing the models into AIAA to actual making requests to the server.

1. Follow Running AIAA to start your server.

2. Follow Convert PyTorch trained network to convert the example PyTorch model.

3. Write your own transforms that are missing from Clara Train API Below are two specific transforms that you need for this tutorial. Put them inside a file called “transforms.py” and put that in /workspace/transforms/ You can also download the file here: transforms.py.

import logging
import numpy as np
import cv2

class MyLabelNPSqueeze(object):
  def __init__(self, label_in='model', label_out='model', dtype='uint8'):
    self.key_label_in = label_in
    self.key_label_out = label_out
    self.dtype = dtype

  def __call__(self, data):
    logger = logging.getLogger(self.__class__.__name__)

    label = data[self.key_label_in]
    logger.debug('Input Label Shape: {}'.format(label.shape))

    label = np.squeeze(label).astype(self.dtype)
    logger.debug('Output Label Shape: {}'.format(label.shape))

    data[self.key_label_out] = label
    return data

class MyOpenCVWriter(object):
  def __init__(self, image_in='model'):
    self.key_image_in = image_in

  def __call__(self, output_file, data):
    logger = logging.getLogger(self.__class__.__name__)
    # convert 0-1 back to 0-255
    image = data[self.key_image_in] * 255

    logger.debug('Saving Image {} to: {}'.format(image.shape, output_file))
    cv2.imwrite(output_file, image)
    return output_file

4. Prepare your configuration:

Now you have models and transforms ready, you should write a config.json file for AIAA to understand. Open a file named segmentation_2d_brain.json and write your configuration as follows. You can also download it here: segmentation_2d_brain.json.

{
  "version": "2",
  "type": "segmentation",
  "labels": [
    "brain"
  ],
  "description": "2D segmentation training and evaluation examples to identify and segment the brain based on Unet model given in https://pytorch.org/hub/mateuszbuda_brain-segmentation-pytorch_unet/",
  "pre_transforms": [
    {
      "name": "LoadPng",
      "args": {
        "fields": "image"
      }
    },
    {
      "name": "ConvertToChannelsFirst",
      "args": {
        "fields": "image"
      }
    },
    {
      "name": "NormalizeNonzeroIntensities",
      "args": {
        "fields": "image"
      }
    }
  ],
  "inference": {
    "image": "image",
    "scanning_window": false,
    "batch_size": 1,
    "roi": [
      256,
      256
    ],
    "trtis": {
      "input_channels": {
        "image": 3
      },
      "output_channels": {
        "model": 1
      },
      "gpu_instance_count": 1,
      "max_batch_size": 8,
      "platform": "pytorch_libtorch"
    }
  },
  "post_transforms": [
    {
      "name":  "ThresholdValues",
      "args": {
        "fields": "model",
        "threshold": 0.5,
        "dtype": "uint8"
      }
    },
    {
      "name": "MyLabelNPSqueeze",
      "path": "transforms.MyLabelNPSqueeze",
      "args": {
        "label_in": "model",
        "label_out": "model"
      }
    }
  ],
  "writer": {
    "name": "MyOpenCVWriter",
    "path": "transforms.MyOpenCVWriter",
    "args": {
      "image_in": "model"
    }
  }
}

5. Upload the model:

Use the command below to load it in AIAA.

# Import Model
curl -X PUT "http://127.0.0.1/admin/model/segmentation_2d_brain" \
     -F "config=@segmentation_2d_brain.json;type=application/json" \
     -F "data=@unet.pt"

6. Call the APIs:

Now the model is loaded in AIAA and ready to serve. Type following commands in your Linux machine and the result will be stored as result.png.

# download example
wget https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png

# call segmentation API
curl -v -X POST "http://127.0.0.1/v1/segmentation?model=segmentation_2d_brain&output=image" \
     -H "accept: multipart/form-data" \
     -H "Content-Type: multipart/form-data" \
     -F "params={}" \
     -F "image=@TCGA_CS_4944.png" \
     -o result.png

7. Verify the result. The input image is on the left while the result is on the right.