Online Self-Supervised Thermal River Segmentation

This repository contains code necessary to run our online self-supervised river segmentation algorithm onboard a low-flying UAV. The ROS nodes were written for and tested on an Nvidia Jetson AGX Orin. It also contains all datasets used for training and validation in our paper.

Getting started

This package subscribes to the following topics:

/boson/thermal/image_raw: This provides a stream of 16 bit thermal imagery.
/imu/imu: Provides the qw, qx, qy, qz orientation of the UAV.

If running on custom UAV platform, please modify:

camera settings (intrinsics, extrinsics, distortion), which are set in startup/config/flir_boson.py
IMU->Camera transformation, which is hardcoded in onr/startup/sync_thermal_imu.py and formatted as [x, y, z, qw, qx, qy, qz]

Environment setup

Setup on Nvidia Jetson AGX Orin

Create your conda environment with Python 3.8 (or whichever Python version that is compatible with both ros and pytorch).
Install ros-noetic from source or via robostack. You'll likely have to install from source due to python version incompatabilities between the Jetson's PyTorch and RoboStack's ros-noetic.
Install pytorch using a compatible wheel from here or build from source (not recommended).
Install onnxruntime from source.
Install torchvision from source.
Install the other requirements via mamba/conda/pip using environment.yaml or requirements.txt.
Install ROS python packages as needed via mamba/conda/pip.

Setup on any other computer

Install ros-noetic via robostack or source. I recommend robostack.
Install all requirements via mamba/conda/pip. No building from source required here.

Running on the online self-supervised algorithm

Start the ROS master node.

roscore

Start the segmentation node. This takes a few seconds to initialize (loading pretrained weights, warming up network for training, etc...) so do it first. See command line arguments in segmentation/online_segmentation.py for more details on modifying training configurations.

python segmentation/online_segmentation.py --use-texture --postprocess --adapt

Start the pseudolabeling node (texture cue in example, but use motion cue if desired).

python pseudolabel/texture_cue.py

Start the node to preprocess data and perform horizon line estimation.

python startup/sync_thermal_imu.py

# or, if camera oriented upside down on UAV:
python startup/sync_thermal_imu.py --rotate-180

# Start the sky segmentation node if no IMU is available for horizon line estimation.
python sky_segmentation/segmentation.py --weights-path weights/fast_scnn.onnx

Play the rosbag file

rosbag play path/to/bagfile.bag

Datasets

Network pretraining datasets

The segmentation network used in this work was pretrained on RGB images from COCO-Stuff, ADE20k, Fuentes river dataset that contained water-related pixels. We supplemented with images scraped from Flickr via water- and aerial-related search terms and labeled them using an ADE20k-pretrained ResNet50dilated + PPM_deepsup convolutional neural network from the semantic-segmentation-pytorch library. The set of scraped Flickr images with their annotations are made available below.

Dataset	Num. Train	Num. Validation	Water-related indices	Link
COCO-Stuff	10977	458	148, 155, 178, 179	Link to download
ADE20K	1743	163	22, 27, 61, 114, 129	Link to download
Fuentes River	300	0	---	Link to download
Flickr	1220	0	---	Link to download

Sky segmentation network training datasets

The sky segmentation network was trained on publicly-available thermal images containing sky pixels. The datasets are listed below:

Dataset	Link
KAIST Pedestrian Segmentation	Link to download
SODA	Link to download
MassMIND	Link to download
FLIR	Link to download

Annotated river/coastal scenes

The set of annotated images used for online segmentation validation contains N annotated images pulled from the Kentucky River, KY; Colorado River, CA; Castaic Lake, CA; and Duck Pier, NC sequences.

ROS bagfiles

Location	Sequence Name	Link
Kentucky River	flight 2-1	Link
Colorado River	flight 1	Link
Colorado River	flight 3	Link
Castaic Lake	flight 2	TODO
Castaic Lake	flight 4	TODO

Nvidia Jetson Tips/Bugs

Place ROS python imports after everything else has been imported. It seems to have conflicts with skimage specifically on the Jetson.

Citation

@INPROCEEDINGS{10342016,
  author={Lee, Connor and Frennert, Jonathan Gustafsson and Gan, Lu and Anderson, Matthew and Chung, Soon-Jo},
  booktitle={2023 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, 
  title={Online Self-Supervised Thermal Water Segmentation for Aerial Vehicles}, 
  year={2023},
  volume={},
  number={},
  pages={7734-7741},
  doi={10.1109/IROS55552.2023.10342016}}

connorlee77 / uav-thermal-water-segmentation