samstars.training.sam_training

samstars.training.sam_training

Train the proposal decoder from prepared training chips.

Upgrade: align training prompts with inference: - Use box prompts (json["box"] if present, else derived from mask bbox) - Use many fg/bg points derived from CHM, like inference - Pad variable-length points to a fixed size with (-10,-10) and label -1

We still: - freeze image_encoder + prompt_encoder - train mask_decoder only - use global RGB scaling to uint8 (via --scale-stats)

Chips

Bases: Dataset

Source code in samstars/training/sam_training.py

class Chips(Dataset):
    def __init__(
        self,
        chip_dir: Path,
        img_size: int = 512,
        enc_size: int = 1024,
        rgb_lo_hi=None,
        *,
        fg_extra: int = FG_EXTRA_DEFAULT,
        neg_max_gnd: int = NEG_MAX_GND_DEFAULT,
        neg_max_can: int = NEG_MAX_CAN_DEFAULT,
        chm_ground: float = CHM_GROUND_DEFAULT,
        box_expand: float = 1.0,
    ):
        """Initialize the instance."""
        self.chip_dir = chip_dir
        self.ids = sorted((self.chip_dir / "images").glob("*.npy"))
        if not self.ids:
            raise FileNotFoundError(f"No chips in {self.chip_dir}")

        self.img_size = int(img_size)
        self.resize = ResizeLongestSide(enc_size)
        self.rgb_lo_hi = rgb_lo_hi

        self.fg_extra = int(fg_extra)
        self.neg_max_gnd = int(neg_max_gnd)
        self.neg_max_can = int(neg_max_can)
        self.chm_ground = float(chm_ground)
        self.box_expand = float(box_expand)

        # Fixed maximum number of points so DataLoader can stack tensors.
        # seed + fg + neg_gnd + neg_can
        self.max_points = 1 + self.fg_extra + self.neg_max_gnd + self.neg_max_can

    def __len__(self):
        """Return the number of items."""
        return len(self.ids)

    def __getitem__(self, i):
        """Return one item by index."""
        name = self.ids[i].stem

        # --- load RGB chip (C,H,W) ---
        img = np.load(self.chip_dir / "images" / f"{name}.npy")
        if img.ndim != 3:
            raise ValueError(f"{name}: expected img CHW, got {img.shape}")
        C, H, W = img.shape
        if H != self.img_size or W != self.img_size:
            # if your chips are padded/cropped, this keeps things safe
            H = min(H, self.img_size)
            W = min(W, self.img_size)

        # Convert RGB -> uint8 (global scaling if available)
        img_hwc = img.transpose(1, 2, 0).astype(np.float32)
        if self.rgb_lo_hi is not None and len(self.rgb_lo_hi) >= img_hwc.shape[2]:
            img_u8 = scale_to_u8_per_channel(img_hwc, self.rgb_lo_hi[: img_hwc.shape[2]])
        else:
            img_u8 = np.stack([stretch_u8(img_hwc[..., c]) for c in range(img_hwc.shape[2])], axis=-1)

        # Resize to encoder scale (1024 long side)
        img_u8 = self.resize.apply_image(img_u8)
        img_t = img_u8.transpose(2, 0, 1).astype(np.float32)  # CHW float, 0..255

        # --- load mask (H,W) ---
        msk = np.load(self.chip_dir / "masks" / f"{name}.npy")
        msk = (msk > 0).astype(np.float32)

        # --- load CHM (H,W) for prompt sampling ---
        chm_path = self.chip_dir / "chm" / f"{name}.npy"
        if not chm_path.exists():
            raise FileNotFoundError(f"{name}: missing CHM chip {chm_path}")
        chm = np.load(chm_path).astype(np.float32)

        # --- load json (point + optional box) ---
        jpath = self.chip_dir / "prompts" / f"{name}.json"
        j = json.load(open(jpath))
        pt = j.get("point", None)

        if pt is None:
            # fallback: centroid of mask
            yy, xx = np.where(msk > 0)
            if yy.size == 0:
                # should not happen, but keep safe
                sx, sy = self.img_size // 2, self.img_size // 2
            else:
                sx = int(np.round(xx.mean()))
                sy = int(np.round(yy.mean()))
        else:
            sx, sy = int(pt[0]), int(pt[1])

        # Box: prefer json["box"], else bbox from mask
        box = j.get("box", None)
        if box is not None and len(box) == 4:
            x0, y0, x1, y1 = map(float, box)
            box_xyxy = _clamp_box_xyxy((x0, y0, x1, y1), self.img_size, self.img_size)
        else:
            bb = _bbox_from_mask(msk)
            if bb is None:
                # fallback: small box around seed
                box_xyxy = _clamp_box_xyxy((sx - 8, sy - 8, sx + 8, sy + 8), self.img_size, self.img_size)
            else:
                box_xyxy = _clamp_box_xyxy(bb, self.img_size, self.img_size)

        if self.box_expand and self.box_expand != 1.0:
            box_xyxy = _expand_box_xyxy(box_xyxy, self.box_expand, self.img_size, self.img_size)

        # Build inference-style points from CHM
        pts, lbl = _sample_points_from_chm(
            seed_xy=(sx, sy),
            box_xyxy=box_xyxy,
            chm=chm,
            fg_extra=self.fg_extra,
            neg_max_gnd=self.neg_max_gnd,
            neg_max_can=self.neg_max_can,
            chm_ground=self.chm_ground,
        )

        # Pad to fixed length (max_points)
        pts_pad = np.full((self.max_points, 2), PAD_COORD, dtype=np.float32)
        lbl_pad = np.full((self.max_points,), -1, dtype=np.int64)
        n = min(self.max_points, pts.shape[0])
        pts_pad[:n] = pts[:n]
        lbl_pad[:n] = lbl[:n]

        # Transform coords/boxes to resized image coords (1024-long-side space)
        pts_pad = self.resize.apply_coords(pts_pad, (self.img_size, self.img_size)).astype(np.float32)
        box_np = np.array([box_xyxy], dtype=np.float32)
        box_rs = self.resize.apply_boxes(box_np, (self.img_size, self.img_size)).astype(np.float32)[0]

        return (
            torch.tensor(img_t).float(),
            torch.tensor(msk[None]).float(),
            torch.tensor(pts_pad).float(),     # (P,2)
            torch.tensor(lbl_pad).long(),      # (P,)
            torch.tensor(box_rs).float(),      # (4,)
        )

