DeepContrail segmentation: Unet++ + CAFormer-B36 (v0.1.0)

This is a contrail segmentation model for DeepContrail. It takes a GOES-16 ABI false-colour composite and returns, for each pixel, the probability that it shows a contrail. It was trained on the OpenContrails dataset (20250305 export).

Usage

Install the DeepContrail package from PyPI (it installs deepcontrail-core too):

pip install deepcontrail-segmentation

import numpy as np
from deepcontrail.core import false_color_from_spectral, read_band_manifest
from deepcontrail.segmentation import load_pretrained, predict

root = "/path/to/dataset"           # a dataset in the DeepContrail format
record_id = "1801360214472457069"    # any scene id of the test split

model = load_pretrained("eurocontrol-asu/deepcontrail-segmentation", revision="v0.1.0")
bands = np.load(f"{root}/spectral_bands/test/{record_id}.npy")
composite = false_color_from_spectral(bands, model.config.false_color, read_band_manifest(root))
probabilities = predict(model, composite, tta=model.config.tta)  # (256, 256), float32 in [0, 1]
mask = probabilities > model.config.threshold

The OpenContrails converter produces a DeepContrail dataset from the raw TFRecords:

gcloud storage cp -r gs://goes_contrails_dataset/20250305 .

predict also takes a batch (N, 256, 256, 3).

Model

Architecture: Unet++ decoder with a CAFormer-B36 encoder (timm caformer_b36.sail_in22k_ft_in1k_384), from segmentation-models-pytorch; 102.6 M parameters.
- CAFormer has no feature map at half the input resolution, which Unet++ needs.
- DeepContrail's build_model adds a small convolutional stem for that level: two 3 × 3 convolutions with 64 channels.
- Loading the weights outside DeepContrail needs the same stem.
Input: a 3-channel composite in [0, 1] (described below) of a 256 × 256 scene. It is resized to 512 px and normalised with mean [0.485, 0.456, 0.406] and std [0.229, 0.224, 0.225]; predict does both.
Output: contrail probability at 256 × 256. A pixel is a contrail when its probability is above 0.5.
Test-time augmentation: recommended. It averages the 8 flips and rotations.

Input composite (standard ash RGB)

Channel	Bands	Bounds (scaled to [0, 1], clipped)
R	data_12um − data_11um (GOES-16 ABI bands 15 − 14)	[−4, 2] K
G	data_11um − cloud_top (GOES-16 ABI bands 14 − 11)	[−4, 5] K
B	data_11um (GOES-16 ABI band 14)	[243, 303] K

The bands are referred to by their keys in the dataset's bands.json.

Training

