Skip to content

Repository files navigation

SAM2-Contrails

SAM2-Contrails logo

Contrail detection, tracking, and attribution from all-sky camera imagery using SAM2.

License Python 3.10+ Hugging Face arXiv coming soon GVCCS Dataset

This repository contains SAM2-Contrails, a fine-tuned version of SAM 2 for dense contrail segmentation, tracking, and attribution in ground-based all-sky camera videos. Given a sequence of camera frames and ADS-B flight trajectories projected as physics-informed prompt masks, the model tracks each contrail through the full video and attributes it to its source flight.

The full methodology, experiments, and results are described in our paper (preprint will be linked here upon publication). All design choices — prompt encoding, architecture modifications, training setup, and evaluation protocol — are explained in detail there.


Pipeline

SAM2-Contrails pipeline

The pipeline has three stages:

  1. Prompt generation — ADS-B positions + ERA5 wind fields → CoCiP or DryAdvection contrail model → projected to image pixel coordinates → age-weighted ternary prompt masks (positive + negative signal).
  2. SAM2 inference — physics-informed dense prompts fed into an adapted SAM2 model; the memory bank propagates each contrail mask through the full video.
  3. Attribution — each predicted mask carries the flight ID of the prompt that seeded it, directly linking every contrail pixel to a source flight.

Key design choices

  • Ternary prompt encoding — a two-channel prompt encodes both the target flight's contrail (positive, age-weighted intensity) and all competing flights (negative signal), reducing cross-attribution errors.
  • Memory attention re-enabled — disabled by default in SAM2 video mode, re-enabling it was critical for long contrail tracks.
  • Dice-dominant loss — handles the extreme foreground/background imbalance of thin contrail masks.

Dataset

Models are trained and evaluated on the GVCCS dataset (Ground-based Video Camera for Contrail Segmentation), recorded with an all-sky camera at EUROCONTROL's experimental centre in Bretigny-sur-Orge, France.

Dataset: GVCCS on Zenodo

The dataset contains annotated all-sky camera videos with per-flight contrail segmentation masks and ADS-B flight trajectories. ERA5 meteorological data is not included and must be obtained separately via the Copernicus Climate Data Store.


Installation

Requirements: Python >= 3.10, uv, and a CUDA-capable GPU (CPU inference is ~50x slower).

git clone https://github.com/ramondalmau/sam2-contrails.git
cd sam2-contrails

# Inference only
uv sync --extra inference

# Full install — adds prompt generation from ADS-B + ERA5 data
uv sync --extra all

The repo includes a uv.lock file for fully reproducible installs. PyTorch is pre-configured for CUDA 12.8 via [tool.uv.sources] in pyproject.toml. If your system uses a different CUDA version, edit the index name (e.g. cu128cu121) before running uv sync. For CPU-only, remove the [tool.uv.sources] block.

The CUDA extension (sam2._C) is optional. If it fails to compile, SAM2 still works correctly — see INSTALL.md.

All commands run through uv run — no manual environment activation needed:

uv run python examples/04_run_inference.py --help

Model Weights

Pre-trained checkpoints are hosted on Hugging Face.

Variant Description mAP Attr. Precision
ternary Age-weighted + negative signal (recommended) 0.389 96.5%
original Binary baseline (1-min window, positive-only) 0.303 89.0%

Auto-download (Python)

from contrailtrack import load_model

# Downloads and caches weights automatically on first call
model = load_model()                    # ternary (recommended)
model = load_model(config="original")   # binary baseline

Manual download

from huggingface_hub import hf_hub_download

path = hf_hub_download("ramondalmau/sam2-contrails", "checkpoints/ternary.pt")
path = hf_hub_download("ramondalmau/sam2-contrails", "checkpoints/original.pt")

Quick Start with GVCCS

This walkthrough runs contrail detection on a GVCCS video end-to-end. It uses video 00001 (the first GVCCS test video) — replace with the integer COCO video ID of any video you want to process.

# Set paths (adjust to your setup)
GVCCS=/path/to/GVCCS/test
VID=00001          # COCO integer video ID (zero-padded)
VIDEO=20230930055430_20230930075430  # timestamp name, for parquet lookup

# Step 1: Extract per-video frame folders (named by COCO video ID)
uv run python examples/01_prepare_data.py $GVCCS

# Step 2: Convert ADS-B traffic parquet to a pycontrails Fleet JSON
uv run python examples/02_fleet_from_traffic.py data/parquet/$VIDEO.parquet

# Step 3: Run CoCiP contrail model and generate prompt masks
uv run contrailtrack run-cocip \
    --fleet-dir  data/fleet/ \
    --out        data/cocip/ \
    --annotations $GVCCS/annotations.json \
    --videos $VID

uv run contrailtrack generate-prompts \
    --contrail-dir  data/cocip/ \
    --annotations   $GVCCS/annotations.json \
    --out           data/prompts/ \
    --images-source $GVCCS/images/ \
    --max-age 5.0 \
    --videos $VID

# Step 4: Run SAM2-Contrails inference (downloads weights on first run)
uv run contrailtrack run \
    --images  data/frames/$VID \
    --prompts data/prompts/ \
    --out     results/$VID.json

# Step 5: Evaluate against GVCCS ground truth
uv run contrailtrack evaluate results/$VID.json \
    --labels $GVCCS/annotations.json \
    --out    evaluation/

# Step 6: Render an overlay video (prompts + predictions)
uv run python examples/06_render_video.py \
    --frames      data/frames/$VID \
    --prompts     data/prompts \
    --predictions results/$VID.json \
    --video-id    $VID

Python API

import contrailtrack as ct

video_id = "00001"  # COCO integer video ID (zero-padded)

# Load model — downloads weights from Hugging Face on first call
model = ct.load_model(config="ternary", device="cuda")

# Load frames (returns float32 tensor, no ImageNet normalisation)
frames, frame_names, H, W = ct.load_frames(f"data/frames/{video_id}/", image_size=1024)

# Read prompt masks (ternary encoding: positive own + negative competing)
prompts = ct.read_prompts("data/prompts/", video_id=video_id, encoding="ternary")

# Run SAM2 video inference
predictions = ct.run_video(
    model=model,
    frames=frames,
    frame_names=frame_names,
    prompts=prompts,
    original_height=H,
    original_width=W,
    score_threshold=0.5,
    max_propagation_frames=0,   # 0 = unlimited (recommended for full tracks)
)

# Export as COCO RLE JSON
ct.export_coco_json(predictions, video_id, frame_names, H, W, f"results/{video_id}.json")

# Evaluate against ground truth (resolves video ID → COCO timestamp name automatically)
results = ct.evaluate(
    predictions_path=f"results/{video_id}.json",
    gt_annotations="path/to/annotations.json",
    video_name=video_id,
)

CLI Reference

All functionality is available through the contrailtrack command:

uv run contrailtrack --help
Command Description
contrailtrack run-cocip Run CoCiP contrail model on fleet JSON files
contrailtrack run-dry-advection Run DryAdvection contrail model on fleet JSON files
contrailtrack generate-prompts Generate prompt PNGs from contrail model output
contrailtrack run Run SAM2 inference: frames + prompts → COCO RLE JSON
contrailtrack evaluate Evaluate a single video's predictions against ground truth
contrailtrack evaluate-dataset Evaluate all videos in one dataset-wide pass (use for paper metrics)
contrailtrack train Fine-tune SAM2 on GVCCS contrail data

Every batch command accepts an optional --videos flag (repeatable) to restrict processing to a subset of videos. Leading zeros are optional — 1 and 00001 are equivalent.

Full pipeline (annotated dataset)

FLEET=/path/to/fleet/       # fleet JSON files (one per video, named by timestamp)
ANNOTATIONS=/path/to/annotations.json
FRAMES=/path/to/img_folder/ # per-video JPEG frame folders
CACHE=/path/to/met_cache/   # ERA5 disk cache (reused across runs)

# 1. Run contrail model
uv run contrailtrack run-cocip \
    --fleet-dir  $FLEET \
    --out        data/cocip/ \
    --annotations $ANNOTATIONS \
    --cache-dir  $CACHE

# 2. Generate prompt + GT mask PNGs
uv run contrailtrack generate-prompts \
    --contrail-dir  data/cocip/ \
    --annotations   $ANNOTATIONS \
    --out           data/prompts/ \
    --images-source $FRAMES \
    --max-age 5.0

# 3. Train (optional — skip to use pre-trained weights)
uv run contrailtrack train \
    sam2/configs/sam2.1_training/sam2.1_hiera_b+_GVCCS_finetune_ternary.yaml \
    --num-gpus 2

# 4. Run inference on all test videos
uv run contrailtrack run \
    --images-dir $FRAMES \
    --prompts    data/prompts/ \
    --out        results/ \
    --checkpoint checkpoints/ternary.pt \
    --object-batch-size 25

