Skip to content

Experimental Detection

Use obia.detection to train a RetinaNet model that locates objects with bounding boxes. This is an experimental workflow, separate from segmenting and classifying polygon objects; its interfaces may change.

Install detection dependencies

Install the detection extra:

1
pip install "obia[detection]"

The detection extra includes PyTorch, Torchvision, Albumentations, and Matplotlib.

Prepare image chips and annotations

TreeDetectionDataset reads raster chips and JSON annotations:

1
2
3
4
5
from obia.detection.dataset import TreeDetectionDataset
dataset = TreeDetectionDataset(
    images_dir="/path/to/chips",
    annotations_path="/path/to/annotations.json",
)

The JSON file maps image IDs to filenames, boxes, and class labels. Box coordinates are [xmin, ymin, xmax, ymax] in pixels. For example:

1
2
3
4
5
6
7
{
  "chip_01": {
    "file_name": "chip_01.tif",
    "boxes": [[20, 30, 80, 95]],
    "labels": [1]
  }
}

Build a model

1
2
3
4
5
6
from obia.detection.models import build_detection_model

model = build_detection_model(
    num_classes=2,
    in_channels=3,
)

Set in_channels to match the number of raster bands used for detection.

Train

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
from torch.utils.data import DataLoader

from obia.detection.train import train_model
from obia.detection.utils import collate_fn

loader = DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    collate_fn=collate_fn,
)

model = train_model(
    model,
    loader,
    num_epochs=20,
    device="cpu",
)

Use device="cuda" for a supported NVIDIA GPU, or device="mps" for a supported Apple Silicon installation. The CPU example works without GPU support.

Predict

1
2
3
4
5
6
7
8
from obia.detection.predict import predict

detections = predict(
    model,
    image_path="/path/to/image.tif",
    device="cpu",
    score_threshold=0.5,
)

The result contains boxes, scores, and labels arrays. Boxes use pixel coordinates; convert them with the source raster's transform if you need map coordinates.