Data: the OpenContrails 20250305 export. Of its 20,088 training scenes, 10% (2,009, seed
1. were held out for validation, so the model was trained on 18,079 scenes.
Labels: the per-pixel fraction of annotators who marked a contrail (soft labels).
Recipe:
- 40 epochs at 512 px, batch 32, bf16 mixed precision;
- AdamW (learning rate 1e-4, weight decay 1e-2), cosine schedule with 3 warm-up epochs;
- BCE + 0.2 × Dice loss;
- augmentations: independent horizontal and vertical flips, 90° rotations and small shift/scale/rotate.
Checkpoint: the one with the best validation Dice.
Configuration: experiment.yaml.

Evaluation

Global Dice on the held-out test split: 1,879 scenes, the OpenContrails validation shards. They were never used for training or model selection.

Inference	Threshold	Test Dice
Single pass	0.5	0.6890
Single pass	tuned on validation	0.6870
Test-time augmentation	0.5	0.6974
Test-time augmentation	tuned on validation	0.6937

GPU training is not bit-for-bit reproducible, so architectures are compared over several runs (TTA, threshold 0.5, same recipe).

Before the augmentation fix: in DeepContrail before PR #47, the two flips always fired together. With that code, two Unet++ trainings scored 0.6992 and 0.6995, and four UPerNet trainings 0.6916-0.6960.
With the fix: this model scored 0.6974 and UPerNet 0.6912.

Instance segmentation

This revision also holds the instance step tuned for this model. It splits the model's probability map into individual contrails, one mask and one score each:

instance_pipeline.yaml: the settings of the deepcontrail-instance pipeline (threshold, Mannstein-style line detection, growth);
score_model.json: a learned score that ranks the detected contrails. It was trained on this model's test-time-augmented predictions for 1,000 of its validation scenes.

deepcontrail-instance is a preview: its API may still change.

pip install "deepcontrail-instance[map]"

from deepcontrail.instance import load_instance_step, to_coco_results

step = load_instance_step(model.directory, tta=model.config.tta)
masks, scores = step(predict(model, composite, tta=model.config.tta))
results = to_coco_results(masks, scores, image_id=record_id)  # COCO results with RLE masks

Results on the 1,879 test scenes, scored against each annotator's contrails (2,700 annotator samples with at least one contrail):

COCO segmentation mAP of 0.280, averaged over IoU 0.25-0.75;
precision 0.532, recall 0.387 and F1 0.448 at IoU 0.5;
against the consensus contrails (drawn by most annotators): mAP 0.435.

Conditions:

The settings and the score were chosen on validation scenes only.
The score model is valid only for this model's test-time-augmented probabilities on 256 × 256 scenes. load_instance_step and the pipeline refuse other uses.

Limitations

Sensor and scene size. The model was trained on GOES-16 ABI scenes of 256 × 256 pixels (about 2 km per pixel). predict refuses other sizes, since resizing would change the pixel scale; tiling larger scenes is not supported yet.
Label convention. The 20250305 masks are aligned with the images. The older 20230419 export and the Kaggle competition data put the masks half a pixel down-right; on those, call predict(..., pixel_shift=0.5).
Other label sets.
- The 20250305 masks overlap the 20230419/Kaggle ones at an IoU of only about 0.78, so scores on the two are not directly comparable.
- This model scored 0.70414 (private) / 0.69328 (public) on the hidden test set of the Kaggle Google Research – Identify Contrails to Reduce Global Warming competition (TTA, threshold 0.5).

Provenance


Training run	`unetplusplus-caformer-b36-512-20250305-augfix`
Checkpoint	`best-val_dice=0.6978.ckpt`
Code	DeepContrail 0.1.0
Data	OpenContrails GOES-16 ABI TFRecords, export `20250305`

Licence and attribution

EUROCONTROL releases the weights under the Apache License 2.0. If you redistribute them, as they are or modified, keep the NOTICE file with them. The DeepContrail code is under the EUPL-1.2.

Training data: OpenContrails (bucket gs://goes_contrails_dataset), by Google Research, under the CC BY 4.0 licence. GOES-16 imagery is courtesy of NOAA and has no use restrictions.
Encoder initialisation: timm/caformer_b36.sail_in22k_ft_in1k_384, under the Apache License 2.0.

If you use this model, please cite DeepContrail (doi:10.5281/zenodo.24130966) and the OpenContrails dataset:

@software{Heffar_DeepContrail,
  author={Heffar, Amine},
  title={DeepContrail: contrail detection and segmentation in satellite imagery},
  publisher={Zenodo},
  version={v0.1.0},
  doi={10.5281/zenodo.24130967},
  url={https://github.com/eurocontrol-asu/deepcontrail}
}


@article{Ng2024ContrailDO,
  title={Contrail Detection on GOES-16 ABI With the OpenContrails Dataset},
  author={Joe Yue-Hei Ng and Kevin McCloskey and Jian Cui and Vincent R. Meijer and Erica Brand and Aaron Sarna and Nita Goyal and Christopher Van Arsdale and Scott Geraedts},
  journal={IEEE Transactions on Geoscience and Remote Sensing},
  year={2024},
  volume={62},
  pages={1-14},
  doi={10.1109/TGRS.2023.3345226}
}

The model continues an earlier study of contrail segmentation on GOES-16 by the same unit:

@inproceedings{jarry2024deep,
  author    = {Jarry, Gabriel and Very, Philippe and Heffar, Amine and Torjman-Levavasseur, Valentin},
  title     = {Deep Semantic Contrails Segmentation of {GOES-16} Satellite Images: A Hyperparameter Exploration},
  booktitle = {Proceedings of the SESAR Innovation Days},
  year      = {2024},
  url       = {https://sesar.eu/sites/default/files/documents/sid/2024/papers/SIDs_2024_paper_028%20final.pdf}
}

Downloads last month: 123

Safetensors

Model size

0.1B params

Tensor type

F32