"""Segment Anything on aerial imagery with segment-geospatial (samgeo), on a laptop CPU.

Streams a 512 x 512 m patch of 2023 NAIP aerial imagery of Ames, Iowa from the Microsoft Planetary Computer and
runs Meta's Segment Anything Model (SAM, the small ViT-B version) on it three ways:
  1. automatic mask generation: segment everything, then vectorise the masks to a GeoPackage and filter
     them by size, shape and NDVI to find buildings;
  2. point prompts: click on an object (given as longitude/latitude) and get a mask back;
  3. box prompts: draw a box around each OpenStreetMap building and let SAM trace its outline.
Results are checked against OpenStreetMap building footprints.
Run with:  python sam_aerial_segmentation.py
"""
import platform
import time
from importlib.metadata import version
from pathlib import Path

import geopandas as gpd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import osmnx as ox
import pandas as pd
import planetary_computer
import pystac_client
import rasterio
import rasterio.features
import requests
import torch
from matplotlib.colors import ListedColormap
from matplotlib.lines import Line2D
from matplotlib.patches import Patch
from pyproj import Transformer
from rasterio.enums import Resampling
from rasterio.transform import from_origin
from rasterio.vrt import WarpedVRT
from samgeo import SamGeo
from scipy import ndimage
from shapely.geometry import box, shape

HERE = Path(__file__).parent
OUT, DATA, CACHE = HERE / "output", HERE / "data", HERE / "cache"
for d in (OUT, DATA, CACHE):
    d.mkdir(exist_ok=True)
plt.rcParams.update({"figure.dpi": 150, "savefig.bbox": "tight", "font.size": 9})
pd.set_option("display.width", 140)
torch.set_num_threads(8)              # a typical laptop; SAM on CPU scales well up to about 8 threads

print("=== versions ===")
print(f"python {platform.python_version()} | segment-geospatial {version('segment-geospatial')} | "
      f"segment-anything {version('segment-anything')} | torch {torch.__version__} | rasterio {rasterio.__version__} | "
      f"geopandas {gpd.__version__} | osmnx {ox.__version__}")
print(f"device: CPU ({torch.get_num_threads()} threads), CUDA available: {torch.cuda.is_available()}")


def section(title: str) -> None:
    print(f"\n=== {title} ===")


# --- Step 1: a patch of NAIP aerial imagery ---------------------------------------------------
STAC_URL = "https://planetarycomputer.microsoft.com/api/stac/v1"
CENTRE = (-93.684, 42.016)        # lon/lat: a school, houses and apartments in west Ames, Iowa
SIZE, RES = 1024, 0.5             # 1024 x 1024 pixels of 0.5 m = 512 x 512 m, exactly one SAM input tile

catalog = pystac_client.Client.open(STAC_URL, modifier=planetary_computer.sign_inplace)
nearby = list(catalog.search(collections=["naip"], intersects=box(*CENTRE, *CENTRE).buffer(0.004)).items())
year = max(i.datetime.year for i in nearby)               # the most recent NAIP survey
items = [i for i in nearby if i.datetime.year == year]    # NAIP tiles are quarter quads: the patch may span several
item = items[0]
section("NAIP imagery")
print(f"{len(items)} NAIP tiles from {year} overlap the patch, e.g. {item.id}")
print(f"flown {item.datetime:%d %b %Y} | {item.properties['gsd']} m | {item.properties['proj:code']} | "
      f"bands: red, green, blue, near-infrared")

rgbn_tif, rgb_tif = DATA / "naip_ames_rgbn.tif", DATA / "naip_ames_rgb.tif"
with rasterio.open(item.assets["image"].href) as src:
    CRS = src.crs
cx, cy = Transformer.from_crs(4326, CRS, always_xy=True).transform(*CENTRE)
TRANSFORM = from_origin(round(cx) - SIZE * RES / 2, round(cy) + SIZE * RES / 2, RES, RES)
if not rgbn_tif.exists():
    rgbn = np.zeros((4, SIZE, SIZE), dtype="uint8")
    for it in items:                                      # mosaic: fill the empty pixels from each tile in turn
        with rasterio.open(it.assets["image"].href) as src, \
                WarpedVRT(src, crs=CRS, transform=TRANSFORM, width=SIZE, height=SIZE, src_nodata=0, nodata=0,
                          resampling=Resampling.average) as vrt:
            part = vrt.read()                             # 4 bands: red, green, blue, near-infrared
        empty = rgbn.max(axis=0) == 0
        rgbn[:, empty] = part[:, empty]
    profile = dict(driver="GTiff", width=SIZE, height=SIZE, dtype="uint8", crs=CRS, transform=TRANSFORM,
                   compress="deflate")
    with rasterio.open(rgbn_tif, "w", count=4, **profile) as dst:
        dst.write(rgbn)
    with rasterio.open(rgb_tif, "w", count=3, **profile) as dst:   # SAM expects a 3-band RGB image
        dst.write(rgbn[:3])
