deepforest_UNet_OAM: tree-crown segmentation for OpenAerialMap chips

A U-Net that labels every pixel of an RGB aerial chip as tree canopy or background. Its encoder is the ResNet-50 body of the DeepForest detector, in the NEON_OAM version (the NEON model fine-tuned by HOT on OpenAerialMap imagery); the decoder is new and was trained on the Guru trees dataset. It is built for the HOT fAIr platform: RGB chips in, tree polygons out.

This card says what was done, what worked, what did not, and where the numbers should be read with care. The test score is much lower than the validation score; most of that gap comes from missing labels in the dataset, not from a different kind of imagery (section "Read the scores with care").

Files

File	What it is
NEON_OAM_UNET.onnx	the deployable model, 130 MB: input image (3, H, W) float32 RGB in [0, 1], dynamic H and W; output probability (H, W) float32 in [0, 1]. ImageNet normalisation, the network and the sigmoid are inside the graph
NEON_OAM_UNET.pt	the PyTorch weights (a TreeUNet state dict, 130 MB), for further fine-tuning with the ts-hot library
results/	validation and test metrics (threshold sweeps), training history and settings, ONNX parity check, the run report and the train / validation / test comparison

Quick start (ONNX, CPU)

The chips are 256 px tiles named OAM-<x>-<y>-<z>.tif; the zoom and tile row give the true ground size of a pixel, which sets the scale.

import math
import numpy as np
import onnxruntime as ort
import rasterio
from PIL import Image

MODEL_GSD_M = 0.05            # ground size of one model pixel the encoder was trained at
WORLD_M = 40075016.68558      # circumference of the Web-Mercator world, metres

def scale_for(y, z, width, height, max_px=1024):
    lat = math.atan(math.sinh(math.pi * (1 - 2 * (y + 0.5) / 2**z)))      # latitude of the tile centre
    chip_gsd_m = WORLD_M * math.cos(lat) / 2**z / width                    # metres per chip pixel
    return min(chip_gsd_m / MODEL_GSD_M, max_px / max(width, height))

session = ort.InferenceSession("NEON_OAM_UNET.onnx", providers=["CPUExecutionProvider"])
with rasterio.open("OAM-485645-499613-20.tif") as src:
    rgb = src.read([1, 2, 3])                                              # (3, H, W) uint8
_, h, w = rgb.shape
s = scale_for(y=499613, z=20, width=w, height=h)
size = (round(w * s), round(h * s))                                        # 256 px at zoom 20 -> 756 px
image = np.stack([np.asarray(Image.fromarray(b.astype(np.float32) / 255.0).resize(size, Image.BILINEAR))
                  for b in rgb]).astype(np.float32)
probability = session.run(None, {"image": image})[0]                       # (756, 756), tree probability
tree = probability >= 0.3                                                   # threshold chosen on validation

This snippet was run against the ts-hot library's own preprocessing on a real chip and gives identical input and output. To get GeoJSON polygons (threshold, tracing, small-polygon removal, simplification, georeferencing to EPSG:4326) use ts_hot.predict_session from the HOT ts-hot helper library.

Preprocessing

• RGB only, 3 bands, uint8 0 to 255 read as float 0 to 1.
• One resize by the zoom rule: scale = chip ground size / 0.05 m, capped so the longest side stays within 1024 px. A 256 px chip at zoom 20 (about 0.148 m per pixel at these latitudes) is scaled by 2.95 to 756 x 756. Feeding the model chips at their original resolution finds far fewer crowns; the scale is the most important input rule. Bilinear, in float.
• No padding or cropping. The network takes any height and width and returns a map of the same size (checked at 756 x 756, 768 x 768 and several small sizes). A test that pads to a multiple of 32 and crops back was not run.
• Normalisation: the ImageNet mean and standard deviation, applied inside the model, so the input is just RGB in [0, 1].

Model

ResNet-50 encoder (stem and layer1 to layer4, strides 2 to 32, loaded from the backbone.body.* weights of DeepForest NEON_OAM) and a new U-Net decoder of five up-blocks (256, 128, 64, 32, 16 channels), each upsampling to the size of the matching skip features and applying two 3x3 convolutions with BatchNorm; a last 1x1 convolution gives the logit. 32.5 million parameters. The exported graph was compared with the PyTorch model on 200 validation and 200 test chips: the largest probability difference on any pixel is 0.005 (validation) and 0.011 (test), the mean difference is about 1e-5 to 2e-5, and the pixel IoU at threshold 0.3 agrees to 1e-4.

Training data and labels

Dataset: Guru trees (hotosm/fAIr-guru-td-trees), OpenAerialMap tiles at zoom 20 with labels from OpenStreetMap. Used as shipped: 38,705 train, 4,158 validation and 5,627 test chips from the dataset's own spatial split (no chip is in two splits).

Labels: one GeoJSON per split, no masks. A tree_group polygon is filled as tree; a natural=tree point becomes a disc of its diameter_crown metres at the chip's true ground size; a point with no usable diameter becomes a 5 m disc that is ignored (neither rewarded nor punished). Tree wins over ignore. The masks were checked against an independent construction from the dataset's own tile maths on 300 chips per split: tree IoU 0.997 to 0.998.

Training recipe (the published model):

