slice and water
This commit is contained in:
randogoth
parent
2edd6ca38c
commit
189f135840
1 changed files with 81 additions and 13 deletions
94
slice.py
94
slice.py
|
|
@ -7,8 +7,14 @@ import json
|
|||
import cbor2
|
||||
from pathlib import Path
|
||||
from shapely import wkb
|
||||
from shapely.geometry import box
|
||||
from typing import List, Dict
|
||||
|
||||
# Configure logging
|
||||
import logging
|
||||
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.INFO)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
def get_geohashes_from_bbox(min_x, min_y, max_x, max_y) -> List[str]:
|
||||
"""Generates all 4-character Geohashes that intersect a bounding box."""
|
||||
geohashes = set()
|
||||
|
|
@ -37,12 +43,12 @@ def append_to_parquet(file_path: Path, new_data: gpd.GeoDataFrame):
|
|||
|
||||
combined_data = combined_data.reset_index(drop=True)
|
||||
combined_data.to_parquet(file_path, index=False)
|
||||
print(f"✅ Appended new data to {file_path}")
|
||||
logger.info(f"✅ Appended new data to {file_path}")
|
||||
except Exception as e:
|
||||
print(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)
|
||||
print(f"✅ Created new file: {file_path}")
|
||||
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."""
|
||||
|
|
@ -51,7 +57,7 @@ def save_as_geojson(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path)
|
|||
|
||||
with open(geojson_file, "w", encoding="utf-8") as f:
|
||||
json.dump(geojson_dict, f)
|
||||
print(f"✅ Saved GeoJSON: {geojson_file}")
|
||||
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)."""
|
||||
|
|
@ -60,16 +66,76 @@ def save_as_cbor(geohash_code: str, gdf: gpd.GeoDataFrame, output_path: Path):
|
|||
|
||||
with open(cbor_file, "wb") as f:
|
||||
cbor2.dump(geojson_dict, f)
|
||||
print(f"✅ Saved CBOR: {cbor_file}")
|
||||
logger.info(f"✅ Saved CBOR: {cbor_file}")
|
||||
|
||||
def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet: bool, export_geojson: bool, export_cbor: bool):
|
||||
"""Splits the GeoParquet file into multiple files based on 4-character Geohash tiles."""
|
||||
input_path, output_path = Path(input_file), Path(output_dir)
|
||||
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 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
|
||||
|
||||
existing_gdf = ensure_crs_consistency(existing_gdf)
|
||||
|
||||
# Get the geohash bounding box
|
||||
bbox = geohash_bbox(geohash_code)
|
||||
|
||||
# Extract water polygons that overlap with this geohash
|
||||
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
|
||||
|
||||
# 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(geohash_file, index=False)
|
||||
logger.info(f"Updated geohash file {geohash_file} with water polygons.")
|
||||
|
||||
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)
|
||||
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)
|
||||
|
||||
print(f"📂 Processing '{input_file}', output will be saved in '{output_dir}'...")
|
||||
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))
|
||||
|
||||
# Connect to DuckDB and load spatial extension
|
||||
con = duckdb.connect()
|
||||
|
|
@ -129,6 +195,7 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet:
|
|||
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)
|
||||
|
|
@ -137,12 +204,13 @@ def slice_and_split_geoparquet(input_file: str, output_dir: str, export_parquet:
|
|||
save_as_cbor(geohash_code, gdf, output_path)
|
||||
|
||||
con.close()
|
||||
print("🎉 Processing complete!")
|
||||
logger.info("🎉 Processing complete!")
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles with export format options.")
|
||||
parser = argparse.ArgumentParser(description="Split a GeoParquet file into Geohash tiles, 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("--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")
|
||||
|
|
@ -153,9 +221,9 @@ def main():
|
|||
args.parquet = True # Default to GeoParquet if no options are given
|
||||
|
||||
try:
|
||||
slice_and_split_geoparquet(args.input, args.output, args.parquet, args.geojson, args.cbor)
|
||||
slice_and_split_geoparquet(args.input, args.output, args.water, args.parquet, args.geojson, args.cbor)
|
||||
except Exception as e:
|
||||
print(f"❌ Error: {e}")
|
||||
logger.error(f"❌ Error: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue