Detecting Objects From Text Prompts with RasterFlow

import os

import shapely
import wkls
from sedona.spark import *

config = SedonaContext.builder().getOrCreate()
sedona = SedonaContext.create(config)

# Generate a WKT geometry for College Park, Maryland using Well-Known Locations (https://github.com/wherobots/wkls)
aoi = wkls.us.md.collegepark.wkt()

min_lon, min_lat, max_lon, max_lat = shapely.from_wkt(aoi).bounds
print(f"Bounds: [{min_lon:.4f}, {min_lat:.4f}, {max_lon:.4f}, {max_lat:.4f}]")

from datetime import datetime

# Date range for imagery to be used by the model
start_date = datetime(2023, 1, 1)
end_date = datetime(2024, 1, 1)

# The SAM3 recipes run on 30cm NAIP imagery from the RGB (visual) product. Wherobots Open Data
# indexes every scene of that product with its footprint, resolution (gsd) and acquisition time.
MODEL_RES = 0.3
sedona.createDataFrame([(aoi,)], ["wkt"]).selectExpr("ST_GeomFromWKT(wkt) AS geom").createOrReplaceTempView("aoi")

# Scenes at the model's resolution touching the AOI (the table and wkls are both EPSG:4326)
sedona.sql(f"""
    SELECT n.geometry, n.year, n.datetime
    FROM wherobots_open_data.imagery.usda_naip_visual_stac n
    JOIN aoi a ON ST_Intersects(n.geometry, a.geom)
    WHERE n.gsd = {MODEL_RES}
""").cache().createOrReplaceTempView("nearby")

# The recipe needs the AOI fully inside the union of the scene footprints in the date range
coverage = sedona.sql(f"""
    SELECT count(*) AS scenes, array_sort(collect_set(n.year)) AS years,
           ST_Contains(ST_Union_Agg(n.geometry), first(a.geom)) AS covered
    FROM nearby n CROSS JOIN aoi a
    WHERE n.datetime >= TIMESTAMP '{start_date.isoformat()}Z' AND n.datetime < TIMESTAMP '{end_date.isoformat()}Z'
""").first()

if not coverage["covered"]:
    # Years that would work, applying the same full-coverage test to each year on its own
    available = [
        row.year
        for row in sedona.sql("""
            SELECT n.year FROM nearby n CROSS JOIN aoi a
            GROUP BY n.year
            HAVING ST_Contains(ST_Union_Agg(n.geometry), first(a.geom))
            ORDER BY n.year
        """).collect()
    ]
    raise ValueError(
        f"{MODEL_RES:g}m NAIP does not fully cover this AOI between "
        f"{start_date:%Y-%m-%d} and {end_date:%Y-%m-%d}. "
        + (
            f"Years with complete {MODEL_RES:g}m coverage over this AOI: {available}."
            if available
            else f"No year has complete {MODEL_RES:g}m NAIP coverage over this AOI."
        )
    )

print(f"NAIP coverage is available. Found {coverage['scenes']} NAIP scenes at {MODEL_RES:g}m covering the AOI")
print(f"Acquisition years: {coverage['years']}")

import time

from wherobots import TaskRunState, Wherobots

# In a Wherobots notebook the SDK uses your Wherobots credentials, elsewhere it reads WHEROBOTS_API_KEY
wb = Wherobots()


def run_task(task_name, inputs, poll_seconds=30):
    """Run a RasterFlow task, wait for it to finish, and return its output value."""
    run = wb.task_runs.get(wb.tasks.run(task_name, inputs, region="aws-us-west-2").task_run_id)
    print(f"Started {task_name} run {run.task_run_id}")
    while run.is_active:
        time.sleep(poll_seconds)
        run = wb.task_runs.get(run.task_run_id)
    if run.state != TaskRunState.SUCCEEDED:
        raise RuntimeError(f"{task_name} run {run.task_run_id} ended {run.state}: {run.error_message}")
    # A successful run can temporarily have no results if the output read fails.
    for _ in range(3):
        if run.results is not None:
            break
        time.sleep(poll_seconds)
        run = wb.task_runs.get(run.task_run_id)
    if run.results is None:
        raise RuntimeError(
            f"Results for successful run {run.task_run_id} are unavailable. "
            "Fetch this run again with wb.task_runs.get(run_id); do not resubmit the task."
        )
    # The single task output can be a URI or a structured recipe result.
    (output,) = run.results.values()
    return output

recipe_output = run_task(
    "predict_mosaic_geometries_recipe",
    {
        # Area of Interest as WKT (EPSG:4326)
        "aoi": aoi,

        # Date range for imagery to be used by the model (set in the coverage check above), as RFC 3339 timestamps
        "start": start_date.isoformat() + "Z",
        "end": end_date.isoformat() + "Z",

        # Coordinate Reference System for the imagery mosaic (output geometries use EPSG:4326)
        "target_crs": "EPSG:3857",

        # The model recipe and text prompt for object detection
        "model_recipe": "sam3-text-geometry",
        # You can also pass multiple prompts to detect several object types at once,
        # e.g. "text_prompt": ["roofs", "roads"]
        "text_prompt": ["roofs"],
        # Lower confidence thresholds surface more (but noisier) detections; you may
        # want to adjust this to trade off recall against precision for your AOI.
        "confidence_threshold": 0.3,
    },
)

