Retrain polygon refinement for one area¶

This notebook reuses saved proposal rasters from an existing persistent one-area training run. It builds a new refinement training set from crowns_training_smoothed_strong.gpkg, trains a new refinement model in a separate folder, and writes new crowns without rerunning proposal generation.

Inputs¶

This notebook assumes:

• the clipped rasters already exist in data/
• the source one-area bundle-training run already exists
• the saved proposal records live under model/_training/training_run/ in that source run

In [ ]:

import json
import subprocess
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.data import build_scale_stats_from_proposal_records, build_single_tile
from samstars.training.records import build_stardist_training_records
from samstars.training.stardist_training import train_stardist_model

import json import subprocess 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.data import build_scale_stats_from_proposal_records, build_single_tile from samstars.training.records import build_stardist_training_records from samstars.training.stardist_training import train_stardist_model

Set paths¶

This notebook reuses proposal rasters from the source one-area training run.

In [ ]:

data_dir = repo_root / "data"

img = data_dir / "img_clipped.tif"
chm = data_dir / "chm_clipped.tif"
den = data_dir / "den_clipped.tif"
crowns = data_dir / "crowns_training_smoothed_strong.gpkg"

source_work_dir = data_dir / "train_and_segment_one_area_run"
source_training_run_dir = source_work_dir / "model" / "_training" / "training_run"
proposal_records = source_training_run_dir / "proposals" / "proposal_records.json"

work_dir = data_dir / "retrain_stardist_one_area_strong_smoothed"
device = "cpu"

labeled_tile_dir = work_dir / "labeled_tile"
records_dir = work_dir / "stardist_training_records"
models_dir = work_dir / "models"
outputs_dir = work_dir / "crowns"
stardist_model = models_dir / "stardist_strong_smoothed"
scale_stats = work_dir / "training_scale_stats.json"

rgb_bands = (1, 2, 3)
stardist_input_key = "index"
stardist_epochs = 30
val_fraction = 0.2
use_aux = True
prob_thresh = 0.3
nms_thresh = 0.4

print("repo_root:", repo_root)
print("source_training_run_dir:", source_training_run_dir)
print("proposal_records:", proposal_records)
print("work_dir:", work_dir)
print("device:", device)

data_dir = repo_root / "data" img = data_dir / "img_clipped.tif" chm = data_dir / "chm_clipped.tif" den = data_dir / "den_clipped.tif" crowns = data_dir / "crowns_training_smoothed_strong.gpkg" source_work_dir = data_dir / "train_and_segment_one_area_run" source_training_run_dir = source_work_dir / "model" / "_training" / "training_run" proposal_records = source_training_run_dir / "proposals" / "proposal_records.json" work_dir = data_dir / "retrain_stardist_one_area_strong_smoothed" device = "cpu" labeled_tile_dir = work_dir / "labeled_tile" records_dir = work_dir / "stardist_training_records" models_dir = work_dir / "models" outputs_dir = work_dir / "crowns" stardist_model = models_dir / "stardist_strong_smoothed" scale_stats = work_dir / "training_scale_stats.json" rgb_bands = (1, 2, 3) stardist_input_key = "index" stardist_epochs = 30 val_fraction = 0.2 use_aux = True prob_thresh = 0.3 nms_thresh = 0.4 print("repo_root:", repo_root) print("source_training_run_dir:", source_training_run_dir) print("proposal_records:", proposal_records) print("work_dir:", work_dir) print("device:", device)

Prepare one labeled tile¶

This tile is only for rebuilding refinement training masks against the strong-smoothed crowns. Proposal rasters are reused from the existing source run.

In [ ]:

build_single_tile(
    img=img,
    chm=chm,
    den=den,
    crowns=crowns,
    rgb_bands=rgb_bands,
    out=labeled_tile_dir,
    skip_existing=True,
)

build_single_tile( img=img, chm=chm, den=den, crowns=crowns, rgb_bands=rgb_bands, out=labeled_tile_dir, skip_existing=True, )

Reuse the saved proposal records¶

This notebook does not rerun proposal generation.

