diff --git a/.gitignore b/.gitignore index 3687264..1df5586 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ *.parquet *.pbf *.png -*.pyc \ No newline at end of file +*.pyc +data/ \ No newline at end of file diff --git a/.justfile b/.justfile index 674aa87..2949dec 100644 --- a/.justfile +++ b/.justfile @@ -4,14 +4,14 @@ extract input output: download region output: quackosm --geom-filter-geocode {{region}} --osm-extract-source Geofabrik --osm-tags-filter-file tags.json --compact --output {{output}} -convert input output export_flags='--parquet': - python3 slice.py --input {{input}} --output {{output}} {{export_flags}} +convert input water='water.parquet' output='geohash' export_flags='--parquet': + python3 slice.py -i {{input}} -w {{water}} -o {{output}} {{export_flags}} -water input output export_flags='--parquet': - python3 water.py -i {{input}} -o {{output}} -c coastline.parquet {{export_flags}} +water coast='coastline.parquet' water='water.parquet' output='geohash': + python3 water.py {{coast}} {{water}} {{output}} shp input output: - ogr2ogr -of 'Parquet' {{output}} {{input}} + ogr2ogr -s_srs EPSG:4326 -t_srs EPSG:4326 -a_srs EPSG:4326 -of 'Parquet' {{output}} {{input}} plot file: #!/usr/bin/env bash @@ -24,7 +24,7 @@ process file: echo "Extracting Data from OpenStreetMaps PBF" quackosm {{file}} --osm-tags-filter-file tags.json --compact --output ${stem}.parquet echo "Generating geohash parquet files" - python3 slice.py --input ${stem}.parquet --output geohash + python3 slice.py --input ${stem}.parquet --water water.parquet --output geohash --parquet echo "Cleaning up" rm ${stem}.parquet diff --git a/readme.md b/readme.md index f93d466..f50e951 100644 --- a/readme.md +++ b/readme.md @@ -126,10 +126,10 @@ just download '' The original PBF file remains saved in the `files/` folder. -### Split Parquet File to GeoHash Parquet Files +### Split Parquet File to GeoHash Parquet Files and add Sea Polygons ```bash -just convert +just convert ``` ### Plot Parquet File as PNG diff --git a/slice.py b/slice.py index 42d4dfc..73f74d0 100644 --- a/slice.py +++ b/slice.py @@ -3,12 +3,11 @@ import duckdb import geohash import geopandas as gpd import pandas as pd -import json -import cbor2 from pathlib import Path from shapely import wkb from shapely.geometry import box -from typing import List, Dict +from tqdm import tqdm # For progress bar +from typing import List, Dict, Optional # Configure logging import logging @@ -31,39 +30,40 @@ def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]: return sorted(geohashes) -def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame): - """Appends new data to an existing Parquet file while ensuring correct merging.""" +def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame, add_water: bool): + """Appends new data to an existing Parquet file while ensuring correct merging and avoiding duplicates.""" if file_path.exists(): try: existing_data = gpd.read_parquet(file_path) - combined_data = pd.concat([existing_data, new_data], ignore_index=True) + + # Remove sea polygons to avoid duplicates (only if we're adding water polygons) + if add_water: + existing_data = existing_data[~existing_data["tags"].apply(lambda tags: isinstance(tags, dict) and tags.get("natural") == "water")] - if "feature_id" in combined_data.columns: - combined_data = combined_data.drop_duplicates(subset="feature_id", keep="last") + # Ensure no duplicates based on feature_id (for land geometries) + if "feature_id" in existing_data.columns and "feature_id" in new_data.columns: + new_unique_data = new_data[~new_data["feature_id"].isin(existing_data["feature_id"])] + combined_data = pd.concat([existing_data, new_unique_data], ignore_index=True) + else: + combined_data = pd.concat([existing_data, new_data], ignore_index=True) - combined_data = combined_data.reset_index(drop=True) combined_data.to_parquet(file_path, index=False) - logger.info(f"✅ Appended new data to {file_path}") except Exception as e: - logger.warning(f"❌ Error while merging {file_path}: {e}") + logger.warning(f"Error while merging {file_path}: {e}") else: new_data.to_parquet(file_path, index=False) - logger.info(f"✅ Created new file: {file_path}") def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a GeoJSON file.""" geojson_file = output_path / f"{geohash_code}.geojson" - geojson_dict = json.loads(gdf.to_json()) - - with open(geojson_file, "w", encoding="utf-8") as f: - json.dump(geojson_dict, f) + gdf.to_file(geojson_file, driver="GeoJSON") logger.info(f"✅ Saved GeoJSON: {geojson_file}") def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path): """Saves a GeoDataFrame as a CBOR file (compact binary format).""" + import cbor2 cbor_file = output_path / f"{geohash_code}.cbor" geojson_dict = json.loads(gdf.to_json()) - with open(cbor_file, "wb") as f: cbor2.dump(geojson_dict, f) logger.info(f"✅ Saved CBOR: {cbor_file}") @@ -81,14 +81,11 @@ def ensure_crs_consistency(gdf): def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.GeoDataFrame): """Adds water polygons to a geohash tile.""" - logger.info(f"Processing {geohash_file}...") - - # Load existing geohash file try: existing_gdf = gpd.read_parquet(geohash_file) except Exception as e: logger.warning(f"Failed to read {geohash_file}. Skipping. Error: {e}") - return + return False existing_gdf = ensure_crs_consistency(existing_gdf) @@ -99,13 +96,13 @@ def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.G water_in_tile = water_gdf[water_gdf.intersects(bbox)].copy() if water_in_tile.empty: - logger.info(f"No water found for {geohash_code}. Skipping water addition.") - return + return False # No water polygons added # Clip water polygons to the geohash tile boundary water_in_tile["geometry"] = water_in_tile.intersection(bbox) - # Assign "water" metadata + # Assign "water" metadata and a unique identifier for water polygons + water_in_tile["feature_id"] = [f"water_{geohash_code}_{i}" for i in range(len(water_in_tile))] water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile) # Ensure columns match before merging @@ -120,22 +117,22 @@ def add_water_to_geohash(geohash_code: str, geohash_file: Path, water_gdf: gpd.G # Save back to parquet updated_gdf.to_parquet(geohash_file, index=False) - logger.info(f"Updated geohash file {geohash_file} with water polygons.") + return True # Water polygons added -def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str, export_parquet: bool, export_geojson: bool, export_cbor: bool): - """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and adds water polygons.""" - input_path, output_path, water_path = Path(input_file), Path(output_dir), Path(water_file) +def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: Optional[str], export_parquet: bool, export_geojson: bool, export_cbor: bool): + """Splits the GeoParquet file into multiple files based on 4-character Geohash tiles and optionally adds water polygons.""" + input_path, output_path = Path(input_file), Path(output_dir) if not input_path.exists(): raise FileNotFoundError(f"File '{input_file}' not found.") - if not water_path.exists(): - raise FileNotFoundError(f"File '{water_file}' not found.") output_path.mkdir(parents=True, exist_ok=True) - logger.info(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...") - - # Load water dataset - logger.info(f"Loading water dataset from {water_file}...") - water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path)) + # Load water dataset if provided + water_gdf = None + if water_file: + water_path = Path(water_file) + if not water_path.exists(): + raise FileNotFoundError(f"File '{water_file}' not found.") + water_gdf = ensure_crs_consistency(gpd.read_parquet(water_path)) # Connect to DuckDB and load spatial extension con = duckdb.connect() @@ -161,56 +158,65 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, water_file: str geohash_mapping[geohash_code].append(rowid) # Process each geohash tile - for geohash_code, rowids in geohash_mapping.items(): - geohash_bbox = geohash.bbox(geohash_code) + total_tiles = len(geohash_mapping) + tiles_with_water = 0 - # Fetch and clip geometries for this geohash tile - filtered_data = con.execute(f""" - SELECT feature_id, tags, - ST_AsWKB(ST_Intersection(geometry, - ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, - {geohash_bbox["e"]}, {geohash_bbox["n"]}))) - AS clipped_geom - FROM geoparquet - WHERE rowid IN ({','.join(map(str, rowids))}) - AND ST_Intersects(geometry, - ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, - {geohash_bbox["e"]}, {geohash_bbox["n"]})); - """).fetchdf() + with tqdm(total=total_tiles, desc="Processing geohash tiles") as pbar: + for geohash_code, rowids in geohash_mapping.items(): + geohash_bbox = geohash.bbox(geohash_code) - if not filtered_data.empty: - # Convert WKB geometries to Shapely geometries - filtered_data["geometry"] = filtered_data["clipped_geom"].apply( - lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None - ) - filtered_data.drop(columns=["clipped_geom"], inplace=True) - - # Drop rows with invalid geometries - filtered_data = filtered_data.dropna(subset=["geometry"]) + # Fetch and clip geometries for this geohash tile + filtered_data = con.execute(f""" + SELECT feature_id, tags, + ST_AsWKB(ST_Intersection(geometry, + ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, + {geohash_bbox["e"]}, {geohash_bbox["n"]}))) + AS clipped_geom + FROM geoparquet + WHERE rowid IN ({','.join(map(str, rowids))}) + AND ST_Intersects(geometry, + ST_MakeEnvelope({geohash_bbox["w"]}, {geohash_bbox["s"]}, + {geohash_bbox["e"]}, {geohash_bbox["n"]})); + """).fetchdf() if not filtered_data.empty: - gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") + # Convert WKB geometries to Shapely geometries + filtered_data["geometry"] = filtered_data["clipped_geom"].apply( + lambda x: wkb.loads(bytes(x)) if isinstance(x, (bytes, bytearray)) else None + ) + filtered_data.drop(columns=["clipped_geom"], inplace=True) - # Export to desired formats - if export_parquet: - geohash_file = output_path / f"{geohash_code}.parquet" - append_to_parquet(geohash_file, gdf) - add_water_to_geohash(geohash_code, geohash_file, water_gdf) - - if export_geojson: - save_as_geojson(geohash_code, gdf, output_path) - - if export_cbor: - save_as_cbor(geohash_code, gdf, output_path) + # Drop rows with invalid geometries + filtered_data = filtered_data.dropna(subset=["geometry"]) + + if not filtered_data.empty: + gdf = gpd.GeoDataFrame(filtered_data, geometry="geometry", crs="EPSG:4326") + + # Export to GeoParquet + if export_parquet: + geohash_file = output_path / f"{geohash_code}.parquet" + append_to_parquet(geohash_file, gdf, add_water=(water_gdf is not None)) + if water_gdf is not None and add_water_to_geohash(geohash_code, geohash_file, water_gdf): + tiles_with_water += 1 + + # Export to GeoJSON + if export_geojson: + save_as_geojson(geohash_code, gdf, output_path) + + # Export to CBOR + if export_cbor: + save_as_cbor(geohash_code, gdf, output_path) + + pbar.update(1) con.close() - logger.info("🎉 Processing complete!") + print(f"🎉 Processing complete! Generated {total_tiles} tiles, {tiles_with_water} of which had water polygons added.") def main(): - parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, add water polygons, and export in desired formats.") + parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, optionally add water polygons, and export in desired formats.") parser.add_argument("-i", "--input", required=True, help="Path to the input GeoParquet file.") parser.add_argument("-o", "--output", required=True, help="Directory to save output files.") - parser.add_argument("-w", "--water", required=True, help="Path to the water.parquet file.") + parser.add_argument("-w", "--water", required=False, help="Path to the water.parquet file.") parser.add_argument("--parquet", action="store_true", help="Export as GeoParquet") parser.add_argument("--geojson", action="store_true", help="Export as GeoJSON") parser.add_argument("--cbor", action="store_true", help="Export as CBOR") diff --git a/u33d.png b/u33d.png deleted file mode 100644 index c4dae55..0000000 Binary files a/u33d.png and /dev/null differ diff --git a/water.py b/water.py deleted file mode 100644 index 49533d3..0000000 --- a/water.py +++ /dev/null @@ -1,110 +0,0 @@ -import os -import argparse -import pandas as pd -import geopandas as gpd -import geohash -import logging -from shapely.geometry import box - -# Configure logging -logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO) -logger = logging.getLogger(__name__) - -def load_geohash_files(folder): - """Finds all existing geohash parquet files in the folder.""" - geohash_files = {} - for filename in os.listdir(folder): - if filename.endswith(".parquet"): - geohash_code = filename[:-8] # Remove '.parquet' extension - geohash_files[geohash_code] = os.path.join(folder, filename) - logger.info(f"Found {len(geohash_files)} geohash parquet files in {folder}.") - return geohash_files - -def geohash_bbox(geohash_code): - """Returns a bounding box polygon for a geohash.""" - bbox = geohash.bbox(geohash_code) - return box(bbox["w"], bbox["s"], bbox["e"], bbox["n"]) - -def ensure_crs_consistency(gdf): - """Ensures the GeoDataFrame is in EPSG:4326 to prevent CRS mismatches.""" - if gdf.crs is None or gdf.crs.to_string() != "EPSG:4326": - gdf = gdf.to_crs("EPSG:4326") - return gdf - -def process_geohash_files(water_gdf, geohash_files): - """Processes each geohash file: adds water polygons from water.parquet.""" - # Precompute bounding boxes for all geohash codes - geohash_bboxes = {geohash_code: geohash_bbox(geohash_code) for geohash_code in geohash_files.keys()} - - # Spatial index for faster spatial queries - water_sindex = water_gdf.sindex - - for geohash_code, file_path in geohash_files.items(): - logger.info(f"Processing {file_path}...") - - # Load existing geohash file - try: - existing_gdf = gpd.read_parquet(file_path) - except Exception as e: - logger.warning(f"Failed to read {file_path}. Skipping. Error: {e}") - continue - - existing_gdf = ensure_crs_consistency(existing_gdf) - - # Get the precomputed bounding box for this geohash - bbox = geohash_bboxes[geohash_code] - - # Use spatial index to find intersecting water polygons - possible_matches_index = list(water_sindex.intersection(bbox.bounds)) - possible_matches = water_gdf.iloc[possible_matches_index] - water_in_tile = possible_matches[possible_matches.intersects(bbox)].copy() - - if water_in_tile.empty: - logger.info(f"No water found for {geohash_code}. Skipping water addition.") - continue - - # Clip water polygons to the geohash tile boundary - water_in_tile["geometry"] = water_in_tile.intersection(bbox) - - # Assign "water" metadata - water_in_tile["tags"] = [{"natural": "water", "water": "sea"}] * len(water_in_tile) - - # Ensure columns match before merging - for col in ["feature_id", "tags"]: - if col not in existing_gdf.columns: - existing_gdf[col] = None - if col not in water_in_tile.columns: - water_in_tile[col] = None - - # Merge updated water polygons into geohash file - updated_gdf = gpd.GeoDataFrame(pd.concat([existing_gdf, water_in_tile], ignore_index=True), crs="EPSG:4326") - - # Save back to parquet - updated_gdf.to_parquet(file_path, index=False) - logger.info(f"Updated geohash file {file_path} with water polygons.") - -def main(): - parser = argparse.ArgumentParser(description="Add water polygons to geohash parquet files from a global water dataset.") - parser.add_argument("water", help="Path to the water.parquet file.") - parser.add_argument("geohash_folder", help="Path to the folder containing .parquet files.") - - args = parser.parse_args() - - # Load water dataset - logger.info(f"Loading water dataset from {args.water}...") - water_gdf = ensure_crs_consistency(gpd.read_parquet(args.water)) - - # Get existing geohash parquet files - geohash_files = load_geohash_files(args.geohash_folder) - - if not geohash_files: - logger.warning("No geohash parquet files found. Exiting.") - return - - # Process geohash files - process_geohash_files(water_gdf, geohash_files) - - logger.info("Processing complete. All water polygons added.") - -if __name__ == "__main__": - main() \ No newline at end of file