# 5. Evaluate all predictions (dataset-wide pass for comparable metrics)
uv run contrailtrack evaluate-dataset \
    --predictions-dir results/ \
    --labels          $ANNOTATIONS \
    --out             evaluation/ \
    --flight-mappings data/prompts/

Processing a subset of videos

All batch commands accept --videos (repeat per video ID):

uv run contrailtrack run-cocip \
    --fleet-dir $FLEET --out data/cocip/ --annotations $ANNOTATIONS \
    --videos 1 --videos 3 --videos 5

uv run contrailtrack generate-prompts \
    --contrail-dir data/cocip/ --annotations $ANNOTATIONS --out data/prompts/ \
    --videos 1 --videos 3 --videos 5

uv run contrailtrack run \
    --images-dir $FRAMES --prompts data/prompts/ --out results/ \
    --videos 1 --videos 3 --videos 5

uv run contrailtrack evaluate-dataset \
    --predictions-dir results/ --labels $ANNOTATIONS --out evaluation/ \
    --flight-mappings data/prompts/ \
    --videos 1 --videos 3 --videos 5

Single video

uv run contrailtrack run \
    --images  $FRAMES/00001 \
    --prompts data/prompts/ \
    --out     results/00001.json

uv run contrailtrack evaluate results/00001.json \
    --labels $ANNOTATIONS \
    --out    evaluation/00001/

Unannotated traffic (no ground truth)

# Run contrail model (no --annotations needed)
uv run contrailtrack run-dry-advection \
    --fleet-dir data/fleet/ \
    --out       data/dry_advection/

# Generate prompts without GT masks
uv run contrailtrack generate-prompts \
    --contrail-dir data/dry_advection/ \
    --out          data/prompts/

# Infer (no --labels → no evaluation)
uv run contrailtrack run \
    --images-dir data/frames/ \
    --prompts    data/prompts/ \
    --out        results/

Data layout produced by generate-prompts: Frames are organised as img_folder/{video_id:05d}/{frame_idx:05d}.jpg and prompts as per_object_data*/{video_id:05d}/{flight_id}/{frame_idx:05d}_prompt.png, where video_id is the COCO integer video ID and flight_id is the hex ADS-B flight identifier.


Examples

Script Description
01_prepare_data.py Extract per-video frame folders from the flat GVCCS image directory
02_fleet_from_traffic.py Convert ADS-B traffic parquet to pycontrails Fleet JSON
03_generate_prompts.py Generate prompt masks (CoCiP or DryAdvection) — for unannotated / custom data
04_run_inference.py Run SAM2-Contrails on a frame sequence
05_evaluate.py Evaluate predictions against GVCCS ground truth
06_render_video.py Render overlay video with prompts and predictions

Prompt Generation

Prompt masks encode where each flight's contrail is expected to appear in each frame. Two contrail models are supported — both produce compatible output and can be used interchangeably in all downstream steps.

CoCiP (recommended for accuracy)

Full thermodynamic contrail lifecycle model. Predicts persistence, lifetime, and contrail width. Requires ERA5 meteorology and radiation data via a CDS API key.

import pandas as pd
from contrailtrack.prompts.cocip import run_cocip
from contrailtrack.prompts.projection import MiniProjector
from contrailtrack.prompts.writer import generate_prompts_video

stem = "20230930055430_20230930075430"

# Step 1: Run CoCiP — saved to disk so ERA5 downloads are not repeated
contrail_path = run_cocip(fleet_json=f"data/fleet/{stem}.json", output_dir="data/cocip/")

# Step 2: Write per-flight prompt PNGs (projection to pixel space happens in-memory)
generate_prompts_video(
    images_dir=f"data/frames/{stem}/",
    contrail_df=pd.read_parquet(contrail_path),
    output_dir="data/prompts_cocip_age5/",
    video_id=stem,
    projector=MiniProjector(),   # defaults to GVCCS camera constants
    max_age_min=5.0,
)

DryAdvection (faster, no radiation data)

Advects flight waypoints using ERA5 wind fields only. No radiation data, no aircraft performance model. Recommended for rapid prototyping or when ERA5 radiation is unavailable.

import pandas as pd
from contrailtrack.prompts.dry_advection import run_dry_advection
from contrailtrack.prompts.projection import MiniProjector
from contrailtrack.prompts.writer import generate_prompts_video

stem = "20230930055430_20230930075430"

contrail_path = run_dry_advection(fleet_json=f"data/fleet/{stem}.json",
                                  output_dir="data/dry_advection/")
