Detecting field boundaries with RasterFlow

import wkls
import geopandas as gpd
import os

# Generate a geometry for Haskell County, Kansas using WKLS (https://github.com/wherobots/wkls)
gdf = gpd.read_file(wkls['us']['ks']['Haskell County'].geojson())

# Save the geometry to a parquet file in the user's S3 path
aoi_path = os.getenv("USER_S3_PATH") + "haskell.parquet"
gdf.to_parquet(aoi_path)

# Make variables for the bounds of our aoi for visualizing results after inference
min_lon, min_lat, max_lon, max_lat = gdf.total_bounds

from datetime import datetime

from rasterflow_remote import RasterflowClient

from rasterflow_remote.data_models import (
    ModelRecipes, 
    VectorizeMethodEnum
)

rf_client = RasterflowClient()

model_outputs = rf_client.build_and_predict_mosaic_recipe(
    # Path to our AOI in GeoParquet or GeoJSON format
    aoi = aoi_path,

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

    # Coordinate Reference System EPSG code for the output mosaic   
    crs_epsg = 3857,

    # The model recipe to be used for inference (FTW in this case)
    model_recipe = ModelRecipes.FTW,
)

print(model_outputs)

# Import libraries for visualization and coordinate transformation
import hvplot.xarray
import xarray as xr
import s3fs 
import zarr
from pyproj import Transformer
from holoviews.element.tiles import EsriImagery 

# Open the Zarr store
fs = s3fs.S3FileSystem(profile="default", asynchronous=True)
zstore = zarr.storage.FsspecStore(fs, path=model_outputs[5:])
ds = xr.open_zarr(zstore)

# Create a transformer to convert from lat/lon to meters
transformer = Transformer.from_crs("EPSG:4326", "EPSG:3857", always_xy=True)

# Transform bounding box coordinates from lat/lon to meters
(min_x, max_x), (min_y, max_y) = transformer.transform(
    [min_lon, max_lon], 
    [min_lat, max_lat]
)

# Select the height variable and slice the dataset to the bounding box
# y=slice(max_y, min_y) handles the standard "North-to-South" image orientation
ds_subset = ds.sel(band="field_boundaries",
    x=slice(min_x, max_x), 
    y=slice(max_y, min_y) 
)

# Select the first time step and extract the variables array
arr_subset = ds_subset.isel(time=0)["variables"]

# Create a base map layer using Esri satellite imagery
base_map = EsriImagery()

# Create an overlay layer from the model outputs with hvplot
output_layer = arr_subset.hvplot(
    x = "x",
    y = "y",
    geo = True,           # Enable geographic plotting
    dynamic = True,       # Enable dynamic rendering for interactivity
    rasterize = True,     # Use datashader for efficient rendering of large datasets
    cmap = "viridis",     # Color map for visualization
    aspect = "equal",     # Maintain equal aspect ratio
    title = "FTW Model Outputs" 
).opts(
    width = 600, 
    height = 600,
    alpha = 0.7           # Set transparency to see the base map underneath
)

# Combine the base map and output layer
final_plot = base_map * output_layer
final_plot

# Determine the classes that are predicted by the model
model_features = xr.open_zarr(zstore)['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 = rf_client.vectorize_mosaic(
        store = model_outputs,
        features = vector_features,
        threshold = 0.5,
        vectorize_method = VectorizeMethodEnum.SEMANTIC_SEGMENTATION_RASTERIO,
        vectorize_config={"stats": True, "medial_skeletonize": False}
    )

print(vectorized_results)

from sedona.spark import *
config = SedonaContext.builder().getOrCreate()
sedona = SedonaContext.create(config)

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

vectorize_result_df = sedona.read.format("geoparquet").load(vectorized_results)
vectorize_result_df = vectorize_result_df.withColumnRenamed("label", "layer")
vectorize_result_df.writeTo("org_catalog.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 sedona.spark.maps.SedonaKepler import SedonaKepler
map = SedonaKepler.create_map(df=df, name="Vectorized results")
map

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 wherobots.tools.utility.s3_utils import get_url

get_url(full_tiles_path)