Train and segment one area¶

This notebook trains a samstars model bundle from one labeled training area, keeps training products on disk for inspection, and segments a small input area with the trained bundle.

Inputs¶

This walkthrough trains a model bundle and then runs segmentation with that model.

The training package is a 150 m tile with:

• data/training_data_strong/img.tif
• data/training_data_strong/chm.tif
• data/training_data_strong/den.tif
• data/training_data_strong/crowns.gpkg

Segmentation uses:

• data/img_clipped.tif
• data/chm_clipped.tif
• data/den_clipped.tif

In [ ]:

import sys
from pathlib import Path


def find_repo_root(start: Path) -> Path:
    start = start.resolve()
    for candidate in [start, *start.parents]:
        if (candidate / "pyproject.toml").exists() and (candidate / "samstars").exists():
            return candidate
    raise RuntimeError("Could not locate the repo root from the current working directory.")


repo_root = find_repo_root(Path.cwd())
if str(repo_root) not in sys.path:
    sys.path.insert(0, str(repo_root))

from samstars import train, segment

import sys from pathlib import Path def find_repo_root(start: Path) -> Path: start = start.resolve() for candidate in [start, *start.parents]: if (candidate / "pyproject.toml").exists() and (candidate / "samstars").exists(): return candidate raise RuntimeError("Could not locate the repo root from the current working directory.") repo_root = find_repo_root(Path.cwd()) if str(repo_root) not in sys.path: sys.path.insert(0, str(repo_root)) from samstars import train, segment

Set paths¶

Set paths for the training rasters, crown labels, base SAM checkpoint, model bundle, segmentation rasters, and output directory.

In [ ]:

data_dir = repo_root / "data"
training_data_dir = data_dir / "training_data_strong"

training_img = training_data_dir / "img.tif"
training_chm = training_data_dir / "chm.tif"
training_den = training_data_dir / "den.tif"
training_crowns = training_data_dir / "crowns.gpkg"

inference_img = data_dir / "img_clipped.tif"
inference_chm = data_dir / "chm_clipped.tif"
inference_den = data_dir / "den_clipped.tif"

sam_base_ckpt = repo_root / "sam/sam_vit_l_0b3195.pth"
work_dir = data_dir / "train_and_segment_one_area_run"
model_dir = work_dir / "model"
inference_run_dir = work_dir / "segmentation_run"

print("repo_root:", repo_root)
print("training_data_dir:", training_data_dir)
print("model_dir:", model_dir)
print("inference_run_dir:", inference_run_dir)

data_dir = repo_root / "data" training_data_dir = data_dir / "training_data_strong" training_img = training_data_dir / "img.tif" training_chm = training_data_dir / "chm.tif" training_den = training_data_dir / "den.tif" training_crowns = training_data_dir / "crowns.gpkg" inference_img = data_dir / "img_clipped.tif" inference_chm = data_dir / "chm_clipped.tif" inference_den = data_dir / "den_clipped.tif" sam_base_ckpt = repo_root / "sam/sam_vit_l_0b3195.pth" work_dir = data_dir / "train_and_segment_one_area_run" model_dir = work_dir / "model" inference_run_dir = work_dir / "segmentation_run" print("repo_root:", repo_root) print("training_data_dir:", training_data_dir) print("model_dir:", model_dir) print("inference_run_dir:", inference_run_dir)

Set tutorial parameters¶

These parameters control seed generation, proposal smoothing, and model training. With storage="persistent", the run keeps derived products under model/_training/ for inspection.

In [ ]:

rgb_bands = (1, 2, 3)

device = "cpu"
storage = "persistent"

chip_size = 512
chip_stride = 384

sam_epochs = 3
stardist_epochs = 10
val_fraction = 0.2

crop_smooth_sigma = 2.0
smooth_sigma = 4.0
post_smooth_sigma = 0.0

thin_seeds = False
thin_dist = 0.8
seed_d_min = 4.0
seed_min_dist_px = 4
seed_gauss_sigma = 1
seed_min_samples = 1
seed_merge_radius = 3.0
seed_min_canopy_height = 2.0

stardist_input_key = "index"
stardist_use_aux = True
stardist_prob_thr = 0.3
stardist_nms_thr = 0.4
stardist_window_size_m = 150.0
stardist_window_overlap_m = 30.0
stardist_dedup_iou_thr = 0.6

rgb_bands = (1, 2, 3) device = "cpu" storage = "persistent" chip_size = 512 chip_stride = 384 sam_epochs = 3 stardist_epochs = 10 val_fraction = 0.2 crop_smooth_sigma = 2.0 smooth_sigma = 4.0 post_smooth_sigma = 0.0 thin_seeds = False thin_dist = 0.8 seed_d_min = 4.0 seed_min_dist_px = 4 seed_gauss_sigma = 1 seed_min_samples = 1 seed_merge_radius = 3.0 seed_min_canopy_height = 2.0 stardist_input_key = "index" stardist_use_aux = True stardist_prob_thr = 0.3 stardist_nms_thr = 0.4 stardist_window_size_m = 150.0 stardist_window_overlap_m = 30.0 stardist_dedup_iou_thr = 0.6

