Python Operator

This example shows how to implement and use a custom python operator.

Defining an operator implementation

def mix_and_split_channels(image1, image2):
    mixed = (image1 / 255. + image2 / 255.) / 2.
    return mixed[:, :, 0], mixed[:, :, 1], mixed[:, :, 2]

The function above defines a simple operation on RGB images given as numpy arrays. Generally a python operator implementation must be a function that consumes any number of numpy arrays as arguments and produces a numpy array or a tuple of those as an output.

Example pipeline

from nvidia.dali.pipeline import Pipeline
import nvidia.dali.ops as ops
import nvidia.dali.types as types
import nvidia.dali.plugin.pytorch as dalitorch
import torch

image_dir = '../images'

class SimplePipeline(Pipeline):
    def __init__(self, batch_size, num_threads, device_id, _seed, image_dir):
        super(SimplePipeline, self).__init__(batch_size, num_threads, device_id, seed=_seed, exec_async=False,
                                             exec_pipelined=False)
        self.input = ops.FileReader(file_root=image_dir)
        self.decode = ops.ImageDecoder(device='cpu', output_type=types.RGB)
        self.flip = ops.Flip()
        self.mix_and_split = ops.PythonFunction(function=mix_and_split_channels, num_outputs=3)
        self.sqrt = dalitorch.TorchPythonFunction(function=lambda tensor: tensor.sqrt())

    def load(self):
        jpegs, labels = self.input()
        decoded = self.decode(jpegs)
        return decoded

    def define_graph(self):
        images1 = self.load()
        images2 = self.flip(images1)
        r, g, b = self.mix_and_split(images1, images2)
        return self.sqrt(r), g, b

Our custom python operator is introduced into pipeline as PythonFunction. The function object is passed as a named argument. The number of outputs must be specified on the operator construction level. Default value for num_outputs is 1. Please note that for now custom python operators work only in pipelines with exec_async=False and exec_pipelined=False specified. There is also TorchPythonFunction operator added. The only thing which differs it from PythonFunction is that it passes data to the called function as PyTorch tensors and not as numpy arrays.

Running the pipeline and visualizing results

import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib import cm
%matplotlib inline

batch_size = 4

def show_images(image_batch, color_map):
    columns = 4
    rows = (batch_size + 1) // columns
    fig = plt.figure(figsize=(32, (32 // columns) * rows))
    gs = gridspec.GridSpec(rows, columns)
    for j in range(rows*columns):
        plt.subplot(gs[j])
        plt.axis("off")
        plt.imshow(image_batch.at(j), cmap=plt.get_cmap(color_map))


pipe = SimplePipeline(batch_size, 4, 0, 12, image_dir)
pipe.build()
r, g, b = pipe.run()

show_images(r, 'Reds_r')
show_images(g, 'Greens_r')
show_images(b, 'Blues_r')