generate_prompts_video(
    images_dir=f"data/frames/{stem}/",
    contrail_df=pd.read_parquet(contrail_path),
    output_dir="data/prompts_dry_advection_age5/",
    video_id=stem,
    projector=MiniProjector(),
    max_age_min=5.0,
)

Fine-tuning

To fine-tune on your own GVCCS split, first generate training prompts with GT masks, then launch training.

1. Prepare training data

# Run contrail model on the train split
uv run contrailtrack run-cocip \
    --fleet-dir  /path/to/fleet/train/ \
    --out        /path/to/cocip/train/ \
    --annotations /path/to/GVCCS/train/annotations.json

# Generate age-weighted prompts + GT masks
uv run contrailtrack generate-prompts \
    --contrail-dir  /path/to/cocip/train/ \
    --annotations   /path/to/GVCCS/train/annotations.json \
    --out           /path/to/GVCCS/train/per_object_data_age_5/ \
    --images-source /path/to/GVCCS/train/images/ \
    --max-age 5.0

This writes:

per_object_data_age_5/
  {video_id:05d}/
    {flight_id}/
      {frame_idx:05d}_prompt.png   # age-weighted prompt
      {frame_idx:05d}_mask.png     # GT segmentation mask
    {frame_idx:05d}_all_prompts_union.png  # union of all flight prompts (ternary)
img_folder/
  {video_id:05d}/
    {frame_idx:05d}.jpg            # symlinks to source images

Update the img_folder and gt_folder paths in your training config YAML to point to these directories.

2. Launch training

# Age-weighted 5-min prompts, 2 GPUs
uv run contrailtrack train \
    sam2/configs/sam2.1_training/sam2.1_hiera_b+_GVCCS_finetune_age.yaml \
    --num-gpus 2

# Ternary prompts (recommended)
uv run contrailtrack train \
    sam2/configs/sam2.1_training/sam2.1_hiera_b+_GVCCS_finetune_ternary.yaml \
    --num-gpus 2

Available training configs (in sam2/configs/sam2.1_training/):

Config Prompts Notes
..._finetune_age.yaml Age-weighted, 5-min window Recommended starting point
..._finetune_ternary.yaml Ternary (positive + negative), 5-min Best attribution precision
..._finetune_original.yaml Binary, 1-min window Baseline

Checkpoints and TensorBoard logs are written to the experiment_log_dir specified in the config (default: ./logs/).

Before training, download the SAM 2.1 base checkpoint:

bash checkpoints/download_ckpts.sh

Evaluation

Predictions are evaluated against GVCCS ground-truth annotations on three axes: segmentation mAP, tracking (detection rate, completeness, temporal IoU), and attribution (is each prediction assigned to the correct flight?).

from contrailtrack import evaluate

results = evaluate(
    predictions_path="results/20230930055430_20230930075430.json",
    gt_annotations="path/to/annotations.json",
    video_name="20230930055430_20230930075430",
    output_dir="evaluation/",
)
print(f"mAP:                {results['segmentation']['mAP']:.3f}")
print(f"Detection rate:     {results['tracking']['metrics']['detection_rate']:.1%}")
print(f"Attr. precision:    {results['attribution']['metrics']['attribution_precision']:.1%}")

Reproducing Paper Results on GVCCS

This section gives the exact commands to reproduce the numbers reported in the paper on the GVCCS test set using the pre-trained ternary checkpoint.

