Spaces:
Sleeping
Sleeping
File size: 3,501 Bytes
eebd638 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Build the global K-means prototype corpus from all 2022 construction tiles.
Run this once after downloading tiles, before training.
Usage::
python scripts/build_prototypes.py \\
--tiles data/tiles/ \\
--k 20 \\
--out checkpoints/prototypes.pt
The script:
1. Loads every (tile_2022, mask_2022) pair from the tile directory.
2. Passes each tile through the frozen SAM2 image encoder.
3. Extracts foreground feature vectors (where mask > 0).
4. Runs KMeans(k) across all sites → k cluster centroids.
5. Saves as a dict: {"prototypes": Tensor[k, 256], "k": int, "inertia": float}.
"""
from __future__ import annotations
import argparse
import logging
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import torch
from omegaconf import OmegaConf
from data.dataset import ConstructionChangeDataset
from models.corpus_prototype import build_corpus_prototypes
from models.detector import ConstructionChangeDetector
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger(__name__)
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser(description="Build SAM2 prototype corpus from 2022 tiles.")
p.add_argument("--tiles", required=True, help="Root tile directory.")
p.add_argument("--k", type=int, default=20, help="Number of prototype clusters.")
p.add_argument("--out", required=True, help="Output path for prototypes.pt.")
p.add_argument("--config", default="configs/base.yaml")
p.add_argument("--max-fg", type=int, default=2000, help="Max foreground features sampled per site.")
p.add_argument("--max-sites", type=int, default=None, help="Randomly sample this many sites (useful on CPU).")
p.add_argument("--sam2-model", default=None, help="Override SAM2 HF model id (e.g. facebook/sam2.1-hiera-tiny).")
return p.parse_args()
def main() -> None:
args = parse_args()
cfg = OmegaConf.load(args.config)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info("Device: %s", device)
# Load dataset (no augmentation, train split = all)
dataset = ConstructionChangeDataset(
tile_dir=args.tiles,
image_size=cfg.data.image_size,
augment=False,
min_fg_pixels=cfg.data.min_fg_pixels,
)
logger.info("Dataset: %d samples", len(dataset))
# Subsample sites if requested (useful on CPU)
if args.max_sites and args.max_sites < len(dataset):
import random, torch.utils.data as tud
indices = random.sample(range(len(dataset)), args.max_sites)
dataset = tud.Subset(dataset, indices)
logger.info("Subsampled to %d sites (--max-sites)", args.max_sites)
# Build model and load SAM2 encoder
sam2_id = args.sam2_model or cfg.model.sam2_hf_id
logger.info("Loading SAM2: %s", sam2_id)
model = ConstructionChangeDetector(k_prototypes=args.k)
model.load_sam2(sam2_id)
model.to(device)
# Build prototypes
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
prototypes = build_corpus_prototypes(
image_encoder=model.image_encoder,
dataset=dataset,
k=args.k,
max_fg_per_site=args.max_fg,
device=device,
output_path=args.out,
)
logger.info("Prototypes shape: %s", prototypes.shape)
logger.info("Saved to %s", args.out)
if __name__ == "__main__":
main()
|