samstars.segmentation.refinement

samstars.segmentation.refinement

Run a trained polygon-refinement model on proposal outputs and write predicted crowns as polygons.

Inputs

• A proposal-records JSON produced by samstars/segmentation/proposals.py.
• A StarDist model directory (Keras) containing a stardist/ subfolder with config.json and weights (e.g., weights_best.h5).

Outputs

• <out_dir>/crowns.gpkg (layer: crowns) with fields: id, source, label, score
• Optional per-input label rasters under <out_dir>/labels/ if --write-labels

This is inference-only (no ground truth masks required).

main()

Run the command-line entry point.

Source code in samstars/segmentation/refinement.py

def main() -> None:
    """Run the command-line entry point."""
    pa = argparse.ArgumentParser(description="Run StarDist refinement from proposal records.")
    pa.add_argument("--proposal-records", type=Path, required=True, help="Proposal-record JSON from the proposal stage")
    pa.add_argument("--model-dir", type=Path, required=True, help="Directory containing the StarDist model subdir")
    pa.add_argument("--model-name", type=str, default="stardist", help="Model subdir name (default: stardist)")
    pa.add_argument("--out-dir", type=Path, required=True)

    pa.add_argument("--prob-thr", type=float, default=None)
    pa.add_argument("--nms-thr", type=float, default=None)

    pa.add_argument(
        "--input-key",
        type=str,
        default="index",
        help="Proposal-record key to use as StarDist input (default: index). Use file_name for raw SAM prob.",
    )
    pa.add_argument("--normalize", action="store_true", help="Enable csbdeep percentile normalization")
    pa.add_argument("--no-normalize", action="store_true", help="Disable csbdeep percentile normalization")

    # Aux
    pa.add_argument("--use-aux", action="store_true", help="Use CHM aux channel (stacks [primary, chm01])")
    pa.add_argument("--scale-stats", type=Path, default=None, help="Global scaling stats JSON (required for --use-aux)")
    pa.add_argument("--chm-key", type=str, default="chm", help="Proposal-record key for CHM path")

    # Outputs
    pa.add_argument("--write-labels", action="store_true", help="Write per-input label rasters under out_dir/labels")
    pa.add_argument("--preds-name", type=str, default="crowns.gpkg", help="Output GeoPackage filename")
    pa.add_argument(
        "--resume",
        action="store_true",
        help=(
            "Resume from per-record checkpoints under out_dir/record_shards. "
            "Completed records are skipped and the final GeoPackage is rebuilt at the end."
        ),
    )
    pa.add_argument(
        "--resume-from-labels",
        action="store_true",
        help=(
            "With --resume and --write-labels, bootstrap missing per-record checkpoints "
            "from existing label rasters. Bootstrapped features have NaN scores."
        ),
    )
    pa.add_argument(
        "--skip-sources",
        type=str,
        default=None,
        help="Comma-separated source names to skip before StarDist inference, for example tile_086_031_0.",
    )
    pa.add_argument(
        "--skip-sources-file",
        type=Path,
        default=None,
        help="Optional text file of source names to skip, one per line. Lines may contain # comments.",
    )
    pa.add_argument(
        "--print-source-before-run",
        action="store_true",
        help="Print each source immediately before running or bootstrapping it; useful for native crashes.",
    )
    pa.add_argument(
        "--window-size-m",
        type=float,
        default=None,
        help="Optional internal StarDist patch size in map units; if omitted, predict on the full tile.",
    )
    pa.add_argument(
        "--window-overlap-m",
        type=float,
        default=0.0,
        help="Overlap between internal StarDist patches in map units (default: 0).",
    )

    # Polygon filtering (useful when inputs overlap)
    pa.add_argument(
        "--edge-buffer-px",
        type=float,
        default=0.5,
        help="Drop polygons within this many pixels of the input raster edge (default: 0.5)",
    )
    pa.add_argument("--min-area-m2", type=float, default=0.0, help="Drop polygons smaller than this area (m^2)")
    pa.add_argument(
        "--dedup-iou-thr",
        type=float,
        default=0.0,
        help="Optional final polygon dedup IoU threshold; keep the higher-score polygon when IoU exceeds this value.",
    )
    pa.add_argument(
        "--log",
        action="store_true",
        help="Show verbose refinement logs, including StarDist internal tile progress.",
    )
    args = pa.parse_args()

    if args.normalize and args.no_normalize:
        raise SystemExit("Use only one of --normalize or --no-normalize.")
    if args.resume_from_labels and not args.resume:
        raise SystemExit("--resume-from-labels requires --resume.")
    if args.resume_from_labels and not args.write_labels:
        raise SystemExit("--resume-from-labels requires --write-labels.")
    if args.window_size_m is not None and args.window_size_m <= 0:
        raise SystemExit("--window-size-m must be > 0")
    if args.window_overlap_m < 0:
        raise SystemExit("--window-overlap-m must be >= 0")
    if args.window_size_m is not None and args.window_overlap_m >= args.window_size_m:
        raise SystemExit("--window-overlap-m must be smaller than --window-size-m")
    if not (0.0 <= args.dedup_iou_thr < 1.0):
        raise SystemExit("--dedup-iou-thr must be in [0, 1).")
    if args.normalize:
        do_normalize = True
    elif args.no_normalize:
        do_normalize = False
    else:
        do_normalize = True

    recs = json.load(open(args.proposal_records))
    if not isinstance(recs, list):
        raise SystemExit(f"{args.proposal_records} is not a list of records.")
    skip_sources = _load_skip_sources(
        skip_sources=args.skip_sources,
        skip_sources_file=args.skip_sources_file,
    )

    scale_stats = load_scale_stats(args.scale_stats) if args.scale_stats else None
    if args.use_aux and scale_stats is None:
        raise SystemExit("--scale-stats is required when --use-aux is set")

    chm_lo = chm_hi = None
    if scale_stats:
        chm_lo = float(scale_stats.get("chm_01", {}).get("lo", 0.0))
        chm_hi = float(scale_stats.get("chm_01", {}).get("hi", 1.0))

    model = StarDist2D(None, name=args.model_name, basedir=str(args.model_dir))

    prob_thr = args.prob_thr
    nms_thr = args.nms_thr
    if prob_thr is None and nms_thr is None:
        p1, n1 = _extract_thresholds(getattr(model, "thresholds", None))
        if p1 is None or n1 is None:
            p1, n1 = _load_thresholds_from_json(args.model_dir, args.model_name)
        prob_thr, nms_thr = p1, n1
    if (prob_thr is None) ^ (nms_thr is None):
        raise SystemExit("Provide both --prob-thr and --nms-thr, or neither to use optimized thresholds.")
    if args.log and prob_thr is not None and nms_thr is not None:
        print(f"Using effective thresholds: prob_thresh={prob_thr}, nms_thresh={nms_thr}.")

    args.out_dir.mkdir(parents=True, exist_ok=True)
    out_gpkg = args.out_dir / args.preds_name
    labels_dir = args.out_dir / "labels"
    shards_dir = args.out_dir / "record_shards"
    if args.resume:
        shards_dir.mkdir(parents=True, exist_ok=True)
        if args.write_labels:
            labels_dir.mkdir(parents=True, exist_ok=True)
    else:
        if out_gpkg.exists():
            out_gpkg.unlink()
        tmp_out_gpkg = out_gpkg.with_suffix(".tmp.gpkg")
        if tmp_out_gpkg.exists():
            tmp_out_gpkg.unlink()
        if shards_dir.exists():
            shutil.rmtree(shards_dir)
        shards_dir.mkdir(parents=True, exist_ok=True)
        if args.write_labels:
            if labels_dir.exists():
                shutil.rmtree(labels_dir)
            labels_dir.mkdir(parents=True, exist_ok=True)

    wrote = 0
    crs_wkt = _first_crs_wkt(recs, args.input_key)
    skipped = 0
    skipped_by_request = 0

    pbar = tqdm(
        recs,
        desc="StarDist refinement",
        unit="record",
        dynamic_ncols=True,
    )
    for rec in pbar:
            x_path = rec.get(args.input_key, "")
            if not x_path or not Path(x_path).exists():
                continue

            source = _guess_source_name(x_path)
            pbar.set_postfix_str(source, refresh=False)
            if source in skip_sources or _safe_source_name(source) in skip_sources:
                skipped_by_request += 1
                continue
            safe_source = _safe_source_name(source)
            shard_path = shards_dir / f"{safe_source}.gpkg"
            done_marker = shards_dir / f"{safe_source}.done.json"
            label_path = labels_dir / f"{source}_labels.tif" if args.write_labels else None
            if args.resume and _checkpoint_complete(shard_path, done_marker, label_path):
                skipped += 1
                continue
            current_source_path = args.out_dir / "current_refinement_source.txt"
            current_source_path.write_text(f"{source}\n")
            if args.print_source_before_run:
                print(f"processing source: {source}", flush=True)
            if args.resume_from_labels and label_path is not None and label_path.exists():
                try:
                    _bootstrap_checkpoint_from_label(
                        label_path=label_path,
                        source=source,
                        shard_path=shard_path,
                        done_marker=done_marker,
                        crs_wkt=crs_wkt,
                        min_area_m2=float(args.min_area_m2),
                    )
                    skipped += 1
                    continue
                except Exception as exc:
                    if args.log:
                        print(f"Could not bootstrap checkpoint from {label_path}: {exc}")

            with rasterio.open(x_path) as src:
                x = src.read(1).astype(np.float32)
                tfm = src.transform
                res_m = float(max(abs(tfm.a), abs(tfm.e)))
                edge_tol_m = float(max(0.0, args.edge_buffer_px) * res_m)
                extent = box(src.bounds.left, src.bounds.bottom, src.bounds.right, src.bounds.top)
                border = extent.boundary

            x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)

            if args.use_aux:
                chm_path = rec.get(args.chm_key, "")
                if not chm_path or not Path(chm_path).exists():
                    continue
                chm = _read_float_raster(chm_path)
                chm = _align_to_shape(chm, x.shape)
                if scale_stats is not None:
                    chm = scale_to_01(chm, chm_lo, chm_hi)
                x_in = np.stack([x, chm.astype(np.float32)], axis=-1).astype(np.float32)
            else:
                x_in = x[..., None].astype(np.float32)

            if do_normalize:
                x_in = normalize(x_in, 1, 99.8, axis=(0, 1)).astype(np.float32)

            predict_kwargs = {
                "show_tile_progress": bool(args.log),
            }
            if prob_thr is not None and nms_thr is not None:
                predict_kwargs["prob_thresh"] = prob_thr
                predict_kwargs["nms_thresh"] = nms_thr

            windowed = args.window_size_m is not None
            window_polys: list[tuple[Polygon, int, float]] = []

            if not windowed:
                predict_kwargs["n_tiles"] = model._guess_n_tiles(x_in)
                labels, details = model.predict_instances(x_in, **predict_kwargs)

                score_by_label = {}
                if isinstance(details, dict) and "prob" in details:
                    probs = np.asarray(details["prob"], dtype=float)
                    score_by_label = {i + 1: float(p) for i, p in enumerate(probs)}

                for lbl_id, poly in _iter_polys_from_labels(
                    labels,
                    transform=tfm,
                    min_area_m2=float(args.min_area_m2),
                ):
                    if edge_tol_m > 0 and poly.distance(border) < edge_tol_m:
                        continue
                    score = float(score_by_label.get(int(lbl_id), float("nan")))
                    window_polys.append((poly, int(lbl_id), score))

                if args.write_labels:
                    with rasterio.open(x_path) as src:
                        meta = src.meta.copy()
                    _write_labels_atomic(label_path, labels, meta=meta, dtype="uint16")
            else:
                patch_size_px = max(1, int(round(float(args.window_size_m) / res_m)))
                overlap_px = max(0, int(round(float(args.window_overlap_m) / res_m)))

                for window in _iter_patch_windows(
                    x.shape,
                    patch_size_px=patch_size_px,
                    overlap_px=overlap_px,
                ):
                    row_off = int(window.row_off)
                    col_off = int(window.col_off)
                    h = int(window.height)
                    w = int(window.width)
                    x_patch = x_in[row_off : row_off + h, col_off : col_off + w]
                    patch_kwargs = dict(predict_kwargs)
                    patch_kwargs["n_tiles"] = model._guess_n_tiles(x_patch)
                    labels, details = model.predict_instances(x_patch, **patch_kwargs)

                    score_by_label = {}
                    if isinstance(details, dict) and "prob" in details:
                        probs = np.asarray(details["prob"], dtype=float)
                        score_by_label = {i + 1: float(p) for i, p in enumerate(probs)}

                    patch_tfm = window_transform(window, tfm)
                    keep_extent = _window_keep_extent(
                        window,
                        full_shape=x.shape,
                        transform=tfm,
                        overlap_m=float(args.window_overlap_m),
                    )

                    for lbl_id, poly in _iter_polys_from_labels(
                        labels,
                        transform=patch_tfm,
                        min_area_m2=float(args.min_area_m2),
                    ):
                        if edge_tol_m > 0 and poly.distance(border) < edge_tol_m:
                            continue
                        if not poly.representative_point().within(keep_extent):
                            continue
                        score = float(score_by_label.get(int(lbl_id), float("nan")))
                        window_polys.append((poly, int(lbl_id), score))

                if args.write_labels:
                    with rasterio.open(x_path) as src:
                        meta = src.meta.copy()
                    shapes = [
                        (poly, int(i + 1))
                        for i, (poly, _lbl_id, _score) in enumerate(window_polys)
                    ]
                    labels_full = features.rasterize(
                        shapes=shapes,
                        out_shape=x.shape,
                        transform=tfm,
                        fill=0,
                        dtype="uint32",
                    )
                    _write_labels_atomic(label_path, labels_full, meta=meta, dtype="uint32")

            record_features = [
                {
                    "geometry": poly,
                    "source": str(source),
                    "label": int(lbl_id),
                    "score": float(score),
                }
                for poly, lbl_id, score in window_polys
            ]
            _write_features_gpkg(shard_path, record_features, crs_wkt=crs_wkt)
            _write_json_atomic(
                done_marker,
                {
                    "source": str(source),
                    "input": str(x_path),
                    "features": len(record_features),
                },
            )

    out_features: list[dict] = []
    for marker in sorted(shards_dir.glob("*.done.json")):
        shard_path = marker.with_suffix("").with_suffix(".gpkg")
        if not shard_path.exists():
            continue
        out_features.extend(_read_features_gpkg(shard_path))

    out_features = _dedup_features_by_iou(
        out_features,
        iou_thr=float(args.dedup_iou_thr),
    )

    _write_features_gpkg(out_gpkg, out_features, crs_wkt=crs_wkt)
    wrote = len(out_features)

    print(f"✓ wrote {wrote} crown polygons to {out_gpkg}")
    if args.resume:
        print(f"✓ skipped {skipped} completed refinement records")
    if skipped_by_request:
        print(f"✓ skipped {skipped_by_request} requested sources")