In [ ]:

assert proposal_records.exists(), proposal_records

proposal_result = {
    "records_root": source_training_run_dir / "run_records",
    "proposals_root": source_training_run_dir / "proposals",
    "crowns_dir": source_training_run_dir / "crowns",
    "proposal_records": proposal_records,
}

proposal_result

assert proposal_records.exists(), proposal_records proposal_result = { "records_root": source_training_run_dir / "run_records", "proposals_root": source_training_run_dir / "proposals", "crowns_dir": source_training_run_dir / "crowns", "proposal_records": proposal_records, } proposal_result

Build refinement training records¶

These records use the saved proposal rasters together with the strong-smoothed training crowns.

In [ ]:

train_records, val_records = build_stardist_training_records(
    proposal_records=proposal_result["proposal_records"],
    tiles_dir=labeled_tile_dir,
    out_dir=records_dir,
    val_fraction=val_fraction,
)

print("train records:", len(json.load(open(train_records))))
print("val records:", len(json.load(open(val_records))))

train_records, val_records = build_stardist_training_records( proposal_records=proposal_result["proposal_records"], tiles_dir=labeled_tile_dir, out_dir=records_dir, val_fraction=val_fraction, ) print("train records:", len(json.load(open(train_records)))) print("val records:", len(json.load(open(val_records))))

Build scaling stats¶

The refinement input uses the saved proposal rasters from the source run.

In [ ]:

build_scale_stats_from_proposal_records(
    proposal_records=proposal_result["proposal_records"],
    out=scale_stats,
    percentiles=(2.0, 98.0),
    chm_key="chm",
    nodata_values=(-9999.0,),
)

build_scale_stats_from_proposal_records( proposal_records=proposal_result["proposal_records"], out=scale_stats, percentiles=(2.0, 98.0), chm_key="chm", nodata_values=(-9999.0,), )

Train refinement model¶

This writes a new refinement model under a separate work directory.

In [ ]:

train_stardist_model(
    train_records=train_records,
    val_records=val_records,
    out_dir=stardist_model.parent,
    model_name=stardist_model.name,
    epochs=stardist_epochs,
    input_key=stardist_input_key,
    use_aux=use_aux,
    scale_stats_path=scale_stats,
)

train_stardist_model( train_records=train_records, val_records=val_records, out_dir=stardist_model.parent, model_name=stardist_model.name, epochs=stardist_epochs, input_key=stardist_input_key, use_aux=use_aux, scale_stats_path=scale_stats, )

Run refinement¶

This writes new crowns from the saved proposal rasters into the separate rerun folder.

In [ ]:

cmd = [
    sys.executable,
    str(repo_root / "samstars" / "segmentation" / "refinement.py"),
    "--proposal-records", str(proposal_result["proposal_records"]),
    "--model-dir", str(stardist_model.parent),
    "--model-name", stardist_model.name,
    "--out-dir", str(outputs_dir),
    "--input-key", stardist_input_key,
    "--scale-stats", str(scale_stats),
    "--preds-name", "crowns.gpkg",
    "--write-labels",
    "--prob-thr", str(prob_thresh),
    "--nms-thr", str(nms_thresh),
]

if use_aux:
    cmd.append("--use-aux")

print(" ".join(cmd))
subprocess.run(cmd, check=True)

cmd = [ sys.executable, str(repo_root / "samstars" / "segmentation" / "refinement.py"), "--proposal-records", str(proposal_result["proposal_records"]), "--model-dir", str(stardist_model.parent), "--model-name", stardist_model.name, "--out-dir", str(outputs_dir), "--input-key", stardist_input_key, "--scale-stats", str(scale_stats), "--preds-name", "crowns.gpkg", "--write-labels", "--prob-thr", str(prob_thresh), "--nms-thr", str(nms_thresh), ] if use_aux: cmd.append("--use-aux") print(" ".join(cmd)) subprocess.run(cmd, check=True)

Final output¶

In [ ]:

outputs_dir / "crowns.gpkg"

outputs_dir / "crowns.gpkg"