Detecting roads with RasterFlow

import wkls
import geopandas as gpd
import os

# Generate a geometry for Montgomery County, Maryland using Well-Known Locations (https://github.com/wherobots/wkls)
gdf = gpd.read_file(wkls.us.md.montgomerycounty.geojson())

gdf

from datetime import datetime

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

# The ChesapeakeRSC recipe runs on 1m NAIP imagery. The index holds one row per NAIP scene, with its
# footprint, resolution (res) and acquisition time.
MODEL_RES = 1.0
naip_index = gpd.read_parquet(
    "s3://wherobots-examples/rasterflow/indexes/naip_index.parquet",
    columns=["geometry", "res", "year", "time"],
)

# Scenes at the model's resolution touching the AOI, compared in the index's CRS
aoi_geom = gdf.geometry.to_crs(naip_index.crs).union_all()
at_res = naip_index.query(f"res == {MODEL_RES}")
nearby = at_res.iloc[at_res.sindex.query(aoi_geom, predicate="intersects")]

# Narrowed to the requested date range
covering = nearby.query(f"time >= '{start_date:%Y-%m-%d}' and time <= '{end_date:%Y-%m-%d}'")

# The recipe needs the AOI fully inside the matching scenes
if not covering.geometry.union_all().contains(aoi_geom):
    # How much of the AOI is covered, measured in an equal-area CRS
    area_crs = gdf.estimate_utm_crs()
    aoi_area = gdf.geometry.to_crs(area_crs).union_all()
    tiles = covering.geometry.to_crs(area_crs).union_all()
    fraction = max(0.0, 1.0 - aoi_area.difference(tiles).area / aoi_area.area)

    # Years that would work, applying the same full-coverage test as the gate above
    available = sorted(
        year
        for year, scenes in nearby.groupby("year")
        if scenes.geometry.union_all().contains(aoi_geom)
    )
    raise ValueError(
        f"{MODEL_RES:g}m NAIP covers {fraction:.1%} of this AOI between "
        f"{start_date:%Y-%m-%d} and {end_date:%Y-%m-%d}; the recipe needs the AOI fully covered. "
        + (
            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 {len(covering)} NAIP scenes at {MODEL_RES:g}m covering the AOI")
print(f"Acquisition years: {sorted(covering['year'].unique().tolist())}")

from rasterflow_remote import RasterflowClient

from rasterflow_remote.data_models import (
    ModelRecipes, 
    VectorizeMethodEnum
)

rf_client = RasterflowClient()

model_output_index = rf_client.predict_mosaic_recipe(
    # Area of Interest (AOI)
    aoi = gdf,

    # Date range for imagery to be used by the model (set in the coverage check above)
    start = start_date,
    end = end_date,

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

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

model_output_index.mosaic_index_gdf

# This example AOI has one geometry, so we select the first (only) mosaic location.
model_output_store = model_output_index.first_row_mosaic

model_output_store

# optimized_store = rf_client.build_zarr_multiscales(source_store=model_output_store)
# 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 'background' class for vectorization
vector_features = [f for f in model_features if f != 'background']  

# Note: this should take about 5 minutes to complete
vectorized_results = rf_client.vectorize_mosaic(
        mosaic = model_output_store,
        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 examples_temp.chesapeake_db")

df = sedona.read.format("geoparquet").load(vectorized_results.uri)
df = df.withColumnRenamed("label", "layer")
df.writeTo("examples_temp.chesapeake_db.chesapeake_vectorized").createOrReplace()

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

# Compute a small bounding box around the center of the AOI
minx, miny, maxx, maxy = gdf.total_bounds
center_lat = (miny + maxy) / 2
center_lon = (minx + maxx) / 2
delta = 0.01

min_lon = center_lon - delta
max_lon = center_lon + delta
min_lat = center_lat - delta
max_lat = center_lat + delta

df = df.filter(
    ST_Intersects(f.col("geometry"), ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat))
).filter("score_mean > 0.7")

from wherobots_gl import Map

# Wherobots-GL Map is URL-based, so write the filtered subset to GeoParquet first.
vectorized_subset_path = os.getenv("USER_S3_PATH") + "chesapeake_vectorized_subset.parquet"
df.write.format("geoparquet").mode("overwrite").save(vectorized_subset_path)

Map(
    layers=[{"type": "geoparquet", "source": vectorized_subset_path, "name": "Vectorized results"}],
    view={"zoom": 12, "lat": center_lat, "lng": center_lon},
    basemap="dark",
)

from wherobots import vtiles

full_tiles_path = os.getenv("USER_S3_PATH") + "tiles.pmtiles"
vtiles.generate_pmtiles(df, full_tiles_path)

vtiles.show_pmtiles(full_tiles_path)

full_tiles_path