Streams extraction using Multi-Task Road Extractor

  • 🔬 Data Science
  • 🥠 Deep Learning and Pixel Classification


In this notebook, we will use 3 different geo-morphological characteristics derived from a 5 m resolution DEM in one of the watersheds of Alaska to extract streams. These 3 characteristics are Topographic position index derived from 3 cells, Geomorphon landform, and topographic wetness index.

We created a composite out of these 3 characteristic rasters and used Export raster tool to convert scale the pixels values to 8 bit unsigned. Subsequently, the images are exported as "Classified Tiles" to train a Multi-Task Road Extractor model provided by ArcGIS API for Python for extracting the streams.

Before proceeding through this notebook, it is advised to go through the API Reference for Multi-Task Road Extractor. It will help in understanding the Multi-Task Road Extractor's workflow in detail.

Necessary imports

import os
import zipfile
from pathlib import Path

from arcgis.gis import GIS
from arcgis.learn import prepare_data, MultiTaskRoadExtractor

Connect to your GIS

ent_gis = GIS('')
gis = GIS("home")

Get the data for analysis

Here is the composite with 3 bands representing the 3 geo-morphological characteristics namely Topographic position index, Geomorphon landform, and Topographic wetness index.

composite_raster = ent_gis.content.get('43c4824dd4bf41ee886be5042262f192')
Image by api_data_owner
Last Modified: March 04, 2021
0 comments, 3 views
BeaverCreek_Flowlines = ent_gis.content.get('2b198075e53748b4ab84197bc6ddd478')
Feature Layer Collection by api_data_owner
Last Modified: March 04, 2021
0 comments, 7 views

Export training data

Export training data using 'Export Training data for deep learning' tool, click here for detailed documentation:

  • Set 'composite_3bands_6BC_8BC_8bitunsigned' as Input Raster.
  • Set a location where you want to export the training data in Output Folder parameter, it can be an existing folder or the tool will create that for you.
  • Set the 'BeaverCreek_Flowlines' as input to the Input Feature Class Or Classified Raster parameter.
  • Set Class Field Value as 'FClass'.
  • Set Image Format as 'TIFF format'
  • Tile Size X & Tile Size Y can be set to 256.
  • Stride X & Stride Y can be set to 128.
  • Select 'Classified Tiles' as the Meta Data Format.
  • In 'Environments' tab set an optimum Cell Size. For this example, as we have performing the analysis on the geo-morphological characteristics with 5 m resolution, so, we used '5' as the cell size.

arcpy.ia.ExportTrainingDataForDeepLearning("composite_3bands_6BC_8BC_8bitunsigned.tif", r"D:\Stream Extraction\Exported_3bands_composite_8bit_unsigned", "BeaverCreek_Flowlines", "TIFF", 256, 256, 128, 128, "ONLY_TILES_WITH_FEATURES", "Classified_Tiles", 0, "FClass", 5, None, 0, "MAP_SPACE", "PROCESS_AS_MOSAICKED_IMAGE", "NO_BLACKEN", "FIXED_SIZE")

Alternatively, we have provided a subset of training data containing a few samples. You can use the data directly to run the experiments.

training_data = gis.content.get('3a95fd7a25d54898bddabf1989eea87d')
Image Collection by api_data_owner
Last Modified: March 08, 2021
0 comments, 1 views
filepath =
with zipfile.ZipFile(filepath, 'r') as zip_ref:
output_path = Path(os.path.join(os.path.splitext(filepath)[0]))

Prepare data

We will specify the path to our training data and a few hyperparameters.

  • path: path of the folder containing training data.
  • batch_size: Number of images your model will train on each step inside an epoch, it directly depends on the memory of your graphic card. 8 worked for us on a 11GB GPU.
data = prepare_data(output_path, chip_size=512, batch_size=4)

Visualize a few samples from your training data

To get a sense of what the training data looks like, arcgis.learn.show_batch() method randomly picks a few training chips and visualizes them.

  • rows: number of rows we want to see the results for.