# The recipe returns the input imagery and a nullable detections URI.
input_mosaic_index_uri = recipe_output["input_mosaic"]["mosaic_index_uri"]
input_mosaic_store = recipe_output["input_mosaic"]["first_row_mosaic"]
print(f"Input imagery: {input_mosaic_store}")

detections_uri = recipe_output["result"]
if detections_uri is None:
    raise ValueError(
        "No geometries were detected. Try another prompt, AOI, or confidence threshold."
    )

detections_df = sedona.read.format("geoparquet").load(detections_uri)
detections_df.show(5, truncate=60)

print(f"Total detections: {detections_df.count()}")
print(f"Columns: {detections_df.columns}")
print("\nConfidence score stats:")
detections_df.select("bbox_score").summary().show()

sedona.sql("CREATE DATABASE IF NOT EXISTS examples_temp.sam3_db")

df = sedona.read.format("geoparquet").load(detections_uri)
df = df.withColumnRenamed("label", "layer")
df.writeTo("examples_temp.sam3_db.sam3_roofs").createOrReplace()

import pyspark.sql.functions as f

# Geodesic area on the WGS84 spheroid, in square metres
df_filtered = df.withColumn("area_m2", ST_AreaSpheroid("geometry")).filter(f.col("area_m2") > 10)

df_filtered.show()

from wherobots_gl import Map

# Write the filtered detections to GeoParquet so Wherobots-GL Map can load them by URL
filtered_path = os.getenv("USER_S3_PATH") + "sam3_roofs_filtered.parquet"
df_filtered.write.format("geoparquet").mode("overwrite").save(filtered_path)

Map(layers=[{"type": "geoparquet", "source": filtered_path, "name": "SAM3 roof detections"}])

from functools import reduce


def join_intersecting_detections(detections):
    # One label at a time: a partition column would bypass the spatial index
    labels = [row.layer for row in detections.select("layer").distinct().collect()]
    joined_by_label = [
        detections.filter(f.col("layer") == label)
        .selectExpr("layer", "geometry", "ST_IntersectsCC(geometry) AS component_id")
        .groupBy("layer", "component_id")
        .agg(
            ST_Union_Aggr("geometry").alias("geometry"),
            f.count("*").alias("fragment_count"),
        )
        for label in labels
    ]
    return reduce(lambda left, right: left.unionByName(right), joined_by_label)

fixture = sedona.createDataFrame([
    ("roofs", "POLYGON ((0 0, 1 0, 1 1, 0 1, 0 0))"),
    ("roofs", "POLYGON ((1 0, 2 0, 2 1, 1 1, 1 0))"),
    ("roofs", "POLYGON ((3 0, 4 0, 4 1, 3 1, 3 0))"),
    ("pools", "POLYGON ((1.5 0, 3.5 0, 3.5 1, 1.5 1, 1.5 0))"),
], ["layer", "wkt"]).selectExpr("layer", "ST_GeomFromWKT(wkt) AS geometry")

fixture_joined = join_intersecting_detections(fixture)
fixture_results = fixture_joined.selectExpr(
    "layer", "fragment_count", "ST_Area(geometry) AS area"
).collect()
fixture_values = sorted((row.layer, row.fragment_count, row.area) for row in fixture_results)
assert fixture_values == [("pools", 1, 2.0), ("roofs", 1, 1.0), ("roofs", 2, 2.0)], fixture_values
fixture_joined.selectExpr("layer", "fragment_count", "ST_AsText(geometry) AS geometry").show(truncate=False)

joined = join_intersecting_detections(df).cache()
print(f"Raw detections: {df.count()}; joined components: {joined.count()}")
joined.select("layer", "fragment_count", "geometry").show(truncate=False)

raw_path = os.getenv("USER_S3_PATH") + "sam3_roofs_raw.parquet"
joined_path = os.getenv("USER_S3_PATH") + "sam3_roofs_joined.parquet"
df.write.format("geoparquet").mode("overwrite").save(raw_path)
joined.write.format("geoparquet").mode("overwrite").save(joined_path)
joined.unpersist()

Map(layers=[
    {"type": "geoparquet", "source": raw_path, "name": "Raw SAM3 detections"},
    {"type": "geoparquet", "source": joined_path, "name": "Joined SAM3 detections"},
])

from wherobots import vtiles

full_tiles_path = os.getenv("USER_S3_PATH") + "sam3_roofs_tiles.pmtiles"
vtiles.generate_pmtiles(df_filtered, full_tiles_path)

vtiles.show_pmtiles(full_tiles_path)