Train model bundle¶

This writes a model bundle at model_dir and keeps training products under model_dir/_training/.

In [ ]:

model = train(
    img=training_img,
    chm=training_chm,
    den=training_den,
    crowns=training_crowns,
    sam_base_ckpt=sam_base_ckpt,
    out=model_dir,
    storage=storage,
    rgb_bands=rgb_bands,
    device=device,
    chip_size=chip_size,
    chip_stride=chip_stride,
    sam_epochs=sam_epochs,
    stardist_epochs=stardist_epochs,
    val_fraction=val_fraction,
    crop_smooth_sigma=crop_smooth_sigma,
    smooth_sigma=smooth_sigma,
    post_smooth_sigma=post_smooth_sigma,
    thin_seeds=thin_seeds,
    thin_dist=thin_dist,
    seed_d_min=seed_d_min,
    seed_min_dist_px=seed_min_dist_px,
    seed_gauss_sigma=seed_gauss_sigma,
    seed_min_samples=seed_min_samples,
    seed_merge_radius=seed_merge_radius,
    seed_min_canopy_height=seed_min_canopy_height,
    stardist_input_key=stardist_input_key,
    stardist_use_aux=stardist_use_aux,
    stardist_window_size_m=stardist_window_size_m,
    stardist_window_overlap_m=stardist_window_overlap_m,
    stardist_dedup_iou_thr=stardist_dedup_iou_thr,
    fast=False,
)

model = train( img=training_img, chm=training_chm, den=training_den, crowns=training_crowns, sam_base_ckpt=sam_base_ckpt, out=model_dir, storage=storage, rgb_bands=rgb_bands, device=device, chip_size=chip_size, chip_stride=chip_stride, sam_epochs=sam_epochs, stardist_epochs=stardist_epochs, val_fraction=val_fraction, crop_smooth_sigma=crop_smooth_sigma, smooth_sigma=smooth_sigma, post_smooth_sigma=post_smooth_sigma, thin_seeds=thin_seeds, thin_dist=thin_dist, seed_d_min=seed_d_min, seed_min_dist_px=seed_min_dist_px, seed_gauss_sigma=seed_gauss_sigma, seed_min_samples=seed_min_samples, seed_merge_radius=seed_merge_radius, seed_min_canopy_height=seed_min_canopy_height, stardist_input_key=stardist_input_key, stardist_use_aux=stardist_use_aux, stardist_window_size_m=stardist_window_size_m, stardist_window_overlap_m=stardist_window_overlap_m, stardist_dedup_iou_thr=stardist_dedup_iou_thr, fast=False, )

Inspect bundle and persistent training outputs¶

In [ ]:

sorted(p.name for p in model.root.iterdir())

sorted(p.name for p in model.root.iterdir())

In [ ]:

training_proposals_dir = model.root / "_training" / "training_run" / "proposals" / "rasters"
sorted(p.name for p in training_proposals_dir.glob("*.tif"))[:20]

training_proposals_dir = model.root / "_training" / "training_run" / "proposals" / "rasters" sorted(p.name for p in training_proposals_dir.glob("*.tif"))[:20]

Segment the larger clipped inference area¶

This uses the trained bundle, prepares the segmentation rasters, and writes the persistent run products under inference_run_dir.

In [ ]:

result = segment(
    model=model,
    out=inference_run_dir,
    storage=storage,
    img=inference_img,
    chm=inference_chm,
    den=inference_den,
    rgb_bands=rgb_bands,
    device=device,
    stardist_input_key=stardist_input_key,
    stardist_use_aux=stardist_use_aux,
    stardist_prob_thr=stardist_prob_thr,
    stardist_nms_thr=stardist_nms_thr,
    stardist_window_size_m=stardist_window_size_m,
    stardist_window_overlap_m=stardist_window_overlap_m,
    stardist_dedup_iou_thr=stardist_dedup_iou_thr,
    crop_smooth_sigma=crop_smooth_sigma,
    smooth_sigma=smooth_sigma,
    post_smooth_sigma=post_smooth_sigma,
    fast=False,
)

result = segment( model=model, out=inference_run_dir, storage=storage, img=inference_img, chm=inference_chm, den=inference_den, rgb_bands=rgb_bands, device=device, stardist_input_key=stardist_input_key, stardist_use_aux=stardist_use_aux, stardist_prob_thr=stardist_prob_thr, stardist_nms_thr=stardist_nms_thr, stardist_window_size_m=stardist_window_size_m, stardist_window_overlap_m=stardist_window_overlap_m, stardist_dedup_iou_thr=stardist_dedup_iou_thr, crop_smooth_sigma=crop_smooth_sigma, smooth_sigma=smooth_sigma, post_smooth_sigma=post_smooth_sigma, fast=False, )

Final output¶

In [ ]:

result.crowns

result.crowns