__getitem__(i)

Return one item by index.

Source code in samstars/training/sam_training.py

def __getitem__(self, i):
    """Return one item by index."""
    name = self.ids[i].stem

    # --- load RGB chip (C,H,W) ---
    img = np.load(self.chip_dir / "images" / f"{name}.npy")
    if img.ndim != 3:
        raise ValueError(f"{name}: expected img CHW, got {img.shape}")
    C, H, W = img.shape
    if H != self.img_size or W != self.img_size:
        # if your chips are padded/cropped, this keeps things safe
        H = min(H, self.img_size)
        W = min(W, self.img_size)

    # Convert RGB -> uint8 (global scaling if available)
    img_hwc = img.transpose(1, 2, 0).astype(np.float32)
    if self.rgb_lo_hi is not None and len(self.rgb_lo_hi) >= img_hwc.shape[2]:
        img_u8 = scale_to_u8_per_channel(img_hwc, self.rgb_lo_hi[: img_hwc.shape[2]])
    else:
        img_u8 = np.stack([stretch_u8(img_hwc[..., c]) for c in range(img_hwc.shape[2])], axis=-1)

    # Resize to encoder scale (1024 long side)
    img_u8 = self.resize.apply_image(img_u8)
    img_t = img_u8.transpose(2, 0, 1).astype(np.float32)  # CHW float, 0..255

    # --- load mask (H,W) ---
    msk = np.load(self.chip_dir / "masks" / f"{name}.npy")
    msk = (msk > 0).astype(np.float32)

    # --- load CHM (H,W) for prompt sampling ---
    chm_path = self.chip_dir / "chm" / f"{name}.npy"
    if not chm_path.exists():
        raise FileNotFoundError(f"{name}: missing CHM chip {chm_path}")
    chm = np.load(chm_path).astype(np.float32)

    # --- load json (point + optional box) ---
    jpath = self.chip_dir / "prompts" / f"{name}.json"
    j = json.load(open(jpath))
    pt = j.get("point", None)

    if pt is None:
        # fallback: centroid of mask
        yy, xx = np.where(msk > 0)
        if yy.size == 0:
            # should not happen, but keep safe
            sx, sy = self.img_size // 2, self.img_size // 2
        else:
            sx = int(np.round(xx.mean()))
            sy = int(np.round(yy.mean()))
    else:
        sx, sy = int(pt[0]), int(pt[1])

    # Box: prefer json["box"], else bbox from mask
    box = j.get("box", None)
    if box is not None and len(box) == 4:
        x0, y0, x1, y1 = map(float, box)
        box_xyxy = _clamp_box_xyxy((x0, y0, x1, y1), self.img_size, self.img_size)
    else:
        bb = _bbox_from_mask(msk)
        if bb is None:
            # fallback: small box around seed
            box_xyxy = _clamp_box_xyxy((sx - 8, sy - 8, sx + 8, sy + 8), self.img_size, self.img_size)
        else:
            box_xyxy = _clamp_box_xyxy(bb, self.img_size, self.img_size)

    if self.box_expand and self.box_expand != 1.0:
        box_xyxy = _expand_box_xyxy(box_xyxy, self.box_expand, self.img_size, self.img_size)

    # Build inference-style points from CHM
    pts, lbl = _sample_points_from_chm(
        seed_xy=(sx, sy),
        box_xyxy=box_xyxy,
        chm=chm,
        fg_extra=self.fg_extra,
        neg_max_gnd=self.neg_max_gnd,
        neg_max_can=self.neg_max_can,
        chm_ground=self.chm_ground,
    )

    # Pad to fixed length (max_points)
    pts_pad = np.full((self.max_points, 2), PAD_COORD, dtype=np.float32)
    lbl_pad = np.full((self.max_points,), -1, dtype=np.int64)
    n = min(self.max_points, pts.shape[0])
    pts_pad[:n] = pts[:n]
    lbl_pad[:n] = lbl[:n]

    # Transform coords/boxes to resized image coords (1024-long-side space)
    pts_pad = self.resize.apply_coords(pts_pad, (self.img_size, self.img_size)).astype(np.float32)
    box_np = np.array([box_xyxy], dtype=np.float32)
    box_rs = self.resize.apply_boxes(box_np, (self.img_size, self.img_size)).astype(np.float32)[0]

    return (
        torch.tensor(img_t).float(),
        torch.tensor(msk[None]).float(),
        torch.tensor(pts_pad).float(),     # (P,2)
        torch.tensor(lbl_pad).long(),      # (P,)
        torch.tensor(box_rs).float(),      # (4,)
    )

