SR-Seg: A vegetation segmentation method based on spectral reconstruction for UAV field images

Introduction

We first propose SRCNet and SRANet for realizing multispectral reconstruction from RGB images, and then apply the reconstruction results for vegetation segmentation. More specifically, the SR model contains an encoder and decoder for image feature extraction and spectral recovery, respectively. The model is used to learn the implied relationship between RGB and spectral data in the first stage and to reconstruct RGB images into multispectral images. Then, the reconstructed data are processed in the second stage based on the reconstructed spectral attributes combined with vegetation index(VI), and thus segmenting the field vegetation phenotypic data.

Environment

• python=3.8
• pytorch=1.8.0
• gdal=3.4.1
• scikit-image=0.19.2
• matplotlib=3.5.2
• opencv-python==4.6.0.66
• einops==0.4.1

Get Start

Spectral reconstruction model training:

Use SRCNet training:

                python train.py --model_name="cnn" 
            

Use SRANet training:

                python train.py --model_name="attn" 
            

After model training, segmentation is performed:

                python batch_predict.py --model_name="cnn" --rgb_dir="./yourfoldername"
            

Results

We apply the reconstructed results to the vegetation segmentation task. As shown in Figure, we fuse the reconstructed spectral data of each channel and further segment them, as well as report the segmentation results of the spectral data reconstructed by the five models. To ensure the consistency of the experimental results we keep the same fusion strategy and segmentation parameter settings, uniformly adopt the fusion strategy based on the VI, and set the segmentation parameter to 0.2.

Code

Get the source code from Github, click here.