Detecting field boundaries 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 Haskell County, Kansas using Well-Known Locations (https://github.com/wherobots/wkls)
aoi = wkls.us.ks.haskellcounty.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}]")

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_recipe",
    {
        # Area of Interest as WKT (EPSG:4326)
        "aoi": aoi,

        # Date range for imagery to be used by the model, as RFC 3339 timestamps
        "start": "2023-01-01T00:00:00Z",
        "end": "2024-01-01T00:00:00Z",

        # Coordinate Reference System for the output mosaic
        "target_crs": "EPSG:3857",

        # The model recipe to be used for inference (FTW in this case)
        "model_recipe": "ftw",
    },
)

# The recipe returns both the input imagery and the prediction mosaic.
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}")

model_output_index_uri = recipe_output["result"]["mosaic_index_uri"]

# Inspect the prediction index (one row per output mosaic location)
model_output_index_df = sedona.read.format("geoparquet").load(model_output_index_uri)
model_output_index_df.show(truncate=60)

# This example AOI has one geometry, so we select the first (only) mosaic location.
model_output_store = model_output_index_df.first()["location"]

model_output_store

# optimized_store = run_task("build_zarr_multiscales", {"source_store": model_output_store, "aoi": aoi})
# optimized_store

import xarray as xr
import s3fs
import zarr
# Determine the classes that are predicted by the model
fs = s3fs.S3FileSystem(profile="default", asynchronous=True)
zstore = zarr.storage.FsspecStore(fs, path=model_output_store[5:])
ds = xr.open_zarr(zstore)
model_features = ds['band'].data.tolist()

# Remove the 'non_field_background' class for vectorization
vector_features = [f for f in model_features if f != 'non_field_background']    

# Note: this should take about 5 minutes to complete
vectorized_results_uri = run_task(
    "vectorize_mosaic",
    {
        "mosaic_uri": model_output_store,
        "features": vector_features,
        "threshold": 0.5,
        "stats": True,
        "medial_skeletonize": False,
    },
)

print(vectorized_results_uri)

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

vectorize_result_df = sedona.read.format("geoparquet").load(vectorized_results_uri)
vectorize_result_df = vectorize_result_df.withColumnRenamed("label", "layer")
vectorize_result_df.writeTo("examples_temp.ftw_db.ftw_vectorized").createOrReplace()

from sedona.spark.sql.st_constructors import ST_GeomFromText
from sedona.spark.sql.st_predicates import ST_Intersects
import pyspark.sql.functions as f

# Bounding box for Plymell, Kansas
plymell_wkt = "POLYGON ((-100.98 37.83, -100.98 37.68, -100.8 37.68, -100.8 37.83, -100.98 37.83))"
vectorize_result_filtered_df = vectorize_result_df.filter("layer == 'field'").filter("score_mean > 0.5") \
    .filter(ST_Intersects(f.col("geometry"), ST_GeomFromText(lit(plymell_wkt))))

tolerance = 0.00001
decimal_places = 6  # roughly 10cm in lat/lon
df = vectorize_result_filtered_df.withColumn("geometry", f.expr("ST_MakeValid(geometry)"))
df = df.withColumn("geometry", f.expr("ST_SetSRID(geometry, 4326)")) # vectorize_mosaic defaults to returning geometries in the CRS EPSG:4326
df = df.withColumn("geometry", f.expr(f"ST_ReducePrecision(geometry, {decimal_places})"))
df = df.withColumn("geometry", f.expr(f"ST_SimplifyPreserveTopology(geometry, {tolerance})"))
df = df.repartition(200)
df.show()

from wherobots_gl import Map

# Wherobots-GL Map is URL-based, so write the post-processed df to GeoParquet first.
vectorized_viz_path = os.getenv("USER_S3_PATH") + "ftw_vectorized_viz.parquet"
df.write.format("geoparquet").mode("overwrite").save(vectorized_viz_path)

Map(layers=[{"type": "geoparquet", "source": vectorized_viz_path, "name": "Vectorized results"}])

fields_df = df.where("layer = 'field'")
# 1266 is the vertex count of the second-largest boundary geometry.
# We use this instead of the absolute largest boundary (an extreme outlier)
# so the subdivision threshold is high enough to keep most geometries intact
# while still forcing that outlier-sized boundary to be subdivided for
# better PMTiles rendering performance.
boundaries_df = df.where("layer = 'field_boundaries'").withColumn("geometry", f.expr("ST_SubDivideExplode(geometry, 1266)"))
tile_features_df = fields_df.withColumn("layer", f.lit("fields")).unionByName(
    boundaries_df.withColumn("layer", f.lit("boundaries"))
)

from wherobots import vtiles
full_tiles_path = os.getenv("USER_S3_PATH") + "tiles.pmtiles"
vtiles.generate_pmtiles(tile_features_df, full_tiles_path)

vtiles.show_pmtiles(full_tiles_path)

from urllib.parse import quote

print(f"View on the map: https://cloud.wherobots.com/map?l0.src={quote(full_tiles_path, safe='')}")