tree-distribution-shift 0.2.2

Tree Distribution Shift benchmark — WILDS-style PyTorch datasets + COCO export for satellite tree-crown detection under distribution shift.

tree-distribution-shift

A WILDS-style benchmark for satellite tree-crown detection under distribution shift.

~30 K COCO-annotated 400 x 400 px satellite image tiles from all states in India and California (US), organised into 8 config pairs that define country-level, biome-level, and region-level distribution shifts.

Data lives on the HF Hub: aadityabuilds/tree-distribution-shift

---

Install

pip install tree-distribution-shift

Verify:

python -c "import tree_shift; print(tree_shift.__version__)"
# Should print 0.2.0

To edit or contribute, install from source:

git clone https://github.com/aadityabuilds/tree-distribution-shift.git
cd tree-distribution-shift
pip install -e .

Quick start — PyTorch Dataset (recommended)

from tree_shift import get_dataset
from tree_shift.data_loaders import get_train_loader
import torchvision.transforms as transforms

# Load a benchmark config (downloads & caches automatically)
dataset = get_dataset(config="intl_train_IN__ood_US", download=True)

# Get the training split as a PyTorch Dataset
train_data = dataset.get_subset(
    "train",
    transform=transforms.Compose(
        [transforms.Resize((448, 448)), transforms.ToTensor()]
    ),
)

# Standard data loader (shuffled, detection-aware collate)
train_loader = get_train_loader("standard", train_data, batch_size=16)

for images, targets, metadata in train_loader:
    # images:   [B, C, H, W] tensor
    # targets:  list of dicts — boxes [N,4] (xyxy), labels [N], area [N]
    # metadata: [B, M] integer-encoded (country, state, zone, region, biome)
    ...

Evaluation

from tree_shift.data_loaders import get_eval_loader

val_data = dataset.get_subset("val", transform=...)
val_loader = get_eval_loader("standard", val_data, batch_size=16)

# Accumulate predictions and ground truths
all_y_pred, all_y_true = [], []
for images, targets, metadata in val_loader:
    preds = model(images)          # your model
    all_y_pred.extend(preds)
    all_y_true.extend(targets)

metrics = dataset.eval(all_y_pred, all_y_true)
# {'mAP': 0.45, 'mAP_50': 0.62, 'mAP_75': 0.41, ...}

Metadata

Each sample includes integer-encoded metadata:

print(dataset.metadata_fields)
# ['country', 'state', 'zone', 'region', 'biome']

# Decode an integer code back to a string
dataset.metadata_map["country"][0]  # e.g. 'IN'

Quick start — COCO export

Use this when you want a standard COCO folder structure for Detectron2, MMDetection, torchvision detection, etc.

CLI

tree-shift export --config intl_train_IN__ood_US --out ./coco_out

Python

from tree_shift import export_coco

export_coco(
    config="intl_train_IN__ood_US",
    out_root="./coco_out",
    splits=["train", "val", "ood_test"],
)

Output:

coco_out/
└── intl_train_IN__ood_US/
    ├── train/
    │   ├── images/*.tiff
    │   └── annotations/instances_train.json
    ├── val/
    │   ├── images/
    │   └── annotations/instances_val.json
    └── ood_test/
        ├── images/
        └── annotations/instances_ood_test.json

Available configs

Configs come in pairs. For every pair, the held-out test portions are constant regardless of whether they are used as val or ood_test.

#	Config name	Shift type	Train on	OOD from
1	intl_train_IN__ood_US	Country	India	US
2	intl_train_US__ood_IN	Country	US	India
3	biome_Rajasthan_train_WET__ood_DRY	Biome	Rajasthan WET	Rajasthan DRY
4	biome_Rajasthan_train_DRY__ood_WET	Biome	Rajasthan DRY	Rajasthan WET
5	biome_Karnataka_train_WET__ood_DRY	Biome	Karnataka WET	Karnataka DRY
6	biome_Karnataka_train_DRY__ood_WET	Biome	Karnataka DRY	Karnataka WET
7	region_train_North__ood_South	Region	North India	South India
8	region_train_South__ood_North	Region	South India	North India

Each config gives you three splits automatically:

Split	Description
train	90 % of the in-distribution pool
val	10 % held-out from the in-distribution pool
ood_test	Held-out portion from the out-of-distribution pool

Dataset contract

Every example contains:

Column	Type	Description
image_bytes	bytes	Raw image file contents
coco_annotations	str	JSON — list of COCO annotation dicts
coco_categories	str	JSON — list of COCO category dicts
image_id	int	Unique image identifier
filename	str	Image filename
width / height	int	Pixel dimensions
country	str	Country code (IN, US)
state	str	State name
zone	str	Geographic zone
region	str	Region
biome	str	Biome classification

When accessed via the PyTorch API, targets are returned in torchvision detection format (boxes in [x1, y1, x2, y2]), and metadata is an integer-encoded tensor.

Dependencies

datasets>=2.19.0
huggingface_hub>=0.23.0
numpy>=1.19.0
orjson>=3.9.0
Pillow>=8.0.0
pycocotools>=2.0.0
torch>=1.9.0
torchvision>=0.10.0
tqdm>=4.66.0

Example scripts

In examples/, we provide scripts that demonstrate the full workflow:

python examples/simple_example.py      # list configs, inspect a sample
python examples/training_example.py    # train/eval loop skeleton
python examples/workflow_example.py    # PyTorch + COCO side-by-side

These scripts are not part of the installed package. To use them, clone the repo and install from source.

Authentication (recommended)

Set an HF token for higher rate limits:

export HF_TOKEN=hf_...

Or pass per-command:

tree-shift list --token hf_...

Resume support

COCO exports are resume-safe. If interrupted, re-run the same command and already-written images are skipped.

API reference

get_dataset(config, download=True, root_dir=None) -> TreeShiftDataset

Load a benchmark config. Returns a dataset object with get_subset() and eval().

TreeShiftDataset.get_subset(split, transform=None) -> TreeShiftSubset

Return a PyTorch Dataset for the given split (train, val, or ood_test).

TreeShiftDataset.eval(y_pred, y_true, metadata=None) -> dict

Compute COCO mAP metrics. Returns mAP, mAP_50, mAP_75, and more.

get_train_loader(loader_type, dataset, batch_size, ...) -> DataLoader

Shuffled data loader with detection-aware collate.

get_eval_loader(loader_type, dataset, batch_size, ...) -> DataLoader

Non-shuffled data loader for evaluation.

list_configs(repo_id=..., revision=None) -> list[str]

List available configs from the Hub (no data download).

export_coco(config, out_root, splits=None, ...) -> Path

Export a config to COCO folders on disk.

Citation

@misc{nalawade2025tree_distribution_shift,
  title={Tree Distribution Shift: A Benchmark for Out-of-Distribution Tree Detection},
  author={Nalawade, Aaditya},
  year={2025},
  howpublished={\url{https://huggingface.co/datasets/aadityabuilds/tree-distribution-shift}}
}

License

MIT