__init__(chip_dir, img_size=512, enc_size=1024, rgb_lo_hi=None, *, fg_extra=FG_EXTRA_DEFAULT, neg_max_gnd=NEG_MAX_GND_DEFAULT, neg_max_can=NEG_MAX_CAN_DEFAULT, chm_ground=CHM_GROUND_DEFAULT, box_expand=1.0)

Initialize the instance.

Source code in samstars/training/sam_training.py

def __init__(
    self,
    chip_dir: Path,
    img_size: int = 512,
    enc_size: int = 1024,
    rgb_lo_hi=None,
    *,
    fg_extra: int = FG_EXTRA_DEFAULT,
    neg_max_gnd: int = NEG_MAX_GND_DEFAULT,
    neg_max_can: int = NEG_MAX_CAN_DEFAULT,
    chm_ground: float = CHM_GROUND_DEFAULT,
    box_expand: float = 1.0,
):
    """Initialize the instance."""
    self.chip_dir = chip_dir
    self.ids = sorted((self.chip_dir / "images").glob("*.npy"))
    if not self.ids:
        raise FileNotFoundError(f"No chips in {self.chip_dir}")

    self.img_size = int(img_size)
    self.resize = ResizeLongestSide(enc_size)
    self.rgb_lo_hi = rgb_lo_hi

    self.fg_extra = int(fg_extra)
    self.neg_max_gnd = int(neg_max_gnd)
    self.neg_max_can = int(neg_max_can)
    self.chm_ground = float(chm_ground)
    self.box_expand = float(box_expand)

    # Fixed maximum number of points so DataLoader can stack tensors.
    # seed + fg + neg_gnd + neg_can
    self.max_points = 1 + self.fg_extra + self.neg_max_gnd + self.neg_max_can

