Metadata-Version: 2.4
Name: tree-distribution-shift
Version: 0.2.2
Summary: Tree Distribution Shift benchmark — WILDS-style PyTorch datasets + COCO export for satellite tree-crown detection under distribution shift.
Author: Aaditya Nalawade
License-Expression: MIT
Project-URL: Homepage, https://huggingface.co/datasets/aadityabuilds/tree-distribution-shift
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: datasets>=2.19.0
Requires-Dist: huggingface_hub>=0.23.0
Requires-Dist: numpy>=1.19.0
Requires-Dist: orjson>=3.9.0
Requires-Dist: Pillow>=8.0.0
Requires-Dist: pycocotools>=2.0.0
Requires-Dist: torch>=1.9.0
Requires-Dist: torchvision>=0.10.0
Requires-Dist: tqdm>=4.66.0
Dynamic: license-file

# 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`](https://huggingface.co/datasets/aadityabuilds/tree-distribution-shift)

---

## Install

```bash
pip install tree-distribution-shift
```

Verify:

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

To edit or contribute, install from source:

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

## Quick start — PyTorch Dataset (recommended)

```python
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

```python
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:

```python
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

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

### Python

```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:

```bash
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:

```bash
export HF_TOKEN=hf_...
```

Or pass per-command:

```bash
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

```bibtex
@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
