Class PythonOperatorPipeline

dali/test/python/operator_2/test_python_function.py:86–104

class PythonOperatorPipeline(CommonPipeline):
    def __init__(
        self, batch_size, num_threads, device_id, seed, image_dir, function, prefetch_queue_depth=2
    ):
        super().__init__(
            batch_size,
            num_threads,
            device_id,
            seed,
            image_dir,
            prefetch_queue_depth=prefetch_queue_depth,
        )
        self.python_function = ops.PythonFunction(function=function)

    def define_graph(self):
        images, labels = self.load()
        processed = self.python_function(images)
        assert isinstance(processed, _DataNode)
        return processed


class FlippingPipeline(CommonPipeline):