__len__()

Return the number of items.

Source code in samstars/training/sam_training.py

def __len__(self):
    """Return the number of items."""
    return len(self.ids)

main()

Run the command-line entry point.

Source code in samstars/training/sam_training.py

def main():
    """Run the command-line entry point."""
    pa = argparse.ArgumentParser(description="Train the proposal decoder from a prepared training-chip directory.")
    pa.add_argument(
        "--training-chips-dir",
        type=Path,
        required=True,
        help="Training-chip directory containing images/, masks/, chm/, and prompts/ subdirectories",
    )
    pa.add_argument("--out-path", type=Path, required=True, help="Path to write the trained decoder checkpoint")
    pa.add_argument("--ckpt", type=Path, required=True, help="Base SAM ViT-L checkpoint")
    pa.add_argument("--epochs", type=int, default=15, help="Number of training epochs")
    pa.add_argument("--min-iters", type=int, default=0, help="Minimum training iterations (converted to epochs)")
    pa.add_argument("--device", default="cpu", help="Training device")
    pa.add_argument("--scale-stats", type=Path, default=None, help="Optional global imagery scaling stats JSON")

    # New knobs (optional)
    pa.add_argument("--fg-extra", type=int, default=FG_EXTRA_DEFAULT, help="Additional positive prompt points sampled from CHM")
    pa.add_argument("--neg-max-gnd", type=int, default=NEG_MAX_GND_DEFAULT, help="Maximum negative ground prompt points")
    pa.add_argument("--neg-max-can", type=int, default=NEG_MAX_CAN_DEFAULT, help="Maximum negative canopy prompt points")
    pa.add_argument("--chm-ground", type=float, default=CHM_GROUND_DEFAULT, help="CHM threshold used when sampling prompts")
    pa.add_argument("--box-expand", type=float, default=1.0, help="Expand bbox about its center (1.0 = no expand)")

    args = pa.parse_args()

    rgb_lo_hi = None
    if args.scale_stats:
        stats = load_scale_stats(args.scale_stats)
        per_channel = stats.get("sam_rgb_u8", {}).get("per_channel", [])
        rgb_lo_hi = [(float(c["lo"]), float(c["hi"])) for c in per_channel]

    chip_dir = args.training_chips_dir
    if not chip_dir.exists():
        raise FileNotFoundError(f"Missing chips dir: {chip_dir}")

    train_decoder(
        chip_dir,
        args.ckpt,
        args.out_path,
        epochs=args.epochs,
        min_iters=args.min_iters,
        device=args.device,
        rgb_lo_hi=rgb_lo_hi,
        fg_extra=args.fg_extra,
        neg_max_gnd=args.neg_max_gnd,
        neg_max_can=args.neg_max_can,
        chm_ground=args.chm_ground,
        box_expand=args.box_expand,
    )

stretch_u8(arr)

Fallback per-array min/max stretch to uint8.

Source code in samstars/training/sam_training.py

def stretch_u8(arr: np.ndarray) -> np.ndarray:
    """Fallback per-array min/max stretch to uint8."""
    arr = np.asarray(arr, np.float32)
    finite = np.isfinite(arr)
    if not finite.any():
        return np.zeros_like(arr, dtype=np.uint8)
    lo = float(arr[finite].min())
    hi = float(arr[finite].max())
    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
        return np.zeros_like(arr, dtype=np.uint8)
    out = (arr - lo) * 255.0 / (hi - lo)
    return np.clip(out, 0, 255).astype(np.uint8)