with rasterio.open(rgbn_tif) as src:
    rgbn = src.read().astype("float32")
rgb = np.moveaxis(rgbn[:3], 0, -1).astype("uint8")
with np.errstate(invalid="ignore", divide="ignore"):
    ndvi = (rgbn[3] - rgbn[0]) / (rgbn[3] + rgbn[0])
x0, y1 = TRANSFORM.c, TRANSFORM.f
extent = (0, SIZE * RES, 0, SIZE * RES)                       # plot in metres from the lower-left corner
print(f"mosaic of {len(items)} tiles; empty pixels: {(rgbn.max(axis=0) == 0).sum()}")
print(f"patch: {SIZE} x {SIZE} pixels at {RES} m = {SIZE * RES:.0f} x {SIZE * RES:.0f} m, resampled from "
      f"{item.properties['gsd']} m; vegetated (NDVI > 0.2): {np.mean(ndvi > 0.2):.0%}")

# --- Step 2: reference building footprints from OpenStreetMap --------------------------------
osm_file = DATA / "osm_ames.osm"
patch_ll = gpd.GeoSeries([box(x0, y1 - SIZE * RES, x0 + SIZE * RES, y1)], crs=CRS).to_crs(4326).total_bounds
if not osm_file.exists():           # the OSM editing API serves small boxes directly: no Overpass server needed
    r = requests.get("https://api.openstreetmap.org/api/0.6/map", params={"bbox": ",".join(map(str, patch_ll))},
                     headers={"User-Agent": "technicalgis-tutorial"}, timeout=120)
    r.raise_for_status()
    osm_file.write_bytes(r.content)
osm = ox.features_from_xml(osm_file, tags={"building": True})
osm = osm[osm.geometry.type.isin(["Polygon", "MultiPolygon"])].to_crs(CRS).reset_index()
patch = box(x0, y1 - SIZE * RES, x0 + SIZE * RES, y1)
osm = osm[osm.within(patch.buffer(-1))].reset_index(drop=True)    # whole buildings only
osm["area_m2"] = osm.area
section("OpenStreetMap buildings")
print(f"{len(osm)} building footprints entirely inside the patch; area median {osm['area_m2'].median():.0f} m², "
      f"range {osm['area_m2'].min():.0f}-{osm['area_m2'].max():,.0f} m²")
print(osm["building"].value_counts().head(5).to_string())


def iou(a, b) -> float:
    inter = a.intersection(b).area
    return inter / (a.area + b.area - inter) if inter > 0 else 0.0


def best_iou(polys: gpd.GeoSeries, ref: gpd.GeoSeries) -> np.ndarray:
    """For each polygon in ref, the best intersection-over-union with any polygon in polys."""
    out = np.zeros(len(ref))
    sidx = polys.sindex
    for k, g in enumerate(ref):
        cand = sidx.query(g, predicate="intersects")
        if len(cand):
            out[k] = max(iou(g, polys.iloc[c]) for c in cand)
    return out


# --- Step 3: automatic mask generation with SAM ViT-B ------------------------------------------
section("automatic mask generation (SAM ViT-B, CPU)")
masks_tif, masks_gpkg = OUT / "sam_masks.tif", OUT / "sam_masks.gpkg"
t0 = time.time()
sam = SamGeo(model_type="vit_b", automatic=True, device="cpu", checkpoint_dir=str(CACHE),
             sam_kwargs={"points_per_side": 32,            # a 32 x 32 grid of point prompts = 1,024 prompts
                         "pred_iou_thresh": 0.86,          # keep masks the model rates as good...
                         "stability_score_thresh": 0.90,   # ...and that do not change much with the threshold
                         "min_mask_region_area": 40})      # remove holes and islands under 40 pixels (10 m²)