Prerequisites

  1. Download the GVCCS dataset from Zenodo. It unpacks as:

    GVCCS/test/
      images/            ← flat folder of all JPEG frames
      annotations.json   ← COCO-Video GT annotations
      parquet/           ← ADS-B trajectories, one .parquet per video
    
  2. Prepare per-video frame folders and fleet JSONs (see examples/01_prepare_data.py and examples/02_fleet_from_traffic.py):

    GVCCS=/path/to/GVCCS/test
    
    # Organise frames into per-video folders (outputs to data/frames/)
    uv run python examples/01_prepare_data.py $GVCCS --output $GVCCS/img_folder
    
    # Convert each parquet to a fleet JSON (one per video)
    mkdir -p $GVCCS/fleet
    for f in $GVCCS/parquet/*.parquet; do
        name=$(basename "$f" .parquet)
        uv run python examples/02_fleet_from_traffic.py "$f" --output "$GVCCS/fleet/${name}.json"
    done
  3. Generate age-weighted ternary prompt masks (requires a CDS API key for ERA5 data):

    uv run contrailtrack run-cocip \
        --fleet-dir   $GVCCS/fleet/ \
        --out         $GVCCS/cocip/ \
        --annotations $GVCCS/annotations.json
    
    uv run contrailtrack generate-prompts \
        --contrail-dir  $GVCCS/cocip/ \
        --annotations   $GVCCS/annotations.json \
        --out           $GVCCS/per_object_data_age_5/ \
        --images-source $GVCCS/images/ \
        --max-age 5.0

    If ERA5 is unavailable, substitute run-dry-advection for run-cocip — results will be close but not identical to the paper.

Step 1 — Run inference on the full test set

uv run contrailtrack run \
    --images-dir $GVCCS/img_folder/ \
    --prompts    $GVCCS/per_object_data_age_5/ \
    --config     ternary \
    --encoding   ternary \
    --out        results/ \
    --max-propagation-frames 0

Two flags are critical for reproducing the paper:

  • --max-propagation-frames 0 removes the propagation cap so each contrail is tracked for the full video duration.
  • --config ternary --encoding ternary selects the best-performing variant.

Weights are downloaded from Hugging Face automatically on the first run.

Step 2 — Evaluate against ground truth

Use evaluate-dataset (not evaluate) to pool all 24 videos into a single dataset-wide pass — this is how the paper numbers were computed:

uv run contrailtrack evaluate-dataset \
    --predictions-dir results/ \
    --labels          $GVCCS/annotations.json \
    --out             evaluation/ \
    --flight-mappings $GVCCS/per_object_data_age_5/

--flight-mappings restricts the ground truth to only the flights that received prompts, matching the evaluation scope described in the paper.

Expected results

Metric ternary original
mAP (IoU 0.25–0.75) 0.389 0.303
Detection rate 92.6 % 93.4 %
Completeness 0.731 0.724
Temporal IoU 0.515 0.500
Attribution precision 96.5 % 89.0 %

The paper reported mAP = 0.380 for the ternary variant; the open-source release produces slightly higher mAP (0.389) due to minor code-path improvements during the public release refactoring. All other metrics are within ± 0.5 % of the paper values.

Minor additional differences are expected if prompts are regenerated from scratch rather than using the exact ERA5 data used in the paper.


Repository Structure

sam2/               SAM2 base model (Meta AI, Apache 2.0)
  configs/            Training and inference YAML configs
contrailtrack/      Contrail-specific package
  model/              Model loading and Hugging Face Hub integration
  data/               Frame loading, prompt reading
  inference/          SAM2 video predictor wrapper
  output/             COCO RLE export
  eval/               Segmentation, tracking, and attribution metrics
  prompts/            CoCiP, DryAdvection, camera projection, prompt writer
training/           Fine-tuning code (dataset, sampler, loss, trainer)
examples/           Step-by-step usage examples (01–06)
assets/             Figures and logo

Citation

If you use SAM2-Contrails in your research, please cite our paper:

@article{dalmau2025contrails,
  title   = {Contrails Cannot Exist Without Flights:
             Physics-Informed Detection, Tracking, and Attribution},
  author  = {Dalmau, Ramon and Jarry, Gabriel and Very, Philippe},
  year    = {2025},
}

Please also cite the original SAM 2 paper:

@article{ravi2024sam2,
  title   = {SAM 2: Segment Anything in Images and Videos},
  author  = {Ravi, Nikhila and Gabeur, Valentin and Hu, Yuan-Ting and others},
  journal = {arXiv preprint arXiv:2408.00714},
  year    = {2024},
}

Acknowledgements

We are grateful to the following organisations and projects:

  • Reuniwatt — the all-sky cameras used to record the GVCCS dataset were purchased from Reuniwatt by EUROCONTROL.
  • Encord — for the annotation platform used to label the contrail segmentation masks.
  • OpenSky Network — for providing the ADS-B flight trajectory data used to generate contrail prompts.
  • pycontrails — for the CoCiP and DryAdvection contrail models.
  • Meta AI — for SAM 2, the foundation model this work builds upon.
  • cc_torch — for the GPU connected-component post-processing (see LICENSE_cctorch).

License

The contrailtrack package is released under the European Union Public Licence v1.2 (EUPL-1.2).

The SAM 2 base model code (in sam2/) retains its original Apache 2.0 License from Meta AI.

About

Fine-tuned SAM2 for contrail detection, tracking, and attribution in all-sky camera video using physics-informed prompt masks from ADS-B and ERA5 data

Topics

Resources

Stars

4 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages