Dali operator for multi-page TIFF's

Dali is an Nvidia library that provides a computational graph for the pre-processing part of your training script. The interesting bit is that operators in the graph can be run on the GPU. You can read more here.

The data pre-processing could be massive on the CPU, and you may end up in a situation where the GPU is starving while waiting for data. In addition to some offline pre-processing, Dali could then alleviate the burden and balance things up.

Test scenario

For a test scenario, I wanted to feed a network with brain scans for a segmentation task. Usually, the images are stored in formats like NIfTI, which are not easy to handle in a typical image pipeline.

I decided to convert the brain scans to multipage gray-scale TIFF’s. I used nibabel to read the original files. Once you have the multidimensional NumPy arrays containing the images, you can save those as TIFF files using tifffile. Saving the TIFF was a matter of one line of code, like the following:

1
tifffile.imwrite(outfile, data, compress='ZSTD', photometric='minisblack')

Notice I used the Zstandand compression algorithm for deflating the TIFF files. It is swift in compression and decompression and worth using.

Having the images as multi-page TIFF’s is also handy because you can use an ordinary image viewer like geeqie.

Some limitations

The next step was to try to write a Dali pipeline to ingest the data. In the endeavor, I soon discovered that some functionalities and features were missing. For example, it was not possible to open compressed TIFF files. I came up with small modifications to the build script for Dali and solved the problem. They merged those modifications with my pull request to the main branch and released with version 0.26.

The second problem was that the FileReader operator didn’t support a segmentation scenario explicitly. If you want to read an image and the corresponding ground truth, you need to use two FileReader ops. Finally, the ImageDecoder didn’t support the multi-page TIFF files I wanted to use.

Custom operators

Luckily you can write a custom operator in C++ (well, not sure how many people will consider that ‘luck’) and implement the features you need by yourself. You may read more here.

I implemented the custom operators I needed and shared the code here. You can find the instructions to build the shared library in the project’s README.

The following is a code snippet that shows how to load and use these two custom operators:

1
import nvidia.dali.plugin_manager as plugin_manager
2
from nvidia.dali.pipeline import Pipeline
3
import nvidia.dali.ops as ops
4
import nvidia.dali.types as types
5

6

7
class TestPipeline(Pipeline):
8
    def __init__(self, batch_size, num_threads, device_id):
9
        super(TestPipeline, self).__init__(batch_size, num_threads, device_id)
10
        self.input = ops.SegFileReader(file_root='data', file_list='data/file_list.txt', random_shuffle=True)
11
        self.decodeMulti = ops.TiffDecoder()
12
        self.rotate = ops.Rotate(device='gpu', interp_type=types.INTERP_NN, keep_size=True)
13
        self.rotate_range = ops.Uniform(range=(-27, 27))
14
        self.transpose = ops.Transpose(perm=[1, 2, 0])
15
        self.transposeBack = ops.Transpose(device='gpu', perm=[2, 0, 1])
16

17
    def define_graph(self):
18
        angle_range = self.rotate_range()
19
        image, mask = self.input()
20
        image = self.decodeMulti(image)
21
        mask = self.decodeMulti(mask)
22
        image = self.transpose(image)
23
        image = self.rotate(image.gpu(), angle=angle_range)
24
        image = self.transposeBack(image)
25
        mask = self.transpose(mask)
26
        mask = self.rotate(mask.gpu(), angle=angle_range)
27
        mask = self.transposeBack(mask)
28

29
        return image, mask
30

31

32
def main():
33
    plugin_manager.load_library('./cmake-build-debug/libCustomOp.so')
34
    pipe = TestPipeline(batch_size=8, num_threads=4, device_id=0)
35
    pipe.build()
36
    pipe_out = pipe.run()
37
    gimage, gmask = pipe_out
38
    print(gmask.as_tensor())
39
    print(gmask.layout())
40

41

42
if __name__ == '__main__':
43
    main()

At line 33, the shared library will be loaded, and the custom operators will be registered to be used by Dali.

At lines 10, the first custom operator, SegFileReader, is created and can be wired in define_graph. This operator differs from the standard one and it is able to read the image and its corresponding segmentation ground truth in one go.

At line 11, we find our custom TiffDecoder, which will read our multi-page TIFFs’ files. The operator will return tensor with “channel first.” So the shape will be CHW. Unfortunately, not all the standard operators can work with channel first tensors, and sometimes you will need to transpose things around.

You will find the above test code together with a jupyter notebook as an extra starting point in the repository. I hope you will find it useful.

Conclusions

Dali proved itself a useful library, with some shortcomings. You can get around some rough edges by writing your own operators. The library allows easy integration of C++ code, and for me, it is a useful feature.