diff --git a/README.md b/README.md index f4f6d43..2a52144 100644 --- a/README.md +++ b/README.md @@ -118,15 +118,19 @@ CABiNet/ │ ├── datasets/ # Data loading and preprocessing │ │ ├── cityscapes.py # Cityscapes dataset loader │ │ ├── uavid.py # UAVid dataset loader +│ │ ├── aeroscapes.py # AeroScapes dataset loader +│ │ ├── vdd.py # VDD (Varied Drone Dataset) loader │ │ ├── registry.py # Shared dataset registry (train.py + evaluate.py) │ │ └── transform.py # Data augmentation pipeline │ ├── scripts/ # Training, evaluation, inference scripts -│ │ ├── train.py # CABiNet training -│ │ ├── evaluate.py # CABiNet standalone checkpoint evaluation -│ │ ├── visualize.py # Prediction visualization (Cityscapes only) -│ │ ├── train_yolo.py # YOLO26-sem training/validation (Hydra CLI) -│ │ ├── infer_yolo.py # YOLO26-sem inference + showcase mosaic -│ │ └── convert_uavid_to_yolo.py # UAVid RGB masks -> single-channel format +│ │ ├── train.py # CABiNet training +│ │ ├── evaluate.py # CABiNet standalone checkpoint evaluation +│ │ ├── visualize.py # Prediction visualization (Cityscapes only) +│ │ ├── train_yolo.py # YOLO26-sem training/validation (Hydra CLI) +│ │ ├── infer_yolo.py # YOLO26-sem inference + showcase mosaic +│ │ ├── convert_uavid_to_yolo.py # UAVid RGB masks -> single-channel format +│ │ ├── convert_aeroscapes_to_yolo.py # AeroScapes -> shared images/+masks/ format +│ │ └── convert_vdd_to_yolo.py # VDD -> shared images/+masks/ format │ └── utils/ # Utility functions │ ├── loss.py # OHEM Cross Entropy + Focal Loss │ ├── optimizer.py # Custom optimizer with warmup @@ -141,14 +145,19 @@ CABiNet/ │ ├── evaluate.yaml # Standalone evaluate.py CLI configuration │ ├── train_yolo.yaml # YOLO26-sem training configuration (Hydra) │ ├── dataset/ # Dataset-specific configs -│ │ ├── cityscapes.yaml / uavid.yaml # CABiNet training configs -│ │ └── uavid_yolo.yaml # Ultralytics dataset YAML +│ │ ├── cityscapes.yaml / uavid.yaml / aeroscapes.yaml / vdd.yaml # CABiNet training configs +│ │ └── uavid_yolo.yaml / aeroscapes_yolo.yaml / vdd_yolo.yaml # Ultralytics dataset YAMLs │ ├── model/ # CABiNet backbone configs (mobilenetv3_{large,small}.yaml) │ ├── yolo/ # Ultralytics YOLO configs │ │ ├── uavid_train.yaml / uavid_val.yaml +│ │ ├── aeroscapes_train.yaml / aeroscapes_val.yaml +│ │ ├── vdd_train.yaml / vdd_val.yaml │ │ └── model/ # Per-variant configs (yolo26{n,s,m,l,x}-sem.yaml, legacy -seg) +│ ├── train_yolo_aeroscapes.yaml # Hydra root config for train_yolo.py on AeroScapes +│ ├── train_yolo_vdd.yaml # Hydra root config for train_yolo.py on VDD │ ├── cityscapes_info.json # Cityscapes label information -│ └── UAVid_info.json # UAVid label information +│ ├── UAVid_info.json # UAVid label information +│ └── AeroScapes_info.json # AeroScapes label information ├── hf_modelcards/ # Hugging Face model card generator (template + per-model metrics) ├── tests/ # Test suite │ ├── unit/ # Unit tests @@ -172,6 +181,8 @@ CABiNet/ - **[`cityscapes.py`](src/datasets/cityscapes.py)**: Cityscapes dataset loader with label remapping and thread-safe preprocessing - **[`uavid.py`](src/datasets/uavid.py)**: UAVid dataset loader — consumes the pre-converted `images/+masks/` format (same as the YOLO pipeline), `RandomCrop`-based training +- **[`aeroscapes.py`](src/datasets/aeroscapes.py)**: AeroScapes dataset loader — same pre-converted `images/+masks/` format and `RandomCrop`-based training as `uavid.py`; source images are uniformly 1280x720, so unlike UAVid there's no mixed-resolution batching constraint +- **[`vdd.py`](src/datasets/vdd.py)**: VDD (Varied Drone Dataset) loader — same pre-converted `images/+masks/` format; source images are uniformly 4000x3000 and ship with a real train/val/test split already - **[`transform.py`](src/datasets/transform.py)**: Comprehensive data augmentation including geometric, photometric, and regularization transforms #### Scripts ([`src/scripts/`](src/scripts/)) @@ -182,6 +193,8 @@ CABiNet/ - **[`train_yolo.py`](src/scripts/train_yolo.py)**: YOLO26 semantic segmentation training/validation (Hydra CLI wrapping Ultralytics) - **[`infer_yolo.py`](src/scripts/infer_yolo.py)**: YOLO26-sem inference over images/videos/folders, plus a `--showcase-videos` 2x2 mosaic mode - **[`convert_uavid_to_yolo.py`](src/scripts/convert_uavid_to_yolo.py)**: Converts UAVid RGB colour masks → single-channel class-ID format shared by both pipelines +- **[`convert_aeroscapes_to_yolo.py`](src/scripts/convert_aeroscapes_to_yolo.py)**: Converts AeroScapes (already single-channel class-ID masks) into the shared `images/+masks/` layout — copies (not symlinks) so the result is redistributable +- **[`convert_vdd_to_yolo.py`](src/scripts/convert_vdd_to_yolo.py)**: Converts VDD (already single-channel class-ID masks, already split into train/val/test) into the shared `images/+masks/` layout — symlinks (VDD is already published on HF, so no redistribution concern) #### Utilities ([`src/utils/`](src/utils/)) @@ -197,13 +210,18 @@ CABiNet/ - **[`train.yaml`](configs/train.yaml)**: Main CABiNet training configuration (Hydra) - **[`evaluate.yaml`](configs/evaluate.yaml)**: Standalone `evaluate.py` CLI configuration - **`dataset/*.yaml`**: Dataset-specific configurations (paths, preprocessing parameters) - - `cityscapes.yaml` / `uavid.yaml` — CABiNet training configs + - `cityscapes.yaml` / `uavid.yaml` / `aeroscapes.yaml` / `vdd.yaml` — CABiNet training configs - `uavid_yolo.yaml` — Ultralytics dataset YAML (8 classes incl. Clutter, single-channel masks, 255=ignore) + - `aeroscapes_yolo.yaml` — Ultralytics dataset YAML (12 classes, single-channel masks, no test split) + - `vdd_yolo.yaml` — Ultralytics dataset YAML (7 classes, single-channel masks, train/val/test) - **`model/*.yaml`**: CABiNet backbone configurations - **`yolo/*.yaml`**: Ultralytics YOLO training/evaluation configs - [`uavid_train.yaml`](configs/yolo/uavid_train.yaml) — full training config (AMP, EMA, grad accum, resume, augmentation) - [`uavid_val.yaml`](configs/yolo/uavid_val.yaml) — validation / benchmark config + - [`aeroscapes_train.yaml`](configs/yolo/aeroscapes_train.yaml) / [`aeroscapes_val.yaml`](configs/yolo/aeroscapes_val.yaml) — same, for AeroScapes + - [`vdd_train.yaml`](configs/yolo/vdd_train.yaml) / [`vdd_val.yaml`](configs/yolo/vdd_val.yaml) — same, for VDD - `model/yolo26{n,s,m,l,x}-sem.yaml` — per-size model configs, selected via `'yolo/model@model=yolo26s-sem'` +- **[`train_yolo_aeroscapes.yaml`](configs/train_yolo_aeroscapes.yaml)** / **[`train_yolo_vdd.yaml`](configs/train_yolo_vdd.yaml)**: Hydra root configs for `train_yolo.py` on AeroScapes/VDD (`--config-name train_yolo_aeroscapes` / `train_yolo_vdd`) — `train_yolo.yaml` remains the UAVid default ## Usage @@ -275,6 +293,85 @@ Key `training_config` knobs (see [`configs/train.yaml`](configs/train.yaml) for `configs/dataset/uavid.yaml`'s `augmentation:` block (degrees, translate, scale, flips, HSV, mixup) mirrors the YOLO26 pipeline's recipe below — see [`transform.py`](src/datasets/transform.py). +#### AeroScapes Dataset + +A second aerial semantic segmentation dataset, wired in the same way as UAVid — a pre-converted `images/+masks/` layout consumed by both CABiNet and the YOLO26 pipeline. Unlike UAVid, AeroScapes source images are already single-channel class-ID masks (no RGB decoding step) and uniformly 1280x720 (no mixed-resolution batching constraint). + +1. **Download** the AeroScapes dataset (`JPEGImages/`, `SegmentationClass/`, `ImageSets/{trn,val}.txt`) + +2. **Convert** the raw dataset (see [AeroScapes → YOLO Format](#aeroscapes--yolo-format) for full details): + + ```bash + python src/scripts/convert_aeroscapes_to_yolo.py \ + --src /path/to/raw/aeroscapes --dst /path/to/converted --workers 8 + ``` + + Unlike `convert_uavid_to_yolo.py`, this always **copies** files (no symlink option) — the converted output is meant to be redistributable as-is. + +3. **Configure and train**: + + ```bash + export AEROSCAPES_YOLO_ROOT=/path/to/converted + # or edit configs/dataset/aeroscapes.yaml and set 'dataset_path:' directly + + python src/scripts/train.py dataset=aeroscapes + ``` + + Source images are uniformly 1280x720, so — unlike UAVid — `validation_config.batch_size` can be left at its default; no special-casing is required. + +There is no source test split for AeroScapes (only `trn.txt`/`val.txt`); the converter only ever produces `train`/`val`. + +#### VDD Dataset + +A third aerial semantic segmentation dataset — [VDD (Varied Drone Dataset)](https://github.com/RussRobin/VDD), already published on [Hugging Face](https://huggingface.co/datasets/RussRobin/VDD). Wired in the same way as UAVid/AeroScapes. Unlike both, VDD ships with a real train/val/test split already defined by its own directory layout (`train/`, `val/`, `test/`, each with `src/` + `gt/`), and its masks are already single-channel class-ID PNGs like AeroScapes (no RGB decoding needed). + +1. **Download** the VDD dataset (e.g. `git clone` the HF dataset repo) — this gives `train/{src,gt}`, `val/{src,gt}`, `test/{src,gt}`. + +2. **Convert** the raw dataset (see [VDD → YOLO Format](#vdd--yolo-format) for full details): + + ```bash + python src/scripts/convert_vdd_to_yolo.py \ + --src /path/to/VDD --dst /path/to/converted + ``` + + Unlike `convert_aeroscapes_to_yolo.py`, this **symlinks** by default (like `convert_uavid_to_yolo.py`) — VDD is already redistributed on HF, so there's no need to produce a standalone copy. + +3. **Configure and train**: + + ```bash + export VDD_YOLO_ROOT=/path/to/converted + # or edit configs/dataset/vdd.yaml and set 'dataset_path:' directly + + python src/scripts/train.py dataset=vdd + ``` + + Source images are uniformly 4000x3000, so — like AeroScapes and unlike UAVid — `validation_config.batch_size` can be left at its default. + +VDD is small (280 train / 80 val / 40 test images) — heavier augmentation (`mosaic`, `mixup`, `copy_paste`) is enabled by default in the YOLO configs to help offset this. + +#### Finetuning CABiNet from a pretrained checkpoint + +All CABiNet aerial datasets (UAVid, AeroScapes, VDD) share the same backbone/CAB/spatial/fusion +architecture and only differ in the final classifier heads (sized by `num_classes`). Rather than +always starting from ImageNet backbone weights, `training_config.pretrained_ckpt_path` warm-starts +the **whole model** from a checkpoint trained on a different aerial dataset — converges +substantially faster than backbone-only init, since the CAB/fusion layers are already adapted to +aerial imagery, not just the backbone. + +```bash +# Finetune on AeroScapes starting from a UAVid-trained checkpoint +python src/scripts/train.py dataset=aeroscapes \ + training_config.pretrained_ckpt_path=experiments/uavid/.../checkpoint_last.pth +``` + +Only parameters whose name **and shape** match are loaded — the two classifier heads +(`ab.b4`, `conv_out.conv_out`) are automatically skipped (left at fresh init) whenever +`num_classes` differs between the source checkpoint and the target dataset; everything else +transfers. This is independent of `resume` (which restores optimizer/EMA/epoch state to continue +an interrupted run of the _same_ experiment) — `pretrained_ckpt_path` always starts a fresh run +(epoch 0, fresh optimizer/EMA) with warm-started weights. Accepts either a full training +checkpoint (`checkpoint_last.pth`) or a raw model state_dict (`*_best.pth`). + ### UAVid → YOLO Format UAVid distributes ground-truth labels as **3-channel RGB colour-coded masks** (in each sequence's `Labels/` directory). YOLO's semantic segmentation task requires **single-channel PNG masks** where each pixel value is a class index (0 – N-1); pixel value `255` is reserved for pixels whose colour doesn't match any known class (corrupted/anti-aliased data) and is excluded from training/eval. Per the original UAVid paper, Clutter is a valid class — it is **not** mapped to the ignore label. @@ -392,6 +489,154 @@ All 8 UAVid classes are valid and active — none are mapped to the ignore label --showcase-videos clip1.mp4 clip2.mp4 clip3.mp4 clip4.mp4 --output outputs/inference ``` +### AeroScapes → YOLO Format + +AeroScapes distributes ground-truth labels as **already single-channel PNG masks** (`SegmentationClass/`, pixel value = class index 0-11) — unlike UAVid, no RGB colour decoding step is needed. `convert_aeroscapes_to_yolo.py` reads the split membership from `ImageSets/{trn,val}.txt` and copies each `(image, mask)` pair into the shared layout, validating that every mask pixel is a known class ID (or `255`) along the way. + +#### Class mapping + +All 12 AeroScapes classes are valid and active — none are mapped to the ignore label: + +| YOLO ID | Class | Original RGB (Visualizations) | +| ------- | ------------ | ----------------------------- | +| 0 | Background | `[ 0, 0, 0]` | +| 1 | Person | `[192, 128, 128]` | +| 2 | Bike | `[ 0, 128, 0]` | +| 3 | Car | `[128, 128, 128]` | +| 4 | Drone | `[128, 0, 0]` | +| 5 | Boat | `[ 0, 0, 128]` | +| 6 | Animal | `[192, 0, 128]` | +| 7 | Obstacle | `[192, 0, 0]` | +| 8 | Construction | `[192, 128, 0]` | +| 9 | Vegetation | `[ 0, 64, 0]` | +| 10 | Road | `[128, 128, 0]` | +| 11 | Sky | `[ 0, 128, 128]` | + +`255` is reserved for genuinely unrecognized pixel values, which should not occur in a clean copy of the dataset — the converter validates and warns on any mask that contains one. + +#### Step-by-step workflow + +1. **Convert** (copies, not symlinks — the result is meant to be redistributable): + + ```bash + python src/scripts/convert_aeroscapes_to_yolo.py \ + --src /path/to/aeroscapes \ + --dst /path/to/aeroscapes_yolo \ + --workers 8 + ``` + + Output layout: + + ``` + /path/to/aeroscapes_yolo/ + ├── images/ + │ ├── train/ ← copies of JPEGImages/*.jpg (2621 images) + │ └── val/ ← copies of JPEGImages/*.jpg (648 images) + └── masks/ + ├── train/ ← copies of SegmentationClass/*.png (already single-channel) + └── val/ + ``` + + There is no source test split — only `train`/`val` are produced. + +2. **Point the dataset YAML** at the converted data: + + ```bash + export AEROSCAPES_YOLO_ROOT=/path/to/aeroscapes_yolo + # or edit configs/dataset/aeroscapes_yolo.yaml and set 'path:' directly + ``` + +3. **Train** — either the direct Ultralytics CLI: + + ```bash + yolo semantic train cfg=configs/yolo/aeroscapes_train.yaml + ``` + + or the repo's Hydra wrapper: + + ```bash + python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes + ``` + +4. **Evaluate**: + + ```bash + yolo semantic val cfg=configs/yolo/aeroscapes_val.yaml model=runs/aeroscapes/yolo26n/weights/best.pt + # or + python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes mode=val + ``` + +### VDD → YOLO Format + +VDD distributes ground-truth labels as **already single-channel PNG masks** (`gt/`, pixel value = class index 0-6) — like AeroScapes and unlike UAVid, no RGB colour decoding step is needed. `convert_vdd_to_yolo.py` discovers stems present in both `/src/` and `/gt/` and symlinks each `(image, mask)` pair into the shared layout (renaming the image extension to lowercase `.jpg` along the way), validating that every mask pixel is a known class ID (or `255`). + +#### Class mapping + +All 7 VDD classes are valid and active — none are mapped to the ignore label: + +| YOLO ID | Class | +| ------- | ---------- | +| 0 | Other | +| 1 | Wall | +| 2 | Road | +| 3 | Vegetation | +| 4 | Vehicle | +| 5 | Roof | +| 6 | Water | + +`255` is reserved for genuinely unrecognized pixel values, which should not occur in a clean copy of the dataset — the converter validates and warns on any mask that contains one. + +#### Step-by-step workflow + +1. **Convert** (symlinks — VDD is already published on HF, no redistribution copy needed): + + ```bash + python src/scripts/convert_vdd_to_yolo.py \ + --src /path/to/VDD \ + --dst /path/to/vdd_yolo + ``` + + Output layout: + + ``` + /path/to/vdd_yolo/ + ├── images/ + │ ├── train/ ← symlinks to src/*.JPG (280 images), renamed to lowercase .jpg + │ ├── val/ ← 80 images + │ └── test/ ← 40 images + └── masks/ + ├── train/ ← symlinks to gt/*.png (already single-channel) + ├── val/ + └── test/ + ``` + +2. **Point the dataset YAML** at the converted data: + + ```bash + export VDD_YOLO_ROOT=/path/to/vdd_yolo + # or edit configs/dataset/vdd_yolo.yaml and set 'path:' directly + ``` + +3. **Train** — either the direct Ultralytics CLI: + + ```bash + yolo semantic train cfg=configs/yolo/vdd_train.yaml + ``` + + or the repo's Hydra wrapper: + + ```bash + python src/scripts/train_yolo.py --config-name train_yolo_vdd + ``` + +4. **Evaluate**: + + ```bash + yolo semantic val cfg=configs/yolo/vdd_val.yaml model=runs/vdd/yolo26n/weights/best.pt + # or + python src/scripts/train_yolo.py --config-name train_yolo_vdd mode=val + ``` + ### Evaluation Evaluate a trained checkpoint standalone, independent of a training run (Hydra config: [`configs/evaluate.yaml`](configs/evaluate.yaml)). Accepts either a raw model state_dict (`*_best.pth` / the final saved `.pth`) or a full training checkpoint (`checkpoint_last.pth`): @@ -407,6 +652,18 @@ python src/scripts/evaluate.py checkpoint_path=... dataset=uavid validation_conf # UAVid test split python src/scripts/evaluate.py checkpoint_path=... dataset=uavid validation_config.batch_size=1 split=test +# AeroScapes — source images are uniformly 1280x720, so no batch_size=1 +# constraint is needed; dataset_path comes from configs/dataset/aeroscapes.yaml +# (no test split exists for this dataset) +python src/scripts/evaluate.py checkpoint_path=... dataset=aeroscapes + +# VDD — source images are uniformly 4000x3000, so no batch_size=1 constraint +# is needed; dataset_path comes from configs/dataset/vdd.yaml +python src/scripts/evaluate.py checkpoint_path=... dataset=vdd + +# VDD test split +python src/scripts/evaluate.py checkpoint_path=... dataset=vdd split=test + # Fast single-scale evaluation (no flip) python src/scripts/evaluate.py checkpoint_path=... \ validation_config.eval_scales=[1.0] validation_config.flip=false diff --git a/configs/AeroScapes_info.json b/configs/AeroScapes_info.json new file mode 100644 index 0000000..ae150fd --- /dev/null +++ b/configs/AeroScapes_info.json @@ -0,0 +1,122 @@ +[ + { + "hasInstances": false, + "category": "void", + "catid": 0, + "name": "Background", + "ignoreInEval": false, + "id": 0, + "color": [0, 0, 0], + "trainId": 0 + }, + { + "hasInstances": false, + "category": "human", + "catid": 1, + "name": "Person", + "ignoreInEval": false, + "id": 1, + "color": [192, 128, 128], + "trainId": 1 + }, + { + "hasInstances": false, + "category": "vehicle", + "catid": 2, + "name": "Bike", + "ignoreInEval": false, + "id": 2, + "color": [0, 128, 0], + "trainId": 2 + }, + { + "hasInstances": false, + "category": "vehicle", + "catid": 2, + "name": "Car", + "ignoreInEval": false, + "id": 3, + "color": [128, 128, 128], + "trainId": 3 + }, + { + "hasInstances": false, + "category": "vehicle", + "catid": 2, + "name": "Drone", + "ignoreInEval": false, + "id": 4, + "color": [128, 0, 0], + "trainId": 4 + }, + { + "hasInstances": false, + "category": "vehicle", + "catid": 2, + "name": "Boat", + "ignoreInEval": false, + "id": 5, + "color": [0, 0, 128], + "trainId": 5 + }, + { + "hasInstances": false, + "category": "animal", + "catid": 3, + "name": "Animal", + "ignoreInEval": false, + "id": 6, + "color": [192, 0, 128], + "trainId": 6 + }, + { + "hasInstances": false, + "category": "void", + "catid": 0, + "name": "Obstacle", + "ignoreInEval": false, + "id": 7, + "color": [192, 0, 0], + "trainId": 7 + }, + { + "hasInstances": false, + "category": "construction", + "catid": 4, + "name": "Construction", + "ignoreInEval": false, + "id": 8, + "color": [192, 128, 0], + "trainId": 8 + }, + { + "hasInstances": false, + "category": "nature", + "catid": 5, + "name": "Vegetation", + "ignoreInEval": false, + "id": 9, + "color": [0, 64, 0], + "trainId": 9 + }, + { + "hasInstances": false, + "category": "flat", + "catid": 6, + "name": "Road", + "ignoreInEval": false, + "id": 10, + "color": [128, 128, 0], + "trainId": 10 + }, + { + "hasInstances": false, + "category": "nature", + "catid": 5, + "name": "Sky", + "ignoreInEval": false, + "id": 11, + "color": [0, 128, 128], + "trainId": 11 + } +] diff --git a/configs/dataset/aeroscapes.yaml b/configs/dataset/aeroscapes.yaml new file mode 100644 index 0000000..28537fc --- /dev/null +++ b/configs/dataset/aeroscapes.yaml @@ -0,0 +1,48 @@ +# Dataset Configuration for CABiNet training on AeroScapes +# +# dataset_path must point at the PRE-CONVERTED dataset produced by +# src/scripts/convert_aeroscapes_to_yolo.py: +# +# / +# ├── images/{train,val}/ ← RGB JPEGs (uniformly 1280x720) +# └── masks/{train,val}/ ← single-channel PNGs, pixel value = class ID +# +# This is NOT the raw AeroScapes distribution (JPEGImages/, SegmentationClass/, +# ImageSets/) — convert it first: +# python src/scripts/convert_aeroscapes_to_yolo.py \ +# --src /path/to/raw/aeroscapes --dst /path/to/converted +# +# All 12 classes are valid trainIds 0-11 — see configs/AeroScapes_info.json. +# No colour palette or lookup table is needed at load time; the conversion +# step's source masks were already single-channel class IDs. +# +# Unlike UAVid, AeroScapes source images are uniformly 1280x720, so there is +# no batch_size=1 validation constraint here — standard batched validation +# works. +# +# NOTE: the source dataset has no test split (only ImageSets/trn.txt and +# val.txt) — the converter only ever produces train/val. + +name: aeroscapes +num_classes: 12 # Background(0) … Sky(11); all 12 classes valid, none ignored +cropsize: [720, 720] # native res is 1280x720 (uniform) — smaller than UAVid's 3840x2160 +dataset_path: ${oc.env:AEROSCAPES_YOLO_ROOT,''} # path to the CONVERTED dataset root +ignore_idx: 255 # label for genuinely unrecognized pixel values (should not occur) +seed: 0 + +# Class weights are computed at training startup from a sample of the +# training set via training_config.cls_pw (see configs/train.yaml) — same +# mechanism as uavid.yaml, not a static per-dataset table. + +# Training augmentation — same knobs/defaults as configs/dataset/uavid.yaml; +# AeroScapes is also UAV-captured imagery so the same recipe applies. +augmentation: + degrees: 10.0 # rotation ±10° (matches RandomRotate) + translate: 0.05 # small translation, fraction of image size + scale: 0.3 # scale jitter ±30% (continuous) + flipud: 0.2 # vertical flip — AeroScapes includes top-down and oblique views + fliplr: 0.5 # horizontal flip + hsv_h: 0.01 # small hue shift (controlled imaging conditions) + hsv_s: 0.4 # saturation jitter + hsv_v: 0.3 # brightness/value jitter + mixup: 0.1 # probability of blending with a second random sample diff --git a/configs/dataset/aeroscapes_yolo.yaml b/configs/dataset/aeroscapes_yolo.yaml new file mode 100644 index 0000000..6509c54 --- /dev/null +++ b/configs/dataset/aeroscapes_yolo.yaml @@ -0,0 +1,72 @@ +# AeroScapes dataset configuration for Ultralytics YOLO semantic segmentation +# +# Format: https://docs.ultralytics.com/datasets/semantic/ +# Masks: single-channel PNG, pixel value = class index (0-11). 255 remains +# reserved for genuinely unrecognized values. +# +# Pre-processing required +# ----------------------- +# Convert the raw AeroScapes distribution (JPEGImages/, SegmentationClass/, +# ImageSets/) to the shared images/+masks/ layout first: +# +# python src/scripts/convert_aeroscapes_to_yolo.py \ +# --src /path/to/aeroscapes \ +# --dst /path/to/aeroscapes_yolo +# +# Then set the ``path`` field below to /path/to/aeroscapes_yolo. +# +# Class mapping (all 12 classes are valid; none are ignored) +# ----------------------------------------------------------- +# 0 Background [ 0, 0, 0] +# 1 Person [192, 128, 128] +# 2 Bike [ 0, 128, 0] +# 3 Car [128, 128, 128] +# 4 Drone [128, 0, 0] +# 5 Boat [ 0, 0, 128] +# 6 Animal [192, 0, 128] +# 7 Obstacle [192, 0, 0] +# 8 Construction [192, 128, 0] +# 9 Vegetation [ 0, 64, 0] +# 10 Road [128, 128, 0] +# 11 Sky [ 0, 128, 128] + +# Root path. Ultralytics parses this file with plain YAML (no Hydra/OmegaConf +# env-var resolution), so this must be a literal path — edit it directly for +# your machine, or point `data=` at a copy of this file with your own path. +path: /home/neural_debugger/Downloads/datasets/AeroScapes/AeroScapes-YOLO + +train: images/train # relative to path +val: images/val # relative to path +# NOTE: no 'test:' key — the source AeroScapes distribution has no test +# split (only ImageSets/trn.txt and val.txt); the converter never produces one. + +# Mask directory (must mirror images/ structure): images/train -> masks/train, etc. +masks_dir: masks # relative to path + +# Number of semantic classes (all 12 AeroScapes classes) +nc: 12 + +# Class names — order must match the integer class IDs above +names: + 0: Background + 1: Person + 2: Bike + 3: Car + 4: Drone + 5: Boat + 6: Animal + 7: Obstacle + 8: Construction + 9: Vegetation + 10: Road + 11: Sky + +# Ultralytics training command example +# ------------------------------------- +# yolo semantic train \ +# model=yolo26n-sem.pt \ +# data=configs/dataset/aeroscapes_yolo.yaml \ +# imgsz=896 \ +# epochs=100 \ +# batch=8 \ +# device=0 diff --git a/configs/dataset/vdd.yaml b/configs/dataset/vdd.yaml new file mode 100644 index 0000000..e0d17e2 --- /dev/null +++ b/configs/dataset/vdd.yaml @@ -0,0 +1,47 @@ +# Dataset Configuration for CABiNet training on VDD (Varied Drone Dataset) +# +# dataset_path must point at the PRE-CONVERTED dataset produced by +# src/scripts/convert_vdd_to_yolo.py: +# +# / +# ├── images/{train,val,test}/ ← RGB JPEGs (uniformly 4000x3000) +# └── masks/{train,val,test}/ ← single-channel PNGs, pixel value = class ID +# +# This is NOT the raw VDD distribution (train/val/test with src/, gt/ +# sub-dirs) — convert it first: +# python src/scripts/convert_vdd_to_yolo.py \ +# --src /path/to/raw/VDD --dst /path/to/converted +# +# All 7 classes are valid trainIds 0-6 — see https://github.com/RussRobin/VDD. +# No colour palette or lookup table is needed at load time; the source masks +# were already single-channel class IDs. +# +# Like AeroScapes and unlike UAVid, VDD source images are uniformly +# 4000x3000, so there is no batch_size=1 validation constraint here — +# standard batched validation works. Unlike AeroScapes, VDD ships with a +# real train/val/test split already defined by the source distribution. + +name: vdd +num_classes: 7 # Other(0), Wall(1), Road(2), Vegetation(3), Vehicle(4), Roof(5), Water(6) +cropsize: [1024, 1024] # native res is 4000x3000 (uniform) — much larger than UAVid/AeroScapes +dataset_path: ${oc.env:VDD_YOLO_ROOT,''} # path to the CONVERTED dataset root +ignore_idx: 255 # label for genuinely unrecognized pixel values (should not occur) +seed: 0 + +# Class weights are computed at training startup from a sample of the +# training set via training_config.cls_pw (see configs/train.yaml) — same +# mechanism as uavid.yaml/aeroscapes.yaml, not a static per-dataset table. + +# Training augmentation — same knobs/defaults as configs/dataset/uavid.yaml +# and aeroscapes.yaml; VDD is also UAV-captured imagery so the same recipe +# applies. +augmentation: + degrees: 10.0 # rotation ±10° (matches RandomRotate) + translate: 0.05 # small translation, fraction of image size + scale: 0.3 # scale jitter ±30% (continuous) + flipud: 0.2 # vertical flip — VDD includes top-down and oblique views + fliplr: 0.5 # horizontal flip + hsv_h: 0.01 # small hue shift (controlled imaging conditions) + hsv_s: 0.4 # saturation jitter + hsv_v: 0.3 # brightness/value jitter + mixup: 0.1 # probability of blending with a second random sample diff --git a/configs/dataset/vdd_yolo.yaml b/configs/dataset/vdd_yolo.yaml new file mode 100644 index 0000000..abaed01 --- /dev/null +++ b/configs/dataset/vdd_yolo.yaml @@ -0,0 +1,61 @@ +# VDD (Varied Drone Dataset) configuration for Ultralytics YOLO semantic segmentation +# +# Format: https://docs.ultralytics.com/datasets/semantic/ +# Masks: single-channel PNG, pixel value = class index (0-6). 255 remains +# reserved for genuinely unrecognized values. +# +# Pre-processing required +# ----------------------- +# Convert the raw VDD distribution (train/val/test, each with src/ and gt/) +# to the shared images/+masks/ layout first: +# +# python src/scripts/convert_vdd_to_yolo.py \ +# --src /path/to/VDD \ +# --dst /path/to/vdd_yolo +# +# Then set the ``path`` field below to /path/to/vdd_yolo. +# +# Class mapping (all 7 classes are valid; none are ignored) +# ----------------------------------------------------------- +# 0 Other (background clutter) +# 1 Wall +# 2 Road +# 3 Vegetation +# 4 Vehicle +# 5 Roof +# 6 Water + +# Root path. Ultralytics parses this file with plain YAML (no Hydra/OmegaConf +# env-var resolution), so this must be a literal path — edit it directly for +# your machine, or point `data=` at a copy of this file with your own path. +path: /home/neural_debugger/Downloads/datasets/VDD/VDD-YOLO + +train: images/train # relative to path +val: images/val # relative to path +test: images/test # relative to path + +# Mask directory (must mirror images/ structure): images/train -> masks/train, etc. +masks_dir: masks # relative to path + +# Number of semantic classes +nc: 7 + +# Class names — order must match the integer class IDs above +names: + 0: Other + 1: Wall + 2: Road + 3: Vegetation + 4: Vehicle + 5: Roof + 6: Water + +# Ultralytics training command example +# ------------------------------------- +# yolo semantic train \ +# model=yolo26n-sem.pt \ +# data=configs/dataset/vdd_yolo.yaml \ +# imgsz=1024 \ +# epochs=120 \ +# batch=4 \ +# device=0 diff --git a/configs/train.yaml b/configs/train.yaml index 4804d47..1c52222 100644 --- a/configs/train.yaml +++ b/configs/train.yaml @@ -18,6 +18,17 @@ training_config: optimizer_power: 0.9 warmup_steps: 4000 warmup_start_lr: 1e-5 + # Warm-start model weights from a full CABiNet checkpoint trained on a + # (possibly different) dataset — e.g. finetune AeroScapes/VDD from a + # UAVid-trained checkpoint, which converges much faster than + # backbone-only (ImageNet) init. Only parameters whose name AND shape + # match are loaded — the two classifier heads (sized by num_classes) are + # skipped when they differ, everything else (backbone/CAB/spatial/fusion) + # transfers. Independent of `resume` below: this never restores + # optimizer/EMA/epoch state, it's a fresh training run with warm-started + # weights. Accepts either a full training checkpoint (checkpoint_last.pth) + # or a raw model state_dict (*_best.pth). + pretrained_ckpt_path: null # max_iterations is in OPTIMIZER STEPS (not batches). # train.py computes this as ceil(epochs * batches_per_epoch / accum_steps). # Setting it here is optional — if absent, train.py derives it automatically. @@ -34,7 +45,7 @@ training_config: max_grad_norm: 1.0 # Gradient clipping threshold; set to null/0 to disable log_iter: 20 # Renamed from 'msg_iterations' — clearer intent - experiments_path: experiments/cabinet_experiments_large_1024x1024 + experiments_path: experiments/${dataset.name}/cabinet_experiments_large_${dataset.cropsize[0]}x${dataset.cropsize[1]} # Output path for checkpoints, logs, etc. model_save_name: cabinet.pth resume: false # Set to true to resume from checkpoint_last.pth in experiments_path diff --git a/configs/train_yolo_aeroscapes.yaml b/configs/train_yolo_aeroscapes.yaml new file mode 100644 index 0000000..c8620a2 --- /dev/null +++ b/configs/train_yolo_aeroscapes.yaml @@ -0,0 +1,113 @@ +# train_yolo_aeroscapes.yaml +# Hydra root config for YOLO semantic segmentation on AeroScapes +# +# Uses Ultralytics' dedicated 'semantic' task (dense per-pixel FCN-style head, +# CE+Dice+aux loss, mIoU/pixel-acc metric) — distinct from 'segment' (instance/ +# polygon-based). Only the yolo26*-sem model family supports it. +# +# Usage (train): +# python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes +# python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes 'yolo/model@model=yolo26s-sem' +# +# Override individual hyperparameters (Hydra dotted-path style): +# python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes training_config.epochs=200 training_config.batch_size=8 +# +# Usage (validate best checkpoint): +# python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes mode=val \ +# validation_config.weights=runs/aeroscapes_yolo/yolo26n/weights/best.pt +# +# Resume interrupted training: +# python src/scripts/train_yolo.py --config-name train_yolo_aeroscapes training_config.resume=true + +defaults: + - _self_ + - yolo/model@model: yolo26n-sem # override: 'yolo/model@model=yolo26x-sem' + +# Mode: train | val | test | predict +mode: train + +# Path to the Ultralytics-format dataset YAML (not Hydra-format) +dataset: + name: aeroscapes_yolo + config_file: configs/dataset/aeroscapes_yolo.yaml + +# ── Training configuration ──────────────────────────────────────────────────── +training_config: + epochs: 120 + batch_size: 8 + imgsz: 896 # native res is 1280x720 (uniform) — smaller than UAVid, so a + # smaller imgsz than UAVid's 1024 is appropriate here + num_workers: 8 + + nbs: 64 + + amp: true + + # Optimizer + optimizer: SGD + optimizer_lr_start: 0.01 + lrf: 0.01 + optimizer_momentum: 0.937 + optimizer_weight_decay: 0.0005 + warmup_epochs: 3.0 + warmup_momentum: 0.8 + warmup_bias_lr: 0.1 + cos_lr: true + + patience: 50 + + save_period: 50 + resume: false + + # Class imbalance (AeroScapes: Person/Bike/Drone/Boat/Animal are rare) + cls_pw: 0.5 + + experiments_path: experiments/${dataset.name} + exist_ok: false + + # ── Augmentation ─────────────────────────────────────────────────────────── + augmentation: + degrees: 10.0 + translate: 0.05 + scale: 0.3 + shear: 0.0 + perspective: 0.0 + flipud: 0.2 # AeroScapes includes top-down and oblique views + fliplr: 0.5 + + hsv_h: 0.01 + hsv_s: 0.4 + hsv_v: 0.3 + + mosaic: 0.8 + mixup: 0.1 + copy_paste: 0.15 + copy_paste_mode: flip + + close_mosaic: 15 + + multi_scale: 0.0 + +# ── Validation configuration ───────────────────────────────────────────────── +validation_config: + # Unlike UAVid (batch_size=1 required — mixed source resolution), + # AeroScapes is uniformly 1280x720, so a batched eval is safe here. + batch_size: 4 + save_json: true + augment: false + weights: ~ + split: val # no test split exists for this dataset + +# ── Runtime ────────────────────────────────────────────────────────────────── +runtime: + seed: 0 + device: 0 + deterministic: true + +# ── Hydra output ───────────────────────────────────────────────────────────── +hydra: + run: + dir: ${training_config.experiments_path}/hydra/${now:%Y-%m-%d_%H-%M-%S} + sweep: + dir: ${training_config.experiments_path}/sweeps + subdir: ${hydra.job.num} diff --git a/configs/train_yolo_vdd.yaml b/configs/train_yolo_vdd.yaml new file mode 100644 index 0000000..c63db38 --- /dev/null +++ b/configs/train_yolo_vdd.yaml @@ -0,0 +1,114 @@ +# train_yolo_vdd.yaml +# Hydra root config for YOLO semantic segmentation on VDD (Varied Drone Dataset) +# +# Uses Ultralytics' dedicated 'semantic' task (dense per-pixel FCN-style head, +# CE+Dice+aux loss, mIoU/pixel-acc metric) — distinct from 'segment' (instance/ +# polygon-based). Only the yolo26*-sem model family supports it. +# +# Usage (train): +# python src/scripts/train_yolo.py --config-name train_yolo_vdd +# python src/scripts/train_yolo.py --config-name train_yolo_vdd 'yolo/model@model=yolo26s-sem' +# +# Override individual hyperparameters (Hydra dotted-path style): +# python src/scripts/train_yolo.py --config-name train_yolo_vdd training_config.epochs=200 training_config.batch_size=8 +# +# Usage (validate best checkpoint): +# python src/scripts/train_yolo.py --config-name train_yolo_vdd mode=val \ +# validation_config.weights=runs/vdd_yolo/yolo26n/weights/best.pt +# +# Resume interrupted training: +# python src/scripts/train_yolo.py --config-name train_yolo_vdd training_config.resume=true + +defaults: + - _self_ + - yolo/model@model: yolo26n-sem # override: 'yolo/model@model=yolo26x-sem' + +# Mode: train | val | test | predict +mode: train + +# Path to the Ultralytics-format dataset YAML (not Hydra-format) +dataset: + name: vdd_yolo + config_file: configs/dataset/vdd_yolo.yaml + +# ── Training configuration ──────────────────────────────────────────────────── +training_config: + epochs: 150 # VDD is tiny (280 train images) → more epochs help + batch_size: 4 + imgsz: 1024 # native res is 4000x3000 (uniform) — a genuine downscale + num_workers: 8 + + nbs: 64 + + amp: true + + # Optimizer + optimizer: SGD + optimizer_lr_start: 0.01 + lrf: 0.01 + optimizer_momentum: 0.937 + optimizer_weight_decay: 0.0005 + warmup_epochs: 3.0 + warmup_momentum: 0.8 + warmup_bias_lr: 0.1 + cos_lr: true + + patience: 50 + + save_period: 50 + resume: false + + # Class imbalance (VDD: Vehicle/Roof/Water are comparatively rare) + cls_pw: 0.5 + + experiments_path: experiments/${dataset.name} + exist_ok: false + + # ── Augmentation ─────────────────────────────────────────────────────────── + augmentation: + degrees: 10.0 + translate: 0.05 + scale: 0.3 + shear: 0.0 + perspective: 0.0 + flipud: 0.2 # VDD includes top-down and oblique views + fliplr: 0.5 + + hsv_h: 0.01 + hsv_s: 0.4 + hsv_v: 0.3 + + # VDD has very few training images (280); heavier augmentation helps + # offset the small dataset size + mosaic: 0.8 + mixup: 0.15 + copy_paste: 0.15 + copy_paste_mode: flip + + close_mosaic: 15 + + multi_scale: 0.0 + +# ── Validation configuration ───────────────────────────────────────────────── +validation_config: + # Unlike UAVid (batch_size=1 required — mixed source resolution), VDD is + # uniformly 4000x3000, so a batched eval is safe here. + batch_size: 4 + save_json: true + augment: false + weights: ~ + split: val # 'val' or 'test' — override: validation_config.split=test + +# ── Runtime ────────────────────────────────────────────────────────────────── +runtime: + seed: 0 + device: 0 + deterministic: true + +# ── Hydra output ───────────────────────────────────────────────────────────── +hydra: + run: + dir: ${training_config.experiments_path}/hydra/${now:%Y-%m-%d_%H-%M-%S} + sweep: + dir: ${training_config.experiments_path}/sweeps + subdir: ${hydra.job.num} diff --git a/configs/yolo/aeroscapes_train.yaml b/configs/yolo/aeroscapes_train.yaml new file mode 100644 index 0000000..a09ba90 --- /dev/null +++ b/configs/yolo/aeroscapes_train.yaml @@ -0,0 +1,116 @@ +# ============================================================================= +# Ultralytics YOLO — AeroScapes Semantic Segmentation Training Config +# ============================================================================= +# Usage: +# yolo semantic train cfg=configs/yolo/aeroscapes_train.yaml +# +# Resume an interrupted run: +# yolo semantic train cfg=configs/yolo/aeroscapes_train.yaml resume=True +# +# Multi-GPU (DDP): +# yolo semantic train cfg=configs/yolo/aeroscapes_train.yaml device=0,1,2,3 +# +# Swap model variant (nano / small / medium / large / xlarge): +# yolo semantic train cfg=configs/yolo/aeroscapes_train.yaml model=yolo26s-sem.pt name=yolo26s +# +# NOTE: 'semantic' is a dedicated Ultralytics task (dense per-pixel FCN-style +# head, CE+Dice+aux loss, mIoU/pixel-acc metric) — distinct from 'segment' +# (instance/polygon-based detection+mask head). Only the yolo26*-sem model +# family supports it; yolov8/9/11-seg checkpoints do NOT apply here. +# ============================================================================= + +# ── Task & Model ───────────────────────────────────────────────────────────── +task: semantic +model: yolo26n-sem.pt # pretrained checkpoint (downloads automatically) +data: configs/dataset/aeroscapes_yolo.yaml + +# ── Output ─────────────────────────────────────────────────────────────────── +project: runs/aeroscapes # results root +name: yolo26n # run sub-dir; change per variant (yolo26s, yolo26m, …) +exist_ok: False # set True to overwrite a previous run with the same name + +# ── Core training schedule ─────────────────────────────────────────────────── +epochs: 120 # total epochs (AeroScapes is small → more epochs help) +patience: 30 # early stop after N epochs without val mIoU improvement +imgsz: 896 # native res is 1280x720 (uniform) — smaller than UAVid, so a +# smaller imgsz than UAVid's 1024 is appropriate here +batch: 8 # per-GPU batch size — resolution is uniform, so unlike UAVid this +# can safely exceed 1 during validation too (no mixed-resolution stacking issue) +seed: 0 # reproducibility + +# ── Gradient accumulation ───────────────────────────────────────────────────── +# YOLO implements accumulation via the nominal batch size (nbs). +# Effective accumulation steps = ceil(nbs / (batch * num_GPUs)). +nbs: 64 + +# ── Mixed precision (AMP) ──────────────────────────────────────────────────── +# NOTE: EMA (Exponential Moving Average) is ALWAYS applied by default. +# 'best.pt' and 'last.pt' both contain the EMA-averaged weights. +amp: True + +# ── Checkpointing ──────────────────────────────────────────────────────────── +save: True # always save best.pt and last.pt +save_period: + 10 # additionally save checkpoint_N.pt every N epochs + # protects against crashes; set -1 to save best/last only + +# ── Resume ─────────────────────────────────────────────────────────────────── +# Set resume=True (or pass on CLI) to continue from last.pt in the same run dir. +# The run directory is: {project}/{name}/weights/last.pt +resume: False # override with resume=True on CLI to resume + +# ── Optimizer ──────────────────────────────────────────────────────────────── +optimizer: SGD # SGD with momentum (matches CABiNet; AdamW also works) +lr0: 0.01 # initial learning rate +lrf: 0.01 # final LR = lr0 * lrf (cosine/linear decay floor) +momentum: 0.937 # SGD momentum +weight_decay: 0.0005 # L2 regularisation (matches CABiNet weight decay) +warmup_epochs: 3.0 # linear LR warmup; helps stability with large imgsz +warmup_momentum: 0.8 # momentum during warmup +warmup_bias_lr: 0.1 # bias LR during warmup (higher than main to stabilise BN) +cos_lr: True # cosine annealing (smoother than step decay for seg) + +# ── Class-imbalance handling ───────────────────────────────────────────────── +# cls_pw applies inverse-frequency class weighting to the classification loss. +# 0.0 = disabled (uniform), 1.0 = full inverse-frequency. +# AeroScapes is heavily skewed toward Background/Vegetation/Sky — Person, +# Bike, Drone, Boat, and Animal are rare; 0.5 is a good start (matches UAVid). +cls_pw: 0.5 + +# ── Augmentation ───────────────────────────────────────────────────────────── +# Geometric +degrees: 10.0 # rotation ±10° (matches CABiNet RandomRotate) +translate: 0.05 # small translation (aerial scenes are large) +scale: 0.3 # scale jitter ±30% +shear: 0.0 # shear disabled (aerial imagery is mostly ortho) +perspective: 0.0 # perspective disabled +flipud: 0.2 # vertical flip (AeroScapes includes top-down and oblique views) +fliplr: 0.5 # horizontal flip + +# Colour / photometric +hsv_h: 0.01 # small hue shift +hsv_s: 0.4 # saturation jitter +hsv_v: 0.3 # value (brightness) jitter + +# Mosaic / mixup — useful for small objects (people, bikes, drones) in aerial view +mosaic: 0.8 # mosaic probability (set 0.0 to disable) +mixup: 0.1 # MixUp blending probability +copy_paste: 0.15 # copy-paste small objects (person, bike, animal) +copy_paste_mode: flip # copy-paste strategy + +# Close mosaic N epochs before end to stabilise fine predictions +close_mosaic: 15 + +# ── Validation ─────────────────────────────────────────────────────────────── +val: True # run validation after every epoch +split: val # validation split name (no test split available for this dataset) + +# ── Compute / logging ──────────────────────────────────────────────────────── +device: 0 # GPU id(s); use 'cpu' for CPU-only, '0,1' for multi-GPU +workers: 8 # dataloader workers per GPU +plots: True # save training curves and confusion matrix +verbose: True +deterministic: True # fully reproducible (may be slower) + +# ── Multi-scale training (optional) ────────────────────────────────────────── +multi_scale: 0.0 # set 0.2 to enable diff --git a/configs/yolo/aeroscapes_val.yaml b/configs/yolo/aeroscapes_val.yaml new file mode 100644 index 0000000..d151080 --- /dev/null +++ b/configs/yolo/aeroscapes_val.yaml @@ -0,0 +1,31 @@ +# ============================================================================= +# Ultralytics YOLO — AeroScapes Validation / Benchmark Config +# ============================================================================= +# Usage (after training): +# Single-scale: +# yolo semantic val cfg=configs/yolo/aeroscapes_val.yaml \ +# model=runs/aeroscapes/yolo26n/weights/best.pt +# +# Test-time augmentation (TTA): +# yolo semantic val cfg=configs/yolo/aeroscapes_val.yaml \ +# model=runs/aeroscapes/yolo26n/weights/best.pt \ +# augment=True +# ============================================================================= + +task: semantic +data: configs/dataset/aeroscapes_yolo.yaml + +imgsz: 896 +# Unlike UAVid (batch=1 required — mixed source resolution), AeroScapes is +# uniformly 1280x720, so a batched eval is safe here. +batch: 8 +split: val # no test split available for this dataset +device: 0 +workers: 4 + +save_json: True # save per-image results for offline analysis +plots: True # confusion matrix + class-level IoU bar chart +verbose: True + +# Test-time augmentation (TTA) — enable with: augment=True on CLI +augment: False diff --git a/configs/yolo/vdd_train.yaml b/configs/yolo/vdd_train.yaml new file mode 100644 index 0000000..ba7adad --- /dev/null +++ b/configs/yolo/vdd_train.yaml @@ -0,0 +1,110 @@ +# ============================================================================= +# Ultralytics YOLO — VDD (Varied Drone Dataset) Semantic Segmentation Training Config +# ============================================================================= +# Usage: +# yolo semantic train cfg=configs/yolo/vdd_train.yaml +# +# Resume an interrupted run: +# yolo semantic train cfg=configs/yolo/vdd_train.yaml resume=True +# +# Multi-GPU (DDP): +# yolo semantic train cfg=configs/yolo/vdd_train.yaml device=0,1,2,3 +# +# Swap model variant (nano / small / medium / large / xlarge): +# yolo semantic train cfg=configs/yolo/vdd_train.yaml model=yolo26s-sem.pt name=yolo26s +# +# NOTE: 'semantic' is a dedicated Ultralytics task (dense per-pixel FCN-style +# head, CE+Dice+aux loss, mIoU/pixel-acc metric) — distinct from 'segment' +# (instance/polygon-based detection+mask head). Only the yolo26*-sem model +# family supports it; yolov8/9/11-seg checkpoints do NOT apply here. +# ============================================================================= + +# ── Task & Model ───────────────────────────────────────────────────────────── +task: semantic +model: yolo26n-sem.pt # pretrained checkpoint (downloads automatically) +data: configs/dataset/vdd_yolo.yaml + +# ── Output ─────────────────────────────────────────────────────────────────── +project: runs/vdd # results root +name: yolo26n # run sub-dir; change per variant (yolo26s, yolo26m, …) +exist_ok: False # set True to overwrite a previous run with the same name + +# ── Core training schedule ─────────────────────────────────────────────────── +epochs: 150 # VDD is tiny (280 train images) → more epochs help +patience: 30 # early stop after N epochs without val mIoU improvement +imgsz: 1024 # native res is 4000x3000 (uniform) — far larger than the model +# input, so imgsz is a genuine downscale here (unlike UAVid/AeroScapes, +# which are already closer to their imgsz) +batch: 4 # per-GPU batch size (adjust to your VRAM — native-res VDD images are large) +seed: 0 # reproducibility + +# ── Gradient accumulation ───────────────────────────────────────────────────── +nbs: 64 + +# ── Mixed precision (AMP) ──────────────────────────────────────────────────── +# NOTE: EMA (Exponential Moving Average) is ALWAYS applied by default. +# 'best.pt' and 'last.pt' both contain the EMA-averaged weights. +amp: True + +# ── Checkpointing ──────────────────────────────────────────────────────────── +save: True +save_period: 10 # additionally save checkpoint_N.pt every N epochs + +# ── Resume ─────────────────────────────────────────────────────────────────── +resume: False # override with resume=True on CLI to resume + +# ── Optimizer ──────────────────────────────────────────────────────────────── +optimizer: SGD +lr0: 0.01 +lrf: 0.01 +momentum: 0.937 +weight_decay: 0.0005 +warmup_epochs: 3.0 +warmup_momentum: 0.8 +warmup_bias_lr: 0.1 +cos_lr: True + +# ── Class-imbalance handling ───────────────────────────────────────────────── +# cls_pw applies inverse-frequency class weighting to the classification loss. +# 0.0 = disabled (uniform), 1.0 = full inverse-frequency. +# VDD is skewed toward Vegetation/Road/Wall — Vehicle, Roof, and Water are +# comparatively rare; 0.5 is a good start (matches UAVid/AeroScapes). +cls_pw: 0.5 + +# ── Augmentation ───────────────────────────────────────────────────────────── +# Geometric +degrees: 10.0 +translate: 0.05 +scale: 0.3 +shear: 0.0 +perspective: 0.0 +flipud: 0.2 # vertical flip (VDD includes top-down and oblique views) +fliplr: 0.5 + +# Colour / photometric +hsv_h: 0.01 +hsv_s: 0.4 +hsv_v: 0.3 + +# Mosaic / mixup — VDD has very few training images (280); heavier +# augmentation helps offset the small dataset size +mosaic: 0.8 +mixup: 0.15 +copy_paste: 0.15 +copy_paste_mode: flip + +close_mosaic: 15 + +# ── Validation ─────────────────────────────────────────────────────────────── +val: True +split: val + +# ── Compute / logging ──────────────────────────────────────────────────────── +device: 0 +workers: 8 +plots: True +verbose: True +deterministic: True + +# ── Multi-scale training (optional) ────────────────────────────────────────── +multi_scale: 0.0 diff --git a/configs/yolo/vdd_val.yaml b/configs/yolo/vdd_val.yaml new file mode 100644 index 0000000..58c26f9 --- /dev/null +++ b/configs/yolo/vdd_val.yaml @@ -0,0 +1,31 @@ +# ============================================================================= +# Ultralytics YOLO — VDD (Varied Drone Dataset) Validation / Benchmark Config +# ============================================================================= +# Usage (after training): +# Single-scale: +# yolo semantic val cfg=configs/yolo/vdd_val.yaml \ +# model=runs/vdd/yolo26n/weights/best.pt +# +# Test-time augmentation (TTA): +# yolo semantic val cfg=configs/yolo/vdd_val.yaml \ +# model=runs/vdd/yolo26n/weights/best.pt \ +# augment=True +# ============================================================================= + +task: semantic +data: configs/dataset/vdd_yolo.yaml + +imgsz: 1024 +# Unlike UAVid (batch=1 required — mixed source resolution), VDD is +# uniformly 4000x3000, so a batched eval is safe here. +batch: 4 +split: val +device: 0 +workers: 4 + +save_json: True # save per-image results for offline analysis +plots: True # confusion matrix + class-level IoU bar chart +verbose: True + +# Test-time augmentation (TTA) — enable with: augment=True on CLI +augment: False diff --git a/src/datasets/aeroscapes.py b/src/datasets/aeroscapes.py new file mode 100644 index 0000000..81c9abf --- /dev/null +++ b/src/datasets/aeroscapes.py @@ -0,0 +1,282 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import os +import os.path as osp +from typing import Any, Dict, Optional, Tuple +import warnings + +from PIL import Image +import numpy as np +import torch +from torch.utils.data import Dataset +import torchvision.transforms as transforms + +from src.datasets.transform import ( + Compose, + RandomColorJitter, + RandomCrop, + RandomCutout, + RandomGamma, + RandomHorizontalFlip, + RandomHSV, + RandomNoise, + RandomRotate, + RandomScale, + RandomTranslate, + RandomVerticalFlip, +) + +# Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION — +# AeroScapes is also UAV-captured imagery, so the same augmentation recipe +# applies unchanged. See uavid.py's module docstring for the rationale +# behind each knob (mosaic/copy_paste unsupported, mixup's hard-label +# simplification, etc.) — identical here. +DEFAULT_AUGMENTATION: Dict[str, float] = { + "degrees": 10.0, + "translate": 0.05, + "scale": 0.3, + "flipud": 0.2, + "fliplr": 0.5, + "hsv_h": 0.01, + "hsv_s": 0.4, + "hsv_v": 0.3, + "mixup": 0.1, +} + + +class AeroScapes(Dataset): + """AeroScapes aerial semantic segmentation dataset. + + Consumes the pre-converted, YOLO-style dataset layout produced by + ``src/scripts/convert_aeroscapes_to_yolo.py``:: + + / + ├── images/ + │ ├── train/ ← RGB JPEGs (uniformly 1280x720) + │ └── val/ + └── masks/ + ├── train/ ← single-channel PNGs, pixel value = class ID + └── val/ + + Mask pixel values are already final trainIds (0-11; Background=0 … + Sky=11 — see configs/AeroScapes_info.json) — no RGB colour palette or + lookup table is needed here, unlike UAVid's raw distribution. Pixel + value 255 is reserved for genuinely unrecognized values, not a real + class. There is no source test split (see the conversion script's + docstring) — only ``"train"``/``"val"`` are populated by the converter, + though ``mode="test"`` is still accepted here for API parity with + :class:`~src.datasets.uavid.UAVid` / :class:`~src.datasets.cityscapes.CityScapes`. + + Parameters + ---------- + ignore_lb: + Label value treated as "ignore" by the loss/metric (255). + rootpth: + Root of the *converted* dataset (i.e. ``convert_aeroscapes_to_yolo.py``'s + ``--dst``) — NOT the raw AeroScapes distribution. + cropsize: + ``(H, W)`` crop applied during training augmentation via + ``RandomCrop``. AeroScapes source images are uniformly 1280x720 + (unlike UAVid's mixed resolutions), so ``RandomCrop`` mainly serves + as a fixed-size training patch here rather than a resolution + normalizer — ``pad_if_needed=True`` is kept for robustness. + mode: + ``"train"``, ``"val"``, or ``"test"``. + augmentation: + Optional dict overriding any subset of ``DEFAULT_AUGMENTATION`` + (degrees/translate/scale/flipud/fliplr/hsv_h/hsv_s/hsv_v/mixup). + + Note on validation batching + ---------------------------- + Unlike UAVid, AeroScapes source images are uniformly 1280x720, so + ``val``/``test`` mode (no crop applied) can still be batched with + ``batch_size > 1`` via the default collate function — no special-casing + is needed in train.py/evaluate.py for this dataset. + """ + + def __init__( + self, + ignore_lb: int, + rootpth: str, + cropsize: Tuple[int, int], + mode: str = "train", + augmentation: Optional[Dict[str, Any]] = None, + ) -> None: + super().__init__() + self.mode = mode + self.ignore_lb = ignore_lb + self.rootpth = rootpth + self.cropsize = tuple(cropsize) + self.aug = {**DEFAULT_AUGMENTATION, **(augmentation or {})} + + if self.mode not in ("train", "val", "test"): + raise ValueError( + f"Mode '{mode}' not supported. Choose 'train', 'val', or 'test'." + ) + if not osp.exists(rootpth): + raise FileNotFoundError(f"Dataset root does not exist: {rootpth}") + + img_dir = osp.join(rootpth, "images", mode) + label_dir = osp.join(rootpth, "masks", mode) + if not osp.exists(img_dir): + raise FileNotFoundError(f"Image directory not found: {img_dir}") + if not osp.exists(label_dir): + raise FileNotFoundError(f"Mask directory not found: {label_dir}") + + # --- Load image and label paths ------------------------------------ + self.imnames = [] + self.imgs = {} + self.labels = {} + + imgnames = sorted(fn for fn in os.listdir(img_dir) if fn.endswith(".jpg")) + for fn in imgnames: + stem = osp.splitext(fn)[0] + self.imgs[stem] = osp.join(img_dir, fn) + label_path = osp.join(label_dir, f"{stem}.png") + if osp.exists(label_path): + self.labels[stem] = label_path + self.imnames.append(stem) + + # Drop any image that has no matching mask + missing = [name for name in self.imnames if name not in self.labels] + if missing: + warnings.warn( + f"{len(missing)} image(s) have no matching mask in {label_dir} " + f"and will be skipped: {sorted(missing)[:5]}…" + ) + self.imnames = [name for name in self.imnames if name in self.labels] + + if len(self.imnames) == 0: + raise RuntimeError( + f"No valid image-mask pairs found for mode='{mode}' in {rootpth}." + ) + + self.len = len(self.imnames) + + # --- Image normalisation -------------------------------------------------- + # Mean/std computed from the AeroScapes train set (see + # src/datasets/compute_aeroscapes_stats.py) — distinct from ImageNet + # and UAVid stats. + self.to_tensor = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize( + mean=(0.439, 0.508, 0.460), + std=(0.176, 0.157, 0.194), + ), + ] + ) + + # --- Training augmentation (Geometric → Photometric → Regularisation) --- + degrees = float(self.aug["degrees"]) + scale = float(self.aug["scale"]) + self.trans_train = ( + Compose( + [ + RandomHorizontalFlip(p=float(self.aug["fliplr"])), + RandomVerticalFlip(p=float(self.aug["flipud"])), + RandomTranslate( + translate=float(self.aug["translate"]), + ignore_label=self.ignore_lb, + ), + RandomRotate( + degrees=(-degrees, degrees), ignore_label=self.ignore_lb + ), + RandomScale((1.0 - scale, 1.0 + scale), continuous=True), + RandomCrop( + size=self.cropsize, + pad_if_needed=True, + ignore_label=self.ignore_lb, + ), + RandomHSV( + hgain=float(self.aug["hsv_h"]), + sgain=float(self.aug["hsv_s"]), + vgain=float(self.aug["hsv_v"]), + ), + RandomColorJitter(contrast=0.5), + RandomGamma(gamma_range=(0.8, 1.2), p=0.3), + RandomNoise(mode="gaussian", sigma=0.03, p=0.3), + RandomCutout(p=0.3, size=64), + ] + ) + if mode == "train" + else None + ) + self.mixup_p = float(self.aug["mixup"]) if mode == "train" else 0.0 + + print( + f"[INFO] AeroScapes dataset loaded: {self.len} samples ({mode}) from {img_dir}" + ) + + def _load_one(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Load, augment, and tensorize a single sample (no MixUp).""" + stem = self.imnames[idx] + img = Image.open(self.imgs[stem]).convert("RGB") + label = Image.open(self.labels[stem]) + if label.mode != "L": + label = label.convert("L") + + if self.mode == "train" and self.trans_train is not None: + im_lb = self.trans_train({"im": img, "lb": label}) + img, label = im_lb["im"], im_lb["lb"] + + img_t = self.to_tensor(img) + label_np = np.array(label, dtype=np.int64) # (H, W), already final class IDs + label_t = torch.from_numpy(label_np).long() + return img_t, label_t + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + img, label = self._load_one(idx) + + if ( + self.mode == "train" + and self.mixup_p > 0 + and np.random.random() < self.mixup_p + ): + other_idx = int(np.random.randint(0, self.len)) + img2, label2 = self._load_one(other_idx) + # Both samples went through the same RandomCrop, so shapes match. + r = float(np.random.beta(32.0, 32.0)) # matches Ultralytics' MixUp exactly + img = img * r + img2 * (1.0 - r) + # Class-ID masks can't be blended the same way (no meaningful + # "average" of two class indices) — take the label from + # whichever image contributed the larger share of the blend. + label = label if r >= 0.5 else label2 + + return img, label + + def __len__(self) -> int: + return self.len + + +# --------------------------------------------------------------------------- +# Smoke test (run: python src/datasets/aeroscapes.py) +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import sys + + rootpth = os.environ.get("AEROSCAPES_YOLO_ROOT", "") + if not rootpth: + print( + "Set AEROSCAPES_YOLO_ROOT=/path/to/converted/aeroscapes_yolo and re-run " + "(this must be the OUTPUT of convert_aeroscapes_to_yolo.py, not raw AeroScapes data)." + ) + sys.exit(0) + + for split in ("train", "val", "test"): + try: + ds = AeroScapes( + ignore_lb=255, + rootpth=rootpth, + cropsize=(640, 640), + mode=split, + ) + except (FileNotFoundError, RuntimeError) as exc: + print(f"{split}: skipped ({exc})") + continue + print(f"{split}: {len(ds)} samples") + img, lb = ds[0] + unique = torch.unique(lb[lb != 255]) + print(f" img shape: {tuple(img.shape)}, label shape: {tuple(lb.shape)}") + print(f" unique class IDs in first item: {sorted(unique.tolist())}") diff --git a/src/datasets/compute_aeroscapes_stats.py b/src/datasets/compute_aeroscapes_stats.py new file mode 100644 index 0000000..586503a --- /dev/null +++ b/src/datasets/compute_aeroscapes_stats.py @@ -0,0 +1,30 @@ +import torch +from torchvision import transforms +from PIL import Image +from pathlib import Path +from tqdm import tqdm # pip install tqdm for a progress bar + +root = Path( + "/home/neural_debugger/Downloads/datasets/AeroScapes/AeroScapes-YOLO/images/train" +) # <-- change to your converted AeroScapes train images folder +to_tensor = transforms.ToTensor() # converts to [0,1] float tensor + +mean = torch.zeros(3) +std = torch.zeros(3) +num_pixels = 0 + +image_files = list(root.glob("*.jpg")) + +for img_path in tqdm(image_files): + img = Image.open(img_path).convert("RGB") + tensor = to_tensor(img) # [C,H,W], values in [0,1] + num_pixels += tensor.numel() / 3 # total pixels per channel + mean += tensor.sum(dim=[1, 2]) + std += (tensor**2).sum(dim=[1, 2]) + +mean /= num_pixels +std = (std / num_pixels - mean**2).sqrt() + +print("\nAeroScapes Train Statistics") +print("Mean :", mean.tolist()) +print("Std :", std.tolist()) diff --git a/src/datasets/compute_vdd_stats.py b/src/datasets/compute_vdd_stats.py new file mode 100644 index 0000000..48306ce --- /dev/null +++ b/src/datasets/compute_vdd_stats.py @@ -0,0 +1,30 @@ +import torch +from torchvision import transforms +from PIL import Image +from pathlib import Path +from tqdm import tqdm # pip install tqdm for a progress bar + +root = Path( + "/home/neural_debugger/Downloads/datasets/VDD/VDD-YOLO/images/train" +) # <-- change to your converted VDD train images folder +to_tensor = transforms.ToTensor() # converts to [0,1] float tensor + +mean = torch.zeros(3) +std = torch.zeros(3) +num_pixels = 0 + +image_files = list(root.glob("*.jpg")) + +for img_path in tqdm(image_files): + img = Image.open(img_path).convert("RGB") + tensor = to_tensor(img) # [C,H,W], values in [0,1] + num_pixels += tensor.numel() / 3 # total pixels per channel + mean += tensor.sum(dim=[1, 2]) + std += (tensor**2).sum(dim=[1, 2]) + +mean /= num_pixels +std = (std / num_pixels - mean**2).sqrt() + +print("\nVDD Train Statistics") +print("Mean :", mean.tolist()) +print("Std :", std.tolist()) diff --git a/src/datasets/registry.py b/src/datasets/registry.py index d810082..ef1ab49 100644 --- a/src/datasets/registry.py +++ b/src/datasets/registry.py @@ -5,18 +5,23 @@ imports MscEvalV0 from evaluate.py. """ +from src.datasets.aeroscapes import AeroScapes from src.datasets.cityscapes import CityScapes from src.datasets.uavid import UAVid +from src.datasets.vdd import VDD DATASET_REGISTRY = { "cityscapes": CityScapes, "uavid": UAVid, + "aeroscapes": AeroScapes, + "vdd": VDD, } -# Per-dataset constructor kwargs. UAVid consumes a pre-converted images/+masks/ -# directory (see convert_uavid_to_yolo.py) and no longer needs a colour-palette -# config_file, unlike CityScapes — the two datasets have genuinely different -# constructors, so they can't share one flat kwargs dict. +# Per-dataset constructor kwargs. UAVid/AeroScapes/VDD consume a +# pre-converted images/+masks/ directory (see convert_uavid_to_yolo.py / +# convert_aeroscapes_to_yolo.py / convert_vdd_to_yolo.py) and no longer need +# a colour-palette config_file, unlike CityScapes — the two datasets have +# genuinely different constructors, so they can't share one flat kwargs dict. DATASET_KWARGS_BUILDERS = { "cityscapes": lambda cfg, ignore_idx, cropsize: dict( config_file=cfg.dataset.config_file, @@ -30,4 +35,16 @@ cropsize=cropsize, augmentation=cfg.dataset.get("augmentation"), ), + "aeroscapes": lambda cfg, ignore_idx, cropsize: dict( + ignore_lb=ignore_idx, + rootpth=cfg.dataset.dataset_path, + cropsize=cropsize, + augmentation=cfg.dataset.get("augmentation"), + ), + "vdd": lambda cfg, ignore_idx, cropsize: dict( + ignore_lb=ignore_idx, + rootpth=cfg.dataset.dataset_path, + cropsize=cropsize, + augmentation=cfg.dataset.get("augmentation"), + ), } diff --git a/src/datasets/vdd.py b/src/datasets/vdd.py new file mode 100644 index 0000000..be66558 --- /dev/null +++ b/src/datasets/vdd.py @@ -0,0 +1,282 @@ +#!/usr/bin/python +# -*- encoding: utf-8 -*- + +import os +import os.path as osp +from typing import Any, Dict, Optional, Tuple +import warnings + +from PIL import Image +import numpy as np +import torch +from torch.utils.data import Dataset +import torchvision.transforms as transforms + +from src.datasets.transform import ( + Compose, + RandomColorJitter, + RandomCrop, + RandomCutout, + RandomGamma, + RandomHorizontalFlip, + RandomHSV, + RandomNoise, + RandomRotate, + RandomScale, + RandomTranslate, + RandomVerticalFlip, +) + +# Same knobs/semantics as src/datasets/uavid.py's DEFAULT_AUGMENTATION — VDD +# is also UAV-captured imagery, so the same augmentation recipe applies +# unchanged. See uavid.py's module docstring for the rationale behind each +# knob (mosaic/copy_paste unsupported, mixup's hard-label simplification, +# etc.) — identical here. +DEFAULT_AUGMENTATION: Dict[str, float] = { + "degrees": 10.0, + "translate": 0.05, + "scale": 0.3, + "flipud": 0.2, + "fliplr": 0.5, + "hsv_h": 0.01, + "hsv_s": 0.4, + "hsv_v": 0.3, + "mixup": 0.1, +} + + +class VDD(Dataset): + """VDD (Varied Drone Dataset) semantic segmentation dataset. + + Consumes the pre-converted, YOLO-style dataset layout produced by + ``src/scripts/convert_vdd_to_yolo.py``:: + + / + ├── images/ + │ ├── train/ ← RGB JPEGs (uniformly 4000x3000) + │ ├── val/ + │ └── test/ + └── masks/ + ├── train/ ← single-channel PNGs, pixel value = class ID + ├── val/ + └── test/ + + Mask pixel values are already final trainIds (0-6; Other=0 … Water=6 — + see https://github.com/RussRobin/VDD) — no RGB colour palette or lookup + table is needed here, unlike UAVid's raw distribution. Pixel value 255 + is reserved for genuinely unrecognized values, not a real class. Unlike + AeroScapes, VDD ships with all three splits (train/val/test) already + defined by the source distribution's own directory layout. + + Parameters + ---------- + ignore_lb: + Label value treated as "ignore" by the loss/metric (255). + rootpth: + Root of the *converted* dataset (i.e. ``convert_vdd_to_yolo.py``'s + ``--dst``) — NOT the raw VDD distribution. + cropsize: + ``(H, W)`` crop applied during training augmentation via + ``RandomCrop``. VDD source images are uniformly 4000x3000 (unlike + UAVid's mixed resolutions), so ``RandomCrop`` mainly serves as a + fixed-size training patch here — ``pad_if_needed=True`` is kept for + robustness. + mode: + ``"train"``, ``"val"``, or ``"test"``. + augmentation: + Optional dict overriding any subset of ``DEFAULT_AUGMENTATION`` + (degrees/translate/scale/flipud/fliplr/hsv_h/hsv_s/hsv_v/mixup). + + Note on validation batching + ---------------------------- + Like AeroScapes and unlike UAVid, VDD source images are uniformly + 4000x3000, so ``val``/``test`` mode (no crop applied) can still be + batched with ``batch_size > 1`` via the default collate function — no + special-casing is needed in train.py/evaluate.py for this dataset. + """ + + def __init__( + self, + ignore_lb: int, + rootpth: str, + cropsize: Tuple[int, int], + mode: str = "train", + augmentation: Optional[Dict[str, Any]] = None, + ) -> None: + super().__init__() + self.mode = mode + self.ignore_lb = ignore_lb + self.rootpth = rootpth + self.cropsize = tuple(cropsize) + self.aug = {**DEFAULT_AUGMENTATION, **(augmentation or {})} + + if self.mode not in ("train", "val", "test"): + raise ValueError( + f"Mode '{mode}' not supported. Choose 'train', 'val', or 'test'." + ) + if not osp.exists(rootpth): + raise FileNotFoundError(f"Dataset root does not exist: {rootpth}") + + img_dir = osp.join(rootpth, "images", mode) + label_dir = osp.join(rootpth, "masks", mode) + if not osp.exists(img_dir): + raise FileNotFoundError(f"Image directory not found: {img_dir}") + if not osp.exists(label_dir): + raise FileNotFoundError(f"Mask directory not found: {label_dir}") + + # --- Load image and label paths ------------------------------------ + self.imnames = [] + self.imgs = {} + self.labels = {} + + imgnames = sorted(fn for fn in os.listdir(img_dir) if fn.endswith(".jpg")) + for fn in imgnames: + stem = osp.splitext(fn)[0] + self.imgs[stem] = osp.join(img_dir, fn) + label_path = osp.join(label_dir, f"{stem}.png") + if osp.exists(label_path): + self.labels[stem] = label_path + self.imnames.append(stem) + + # Drop any image that has no matching mask + missing = [name for name in self.imnames if name not in self.labels] + if missing: + warnings.warn( + f"{len(missing)} image(s) have no matching mask in {label_dir} " + f"and will be skipped: {sorted(missing)[:5]}…" + ) + self.imnames = [name for name in self.imnames if name in self.labels] + + if len(self.imnames) == 0: + raise RuntimeError( + f"No valid image-mask pairs found for mode='{mode}' in {rootpth}." + ) + + self.len = len(self.imnames) + + # --- Image normalisation -------------------------------------------------- + # Mean/std computed from the VDD train set (see + # src/datasets/compute_vdd_stats.py) — distinct from ImageNet, UAVid, + # and AeroScapes stats. + self.to_tensor = transforms.Compose( + [ + transforms.ToTensor(), + transforms.Normalize( + mean=(0.486, 0.487, 0.441), + std=(0.190, 0.178, 0.214), + ), + ] + ) + + # --- Training augmentation (Geometric → Photometric → Regularisation) --- + degrees = float(self.aug["degrees"]) + scale = float(self.aug["scale"]) + self.trans_train = ( + Compose( + [ + RandomHorizontalFlip(p=float(self.aug["fliplr"])), + RandomVerticalFlip(p=float(self.aug["flipud"])), + RandomTranslate( + translate=float(self.aug["translate"]), + ignore_label=self.ignore_lb, + ), + RandomRotate( + degrees=(-degrees, degrees), ignore_label=self.ignore_lb + ), + RandomScale((1.0 - scale, 1.0 + scale), continuous=True), + RandomCrop( + size=self.cropsize, + pad_if_needed=True, + ignore_label=self.ignore_lb, + ), + RandomHSV( + hgain=float(self.aug["hsv_h"]), + sgain=float(self.aug["hsv_s"]), + vgain=float(self.aug["hsv_v"]), + ), + RandomColorJitter(contrast=0.5), + RandomGamma(gamma_range=(0.8, 1.2), p=0.3), + RandomNoise(mode="gaussian", sigma=0.03, p=0.3), + RandomCutout(p=0.3, size=64), + ] + ) + if mode == "train" + else None + ) + self.mixup_p = float(self.aug["mixup"]) if mode == "train" else 0.0 + + print( + f"[INFO] VDD dataset loaded: {self.len} samples ({mode}) from {img_dir}" + ) + + def _load_one(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """Load, augment, and tensorize a single sample (no MixUp).""" + stem = self.imnames[idx] + img = Image.open(self.imgs[stem]).convert("RGB") + label = Image.open(self.labels[stem]) + if label.mode != "L": + label = label.convert("L") + + if self.mode == "train" and self.trans_train is not None: + im_lb = self.trans_train({"im": img, "lb": label}) + img, label = im_lb["im"], im_lb["lb"] + + img_t = self.to_tensor(img) + label_np = np.array(label, dtype=np.int64) # (H, W), already final class IDs + label_t = torch.from_numpy(label_np).long() + return img_t, label_t + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + img, label = self._load_one(idx) + + if ( + self.mode == "train" + and self.mixup_p > 0 + and np.random.random() < self.mixup_p + ): + other_idx = int(np.random.randint(0, self.len)) + img2, label2 = self._load_one(other_idx) + # Both samples went through the same RandomCrop, so shapes match. + r = float(np.random.beta(32.0, 32.0)) # matches Ultralytics' MixUp exactly + img = img * r + img2 * (1.0 - r) + # Class-ID masks can't be blended the same way (no meaningful + # "average" of two class indices) — take the label from + # whichever image contributed the larger share of the blend. + label = label if r >= 0.5 else label2 + + return img, label + + def __len__(self) -> int: + return self.len + + +# --------------------------------------------------------------------------- +# Smoke test (run: python src/datasets/vdd.py) +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import sys + + rootpth = os.environ.get("VDD_YOLO_ROOT", "") + if not rootpth: + print( + "Set VDD_YOLO_ROOT=/path/to/converted/vdd_yolo and re-run " + "(this must be the OUTPUT of convert_vdd_to_yolo.py, not raw VDD data)." + ) + sys.exit(0) + + for split in ("train", "val", "test"): + try: + ds = VDD( + ignore_lb=255, + rootpth=rootpth, + cropsize=(1024, 1024), + mode=split, + ) + except (FileNotFoundError, RuntimeError) as exc: + print(f"{split}: skipped ({exc})") + continue + print(f"{split}: {len(ds)} samples") + img, lb = ds[0] + unique = torch.unique(lb[lb != 255]) + print(f" img shape: {tuple(img.shape)}, label shape: {tuple(lb.shape)}") + print(f" unique class IDs in first item: {sorted(unique.tolist())}") diff --git a/src/scripts/convert_aeroscapes_to_yolo.py b/src/scripts/convert_aeroscapes_to_yolo.py new file mode 100644 index 0000000..9913f50 --- /dev/null +++ b/src/scripts/convert_aeroscapes_to_yolo.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python3 +"""Convert the AeroScapes dataset into the shared CABiNet/YOLO layout. + +AeroScapes original format (as distributed) +-------------------------------------------- + aeroscapes/ ← *--src* argument + ├── JPEGImages/*.jpg ← RGB input images (uniformly 1280x720) + ├── SegmentationClass/*.png ← single-channel (mode "L") masks, pixel + │ value = class ID already (0-11) — NOT + │ RGB colour-coded like UAVid's Labels/, so + │ no colour lookup table is needed here + ├── Visualizations/*.png ← colourized masks (not used for training) + └── ImageSets/{trn,val}.txt ← one filename stem per line; together the + two files cover every image with zero + overlap. There is no source test split. + +Class mapping (standard published AeroScapes 12-class set, see +configs/AeroScapes_info.json) +----------------------------------------------------------------- + 0 Background 6 Animal + 1 Person 7 Obstacle + 2 Bike 8 Construction + 3 Car 9 Vegetation + 4 Drone 10 Road + 5 Boat 11 Sky + +Pixel value 255 is reserved for genuinely unrecognized values (there should +be none in a clean copy of the dataset); this script validates that every +mask pixel is in {0..11, 255} and warns on any file that violates this. + +Output layout produced +------------------------ + / + images/ + train/ ← copies (NOT symlinks — see below) of JPEGImages/*.jpg + val/ + masks/ + train/ ← copies of SegmentationClass/*.png (already single-channel) + val/ + +Unlike convert_uavid_to_yolo.py (which symlinks by default), this script +always COPIES files: the user intends to redistribute the converted dataset +on Hugging Face, and a directory of symlinks is not redistributable. + +Usage +------ + python src/scripts/convert_aeroscapes_to_yolo.py \\ + --src /path/to/aeroscapes \\ + --dst /path/to/aeroscapes_yolo \\ + --workers 8 + + # Dry-run to check counts before writing + python src/scripts/convert_aeroscapes_to_yolo.py \\ + --src /path/to/aeroscapes --dst /tmp/out --dry-run +""" + +from __future__ import annotations + +import argparse +import os +import shutil +from concurrent.futures import ProcessPoolExecutor, as_completed +from pathlib import Path +from typing import Dict, List, Tuple + +import numpy as np +from PIL import Image + +VALID_CLASS_IDS = frozenset(range(12)) +IGNORE_LABEL: int = 255 + + +# --------------------------------------------------------------------------- +# Split discovery +# --------------------------------------------------------------------------- + + +def load_split_stems(imagesets_dir: Path, split_file: str) -> List[str]: + """Return the sorted list of filename stems listed in an ImageSets file.""" + path = imagesets_dir / split_file + if not path.exists(): + return [] + with open(path) as f: + return sorted(line.strip() for line in f if line.strip()) + + +def discover_splits(src_root: Path) -> Dict[str, List[str]]: + """Return {"train": [...], "val": [...]} stems from ImageSets/{trn,val}.txt. + + There is no source test split — a caller that needs one must carve it + out of val/train themselves; this script only ever produces train/val. + """ + imagesets_dir = src_root / "ImageSets" + return { + "train": load_split_stems(imagesets_dir, "trn.txt"), + "val": load_split_stems(imagesets_dir, "val.txt"), + } + + +# --------------------------------------------------------------------------- +# Per-file copy + validation +# --------------------------------------------------------------------------- + + +def validate_mask(mask_path: Path) -> Tuple[bool, str]: + """Check that every pixel in *mask_path* is a known class ID or IGNORE_LABEL. + + Returns (is_valid, message). + """ + arr = np.array(Image.open(mask_path)) + unique = set(np.unique(arr).tolist()) + bad = unique - VALID_CLASS_IDS - {IGNORE_LABEL} + if bad: + return False, f"unexpected pixel value(s) {sorted(bad)} in {mask_path.name}" + return True, "ok" + + +def _copy_one(args: Tuple[Path, Path, Path, Path, bool]) -> str: + """Copy one (image, mask) pair; validate the mask first. Returns a status string.""" + img_src, img_dst, mask_src, mask_dst, dry_run = args + + valid, msg = validate_mask(mask_src) + if not valid: + return f"warn:{msg}" + + if not dry_run: + img_dst.parent.mkdir(parents=True, exist_ok=True) + mask_dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(img_src, img_dst) + shutil.copy2(mask_src, mask_dst) + return f"ok:{img_src.name}" + + +# --------------------------------------------------------------------------- +# Split conversion +# --------------------------------------------------------------------------- + + +def convert_split( + src_root: Path, + dst_root: Path, + split: str, + stems: List[str], + workers: int = 1, + dry_run: bool = False, +) -> int: + """Copy all (image, mask) pairs for *stems* into /images|masks//. + + Returns the number of pairs successfully copied (or counted in dry-run). + Stems with a missing image or mask counterpart are skipped with a warning. + """ + img_src_dir = src_root / "JPEGImages" + mask_src_dir = src_root / "SegmentationClass" + img_dst_dir = dst_root / "images" / split + mask_dst_dir = dst_root / "masks" / split + + tasks: List[Tuple[Path, Path, Path, Path, bool]] = [] + for stem in stems: + img_src = img_src_dir / f"{stem}.jpg" + mask_src = mask_src_dir / f"{stem}.png" + if not img_src.exists(): + print(f"[WARN] Missing image for '{stem}' — skipping") + continue + if not mask_src.exists(): + print(f"[WARN] Missing mask for '{stem}' — skipping") + continue + tasks.append( + (img_src, img_dst_dir / f"{stem}.jpg", mask_src, mask_dst_dir / f"{stem}.png", dry_run) + ) + + n_ok = 0 + if workers <= 1 or len(tasks) <= 1: + for task in tasks: + result = _copy_one(task) + if result.startswith("ok"): + n_ok += 1 + else: + print(f"[WARN] {result}") + else: + with ProcessPoolExecutor(max_workers=workers) as pool: + futures = {pool.submit(_copy_one, t): t for t in tasks} + for future in as_completed(futures): + try: + result = future.result() + if result.startswith("ok"): + n_ok += 1 + else: + print(f"[WARN] {result}") + except Exception as exc: + print(f"[ERROR] {futures[future][0].name}: {exc}") + + return n_ok + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Convert AeroScapes (JPEGImages/ + SegmentationClass/) -> " + "the shared CABiNet/YOLO images/+masks/ layout", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--src", + required=True, + type=Path, + help="Root directory containing JPEGImages/, SegmentationClass/, " + "and ImageSets/ (e.g. .../aeroscapes/)", + ) + p.add_argument( + "--dst", + required=True, + type=Path, + help="Output directory for the converted dataset", + ) + p.add_argument( + "--workers", + type=int, + default=os.cpu_count() or 1, + help="Number of parallel worker processes", + ) + p.add_argument( + "--dry-run", + action="store_true", + help="Walk the dataset, validate masks, and report counts without writing any files", + ) + return p.parse_args() + + +def main() -> None: + args = _parse_args() + + if not args.src.exists(): + raise FileNotFoundError(f"--src does not exist: {args.src}") + + splits = discover_splits(args.src) + for split, stems in splits.items(): + if not stems: + print(f"[WARN] No stems found for split '{split}' — skipping") + + print(f"[INFO] Source: {args.src}") + print(f"[INFO] train: {len(splits['train'])} images, val: {len(splits['val'])} images") + print("[INFO] No source test split exists — only train/val will be produced.") + + total = 0 + for split, stems in splits.items(): + if not stems: + continue + n = convert_split( + src_root=args.src, + dst_root=args.dst, + split=split, + stems=stems, + workers=args.workers, + dry_run=args.dry_run, + ) + verb = "would copy" if args.dry_run else "copied" + print(f"[INFO] {split}: {verb} {n} pairs") + total += n + + print(f"\n[DONE] Total pairs {'scanned' if args.dry_run else 'written'}: {total}") + if not args.dry_run: + print(f" Output: {args.dst}") + print(f" Next: set AEROSCAPES_YOLO_ROOT={args.dst}") + + +if __name__ == "__main__": + main() diff --git a/src/scripts/convert_vdd_to_yolo.py b/src/scripts/convert_vdd_to_yolo.py new file mode 100644 index 0000000..b2c9e28 --- /dev/null +++ b/src/scripts/convert_vdd_to_yolo.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +"""Convert VDD (Varied Drone Dataset) into the shared CABiNet/YOLO layout. + +VDD original format (as distributed / downloaded from HF) +------------------------------------------------------------ + VDD/ ← *--src* argument + ├── train/ + │ ├── src/*.JPG ← RGB input images (uniformly 4000x3000) + │ └── gt/*.png ← single-channel masks, pixel value = class ID + ├── val/ + │ ├── src/*.JPG + │ └── gt/*.png + ├── test/ + │ ├── src/*.JPG + │ └── gt/*.png + └── metadata/{train,val,test}.txt ← not used by this script (the + train/val/test directory split + already IS the split; test.txt + also has a typo'd .JPG extension + for gt/ paths that doesn't match + the actual .png files on disk) + +Masks are already single-channel (mode "L") class-ID PNGs — like AeroScapes +and unlike UAVid's RGB colour-coded masks — so no colour lookup table is +needed here, just a straight symlink. + +Class mapping (VDD, see https://github.com/RussRobin/VDD) +----------------------------------------------------------- + 0 Other 4 Vehicle + 1 Wall 5 Roof + 2 Road 6 Water + 3 Vegetation + +Output layout produced +------------------------ + / + images/ + train/ ← symlinks to *.JPG, renamed to lowercase *.jpg + val/ + test/ + masks/ + train/ ← symlinks to *.png (already single-channel) + val/ + test/ + +Unlike convert_aeroscapes_to_yolo.py, this script symlinks by default (like +convert_uavid_to_yolo.py) — VDD is already published on HF, so there's no +need to produce a redistributable standalone copy. + +Usage +------ + python src/scripts/convert_vdd_to_yolo.py \\ + --src /path/to/VDD --dst /path/to/vdd_yolo + + # Dry-run to check counts before writing + python src/scripts/convert_vdd_to_yolo.py \\ + --src /path/to/VDD --dst /tmp/out --dry-run +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from typing import List, Tuple + +import numpy as np +from PIL import Image + +VALID_CLASS_IDS = frozenset(range(7)) +IGNORE_LABEL: int = 255 + +SPLITS = ("train", "val", "test") + + +# --------------------------------------------------------------------------- +# Per-file validation + link +# --------------------------------------------------------------------------- + + +def validate_mask(mask_path: Path) -> Tuple[bool, str]: + """Check that every pixel in *mask_path* is a known class ID or IGNORE_LABEL.""" + arr = np.array(Image.open(mask_path)) + unique = set(np.unique(arr).tolist()) + bad = unique - VALID_CLASS_IDS - {IGNORE_LABEL} + if bad: + return False, f"unexpected pixel value(s) {sorted(bad)} in {mask_path.name}" + return True, "ok" + + +def _link_one(img_src: Path, img_dst: Path, mask_src: Path, mask_dst: Path, dry_run: bool) -> str: + valid, msg = validate_mask(mask_src) + if not valid: + return f"warn:{msg}" + + if not dry_run: + img_dst.parent.mkdir(parents=True, exist_ok=True) + mask_dst.parent.mkdir(parents=True, exist_ok=True) + if not img_dst.exists(): + img_dst.symlink_to(img_src.resolve()) + if not mask_dst.exists(): + mask_dst.symlink_to(mask_src.resolve()) + return f"ok:{img_src.name}" + + +# --------------------------------------------------------------------------- +# Split conversion +# --------------------------------------------------------------------------- + + +def discover_stems(split_dir: Path) -> List[str]: + """Return sorted stems present in both /src/*.JPG and gt/*.png.""" + src_dir = split_dir / "src" + gt_dir = split_dir / "gt" + if not src_dir.exists() or not gt_dir.exists(): + return [] + src_stems = {p.stem for p in src_dir.iterdir() if p.suffix.lower() == ".jpg"} + gt_stems = {p.stem for p in gt_dir.iterdir() if p.suffix.lower() == ".png"} + return sorted(src_stems & gt_stems) + + +def convert_split( + src_root: Path, + dst_root: Path, + split: str, + dry_run: bool = False, +) -> int: + """Symlink all (image, mask) pairs for *split* into /images|masks//. + + Returns the number of pairs successfully linked (or counted in dry-run). + Stems present on only one side are skipped with a warning. + """ + split_dir = src_root / split + src_dir = split_dir / "src" + gt_dir = split_dir / "gt" + img_dst_dir = dst_root / "images" / split + mask_dst_dir = dst_root / "masks" / split + + src_stems = {p.stem: p for p in src_dir.iterdir() if p.suffix.lower() == ".jpg"} if src_dir.exists() else {} + gt_stems = {p.stem: p for p in gt_dir.iterdir() if p.suffix.lower() == ".png"} if gt_dir.exists() else {} + + all_stems = sorted(set(src_stems) | set(gt_stems)) + n_ok = 0 + for stem in all_stems: + img_src = src_stems.get(stem) + mask_src = gt_stems.get(stem) + if img_src is None: + print(f"[WARN] Missing image for '{stem}' — skipping") + continue + if mask_src is None: + print(f"[WARN] Missing mask for '{stem}' — skipping") + continue + result = _link_one( + img_src, img_dst_dir / f"{stem}.jpg", mask_src, mask_dst_dir / f"{stem}.png", dry_run + ) + if result.startswith("ok"): + n_ok += 1 + else: + print(f"[WARN] {result}") + + return n_ok + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Convert VDD (train/val/test/{src,gt}) -> the shared " + "CABiNet/YOLO images/+masks/ layout", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--src", + required=True, + type=Path, + help="Root directory containing train/, val/, test/ (each with src/ and gt/)", + ) + p.add_argument( + "--dst", + required=True, + type=Path, + help="Output directory for the converted dataset", + ) + p.add_argument( + "--dry-run", + action="store_true", + help="Walk the dataset, validate masks, and report counts without writing any files", + ) + return p.parse_args() + + +def main() -> None: + args = _parse_args() + + if not args.src.exists(): + raise FileNotFoundError(f"--src does not exist: {args.src}") + + print(f"[INFO] Source: {args.src}") + + total = 0 + for split in SPLITS: + stems = discover_stems(args.src / split) + if not stems: + print(f"[WARN] No stems found for split '{split}' — skipping") + continue + n = convert_split(args.src, args.dst, split, dry_run=args.dry_run) + verb = "would link" if args.dry_run else "linked" + print(f"[INFO] {split}: {verb} {n} pairs") + total += n + + print(f"\n[DONE] Total pairs {'scanned' if args.dry_run else 'written'}: {total}") + if not args.dry_run: + print(f" Output: {args.dst}") + print(f" Next: set VDD_YOLO_ROOT={args.dst}") + + +if __name__ == "__main__": + main() diff --git a/src/scripts/train.py b/src/scripts/train.py index db66a45..b0eb8da 100644 --- a/src/scripts/train.py +++ b/src/scripts/train.py @@ -123,6 +123,59 @@ def _load_checkpoint( return start_epoch, best_miou, best_loss +def _load_pretrained_checkpoint(net: torch.nn.Module, path: Path, device: torch.device) -> None: + """Warm-start model weights from a checkpoint trained on a (possibly + different) dataset, e.g. finetuning on AeroScapes/VDD/etc. from a + UAVid-trained checkpoint — converges much faster than backbone-only + (ImageNet) initialization since the CAB/spatial/fusion layers and the + backbone are all already adapted to aerial imagery. + + Unlike ``_load_checkpoint`` (used by ``resume``), this only ever touches + model weights: the optimizer, scaler, EMA shadow, and epoch counter all + start fresh — this is a new training run, not a continuation of the + checkpoint's own run. + + Accepts either a raw model state_dict or a full training checkpoint + (``{"model_state": ..., ...}``, as produced by ``_save_checkpoint``). + Only parameters whose name AND shape match the current model are loaded + — this is what actually enables cross-dataset finetuning, since the two + classifier heads (``ab.b4``, ``conv_out.conv_out``) are sized by + ``n_classes`` and virtually never match between datasets (UAVid=8, + AeroScapes=12, VDD=7, Cityscapes=19, …). Those are silently left at + their freshly-initialized values; everything else (backbone, CAB, + spatial branch, feature fusion) transfers. + """ + ckpt = torch.load(path, map_location=device, weights_only=True) + pretrained_state = ckpt["model_state"] if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt + + model_state = _model_state_dict(net) + compatible = { + k: v + for k, v in pretrained_state.items() + if k in model_state and v.shape == model_state[k].shape + } + skipped_shape_mismatch = [ + k for k in pretrained_state if k in model_state and k not in compatible + ] + skipped_unknown = [k for k in pretrained_state if k not in model_state] + + model_state.update(compatible) + net.load_state_dict(model_state) + + logger.info( + "Loaded %d/%d pretrained tensors from %s (%d skipped: shape mismatch — " + "likely the classifier heads, expected when n_classes differs)", + len(compatible), + len(pretrained_state), + path, + len(skipped_shape_mismatch), + ) + if skipped_shape_mismatch: + logger.info(" Shape-mismatched (left at fresh init): %s", skipped_shape_mismatch) + if skipped_unknown: + logger.warning(" Unknown keys in checkpoint (ignored): %s", skipped_unknown) + + def _run_miou_eval( net: torch.nn.Module, dl: DataLoader, @@ -242,6 +295,27 @@ def train_and_evaluate(cfg: DictConfig) -> None: net.to(device) console.log("Model moved to device!", style="info") + # Warm-start from a full checkpoint trained on a (possibly different) + # dataset — e.g. finetune AeroScapes/VDD from a UAVid-trained checkpoint, + # which converges much faster than backbone-only (ImageNet) init since + # the CAB/spatial/fusion layers are already adapted to aerial imagery. + # Independent of `resume` (below): this only ever warm-starts weights — + # optimizer/EMA/epoch always start fresh. If `resume=True` also finds an + # existing checkpoint_last.pth in THIS run's experiments_path, that full + # state (including further-trained weights) takes priority and overwrites + # what's loaded here. + pretrained_ckpt = cfg.training_config.get("pretrained_ckpt_path") + if pretrained_ckpt: + pretrained_ckpt_path = Path(pretrained_ckpt) + if not pretrained_ckpt_path.exists(): + raise ConfigurationError( + f"training_config.pretrained_ckpt_path does not exist: {pretrained_ckpt_path}" + ) + _load_pretrained_checkpoint(net, pretrained_ckpt_path, device) + console.print( + f"✅ Warm-started model weights from {pretrained_ckpt_path}", style="info" + ) + # EMA shadow model — constructed right after net.to(device) so the # internal deepcopy inherits the correct device with no extra .to() call. ema = ModelEMA( diff --git a/src/scripts/train_yolo.py b/src/scripts/train_yolo.py index 63a83e2..3690d49 100644 --- a/src/scripts/train_yolo.py +++ b/src/scripts/train_yolo.py @@ -6,27 +6,13 @@ from pathlib import Path import hydra +import yaml from omegaconf import DictConfig, OmegaConf from src.utils.logger import RichConsoleManager logger = logging.getLogger(__name__) -# UAVid classes in trainId order (0-7) — matches configs/UAVid_info.json and -# hf_modelcards/generate_hf_model_zoo.py's UAVID_CLASSES. Used only to label -# mode=val's per-class IoU printout in a form directly pastable into a -# hf_modelcards/model_metrics//metrics.json. -UAVID_CLASS_NAMES = [ - "Clutter", - "Building", - "Road", - "Static Car", - "Tree", - "Vegetation", - "Human", - "Moving Car", -] - # Repo root: src/scripts/train_yolo.py -> src/scripts -> src -> REPO_ROOT = Path(__file__).resolve().parents[2] @@ -61,6 +47,20 @@ def _resolve_dataset_path(config_file: str) -> Path: return p.resolve() +def _load_dataset_class_names(dataset_path: Path) -> list[str]: + """Read the ``names:`` mapping out of the Ultralytics dataset YAML. + + Used only to label mode=val's per-class IoU printout in a form directly + pastable into a hf_modelcards/model_metrics//metrics.json — reads + from the dataset config itself rather than a hardcoded per-dataset list + so this works for any dataset (UAVid, AeroScapes, …) without edits here. + """ + with open(dataset_path) as f: + data = yaml.safe_load(f) + names = data.get("names", {}) + return [names[i] for i in sorted(names, key=int)] + + def _resolve_experiments_path(experiments_path: str) -> Path: """Resolve experiments_path to an absolute path anchored to this repo. @@ -263,9 +263,10 @@ def main(cfg: DictConfig) -> None: console.print(f"mIoU: {results.miou:.4f}", style="info") console.print(f"Pixel accuracy: {results.pixel_accuracy:.4f}", style="info") + class_names = _load_dataset_class_names(dataset_path) per_class_iou = { name: round(float(iou), 4) - for name, iou in zip(UAVID_CLASS_NAMES, results.per_class_iou) + for name, iou in zip(class_names, results.per_class_iou) } console.print("Per-class IoU:", style="info") for name, iou in per_class_iou.items(): diff --git a/tests/integration/test_training_pipeline.py b/tests/integration/test_training_pipeline.py index 9969e45..c9de203 100644 --- a/tests/integration/test_training_pipeline.py +++ b/tests/integration/test_training_pipeline.py @@ -10,7 +10,7 @@ from torch.utils.data import DataLoader, TensorDataset from src.scripts.evaluate import MscEvalV0 -from src.scripts.train import _load_checkpoint, _save_checkpoint +from src.scripts.train import _load_checkpoint, _load_pretrained_checkpoint, _save_checkpoint from src.utils.early_stopping import EarlyStopping from src.utils.ema import ModelEMA from src.utils.loss import OhemCELoss @@ -833,3 +833,114 @@ def test_lr_decays_to_near_zero_at_max_iter(self, mock_small_model, num_classes) f"LR at max_iter-1 ({lr_near_end:.6f}) should be < 10% of initial " f"({lr_start:.6f}). The schedule may not be using optimizer steps." ) + + +class TestPretrainedCheckpointWarmStart: + """Tests for _load_pretrained_checkpoint — cross-dataset finetuning. + + E.g. warm-starting an AeroScapes (num_classes=12) or VDD (num_classes=7) + run from a checkpoint trained on UAVid (num_classes=8): the two + classifier heads (ab.b4, conv_out.conv_out) are sized by num_classes and + essentially never match across datasets, so they must be skipped while + everything else (backbone/CAB/spatial/fusion) transfers. + """ + + def _save_raw_state_dict(self, model, path): + torch.save(model.state_dict(), path) + + def test_cross_dataset_transfers_backbone_skips_classifier_heads( + self, mock_small_model + ): + source = mock_small_model(num_classes=8) # e.g. UAVid-trained + target = mock_small_model(num_classes=12) # e.g. AeroScapes, fresh init + device = torch.device("cpu") + + # A matching-shape param (backbone) to prove transfer happened. + backbone_key = next( + k for k in source.state_dict() if k.startswith("mobile.") + ) + # Perturb the source's backbone weight so it's distinguishable from + # target's independent random init. + with torch.no_grad(): + source.state_dict()[backbone_key].add_(1.0) + source_backbone_before = source.state_dict()[backbone_key].clone() + target_head_shape_before = target.conv_out.conv_out.weight.shape + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = Path(tmpdir) / "ckpt.pth" + optim = self._make_optimizer_static(source) + scaler = torch.amp.GradScaler(device="cpu", enabled=False) + ema = ModelEMA(source) + stopper = EarlyStopping(patience=10) + _save_checkpoint( + ckpt, + epoch=0, + net=source, + optim=optim, + scaler=scaler, + ema=ema, + stopper=stopper, + best_miou=0.5, + best_loss=0.1, + ) + + _load_pretrained_checkpoint(target, ckpt, device) + + # Matching-shape backbone param transferred exactly. + assert torch.equal(target.state_dict()[backbone_key], source_backbone_before) + # Classifier heads keep the TARGET's own shape (num_classes=12), + # proving they were not overwritten with the source's (8,...) shape. + assert target.conv_out.conv_out.weight.shape == target_head_shape_before + assert target.conv_out.conv_out.weight.shape[0] == 12 + + def test_same_num_classes_full_transfer(self, mock_small_model): + """When num_classes matches, every parameter (including the + classifier heads) should transfer — no shape mismatches to skip.""" + source = mock_small_model(num_classes=8) + target = mock_small_model(num_classes=8) + device = torch.device("cpu") + + with torch.no_grad(): + for p in source.parameters(): + p.add_(1.0) + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = Path(tmpdir) / "ckpt.pth" + self._save_raw_state_dict(source, ckpt) + _load_pretrained_checkpoint(target, ckpt, device) + + for k, v in source.state_dict().items(): + assert torch.equal(target.state_dict()[k], v), f"{k} did not transfer" + + def test_accepts_raw_state_dict_not_just_full_checkpoint(self, mock_small_model): + """_load_pretrained_checkpoint must accept a bare model.state_dict() + (e.g. a *_best.pth save), not only the {"model_state": ...} shape + produced by _save_checkpoint.""" + source = mock_small_model(num_classes=8) + target = mock_small_model(num_classes=8) + device = torch.device("cpu") + + with torch.no_grad(): + for p in source.parameters(): + p.add_(2.0) + + with tempfile.TemporaryDirectory() as tmpdir: + ckpt = Path(tmpdir) / "raw_state_dict.pth" + self._save_raw_state_dict(source, ckpt) + # Must not raise, and must actually transfer weights. + _load_pretrained_checkpoint(target, ckpt, device) + + sample_key = next(iter(source.state_dict())) + assert torch.equal( + target.state_dict()[sample_key], source.state_dict()[sample_key] + ) + + def _make_optimizer_static(self, model): + return Optimizer( + model=model, + lr0=1e-3, + momentum=0.9, + wd=5e-4, + warmup_steps=0, + max_iter=1000, + ) diff --git a/tests/unit/test_aeroscapes_dataset.py b/tests/unit/test_aeroscapes_dataset.py new file mode 100644 index 0000000..22bdfab --- /dev/null +++ b/tests/unit/test_aeroscapes_dataset.py @@ -0,0 +1,380 @@ +"""Unit tests for src/datasets/aeroscapes.py + +Tests the pre-converted images/+masks/ dataset loader: + - __init__ error paths: bad mode, missing root/images/masks dirs + - file discovery: matching image/mask pairs, missing-mask skip-with-warning + - __getitem__: 2-tuple (img, label) output, mask pixel values pass through + untouched (no LUT translation — the conversion step already produced + final class IDs), train-mode crop size, val-mode native size + - DataLoader integration with default collation, including a batch_size>1 + val-mode batch (AeroScapes source images are uniformly 1280x720, unlike + UAVid's mixed resolutions, so this is safe here) +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image +from torch.utils.data import DataLoader + +from src.datasets.aeroscapes import AeroScapes + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_dataset_root( + root: Path, mode: str, names_and_values: dict, size: tuple = (8, 8) +) -> None: + """Build a minimal /images// + /masks// tree. + + names_and_values: {stem: fill_value} — each entry becomes one RGB image + (content irrelevant) and one single-channel mask filled with fill_value. + """ + h, w = size + img_dir = root / "images" / mode + mask_dir = root / "masks" / mode + img_dir.mkdir(parents=True, exist_ok=True) + mask_dir.mkdir(parents=True, exist_ok=True) + + for stem, value in names_and_values.items(): + rgb = np.zeros((h, w, 3), dtype=np.uint8) + Image.fromarray(rgb).save(img_dir / f"{stem}.jpg") + mask = np.full((h, w), value, dtype=np.uint8) + Image.fromarray(mask).save(mask_dir / f"{stem}.png") + + +# --------------------------------------------------------------------------- +# __init__ error paths +# --------------------------------------------------------------------------- + + +class TestAeroScapesInitErrors: + def test_invalid_mode_raises(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"000001_001": 1}) + with pytest.raises(ValueError, match="not supported"): + AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="bogus" + ) + + def test_missing_root_raises(self): + with pytest.raises(FileNotFoundError, match="does not exist"): + AeroScapes(ignore_lb=255, rootpth="/nonexistent/path", cropsize=(4, 4)) + + def test_missing_images_dir_raises(self, tmp_path): + (tmp_path / "masks" / "train").mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="Image directory not found"): + AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train" + ) + + def test_missing_masks_dir_raises(self, tmp_path): + (tmp_path / "images" / "train").mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="Mask directory not found"): + AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train" + ) + + def test_no_valid_pairs_raises(self, tmp_path): + """An images/ dir with files but a completely empty masks/ dir (no + matches at all) must raise, not silently produce an empty dataset.""" + (tmp_path / "images" / "train").mkdir(parents=True) + (tmp_path / "masks" / "train").mkdir(parents=True) + Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).save( + tmp_path / "images" / "train" / "000001_001.jpg" + ) + with pytest.raises(RuntimeError, match="No valid image-mask pairs"): + AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train" + ) + + +# --------------------------------------------------------------------------- +# File discovery +# --------------------------------------------------------------------------- + + +class TestFileDiscovery: + def test_finds_all_matching_pairs(self, tmp_path): + _make_dataset_root( + tmp_path, + "train", + {"000001_001": 0, "000001_002": 1, "000002_001": 2}, + ) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train" + ) + assert len(ds) == 3 + + def test_image_with_no_matching_mask_is_skipped_with_warning(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"000001_001": 0}) + # Add an orphan image with no mask counterpart + Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).save( + tmp_path / "images" / "train" / "000001_002.jpg" + ) + with pytest.warns(UserWarning, match="no matching mask"): + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train" + ) + assert len(ds) == 1 + + def test_train_val_test_are_independent(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 0, "b": 1}) + _make_dataset_root(tmp_path, "val", {"c": 0}) + _make_dataset_root(tmp_path, "test", {"d": 0, "e": 0, "f": 0}) + + assert len(AeroScapes(255, str(tmp_path), (4, 4), mode="train")) == 2 + assert len(AeroScapes(255, str(tmp_path), (4, 4), mode="val")) == 1 + assert len(AeroScapes(255, str(tmp_path), (4, 4), mode="test")) == 3 + + +# --------------------------------------------------------------------------- +# __getitem__ +# --------------------------------------------------------------------------- + + +class TestGetItem: + def test_returns_two_tuple(self, tmp_path): + _make_dataset_root(tmp_path, "val", {"a": 3}, size=(16, 16)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val" + ) + item = ds[0] + assert isinstance(item, tuple) and len(item) == 2 + img, label = item + assert img.ndim == 3 # (C, H, W) + assert label.ndim == 2 # (H, W) — no LUT translation, no patch dict + + def test_mask_pixel_values_pass_through_untouched(self, tmp_path): + """Pixel value 6 (a real class ID) and 255 (ignore) must survive + exactly — masks are already final class IDs, no LUT translation.""" + _make_dataset_root(tmp_path, "val", {"class6": 6, "ignore": 255}, size=(8, 8)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val" + ) + + labels_by_name = {} + for i, name in enumerate(ds.imnames): + _, label = ds[i] + labels_by_name[name] = label + + assert (labels_by_name["class6"] == 6).all() + assert (labels_by_name["ignore"] == 255).all() + + def test_train_mode_output_matches_cropsize(self, tmp_path): + """RandomCrop must fire on the full image — output spatial size must + equal cropsize regardless of source size. + + Note: cropsize is kept >= 64 here because RandomCutout has a + hardcoded size=64 with no bounds check against the post-crop image + size — harmless in production (real cropsize is 640x640) but not + exercised by a smaller test cropsize. + """ + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(128, 160)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(80, 80), mode="train" + ) + img, label = ds[0] + assert tuple(img.shape[-2:]) == (80, 80) + assert tuple(label.shape) == (80, 80) + + def test_val_mode_preserves_native_size_no_crop(self, tmp_path): + _make_dataset_root(tmp_path, "val", {"a": 1}, size=(20, 30)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val" + ) + img, label = ds[0] + assert tuple(img.shape[-2:]) == (20, 30) + assert tuple(label.shape) == (20, 30) + + def test_train_mode_handles_source_smaller_than_cropsize(self, tmp_path): + """RandomCrop(pad_if_needed=True) must handle a source image smaller + than cropsize without erroring.""" + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(8, 8)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(64, 64), mode="train" + ) + img, label = ds[0] + assert tuple(img.shape[-2:]) == (64, 64) + assert tuple(label.shape) == (64, 64) + + +# --------------------------------------------------------------------------- +# DataLoader integration +# --------------------------------------------------------------------------- + + +class TestDataLoaderIntegration: + def test_default_collate_produces_correct_batch_shape(self, tmp_path): + _make_dataset_root( + tmp_path, + "train", + {f"000001_{i:03d}": 0 for i in range(6)}, + size=(128, 128), + ) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(64, 64), mode="train" + ) + dl = DataLoader(ds, batch_size=4, shuffle=False) + imgs, labels = next(iter(dl)) + assert imgs.shape == (4, 3, 64, 64) + assert labels.shape == (4, 64, 64) + + def test_val_batch_size_greater_than_one_works_for_uniform_resolution( + self, tmp_path + ): + """Unlike UAVid (mixed source resolution → batch_size=1 required), + AeroScapes source images are uniformly 1280x720, so a batched val + DataLoader must work with the default collate function.""" + _make_dataset_root( + tmp_path, + "val", + {"a": 0, "b": 1, "c": 2, "d": 3}, + size=(20, 20), + ) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val" + ) + dl = DataLoader(ds, batch_size=4, shuffle=False) + imgs, labels = next(iter(dl)) + assert imgs.shape == (4, 3, 20, 20) + assert labels.shape == (4, 20, 20) + + +# --------------------------------------------------------------------------- +# Augmentation config +# --------------------------------------------------------------------------- + + +class TestAugmentationConfig: + def test_default_augmentation_used_when_none_passed(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(32, 32)) + ds = AeroScapes( + ignore_lb=255, rootpth=str(tmp_path), cropsize=(8, 8), mode="train" + ) + assert ds.aug["fliplr"] == 0.5 + assert ds.aug["flipud"] == 0.2 + assert ds.aug["mixup"] == 0.1 + + def test_partial_override_falls_back_to_defaults(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(32, 32)) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(8, 8), + mode="train", + augmentation={"mixup": 0.0}, + ) + assert ds.aug["mixup"] == 0.0 + assert ds.aug["degrees"] == 10.0 # untouched default + + def test_mixup_disabled_in_non_train_modes(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(16, 16)) + _make_dataset_root(tmp_path, "val", {"b": 1}, size=(16, 16)) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(8, 8), + mode="val", + augmentation={"mixup": 1.0}, + ) + assert ds.mixup_p == 0.0 + + +# --------------------------------------------------------------------------- +# MixUp +# --------------------------------------------------------------------------- + + +class TestMixUp: + # cropsize is kept >= 64 throughout this class for the same reason + # noted in TestGetItem: RandomCutout has a hardcoded size=64 with no + # bounds check. + + def test_mixup_probability_zero_never_blends(self, tmp_path, monkeypatch): + _make_dataset_root(tmp_path, "train", {"a": 1, "b": 2}, size=(128, 128)) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 0.0}, + ) + call_count = 0 + original = ds._load_one + + def counting_load_one(idx): + nonlocal call_count + call_count += 1 + return original(idx) + + monkeypatch.setattr(ds, "_load_one", counting_load_one) + for i in range(len(ds)): + ds[i] + assert call_count == len(ds), "mixup=0.0 must never trigger a second load" + + def test_mixup_probability_one_always_blends(self, tmp_path, monkeypatch): + _make_dataset_root( + tmp_path, "train", {"a": 1, "b": 2, "c": 3}, size=(128, 128) + ) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 1.0}, + ) + call_count = 0 + original = ds._load_one + + def counting_load_one(idx): + nonlocal call_count + call_count += 1 + return original(idx) + + monkeypatch.setattr(ds, "_load_one", counting_load_one) + ds[0] + assert call_count == 2 + + def test_mixup_blends_image_but_label_stays_hard(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1, "b": 9}, size=(128, 128)) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={ + "mixup": 1.0, + "degrees": 0.0, + "translate": 0.0, + "scale": 0.0, + "flipud": 0.0, + "fliplr": 0.0, + "hsv_h": 0.0, + "hsv_s": 0.0, + "hsv_v": 0.0, + }, + ) + _, label = ds[0] + unique = set(label.unique().tolist()) + assert unique.issubset({1, 9}), ( + f"Blended label must be a hard copy of one source label, got {unique}" + ) + + def test_mixup_never_crashes_on_single_sample_dataset(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(128, 128)) + ds = AeroScapes( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 1.0}, + ) + img, label = ds[0] + assert img.shape[-2:] == (64, 64) + assert label.shape == (64, 64) diff --git a/tests/unit/test_aeroscapes_yolo_convert.py b/tests/unit/test_aeroscapes_yolo_convert.py new file mode 100644 index 0000000..f20672e --- /dev/null +++ b/tests/unit/test_aeroscapes_yolo_convert.py @@ -0,0 +1,212 @@ +"""Tests for src/scripts/convert_aeroscapes_to_yolo.py + +Unlike UAVid's converter, AeroScapes masks are already single-channel class +IDs (no RGB colour decoding), so this converter is much simpler: it just +copies files while validating mask pixel values. Coverage: + - load_split_stems / discover_splits: ImageSets/{trn,val}.txt parsing + - validate_mask: accepts known class IDs + IGNORE_LABEL, rejects anything else + - convert_split: copies (not symlinks) matching pairs, skips missing + counterparts with a warning, dry-run writes nothing, train/val separation +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from PIL import Image + +from src.scripts.convert_aeroscapes_to_yolo import ( + IGNORE_LABEL, + VALID_CLASS_IDS, + convert_split, + discover_splits, + load_split_stems, + validate_mask, +) + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _make_aeroscapes_tree(root: Path, stems_and_values: dict, size=(4, 4)) -> None: + """Create a minimal AeroScapes-style source tree. + + stems_and_values: {stem: mask_fill_value} — each entry becomes one JPEG + image and one single-channel PNG mask filled with mask_fill_value. + """ + h, w = size + img_dir = root / "JPEGImages" + mask_dir = root / "SegmentationClass" + img_dir.mkdir(parents=True, exist_ok=True) + mask_dir.mkdir(parents=True, exist_ok=True) + + for stem, value in stems_and_values.items(): + rgb = np.zeros((h, w, 3), dtype=np.uint8) + Image.fromarray(rgb).save(img_dir / f"{stem}.jpg") + mask = np.full((h, w), value, dtype=np.uint8) + Image.fromarray(mask).save(mask_dir / f"{stem}.png") + + +def _write_imagesets(root: Path, trn: list[str], val: list[str]) -> None: + imagesets_dir = root / "ImageSets" + imagesets_dir.mkdir(parents=True, exist_ok=True) + (imagesets_dir / "trn.txt").write_text("\n".join(trn) + "\n" if trn else "") + (imagesets_dir / "val.txt").write_text("\n".join(val) + "\n" if val else "") + + +# --------------------------------------------------------------------------- +# load_split_stems / discover_splits +# --------------------------------------------------------------------------- + + +class TestLoadSplitStems: + def test_parses_stems_sorted(self, tmp_path): + _write_imagesets(tmp_path, trn=["b", "a", "c"], val=[]) + stems = load_split_stems(tmp_path / "ImageSets", "trn.txt") + assert stems == ["a", "b", "c"] + + def test_missing_file_returns_empty(self, tmp_path): + (tmp_path / "ImageSets").mkdir() + assert load_split_stems(tmp_path / "ImageSets", "trn.txt") == [] + + def test_blank_lines_ignored(self, tmp_path): + imagesets_dir = tmp_path / "ImageSets" + imagesets_dir.mkdir() + (imagesets_dir / "trn.txt").write_text("a\n\nb\n\n") + assert load_split_stems(imagesets_dir, "trn.txt") == ["a", "b"] + + +class TestDiscoverSplits: + def test_returns_train_and_val(self, tmp_path): + _write_imagesets(tmp_path, trn=["a", "b"], val=["c"]) + splits = discover_splits(tmp_path) + assert splits["train"] == ["a", "b"] + assert splits["val"] == ["c"] + + def test_no_test_key_produced(self, tmp_path): + _write_imagesets(tmp_path, trn=["a"], val=["b"]) + splits = discover_splits(tmp_path) + assert "test" not in splits + + +# --------------------------------------------------------------------------- +# validate_mask +# --------------------------------------------------------------------------- + + +class TestValidateMask: + def test_valid_class_id_passes(self, tmp_path): + mask = np.full((4, 4), 7, dtype=np.uint8) + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid + + def test_ignore_label_passes(self, tmp_path): + mask = np.full((4, 4), IGNORE_LABEL, dtype=np.uint8) + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid + + def test_out_of_range_value_fails(self, tmp_path): + mask = np.full((4, 4), 200, dtype=np.uint8) # not a class ID, not IGNORE_LABEL + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, msg = validate_mask(p) + assert not valid + assert "200" in msg + + def test_all_twelve_classes_valid(self, tmp_path): + assert VALID_CLASS_IDS == frozenset(range(12)) + for c in range(12): + mask = np.full((2, 2), c, dtype=np.uint8) + p = tmp_path / f"mask_{c}.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid, f"class id {c} should be valid" + + +# --------------------------------------------------------------------------- +# convert_split +# --------------------------------------------------------------------------- + + +class TestConvertSplit: + def test_copies_images_and_masks(self, tmp_path): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 0, "000001_002": 5}) + n = convert_split(src, dst, "train", ["000001_001", "000001_002"], workers=1) + assert n == 2 + assert (dst / "images" / "train" / "000001_001.jpg").exists() + assert (dst / "masks" / "train" / "000001_001.png").exists() + + def test_copies_are_not_symlinks(self, tmp_path): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 0}) + convert_split(src, dst, "train", ["000001_001"], workers=1) + out_img = dst / "images" / "train" / "000001_001.jpg" + out_mask = dst / "masks" / "train" / "000001_001.png" + assert not out_img.is_symlink() + assert not out_mask.is_symlink() + + def test_missing_image_skipped_with_warning(self, tmp_path, capsys): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 0}) + n = convert_split(src, dst, "train", ["000001_001", "ghost"], workers=1) + assert n == 1 + assert "Missing image" in capsys.readouterr().out + + def test_missing_mask_skipped_with_warning(self, tmp_path, capsys): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 0}) + (src / "SegmentationClass" / "000001_001.png").unlink() + n = convert_split(src, dst, "train", ["000001_001"], workers=1) + assert n == 0 + assert "Missing mask" in capsys.readouterr().out + + def test_invalid_mask_pixel_skipped_with_warning(self, tmp_path, capsys): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 200}) # out-of-range value + n = convert_split(src, dst, "train", ["000001_001"], workers=1) + assert n == 0 + assert "unexpected pixel value" in capsys.readouterr().out + assert not (dst / "images" / "train" / "000001_001.jpg").exists() + + def test_dry_run_writes_nothing(self, tmp_path): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"000001_001": 0}) + n = convert_split(src, dst, "train", ["000001_001"], workers=1, dry_run=True) + assert n == 1 + assert not dst.exists() + + def test_train_val_separation(self, tmp_path): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + _make_aeroscapes_tree(src, {"a": 0, "b": 1}) + + convert_split(src, dst, "train", ["a"], workers=1) + convert_split(src, dst, "val", ["b"], workers=1) + + train_imgs = list((dst / "images" / "train").glob("*.jpg")) + val_imgs = list((dst / "images" / "val").glob("*.jpg")) + assert [p.stem for p in train_imgs] == ["a"] + assert [p.stem for p in val_imgs] == ["b"] + + def test_parallel_workers_produce_same_result(self, tmp_path): + src = tmp_path / "aeroscapes" + dst = tmp_path / "aeroscapes_yolo" + stems = [f"000001_{i:03d}" for i in range(6)] + _make_aeroscapes_tree(src, {s: i % 12 for i, s in enumerate(stems)}) + n = convert_split(src, dst, "train", stems, workers=4) + assert n == 6 + assert len(list((dst / "images" / "train").glob("*.jpg"))) == 6 diff --git a/tests/unit/test_vdd_dataset.py b/tests/unit/test_vdd_dataset.py new file mode 100644 index 0000000..d4d935b --- /dev/null +++ b/tests/unit/test_vdd_dataset.py @@ -0,0 +1,331 @@ +"""Unit tests for src/datasets/vdd.py + +Tests the pre-converted images/+masks/ dataset loader: + - __init__ error paths: bad mode, missing root/images/masks dirs + - file discovery: matching image/mask pairs, missing-mask skip-with-warning + - __getitem__: 2-tuple (img, label) output, mask pixel values pass through + untouched (no LUT translation — the source masks are already final + class IDs), train-mode crop size, val-mode native size + - DataLoader integration with default collation, including a batch_size>1 + val-mode batch (VDD source images are uniformly 4000x3000, unlike + UAVid's mixed resolutions, so this is safe here) +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +from PIL import Image +from torch.utils.data import DataLoader + +from src.datasets.vdd import VDD + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_dataset_root( + root: Path, mode: str, names_and_values: dict, size: tuple = (8, 8) +) -> None: + """Build a minimal /images// + /masks// tree. + + names_and_values: {stem: fill_value} — each entry becomes one RGB image + (content irrelevant) and one single-channel mask filled with fill_value. + """ + h, w = size + img_dir = root / "images" / mode + mask_dir = root / "masks" / mode + img_dir.mkdir(parents=True, exist_ok=True) + mask_dir.mkdir(parents=True, exist_ok=True) + + for stem, value in names_and_values.items(): + rgb = np.zeros((h, w, 3), dtype=np.uint8) + Image.fromarray(rgb).save(img_dir / f"{stem}.jpg") + mask = np.full((h, w), value, dtype=np.uint8) + Image.fromarray(mask).save(mask_dir / f"{stem}.png") + + +# --------------------------------------------------------------------------- +# __init__ error paths +# --------------------------------------------------------------------------- + + +class TestVDDInitErrors: + def test_invalid_mode_raises(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"DJI_0008": 1}) + with pytest.raises(ValueError, match="not supported"): + VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="bogus") + + def test_missing_root_raises(self): + with pytest.raises(FileNotFoundError, match="does not exist"): + VDD(ignore_lb=255, rootpth="/nonexistent/path", cropsize=(4, 4)) + + def test_missing_images_dir_raises(self, tmp_path): + (tmp_path / "masks" / "train").mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="Image directory not found"): + VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train") + + def test_missing_masks_dir_raises(self, tmp_path): + (tmp_path / "images" / "train").mkdir(parents=True) + with pytest.raises(FileNotFoundError, match="Mask directory not found"): + VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train") + + def test_no_valid_pairs_raises(self, tmp_path): + (tmp_path / "images" / "train").mkdir(parents=True) + (tmp_path / "masks" / "train").mkdir(parents=True) + Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).save( + tmp_path / "images" / "train" / "DJI_0008.jpg" + ) + with pytest.raises(RuntimeError, match="No valid image-mask pairs"): + VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train") + + +# --------------------------------------------------------------------------- +# File discovery +# --------------------------------------------------------------------------- + + +class TestFileDiscovery: + def test_finds_all_matching_pairs(self, tmp_path): + _make_dataset_root( + tmp_path, "train", {"DJI_0008": 0, "DJI_0010": 1, "DJI_0011": 2} + ) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train") + assert len(ds) == 3 + + def test_image_with_no_matching_mask_is_skipped_with_warning(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"DJI_0008": 0}) + Image.fromarray(np.zeros((4, 4, 3), dtype=np.uint8)).save( + tmp_path / "images" / "train" / "DJI_0099.jpg" + ) + with pytest.warns(UserWarning, match="no matching mask"): + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="train") + assert len(ds) == 1 + + def test_train_val_test_are_independent(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 0, "b": 1}) + _make_dataset_root(tmp_path, "val", {"c": 0}) + _make_dataset_root(tmp_path, "test", {"d": 0, "e": 0, "f": 0}) + + assert len(VDD(255, str(tmp_path), (4, 4), mode="train")) == 2 + assert len(VDD(255, str(tmp_path), (4, 4), mode="val")) == 1 + assert len(VDD(255, str(tmp_path), (4, 4), mode="test")) == 3 + + +# --------------------------------------------------------------------------- +# __getitem__ +# --------------------------------------------------------------------------- + + +class TestGetItem: + def test_returns_two_tuple(self, tmp_path): + _make_dataset_root(tmp_path, "val", {"a": 3}, size=(16, 16)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val") + item = ds[0] + assert isinstance(item, tuple) and len(item) == 2 + img, label = item + assert img.ndim == 3 # (C, H, W) + assert label.ndim == 2 # (H, W) + + def test_mask_pixel_values_pass_through_untouched(self, tmp_path): + """Pixel value 6 (Water) and 255 (ignore) must survive exactly — masks + are already final class IDs, no LUT translation.""" + _make_dataset_root(tmp_path, "val", {"class6": 6, "ignore": 255}, size=(8, 8)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val") + + labels_by_name = {} + for i, name in enumerate(ds.imnames): + _, label = ds[i] + labels_by_name[name] = label + + assert (labels_by_name["class6"] == 6).all() + assert (labels_by_name["ignore"] == 255).all() + + def test_train_mode_output_matches_cropsize(self, tmp_path): + """cropsize is kept >= 64 here because RandomCutout has a hardcoded + size=64 with no bounds check — harmless in production (real cropsize + is 1024x1024) but not exercised by a smaller test cropsize.""" + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(128, 160)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(80, 80), mode="train") + img, label = ds[0] + assert tuple(img.shape[-2:]) == (80, 80) + assert tuple(label.shape) == (80, 80) + + def test_val_mode_preserves_native_size_no_crop(self, tmp_path): + _make_dataset_root(tmp_path, "val", {"a": 1}, size=(20, 30)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val") + img, label = ds[0] + assert tuple(img.shape[-2:]) == (20, 30) + assert tuple(label.shape) == (20, 30) + + def test_train_mode_handles_source_smaller_than_cropsize(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(8, 8)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(64, 64), mode="train") + img, label = ds[0] + assert tuple(img.shape[-2:]) == (64, 64) + assert tuple(label.shape) == (64, 64) + + +# --------------------------------------------------------------------------- +# DataLoader integration +# --------------------------------------------------------------------------- + + +class TestDataLoaderIntegration: + def test_default_collate_produces_correct_batch_shape(self, tmp_path): + _make_dataset_root( + tmp_path, + "train", + {f"DJI_{i:04d}": 0 for i in range(6)}, + size=(128, 128), + ) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(64, 64), mode="train") + dl = DataLoader(ds, batch_size=4, shuffle=False) + imgs, labels = next(iter(dl)) + assert imgs.shape == (4, 3, 64, 64) + assert labels.shape == (4, 64, 64) + + def test_val_batch_size_greater_than_one_works_for_uniform_resolution( + self, tmp_path + ): + """Unlike UAVid (mixed source resolution → batch_size=1 required), + VDD source images are uniformly 4000x3000, so a batched val + DataLoader must work with the default collate function.""" + _make_dataset_root( + tmp_path, "val", {"a": 0, "b": 1, "c": 2, "d": 3}, size=(20, 20) + ) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(4, 4), mode="val") + dl = DataLoader(ds, batch_size=4, shuffle=False) + imgs, labels = next(iter(dl)) + assert imgs.shape == (4, 3, 20, 20) + assert labels.shape == (4, 20, 20) + + +# --------------------------------------------------------------------------- +# Augmentation config +# --------------------------------------------------------------------------- + + +class TestAugmentationConfig: + def test_default_augmentation_used_when_none_passed(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(32, 32)) + ds = VDD(ignore_lb=255, rootpth=str(tmp_path), cropsize=(8, 8), mode="train") + assert ds.aug["fliplr"] == 0.5 + assert ds.aug["flipud"] == 0.2 + assert ds.aug["mixup"] == 0.1 + + def test_partial_override_falls_back_to_defaults(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(32, 32)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(8, 8), + mode="train", + augmentation={"mixup": 0.0}, + ) + assert ds.aug["mixup"] == 0.0 + assert ds.aug["degrees"] == 10.0 # untouched default + + def test_mixup_disabled_in_non_train_modes(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(16, 16)) + _make_dataset_root(tmp_path, "val", {"b": 1}, size=(16, 16)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(8, 8), + mode="val", + augmentation={"mixup": 1.0}, + ) + assert ds.mixup_p == 0.0 + + +# --------------------------------------------------------------------------- +# MixUp +# --------------------------------------------------------------------------- + + +class TestMixUp: + def test_mixup_probability_zero_never_blends(self, tmp_path, monkeypatch): + _make_dataset_root(tmp_path, "train", {"a": 1, "b": 2}, size=(128, 128)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 0.0}, + ) + call_count = 0 + original = ds._load_one + + def counting_load_one(idx): + nonlocal call_count + call_count += 1 + return original(idx) + + monkeypatch.setattr(ds, "_load_one", counting_load_one) + for i in range(len(ds)): + ds[i] + assert call_count == len(ds), "mixup=0.0 must never trigger a second load" + + def test_mixup_probability_one_always_blends(self, tmp_path, monkeypatch): + _make_dataset_root(tmp_path, "train", {"a": 1, "b": 2, "c": 3}, size=(128, 128)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 1.0}, + ) + call_count = 0 + original = ds._load_one + + def counting_load_one(idx): + nonlocal call_count + call_count += 1 + return original(idx) + + monkeypatch.setattr(ds, "_load_one", counting_load_one) + ds[0] + assert call_count == 2 + + def test_mixup_blends_image_but_label_stays_hard(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1, "b": 6}, size=(128, 128)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={ + "mixup": 1.0, + "degrees": 0.0, + "translate": 0.0, + "scale": 0.0, + "flipud": 0.0, + "fliplr": 0.0, + "hsv_h": 0.0, + "hsv_s": 0.0, + "hsv_v": 0.0, + }, + ) + _, label = ds[0] + unique = set(label.unique().tolist()) + assert unique.issubset({1, 6}), ( + f"Blended label must be a hard copy of one source label, got {unique}" + ) + + def test_mixup_never_crashes_on_single_sample_dataset(self, tmp_path): + _make_dataset_root(tmp_path, "train", {"a": 1}, size=(128, 128)) + ds = VDD( + ignore_lb=255, + rootpth=str(tmp_path), + cropsize=(64, 64), + mode="train", + augmentation={"mixup": 1.0}, + ) + img, label = ds[0] + assert img.shape[-2:] == (64, 64) + assert label.shape == (64, 64) diff --git a/tests/unit/test_vdd_yolo_convert.py b/tests/unit/test_vdd_yolo_convert.py new file mode 100644 index 0000000..ade81bd --- /dev/null +++ b/tests/unit/test_vdd_yolo_convert.py @@ -0,0 +1,185 @@ +"""Tests for src/scripts/convert_vdd_to_yolo.py + +Like AeroScapes, VDD masks are already single-channel class IDs (no RGB +colour decoding), so this converter just symlinks files while validating +mask pixel values. Coverage: + - discover_stems: finds stems present in both src/ and gt/ + - validate_mask: accepts known class IDs + IGNORE_LABEL, rejects anything else + - convert_split: symlinks (not copies) matching pairs, skips missing + counterparts with a warning, dry-run writes nothing, extension is + normalized to lowercase .jpg +""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +from PIL import Image + +from src.scripts.convert_vdd_to_yolo import ( + IGNORE_LABEL, + VALID_CLASS_IDS, + convert_split, + discover_stems, + validate_mask, +) + + +# --------------------------------------------------------------------------- +# Fixtures / helpers +# --------------------------------------------------------------------------- + + +def _make_vdd_split(root: Path, split: str, stems_and_values: dict, size=(4, 4)) -> None: + """Create a minimal VDD-style //{src,gt}/ tree. + + stems_and_values: {stem: mask_fill_value} — each entry becomes one + uppercase-extension JPEG image and one single-channel PNG mask. + """ + h, w = size + src_dir = root / split / "src" + gt_dir = root / split / "gt" + src_dir.mkdir(parents=True, exist_ok=True) + gt_dir.mkdir(parents=True, exist_ok=True) + + for stem, value in stems_and_values.items(): + rgb = np.zeros((h, w, 3), dtype=np.uint8) + Image.fromarray(rgb).save(src_dir / f"{stem}.JPG") + mask = np.full((h, w), value, dtype=np.uint8) + Image.fromarray(mask).save(gt_dir / f"{stem}.png") + + +# --------------------------------------------------------------------------- +# discover_stems +# --------------------------------------------------------------------------- + + +class TestDiscoverStems: + def test_finds_common_stems(self, tmp_path): + _make_vdd_split(tmp_path, "train", {"DJI_0008": 0, "DJI_0010": 1}) + stems = discover_stems(tmp_path / "train") + assert stems == ["DJI_0008", "DJI_0010"] + + def test_missing_split_dir_returns_empty(self, tmp_path): + assert discover_stems(tmp_path / "nonexistent") == [] + + def test_orphan_stems_excluded(self, tmp_path): + _make_vdd_split(tmp_path, "train", {"DJI_0008": 0}) + (tmp_path / "train" / "src" / "DJI_0099.JPG").write_bytes(b"") + stems = discover_stems(tmp_path / "train") + assert stems == ["DJI_0008"] + + +# --------------------------------------------------------------------------- +# validate_mask +# --------------------------------------------------------------------------- + + +class TestValidateMask: + def test_valid_class_id_passes(self, tmp_path): + mask = np.full((4, 4), 5, dtype=np.uint8) + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid + + def test_ignore_label_passes(self, tmp_path): + mask = np.full((4, 4), IGNORE_LABEL, dtype=np.uint8) + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid + + def test_out_of_range_value_fails(self, tmp_path): + mask = np.full((4, 4), 200, dtype=np.uint8) + p = tmp_path / "mask.png" + Image.fromarray(mask).save(p) + valid, msg = validate_mask(p) + assert not valid + assert "200" in msg + + def test_all_seven_classes_valid(self, tmp_path): + assert VALID_CLASS_IDS == frozenset(range(7)) + for c in range(7): + mask = np.full((2, 2), c, dtype=np.uint8) + p = tmp_path / f"mask_{c}.png" + Image.fromarray(mask).save(p) + valid, _ = validate_mask(p) + assert valid, f"class id {c} should be valid" + + +# --------------------------------------------------------------------------- +# convert_split +# --------------------------------------------------------------------------- + + +class TestConvertSplit: + def test_links_images_and_masks(self, tmp_path): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 0, "DJI_0010": 5}) + n = convert_split(src, dst, "train") + assert n == 2 + assert (dst / "images" / "train" / "DJI_0008.jpg").exists() + assert (dst / "masks" / "train" / "DJI_0008.png").exists() + + def test_extension_normalized_to_lowercase_jpg(self, tmp_path): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 0}) + convert_split(src, dst, "train") + assert (dst / "images" / "train" / "DJI_0008.jpg").exists() + assert not list((dst / "images" / "train").glob("*.JPG")) + + def test_links_are_symlinks(self, tmp_path): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 0}) + convert_split(src, dst, "train") + out_img = dst / "images" / "train" / "DJI_0008.jpg" + out_mask = dst / "masks" / "train" / "DJI_0008.png" + assert out_img.is_symlink() + assert out_mask.is_symlink() + assert out_img.resolve() == (src / "train" / "src" / "DJI_0008.JPG").resolve() + + def test_missing_mask_skipped_with_warning(self, tmp_path, capsys): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 0}) + (src / "train" / "gt" / "DJI_0008.png").unlink() + n = convert_split(src, dst, "train") + assert n == 0 + assert "Missing mask" in capsys.readouterr().out + + def test_invalid_mask_pixel_skipped_with_warning(self, tmp_path, capsys): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 200}) # out-of-range value + n = convert_split(src, dst, "train") + assert n == 0 + assert "unexpected pixel value" in capsys.readouterr().out + assert not (dst / "images" / "train" / "DJI_0008.jpg").exists() + + def test_dry_run_writes_nothing(self, tmp_path): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"DJI_0008": 0}) + n = convert_split(src, dst, "train", dry_run=True) + assert n == 1 + assert not dst.exists() + + def test_train_val_test_separation(self, tmp_path): + src = tmp_path / "vdd" + dst = tmp_path / "vdd_yolo" + _make_vdd_split(src, "train", {"a": 0}) + _make_vdd_split(src, "val", {"b": 1}) + _make_vdd_split(src, "test", {"c": 2}) + + convert_split(src, dst, "train") + convert_split(src, dst, "val") + convert_split(src, dst, "test") + + assert [p.stem for p in (dst / "images" / "train").glob("*.jpg")] == ["a"] + assert [p.stem for p in (dst / "images" / "val").glob("*.jpg")] == ["b"] + assert [p.stem for p in (dst / "images" / "test").glob("*.jpg")] == ["c"]