ckpt = next(CACHE.glob("sam_vit_b*.pth"))
print(f"model loaded in {time.time() - t0:.1f} s: {ckpt.name}, {ckpt.stat().st_size / 1e6:.0f} MB, "
      f"{sum(p.numel() for p in sam.sam.parameters()) / 1e6:.0f} M parameters")
t0 = time.time()
sam.generate(str(rgb_tif), output=str(masks_tif), foreground=True, unique=True)
t_auto = time.time() - t0
print(f"generate(): {len(sam.masks)} masks in {t_auto:.0f} s")

sam.tiff_to_gpkg(str(masks_tif), str(masks_gpkg), simplify_tolerance=None)
segs = gpd.read_file(masks_gpkg).dissolve(by="value", as_index=False)   # one (multi)polygon per mask
segs = segs.set_crs(CRS, allow_override=True)
with rasterio.open(masks_tif) as src:
    mask_ids = src.read(1)
covered = mask_ids > 0
print(f"vectorised to {len(segs)} polygons in {masks_gpkg.name}; {covered.mean():.0%} of the patch is inside a mask")

# Shape and spectral attributes for every segment
segs["area_m2"] = segs.area
segs["compactness"] = 4 * np.pi * segs.area / segs.length ** 2         # 1 for a circle, ~0.785 for a square
mrr = segs.minimum_rotated_rectangle()
segs["rectangularity"] = segs.area / mrr.area                          # 1 for a perfect rectangle
edges = mrr.apply(lambda g: sorted({round(np.hypot(*np.diff(np.array(g.exterior.coords[i:i + 2]), axis=0)[0]), 2)
                                    for i in range(2)}))
segs["elongation"] = [e[-1] / max(e[0], 0.1) for e in edges]
segs["ndvi"] = ndimage.mean(ndvi, mask_ids, segs["value"].to_numpy())
print(segs[["area_m2", "compactness", "rectangularity", "elongation", "ndvi"]].describe(
    percentiles=[0.1, 0.5, 0.9]).loc[["count", "10%", "50%", "90%"]].round(2).to_string())

# --- Step 4: filter the masks by area, shape and NDVI to find buildings -------------------------
section("filter: building-like segments")
RULES = {"40 m² <= area <= 12,000 m²": segs["area_m2"].between(40, 12000),
         "NDVI < 0.15 (not vegetation)": segs["ndvi"] < 0.15,
         "rectangularity >= 0.75": segs["rectangularity"] >= 0.75,
         "elongation <= 4 (not a road)": segs["elongation"] <= 4}
keep = pd.Series(True, index=segs.index)
for rule, ok in RULES.items():
    keep &= ok
    print(f"  + {rule:<30} -> {keep.sum():4d} segments left")
bld = segs[keep].copy()
bld["best_iou"] = [max((iou(g, o) for o in osm.geometry[osm.intersects(g)]), default=0.0) for g in bld.geometry]
bld.to_file(OUT / "sam_buildings.gpkg", driver="GPKG")
osm["auto_iou_any"] = best_iou(segs.geometry, osm.geometry)              # did SAM find it at all?
osm["auto_iou_kept"] = best_iou(bld.geometry, osm.geometry)              # ...and did the filter keep it?
tp = (bld["best_iou"] >= 0.5).sum()
print(f"precision: {tp} of {len(bld)} kept segments match an OSM building with IoU >= 0.5 ({tp / len(bld):.0%})")
print(f"recall before filtering: {(osm['auto_iou_any'] >= 0.5).sum()} of {len(osm)} OSM buildings have a SAM mask with "
      f"IoU >= 0.5 ({(osm['auto_iou_any'] >= 0.5).mean():.0%})")
print(f"recall after filtering:  {(osm['auto_iou_kept'] >= 0.5).sum()} of {len(osm)} "
      f"({(osm['auto_iou_kept'] >= 0.5).mean():.0%}); median IoU of matched buildings "
      f"{osm.loc[osm['auto_iou_kept'] >= 0.5, 'auto_iou_kept'].median():.2f}")
small = osm["area_m2"] < 100
print(f"recall after filtering, buildings < 100 m² (garages, sheds): {(osm.loc[small, 'auto_iou_kept'] >= 0.5).mean():.0%} "
      f"of {small.sum()}; >= 100 m²: {(osm.loc[~small, 'auto_iou_kept'] >= 0.5).mean():.0%} of {(~small).sum()}")