Start	encoder from NEON_OAM, decoder random
Encoder	frozen, except layer4, trained at a tenth of the learning rate (BatchNorm layers never train)
Optimiser	AdamW, learning rate 1e-3, weight decay 1e-4, cosine schedule per batch, gradient clipping 1.0
Loss	binary cross-entropy + Dice (weight 1.0) over the pixels that are not ignored
Length	10 epochs, each a fresh random tenth of the training chips (one pass over the data in all), batch size 4, whole chips (756 x 756)
Selection	the epoch with the lowest validation loss on all 4,158 validation chips (the last, epoch 10)
Seed	42, a single run
Compute	one NVIDIA RTX 4060 laptop GPU, about 119 minutes, 3.3 GB of GPU memory

Augmentation

Used for this model: a random quarter turn (0, 90, 180 or 270 degrees) plus a random mirror, on the training chips only, applied to the image and its mask together. Nothing else. No brightness or contrast change, no scale jitter, no blur or JPEG.

What the augmentation experiments showed: augmentation did not help on this dataset. In short runs (3 epochs of a tenth of the chips, 500 validation chips, one seed) the validation IoU against the reference of 0.5654 was:

Augmentation	Val IoU
quarter turn + mirror (reference)	0.5654
mirror only	0.5664
none	0.5676
+ brightness / contrast / gamma 0.2	0.5657
+ scale jitter 0.15	0.5669
+ blur or JPEG, probability 0.3	0.5641

All six are within 0.003 of each other, which is inside the noise of a single seed. The worry that rotating the chip would turn the shadows into a direction the sun never cast showed no measurable effect. The ts-hot library therefore now defaults to no augmentation; this published model was trained earlier with the quarter turn and mirror, and a full-length run without any augmentation has not been done.

What gave the best result

Change (short runs, validation IoU, 500 chips)	Result
Dice weight 0 / 0.5 / 1.0	0.564 / 0.566 / 0.565: no difference
Learning rate 3e-4 / 1e-3 / 3e-3	0.564 / 0.565 / 0.565: no difference
Unfreeze layer4, encoder rate 0.1 / 0.3	0.574 / 0.560
Post-processing: threshold 0.2 to 0.4, minimum area, simplification, smoothing, opening / closing, hole fill	within 0.001 of each other (IoU 0.565); threshold 0.1 or 0.5 and above is worse
Averaging 4 mirrored passes at inference	+0.006 IoU for 4 times the CPU time (not adopted)

The one change that helped was training layer4 of the encoder at a low rate. At full length (10 epochs, all validation and test chips):

Run	Val IoU / F1	Test IoU / F1 (threshold 0.3)
Frozen encoder (decoder only)	0.559 / 0.718	0.330 / 0.496
layer4 trained at 0.1 (this model)	0.570 / 0.726	0.404 / 0.576

For reference, the HOT td-hot detector with its boxes filled in as masks scores IoU 0.462 on 500 validation chips, and an untrained decoder on the same encoder at most 0.107.

Results of this model (threshold 0.3, chosen on the validation chips; precision / recall / F1 / IoU of the tree class)

Split	Chips	Precision	Recall	F1	IoU
Validation	4,158	0.697	0.759	0.726	0.570
Test	5,627	0.440	0.833	0.576	0.404

At the test-optimal threshold (0.8) the test IoU is 0.427; this is shown only to explain the gap below, the operating point is the validation threshold. One seed: differences under about 0.01 IoU should be treated as noise.

Read the scores with care

The labels come from OpenStreetMap, so they hold only the trees somebody chose to map.

• About 25% of the train and test chips have no label at all (8% of the validation chips), and many of those show dense canopy. The model marks such canopy as tree, which is counted as a false positive. In the 8 unlabelled test chips with the most predicted trees, all were forest canopy that looks like real tree cover. This was judged from pictures and not counted, and it was not checked on the ground.
• Recall is high on every split (0.75 to 0.87); precision falls on the test split (0.44 against 0.70). The same effect shows on a train sample, so it is a label gap and not overfitting.
• The validation chips are the most completely mapped (5.0 labels per chip against 2.9 in test), so the validation score is the optimistic one and the test score is the pessimistic one. Neither is a clean accuracy.
• A point tree is drawn as a disc of its crown-diameter tag. Real crowns are irregular, and the model often outlines the crown larger than the disc, so part of the pixel error is the label, not the model.
• Over continuous forest the model returns one large polygon. Whether that is wanted is a product decision, not something the score decides.

Full tables and 100-chip overlay pictures per split are in results/split_comparison.md.

Limitations

• One dataset, one region, one zoom (20), one seed. Other cities, sensors, seasons and zooms were not tested.
• Pixel metrics only (IoU, F1 of the tree class). Object-level scores (per-crown precision and recall) and a point hit rate were not measured. The model does not separate touching crowns; they merge into one polygon.
• Polygons are traced per chip, so a crown cut by a chip border comes back as two polygons. predict_session returned valid, closed, finite polygons on 20 real chips (72 polygons), but polygons were not tested for self-intersection.
• Not suited to counting individual trees, or to forest-cover statistics without checking the labels of the area.
• Licence: none is set on this repository yet. DeepForest is MIT licensed; the licence of the dataset was not checked here, so check it before reuse.

Reproducing

Everything above was produced with the HOT ts-hot helper library and research scripts (training, trials, the post-processing sweep, the split comparison). The training settings are in results/training_settings.json and the first line of the run report; results/run_report.md holds the loss per epoch, and results/validation_metrics.json and results/test_metrics.json the threshold sweeps.