unet_segmentation: swap tree-crown base checkpoint for a building-pre…

models/unet_segmentation/README.md

@@ -12,7 +12,40 @@ Semantic segmentation model for building footprint extraction from aerial imager
 
 ## Pretrained Source
 
-OAM-TCD via torchgeo `Unet_Weights.OAM_RGB_RESNET34_TCD` (NeurIPS 2024, [arxiv.org/abs/2407.11743](https://arxiv.org/abs/2407.11743))
+Building-segmentation base checkpoint trained on
+[hotosm/vhr-building-segmentation](https://huggingface.co/datasets/hotosm/vhr-building-segmentation)
+(57,890 train / 7,237 val chips).
+
+### Base checkpoint provenance
+
+The previous base (`unet_resnet34_oam_rgb_tcd`) is a **tree-crown delineation**
+model: it scores F1 ≈ 0 on buildings zero-shot, and fine-tuning from it reached
+50.0 pooled test F1 at 32 chips — below training the decoder from random
+initialization (52.4; 6 regions × 5 folds). The building-pretrained base of
+identical architecture reaches 59.3 pooled (+9.3) under the same recipe.
+
+Training config of the replacement base (mirrors this pipeline's fine-tune
+recipe, seed 1337):
+
+- Architecture: `torchgeo`/`segmentation_models_pytorch` UNet, ResNet34
+  encoder (ImageNet init, frozen); decoder + 2-class head from random init
+- Data: `hotosm/vhr-building-segmentation`, 256×256 chips, inputs scaled `/255`
+- Loss/optim: 2-class cross-entropy, AdamW lr 1e-3 (sweep-selected for
+  from-scratch decoder training), weight decay 1e-4, batch 32, cosine schedule
+  (T_max 50), gradient clip 1.0, fp16 autocast
+- Early stopping patience 5 on validation cross-entropy; best checkpoint kept
+  (val CE 0.2351)
+
+Zero-shot on the dataset's own held-out splits (building-class pixel metrics,
+argmax operating point; val n=7,237 / test n=7,236 chips):
+
+| Base | Split | Building IoU | Building F1 |
+|---|---|---|---|
+| tree-crown (previous) | val / test | 0.000 / 0.000 | 0.000 / 0.000 |
+| buildings (this) | val / test | 0.252 / 0.342 | 0.402 / 0.510 |
+
+Both checkpoints load into the pipeline's model with 0 missing / 0 unexpected
+keys (`scripts/eval_bases_dataset_val.py`).
 
 ## Pipeline
 

models/unet_segmentation/pipeline.py

@@ -1,7 +1,8 @@
 """ZenML pipeline for UNet building segmentation.
 
 Entrypoints referenced by models/unet_segmentation/stac-item.json.
-Pretrained weights: OAM-TCD (arxiv.org/abs/2407.11743).
+Pretrained weights: building-segmentation base trained on
+hotosm/vhr-building-segmentation (see README for provenance).
 """
 
 import tempfile

models/unet_segmentation/stac-item.json

@@ -48,7 +48,7 @@
     "created": "2024-01-01T00:00:00Z",
     "updated": "2024-01-01T00:00:00Z",
     "title": "UNet Building Segmentation",
-    "description": "UNet semantic segmentation model for building footprint extraction from very high resolution RGB aerial imagery. Pretrained on the OAM-TCD dataset (NeurIPS 2024) via torchgeo, intended for finetuning on OpenAerialMap + OpenStreetMap chip-label pairs.",
+    "description": "UNet semantic segmentation model for building footprint extraction from very high resolution RGB aerial imagery. Pretrained on the hotosm/vhr-building-segmentation dataset (frozen ImageNet encoder, decoder + 2-class head trained on 57,890 building chips), intended for finetuning on OpenAerialMap + OpenStreetMap chip-label pairs.",
     "mlm:name": "unet-segmentation",
     "mlm:architecture": "UNet",
     "mlm:tasks": [
@@ -57,7 +57,7 @@
     "mlm:framework": "PyTorch",
     "mlm:framework_version": "2.10.0",
     "mlm:pretrained": true,
-    "mlm:pretrained_source": "https://arxiv.org/abs/2407.11743",
+    "mlm:pretrained_source": "https://huggingface.co/datasets/hotosm/vhr-building-segmentation",
     "mlm:accelerator": "cuda",
     "mlm:accelerator_count": 1,
     "keywords": [
@@ -323,8 +323,8 @@
   },
   "assets": {
     "checkpoint": {
-      "href": "https://hf.co/isaaccorley/unet_resnet34_oam_rgb_tcd/resolve/40c914bbcbe43a6a87c81adb0a22ff2d4a53204d/unet_resnet34_oam_rgb_tcd-9472042e.pth",
-      "title": "Pretrained weights (OAM-TCD UNet ResNet34)",
+      "href": "https://hf.co/nilsho01/unet-resnet34-vhr-buildings/resolve/289e829883fc17b253cf9892946f6502199998a8/unet_bldg_base.pth",
+      "title": "Pretrained weights (VHR buildings UNet ResNet34)",
       "type": "application/octet-stream; framework=pytorch",
       "roles": [
         "mlm:model",
@@ -406,9 +406,9 @@
     },
     {
       "rel": "cite-as",
-      "href": "https://arxiv.org/abs/2407.11743",
+      "href": "https://huggingface.co/datasets/hotosm/vhr-building-segmentation",
       "type": "text/html",
-      "title": "OAM-TCD Pretrained Weights (NeurIPS 2024)"
+      "title": "VHR Building Segmentation Dataset (base-model training data)"
     }
   ]
 }