# --- Step 5: point prompts ---------------------------------------------------------------------
section("point prompts")
t0 = time.time()
sp = SamGeo(model_type="vit_b", automatic=False, device="cpu", checkpoint_dir=str(CACHE))
sp.set_image(str(rgb_tif))                       # runs the image encoder once; every prompt after this is fast
print(f"set_image(): image embedding computed in {time.time() - t0:.1f} s (model load included)")
POINTS = {"school roof": (-93.685575, 42.015863), "playing field": (-93.685261, 42.014694),
          "playground": (-93.684486, 42.015689), "apartment block": (-93.682507, 42.015250)}
to_px = ~TRANSFORM
to_utm = Transformer.from_crs(4326, CRS, always_xy=True)
point_results = {}
for name, lonlat in POINTS.items():
    t0 = time.time()
    m, scores, _ = sp.predict(point_coords=[list(lonlat)], point_labels=1, point_crs="EPSG:4326",
                              multimask_output=True, return_results=True)
    dt = time.time() - t0
    areas = m.reshape(3, -1).sum(axis=1) * RES * RES
    point_results[name] = (m, scores, to_px * to_utm.transform(*lonlat))
    print(f"{name:<16} {dt * 1000:4.0f} ms | 3 candidate masks: " +
          ", ".join(f"{a:,.0f} m² (score {s:.2f})" for a, s in zip(areas, scores)))
sp.predict(point_coords=[list(POINTS["school roof"])], point_labels=1, point_crs="EPSG:4326",
           multimask_output=True, output=str(OUT / "school_point_mask.tif"))   # best mask, as a GeoTIFF
print("best mask for the school point written to school_point_mask.tif")

# --- Step 6: box prompts from OpenStreetMap ---------------------------------------------------
section("box prompts from OpenStreetMap footprints")
PAD = 2.0                                       # metres of margin around each footprint's bounding box
osm_ids = rasterio.features.rasterize(zip(osm.geometry, osm.index + 1), out_shape=(SIZE, SIZE), transform=TRANSFORM)
box_masks = np.zeros((SIZE, SIZE), dtype="int32")
rows = []
t0 = time.time()
for k, g in enumerate(osm.geometry):
    minx, miny, maxx, maxy = g.bounds
    c0, r0 = to_px * (minx - PAD, maxy + PAD)
    c1, r1 = to_px * (maxx + PAD, miny - PAD)
    m, score, _ = sp.predictor.predict(box=np.array([c0, r0, c1, r1]), multimask_output=False)
    m = m[0]
    ref = osm_ids == k + 1
    inter, union = np.logical_and(m, ref).sum(), np.logical_or(m, ref).sum()
    rows.append({"osm_id": osm.loc[k, "id"], "area_m2": g.area, "iou": inter / union, "score": float(score[0])})
    box_masks[m & (box_masks == 0)] = k + 1
t_box = time.time() - t0
res = pd.DataFrame(rows)
print(f"{len(res)} box prompts in {t_box:.1f} s ({t_box / len(res) * 1000:.0f} ms per box)")
print(f"IoU with the OSM footprint: median {res['iou'].median():.2f}, mean {res['iou'].mean():.2f}; "
      f">= 0.5: {(res['iou'] >= 0.5).mean():.0%}, >= 0.75: {(res['iou'] >= 0.75).mean():.0%}")
print(res.groupby(pd.cut(res["area_m2"], [0, 50, 100, 300, 1000, 1e5], labels=["<50", "50-100", "100-300",
                                                                               "300-1000", ">1000"]),
                  observed=True)["iou"].agg(buildings="size", median_iou="median").round(2).to_string())
shapes = [(shape(geom), int(v)) for geom, v in rasterio.features.shapes(box_masks, box_masks > 0, transform=TRANSFORM)]
box_gdf = gpd.GeoDataFrame({"k": [v for _, v in shapes]}, geometry=[g for g, _ in shapes], crs=CRS).dissolve(
    "k", as_index=False)
box_gdf["geometry"] = box_gdf.simplify(0.5)                  # remove the pixel staircase (0.5 m = one pixel)
box_gdf = box_gdf.merge(res.reset_index(names="i").assign(k=lambda d: d["i"] + 1), on="k").drop(columns=["i", "k"])
box_gdf.to_file(OUT / "sam_box_buildings.gpkg", driver="GPKG")
print(f"wrote {len(box_gdf)} traced building polygons to sam_box_buildings.gpkg")

# --- Figures ---------------------------------------------------------------------------------
def to_plot(gs):
    """Shift geometries so the patch's lower-left corner is (0, 0), in metres."""
    return gs.translate(-x0, -(y1 - SIZE * RES))


# Figure 1: the image, colour-infrared, and the OSM buildings
cir = np.clip(np.dstack([rgbn[3], rgbn[0], rgbn[1]]) / 255 * 1.1, 0, 1)
fig, axes = plt.subplots(1, 2, figsize=(11, 5.6), sharey=True)
axes[0].imshow(rgb, extent=extent)
to_plot(osm.geometry).boundary.plot(ax=axes[0], color="#fee08b", lw=0.8)
axes[0].set_title(f"NAIP {item.datetime:%d %b %Y}, true colour, with {len(osm)} OSM buildings")
axes[1].imshow(cir, extent=extent)
axes[1].set_title("Colour infrared (NIR, red, green): vegetation is red")
for ax in axes:
    ax.set_xlabel("metres", fontsize=10)
axes[0].set_ylabel("metres", fontsize=10)
fig.tight_layout()
fig.savefig(OUT / "01-naip-ames-osm-buildings.png")
plt.close(fig)

# Figure 2: every automatic mask in a random colour
rng = np.random.default_rng(3)
lut = np.vstack([[0, 0, 0, 0], np.c_[rng.uniform(0.15, 1, (mask_ids.max(), 3)), np.full(mask_ids.max(), 0.55)]])
fig, axes = plt.subplots(1, 2, figsize=(11, 5.6), sharey=True)
axes[0].imshow(rgb, extent=extent)
axes[0].imshow(lut[mask_ids], extent=extent)
axes[0].set_title(f"Automatic: {len(sam.masks)} masks in {t_auto:.0f} s on CPU")
axes[1].imshow(rgb, extent=extent)
to_plot(segs.geometry).boundary.plot(ax=axes[1], color="white", lw=0.5)
axes[1].set_title(f"Vectorised: {len(segs)} polygons in sam_masks.gpkg")
for ax in axes:
    ax.set_xlabel("metres", fontsize=10)
axes[0].set_ylabel("metres", fontsize=10)
axes[1].set_ylabel("")
fig.tight_layout()
fig.savefig(OUT / "02-automatic-masks.png")
plt.close(fig)

# Figure 3: area versus rectangularity, coloured by NDVI, with the filter
fig, ax = plt.subplots(figsize=(7.5, 4.6))
sc = ax.scatter(segs["area_m2"], segs["rectangularity"], c=segs["ndvi"], cmap="Spectral_r", vmin=-0.2, vmax=0.6,
                s=np.where(keep, 26, 12), edgecolor=np.where(keep, "black", "none"), lw=0.6)
ax.axvspan(40, 12000, ymin=(0.75 - 0.2) / 0.82, ymax=1, color="#3288bd", alpha=0.08)
ax.set_xscale("log")
ax.set_ylim(0.2, 1.02)
ax.set_xlabel("Segment area (m², log scale)")
ax.set_ylabel("Rectangularity (area / minimum rotated rectangle)")
ax.set_title(f"{len(segs)} SAM segments; {keep.sum()} pass all four rules (black outline)")
fig.colorbar(sc, ax=ax, label="Mean NDVI")
fig.tight_layout()
fig.savefig(OUT / "03-segment-area-shape-ndvi.png")
plt.close(fig)

# Figure 4: filtered segments against OSM
fig, ax = plt.subplots(figsize=(7.2, 7.2))
ax.imshow(rgb, extent=extent, alpha=0.75)
to_plot(bld[bld["best_iou"] >= 0.5].geometry).plot(ax=ax, color="#66c2a5", edgecolor="#1a5e4a", lw=0.6)
to_plot(bld[bld["best_iou"] < 0.5].geometry).plot(ax=ax, color="#f46d43", edgecolor="#8c2d04", lw=0.6)
to_plot(osm[osm["auto_iou_kept"] < 0.5].geometry).boundary.plot(ax=ax, color="#fee08b", lw=1.2)
ax.legend(handles=[Patch(color="#66c2a5", label=f"Kept, matches OSM (IoU >= 0.5): {tp}"),
                   Patch(color="#f46d43", label=f"Kept, no OSM match: {len(bld) - tp}"),
                   Line2D([], [], color="#fee08b", lw=2,
                          label=f"OSM building missed: {(osm['auto_iou_kept'] < 0.5).sum()}")],
          loc="lower left", fontsize=8, framealpha=0.9)