Train the model

Load model architecture

arcgis.learn provides the MultiTaskRoadExtractor model for classifying linear features, which is based on multi-task learning mechanism. More details about multi-task learning can be found here.

model = MultiTaskRoadExtractor(data, mtl_model="hourglass")

Find an optimal learning rate

Learning rate is one of the most important hyperparameters in model training. Here, we explore a range of learning rate to guide us to choose the best one. We will use the lr_find() method to find an optimum learning rate at which we can train a robust model.


Fit the model

To train the model, we use the fit() method. To start, we will train our model for 10 epochs. Epoch defines how many times model is exposed to entire training set. We have passes three parameters to fit() method:

  • epochs: Number of cycles of training on the data.
  • lr: Learning rate to be used for training the model.
  • wd: Weight decay to be used.
Input, lr)
epoch train_loss valid_loss accuracy miou dice time
0 2.888322 1.703418 0.842602 0.421301 0.000000 07:29
1 1.175394 0.589130 0.842602 0.421301 0.000000 07:21
2 0.728384 0.544955 0.918369 0.717106 0.614468 07:23
3 0.571631 0.616916 0.914180 0.692409 0.568880 07:24
4 0.502315 0.556234 0.924240 0.739906 0.650921 07:26
5 0.456098 0.480777 0.935033 0.772896 0.697409 07:14
6 0.397026 0.378353 0.952222 0.839674 0.804771 07:25
7 0.347037 0.330573 0.963398 0.874130 0.852213 07:20
8 0.311752 0.304045 0.969619 0.891718 0.871019 06:53
9 0.291295 0.296276 0.971693 0.900430 0.887195 07:15

As you can see, both the losses (valid_loss and train_loss) started from a higher value and ended up to a lower value, that tells our model has learnt well. Let us do an accuracy assessment to validate our observation.

Accuracy Assessment

We can compute the mIOU (Mean Intersection Over Union) for the model we just trained in order to do the accuracy assessment. We can compute the mIOU by calling model.mIOU. It takes the following parameters:

  • mean: If False returns class-wise mean IOU, otherwise returns mean IOU of all classes combined.
#model.mIOU()    #available in 1.8.5 version
100.00% [7/7 00:37<00:00]
{'0': 0.9668936989721181, 1: 0.8336953463440812}

The model has a mean IOU of 0.96 which proves that the model has learnt well. Let us now see it's results on validation set.

Visualize results in validation set

The code below will pick a few random samples and show us ground truth and respective model predictions side by side. This allows us to validate the results of your model in the notebook itself. Once satisfied we can save the model and use it further in our workflow. The model.show_results() method can be used to display the detected streams.


Save the model

We would now save the model which we just trained as a 'Deep Learning Package' or '.dlpk' format. Deep Learning package is the standard format used to deploy deep learning models on the ArcGIS platform.

We will use the save() method to save the model and by default it will be saved to a folder 'models' inside our training data folder itself.


Model inference

The saved model can be used to classify streams using the 'Classify Pixels Using Deep Learning' tool available in ArcGIS Pro, or ArcGIS Enterprise.

arcpy.ia.ClassifyPixelsUsingDeepLearning("composite_3bands_6BC_8BC_8bitunsigned.tif", r"D:\Stream Extraction\Exported_3bands_composite_8bit\models\MTRE_10e\MTRE_10e.dlpk", "padding 64;batch_size 4;predict_background True;return_probability_raster False;threshold 0.5", "PROCESS_AS_MOSAICKED_IMAGE", None);"\Documents\ArcGIS\Packages\Stream Extraction demo_9ee52d\p20\stream_identification.gdb\detected_streams")

The output of the model is a layer of detected streams which is shown below:


The notebook presents the workflow showing how easily you can combine traditional morphological landform characteristics with deep learning methodologies using arcgis.learn to detect streams.

Your browser is no longer supported. Please upgrade your browser for the best experience. See our browser deprecation post for more details.