train_decoder(chip_dir, ckpt, out_path, img_size=512, enc_size=1024, epochs=15, min_iters=0, lr=5e-05, device='cpu', rgb_lo_hi=None, *, fg_extra=FG_EXTRA_DEFAULT, neg_max_gnd=NEG_MAX_GND_DEFAULT, neg_max_can=NEG_MAX_CAN_DEFAULT, chm_ground=CHM_GROUND_DEFAULT, box_expand=1.0)

Train decoder.

Source code in samstars/training/sam_training.py

def train_decoder(
    chip_dir: Path,
    ckpt: Path,
    out_path: Path,
    img_size: int = 512,
    enc_size: int = 1024,
    epochs: int = 15,
    min_iters: int = 0,
    lr: float = 5e-5,
    device: str = "cpu",
    rgb_lo_hi=None,
    *,
    fg_extra: int = FG_EXTRA_DEFAULT,
    neg_max_gnd: int = NEG_MAX_GND_DEFAULT,
    neg_max_can: int = NEG_MAX_CAN_DEFAULT,
    chm_ground: float = CHM_GROUND_DEFAULT,
    box_expand: float = 1.0,
):
    """Train decoder."""
    ds = Chips(
        chip_dir,
        img_size=img_size,
        enc_size=enc_size,
        rgb_lo_hi=rgb_lo_hi,
        fg_extra=fg_extra,
        neg_max_gnd=neg_max_gnd,
        neg_max_can=neg_max_can,
        chm_ground=chm_ground,
        box_expand=box_expand,
    )
    loader = DataLoader(ds, batch_size=1, shuffle=True)

    iters_per_epoch = max(1, len(loader))
    if min_iters > 0:
        min_epochs = int((min_iters + iters_per_epoch - 1) // iters_per_epoch)
        epochs = max(epochs, min_epochs)

    print(f"{chip_dir.name}: epochs={epochs} (iters/epoch={iters_per_epoch})")

    sam = sam_model_registry["vit_l"](checkpoint=str(ckpt)).to(device)
    sam.image_encoder.requires_grad_(False)
    sam.prompt_encoder.requires_grad_(False)

    opt = torch.optim.AdamW(sam.mask_decoder.parameters(), lr=lr)
    base_pe = sam.prompt_encoder.get_dense_pe().to(device)

    def loss_fn(pred, truth):
        """Loss fn."""
        score = torch.sigmoid(pred)
        dice = 1 - (2 * (score * truth).sum((2, 3)) + 1) / (score.sum((2, 3)) + truth.sum((2, 3)) + 1)
        return (dice + F.binary_cross_entropy_with_logits(pred, truth)).mean()

    for ep in range(1, epochs + 1):
        tot = 0.0
        for img, msk, pts, lbl, box in loader:
            img = img.to(device)
            msk = msk.to(device)
            pts = pts.to(device)          # (B,P,2) after batching
            lbl = lbl.to(device)          # (B,P)
            box = box.to(device)          # (B,4)

            with torch.no_grad():
                emb = sam.image_encoder(sam.preprocess(img))

            # IMPORTANT: now we train with the same prompt types as inference (points + boxes)
            sparse, dense = sam.prompt_encoder(
                points=(pts, lbl),
                boxes=box,
                masks=None,
            )

            logit_low, _ = sam.mask_decoder(
                image_embeddings=emb,
                image_pe=base_pe,
                sparse_prompt_embeddings=sparse,
                dense_prompt_embeddings=dense,
                multimask_output=False,
            )

            # supervise against the original 512x512 mask
            logit = F.interpolate(logit_low, size=(img_size, img_size), mode="bilinear", align_corners=False)
            loss = loss_fn(logit, msk)

            opt.zero_grad()
            loss.backward()
            opt.step()

            tot += float(loss.item())

        print(f"{chip_dir.name} epoch {ep}/{epochs} loss={tot/len(loader):.4f}")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(sam.state_dict(), out_path)
    print(f"✓ weights → {out_path}")