ax.set_title("Building-like SAM segments after filtering, against OpenStreetMap")
ax.set_xlabel("metres", fontsize=10)
ax.set_ylabel("metres", fontsize=10)
fig.tight_layout()
fig.savefig(OUT / "04-filtered-buildings-vs-osm.png")
plt.close(fig)

# Figure 5: point prompts, three candidate masks each
cand_colours = ["#3288bd", "#fdae61", "#d53e4f"]
fig, axes = plt.subplots(1, 4, figsize=(13, 3.9))
for ax, (name, (m, scores, (pc, pr))) in zip(axes, point_results.items()):
    half = 110
    r0_, c0_ = int(np.clip(pr - half, 0, SIZE - 2 * half)), int(np.clip(pc - half, 0, SIZE - 2 * half))
    win = (slice(r0_, r0_ + 2 * half), slice(c0_, c0_ + 2 * half))
    ax.imshow(rgb[win])
    for i in np.argsort(scores):
        ax.contour(m[i][win].astype(float), levels=[0.5], colors=[cand_colours[i]], linewidths=1.6)
    ax.plot(pc - c0_, pr - r0_, marker="*", ms=14, color="white", mec="black")
    ax.set_title(name)
    ax.legend(handles=[Line2D([], [], color=cand_colours[i], lw=2, label=f"mask {i + 1}: score {scores[i]:.2f}")
                       for i in range(3)], loc="lower left", fontsize=7, framealpha=0.9)
    ax.set_xticks([])
    ax.set_yticks([])
fig.suptitle("Point prompts: one click (star), three candidate masks (110 x 110 m windows)", y=1.0)
fig.tight_layout()
fig.savefig(OUT / "05-point-prompts.png")
plt.close(fig)

# Figure 6: box prompts against OSM, map and IoU histogram
fig, axes = plt.subplots(1, 2, figsize=(11.5, 5.2), gridspec_kw={"width_ratios": [1.15, 1]})
zoom = (slice(0, 300), slice(0, 420))
axes[0].imshow(rgb[zoom], extent=(0, 420 * RES, 0, 300 * RES))
shift = lambda gs: gs.translate(-x0, -(y1 - 300 * RES))
shift(box_gdf.geometry).boundary.plot(ax=axes[0], color="#d53e4f", lw=1.2)
shift(osm.geometry).boundary.plot(ax=axes[0], color="#fee08b", lw=1.0, linestyle="--")
axes[0].set_xlim(0, 420 * RES)
axes[0].set_ylim(0, 300 * RES)
axes[0].legend(handles=[Line2D([], [], color="#d53e4f", lw=2, label="SAM, prompted with the OSM bounding box"),
                        Line2D([], [], color="#fee08b", lw=2, ls="--", label="OSM footprint")],
               loc="lower left", fontsize=8, framealpha=0.9)
axes[0].set_title("Box prompts: north-west corner of the patch")
axes[0].set_xlabel("metres", fontsize=10)
axes[0].set_ylabel("metres", fontsize=10)
axes[1].hist(res["iou"], bins=np.linspace(0, 1, 21), color="#3288bd", edgecolor="white")
axes[1].axvline(res["iou"].median(), color="black", ls="--", lw=1)
axes[1].annotate(f"median {res['iou'].median():.2f}", xy=(res["iou"].median(), axes[1].get_ylim()[1] * 0.6),
                 xytext=(0.4, axes[1].get_ylim()[1] * 0.6), ha="right", va="center", fontsize=9,
                 arrowprops=dict(arrowstyle="->", lw=0.8))
axes[1].set_xlabel("IoU between SAM mask and OSM footprint")
axes[1].set_ylabel("Buildings")
axes[1].set_title(f"{len(res)} buildings, {t_box / len(res) * 1000:.0f} ms per box on CPU")
fig.tight_layout()
fig.savefig(OUT / "06-box-prompts-iou.png")
plt.close(fig)
print("\nfigures and outputs written to output/")