How to create a DataBlock for Multispectral Satellite Image Segmentation with the Fastai-v2 | by Maurício Cordeiro | Towards Data Science

For information about the course Introduction to Python for Scientists (available on YouTube) and other articles like this, please visit my website cordmaur.carrd.co.
Fastai is an open source deep learning library that adds higher level functionalities to PyTorch and makes it easier to achieve state-of-the-art results with little coding.
The vision module is really handy when we need to quickly create an image dataset, apply data augmentations, resize, crop or even overlay a segmentation mask (Figure 1). However, all this simplification comes with a cost. Most of these higher level APIs for computer vision are optimized to RGB images, and these image libraries don’t support multispectral or multichannel images.
In recent years, due to an increase in data accessibility, Earth Observation researchers have been paying a lot of attention on deep learning techniques, like image recognition, image segmentation, object detection, among others. [1]. The problem is that satellite imagery, usually composed by many different spectral bands (wavelengths), doesn’t fit in most vision libraries used by the deep learning community. For that reason, I have been working directly with PyTorch to create the dataset (here) and train the “home-made” U-Net architecture (here).
With the upcoming of Fastai-v2 (promised to be released in next weeks)[2], I would like to test if it was possible to use it’s Data Block structure to create a multispectral image dataset to train a U-Net model. That’s more advanced than my previous stories as we have to create some custom subclasses, but I tried to make it as simple as possible. The notebook with all the code is available in the GitHub project (notebook here).
Before we start, it’s necessary to install fastai2 library, following the installation guide at https://dev.fast.ai.
To continue, we will need some training patches as multispectral images. That could be a Kaggle dataset, as the 38-cloud dataset, used in this story, or a completely new one. In the story Creating training patches for Deep Learning Image Segmentation of Satellite (Sentinel 2) Imagery using the Google Earth Engine (GEE), I show how to create training patches from Google Earth Engine and consume them as NumPy arrays.
For this tutorial, I made it available in the GitHub repository (https://github.com/cordmaur/Fastai2-Medium), under /data folder, a few sample patches from Orós reservoir in Brazil. Data is saved as NumPy’s arrays and separated in two folders /data/images for the multispectral patches and /data/labels for the target masks.
Checking number of files - images:40      masks:40Checking shapes - image: (13, 366, 366) mask: (366, 366)
As we can see from the output, the images were cropped in 366x366 pixels and the training patches have 13 channels. These are the 12 Sentinel-2 bands, in order, and I have added an additional band to indicate “no data” flag. Channels have been put in the first axis, as it is the standard for PyTorch and Fastai, differently from other image libraries. Mask values are 0-No water, 1-Water and 2-Water Shadow. We will check an image sample using Matplotlib. For this, we need to move the channels axes to the last position using Numpy.transpose and select the corresponding Red, Green and Blue axis.
Now that we have checked our images, we need a way of opening them in Fastai. As Fastai uses the Python Imaging Library (PIL), it is not able to handle multispectral images like these. We will, then, inherit a new class called MSTensorImage from the original TensorImage that is capable of opening .npy extensions and showing it visually. Let’s get to the code.
First, we define a function that opens Numpy arrays and returns it as a given class (open_npy function).
Then, we define the MSTensorImage class with a create method that receives the filename (or the numpy array, or a tensor), the desired channels (if we wouldn’t want to load all 13 channels) and a flag indicating whether the channels are in the first axis or not.
At last, we define a show method to correctly display this image. We can pass to the show method different channels to create false color composites and Matplotlib axes for multiple display. In the following example, we can see that the image was correctly loaded and displayed.
To handle the mask, instead of creating a new class, as we did previously, we will use the TensorMask class already defined in Fastai2. The only difference, in this case, is that the TensorMask, doesn’t know how to open .npy extensions. In this case, we will create it passing the result of our newly defined open_npy function, like so.
The DataBlock is an abstraction to provide transformations to your source data to fit your model. Remember that in order to train a neural net, we have to provide the model a set of inputs (Xs) and corresponding targets (Ys). The cost and the optimization functions will take care of the rest. More complex architectures could have multiple Xs as inputs and even multiple targets, so the DataBlock accepts multiple pipelines (Figure 2).
In our case, that is the most common, we need to create a DataBlock that provides only two blocks, Xs and Ys. For each block we can create a TransformBlock (that’s like a processing pipe) that manipulates each group (Xs or Ys) until they reach the desired format for loading into the model.
To create a DataBlock, besides the TransformBlocks, we have to provide a function responsible to get the items, given a source. In our example, as the files are already saved in disk, the get_items function will just read the directory and return the file_names. The list of files will be the starting point of our transformations pipes.
If the mask is created from the same original file (for example, if the mask was the 14th channel), the X transformation pipeline would be responsible to exclude this channel and the Y transformation pipeline would be responsible to extract this information from the original image. As we have the files saved separately, we will provide a new function called get_lbl_fn that will receive the name of the image file and then return the name of the corresponding target file. Now that we have some understanding of the DataBlock, let’s go to the code.
The last command, DataBlock.summary, will check if everything is working fine and give us a data sample:
Setting-up type transforms pipelinesCollecting items from Data\imagesFound 40 items2 datasets of sizes 36,4Setting up Pipeline: partialSetting up Pipeline: get_lbl_fn -> partialBuilding one sample  Pipeline: partial    starting from      Data\images\Oros_1_19.npy    applying partial gives      MSTensorImage of size 13x366x366  Pipeline: get_lbl_fn -> partial    starting from      Data\images\Oros_1_19.npy    applying get_lbl_fn gives      Data\labels\Oros_1_19.npy    applying partial gives      TensorMask of size 366x366Final sample: (MSTensorImage: torch.Size([13, 366, 366]), TensorMask([[0., 0., 0.,  ..., 0., 0., 0.],        [0., 0., 0.,  ..., 0., 0., 0.],        [0., 0., 0.,  ..., 0., 0., 0.],        ...,        [0., 0., 0.,  ..., 0., 0., 0.],        [0., 0., 0.,  ..., 0., 0., 0.],        [0., 0., 0.,  ..., 0., 0., 0.]]))Setting up after_item: Pipeline: AddMaskCodes -> ToTensorSetting up before_batch: Pipeline: Setting up after_batch: Pipeline: Building one batchApplying item_tfms to the first sample:  Pipeline: AddMaskCodes -> ToTensor    starting from      (MSTensorImage of size 13x366x366, TensorMask of size 366x366)    applying AddMaskCodes gives      (MSTensorImage of size 13x366x366, TensorMask of size 366x366)    applying ToTensor gives      (MSTensorImage of size 13x366x366, TensorMask of size 366x366)Adding the next 3 samplesNo before_batch transform to applyCollating items in a batchNo batch_tfms to apply
As we can see in the results, the final sample, after processing all the transformations, is composed of a tuple with 2 items (X, Y). In our case the (MSTensorImage, TensorMask).
Once we are done creating the DataBlock the creation of a DataSet or a DataLoader are very straightforward. The datasets will be splitted, considering the splitter function (in our case, a random splitter with 10% for validation), into training and validation datasets, and the data can be accessed directly by subscripting:
The same is valid for DataLoaders. Additionally the show_batch method automatically understands that we are performing a segmentation task (because of the TensorMask class) and overlays Xs and Ys:
In today’s story we saw the basic concepts and how to use the Fastai 2’s DataBlock to prepare the multispectral satellite data to fit a deep learning model. The subclassing of the TensorImage class and the use of the open_npy function were enough to handle the multiple channels. As one can note, the basic functionalities of Fastai to create a dataloader and show a batch sample were successfully preserved.
The fact of using Fastai 2 for these Earth Observation tasks is really handy, as differently from Fastai 1 and other high level APIs, Fastai 2’s built-in vision architectures accepts input channels as a parameter. But that’s a subject for a next story.
The notebook and sample data to follow this “how-to” is available at the GitHub repository https://github.com/cordmaur/Fastai2-Medium.
See you in the next story.
[1] Hoeser, T., Kuenzer, C., 2020. Object Detection and Image Segmentation with Deep Learning on Earth Observation Data: A Review-Part I: Evolution and Recent Trends. Remote Sensing 12, 1667. https://doi.org/10.3390/rs12101667
[2] Howard, J., Gugger, S., 2020. fastai: A Layered API for Deep Learning. Information 11, 108. https://doi.org/10.3390/info11020108