Bring your own InferencePipeline

You can also cascade or chain you inference in any order by defining your own InferencePipeline. You need to put your custom code in <AIAA workspace>/lib folder.

The following interfaces is provided to get/run other inferences.


class InferencePipeline:

    def get_model_config(name):
        """Gets model configuration given name."""

    def init_inference(config: ModelConfig, backend: str):
        """Initializes an inference object"""

    def run_inference(name, inference, data, config: ModelConfig, triton_config):
        """Runs inference to get result"""

    def run_transformer(data, transformers_config, log_prefix='POST'):
        """Runs transforms given transform configs"""

    def pipeline(self, name, data, config: ModelConfig, triton_config):
        """Defines a pipeline"""

Write a Custom InferencePipeline

This example pipeline will run all the segmentation models passed into it and union the segmentation masks.

import logging
from typing import List, Optional

import numpy as np

from aiaa.configs.modelconfig import ModelConfig, ModelType
from aiaa.inference.inference import Inference
from aiaa.inference.inference_pipeline import InferencePipeline
from aiaa.utils.aiaa_exception import AIAAError, AIAAException

class CustomInferencePipeline(InferencePipeline):
    def __init__(self, models: Optional[List] = None, backend='TRITON'):
        # HINT:: You can pass any additional init params through config_aiaa.json
        self._inferences = dict()
        self.models = models
        self.backend = backend

    def _init_context(self, models, backend):
        logger = logging.getLogger(__name__)'Init Context for: {}'.format(models))

        # Create and cache inference object for loaded models
        for name in models:
            if self._inferences.get(name):
      'Inference for {} is pre-loaded;  Using the existing one!')

            config = InferencePipeline.get_model_config(name)
            if config is None:
                raise AIAAException(AIAAError.INFERENCE_ERROR, f'Model {name} does not exists/loaded in AIAA.')
            if config.get_type() != ModelType.SEGMENTATION:
                raise AIAAException(AIAAError.INFERENCE_ERROR, f'Model {name} is not type {ModelType.SEGMENTATION}')

            self._inferences[name] = (config, InferencePipeline.init_inference(config, backend))'PreLoaded Inferences: {}'.format(self._inferences.keys()))

    def pipeline(self, name, data, config: ModelConfig, triton_config):
        logger = logging.getLogger(__name__)'Run Inference Pipeline for: {}'.format(name))

        if config.get_type() != ModelType.PIPELINE:
            raise AIAAException(AIAAError.INFERENCE_ERROR,
                                f'Config need to specify model type as "{ModelType.PIPELINE}"')

        input_params = data.get('params')
        if input_params is None:
            raise AIAAException(AIAAError.INFERENCE_ERROR, 'Missing input params.')'Input Params: {}'.format(input_params))

        # Get backend from input parameters
        backend = input_params.get('backend', self.backend)'Backend: {}'.format(backend))
        if backend == 'TRITON' and triton_config is None:
            raise AIAAException(AIAAError.INFERENCE_ERROR,
                                'Triton config is none. Please make sure you run AIAA with Triton backend.')

        # Get model names from input parameters
        models = input_params.get('models', self.models)'Input Models: {}'.format(models))
        if models is None or len(models) == 0:
            raise AIAAException(AIAAError.INFERENCE_ERROR, 'Input params do not have any models to run')

        # Init all referred Inferences
        self._init_context(models, backend)

        results = []
        count = 0
        for model in models:
            i = self._inferences[model]
            model_config: ModelConfig = i[0]
            inference: Inference = i[1]

            pre_transformers = model_config.get_pre_transforms()
            post_transformers = model_config.get_post_transforms()

            # We exclude common pre-transforms
            common_pre_transformers = config.get_pre_transforms()
            pre_transformers = [t for t in pre_transformers if t not in common_pre_transformers]

            # HINT:: You can run pre-transforms or post-transforms for an existing model
            result = InferencePipeline.run_transformer(data.copy(), pre_transformers, log_prefix='PRE-{}'.format(count))

            # HINT:: You can also chain previous model input as input to next model (depending on use-case)
            result = InferencePipeline.run_inference(model, inference, result, model_config, triton_config)

            # We exclude common post-transforms
            common_post_transformers = config.get_post_transforms()
            post_transformers = [t for t in post_transformers if t not in common_post_transformers]
            result = InferencePipeline.run_transformer(result, post_transformers, log_prefix='POST-{}'.format(count))

            # here we assume those models only have one output
            count += 1

        result_mask = None
        for result in results:
            # Union result_mask
            if result_mask is None:
                result_mask = result
                result_mask = np.logical_or(result_mask, result)'Result Mask: {}'.format(result_mask.shape))

        # Store result back into data (we assume this pipeline has only one output)
        data[config.get_inference_output()] = result_mask.astype(np.float)
        return data

    def close(self):
        # Release any resources

Let’s save this custom inference pipeline in and copy it into <AIAA workspace>/lib folder.

An example config_aiaa.json for this pipeline is as below:

  "version": 1,
  "type": "pipeline",
  "labels": [
  "description": "This pipeline will run multiple models on the input image and union their results",
  "pre_transforms": [
      "name": "monai.transforms.LoadImaged",
      "args": {
        "keys": "image"
      "name": "monai.transforms.AddChanneld",
      "args": {
        "keys": "image"
      "name": "monai.transforms.Spacingd",
      "args": {
        "keys": "image",
        "pixdim": [
  "inference": {
    "input": "image",
    "output": "pred",
    "AIAA": {
      "name": "custom_inference_pipeline.CustomInferencePipeline"
    "TRITON": {
      "name": "custom_inference_pipeline.CustomInferencePipeline"
  "post_transforms": [
      "name": "aiaa.transforms.Restored",
      "args": {
        "keys": "pred",
        "ref_image": "image"
      "name": "aiaa.transforms.ExtremePointsd",
      "args": {
        "keys": "pred",
        "result": "result",
        "points": "points"
      "name": "aiaa.transforms.BoundingBoxd",
      "args": {
        "keys": "pred",
        "result": "result",
        "bbox": "bbox"
  "writer": {
    "name": "aiaa.transforms.Writer",
    "args": {
      "image": "pred",
      "json": "result"

Use a Custom InferencePipeline

We will need to first load the custom pipeline into AIAA. It is the same as you would load an AIAA model but without “data” specified. Example:

curl -X PUT "$AIAA_PORT/admin/model/custom_pipeline" -F "config=@config_aiaa.json"

Use /inference API to run the custom inference pipeline. Note that for this cascade pipeline to run, you need to load those inferences to AIAA first. For example if you want to run both spleen model and your own model you can do the following:

# Load spleen model to AIAA from NGC
curl -X PUT "$AIAA_PORT/admin/model/clara_pt_spleen_ct_segmentation" \
     -H "accept: application/json" \
     -H "Content-Type: application/json" \
     -d '{"path":"<NGC Model Path>","version":"1"}'

# Load your own model to AIAA
curl -X PUT "$AIAA_PORT/admin/model/my_own_model" \
     -F "config=@my_own_model/config_aiaa.json;type=application/json" \
     -F "data=@my_own_model/model.ts"

# Run both models
curl -X POST "$AIAA_PORT/v1/inference?model=custom_pipeline&output=image" \
     -H "accept: multipart/form-data" \
     -H "Content-Type: multipart/form-data" \
     -F 'params={"models": ["clara_pt_spleen_ct_segmentation", "my_own_model"]}' \
     -F "image=@image.nii.gz;type=application/gzip" -o output.nii.gz -f