Files
RadioPropagationApi/scripts/bootstrap_canopy.py
T
2026-06-24 09:37:09 +03:00

322 lines
11 KiB
Python

"""Bootstrap canopy height rasters for an AOI."""
from __future__ import annotations
import argparse
import json
import os
import socket
import time
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import unquote, urlparse
from urllib.request import Request, urlopen
META_CHM_V2_INDEX_BASE_URL = (
"https://dataforgood-fb-data.s3.amazonaws.com/"
"forests/v2/global/dinov3_global_chm_v2_ml3"
)
META_CHM_V2_BASE_URL = "https://data.source.coop/tge-labs/meta-chm-v2/chm"
META_CHM_V2_INDEX_URL = f"{META_CHM_V2_INDEX_BASE_URL}/tiles.geojson"
USER_AGENT = "RadioApi canopy bootstrap/0.1"
@dataclass(frozen=True)
class Socks5Proxy:
scheme: str
host: str
port: int
username: str | None
password: str | None
@property
def endpoint(self) -> str:
return f"{self.scheme}://{self.host}:{self.port}"
def parse_socks5_proxy(proxy_url: str) -> Socks5Proxy:
parsed = urlparse(proxy_url)
if parsed.scheme not in {"socks5", "socks5h"}:
raise SystemExit(
f"Unsupported proxy scheme: {parsed.scheme!r}. Use socks5:// or socks5h://"
)
if not parsed.hostname:
raise SystemExit("Proxy URL must include hostname, e.g. socks5://user:pass@host:port")
return Socks5Proxy(
scheme=parsed.scheme,
host=parsed.hostname,
port=parsed.port or 1080,
username=unquote(parsed.username) if parsed.username else None,
password=unquote(parsed.password) if parsed.password else None,
)
def configure_socks5_proxy(proxy_url: str) -> Socks5Proxy:
try:
import socks
except ImportError as exc:
raise SystemExit(
"SOCKS5 proxy requires PySocks. Install with: pip install PySocks"
) from exc
proxy = parse_socks5_proxy(proxy_url)
socks.set_default_proxy(
socks.SOCKS5,
proxy.host,
proxy.port,
username=proxy.username,
password=proxy.password,
)
socket.socket = socks.socksocket
return proxy
def parse_bbox(raw: str) -> tuple[float, float, float, float]:
parts = [float(part.strip()) for part in raw.split(",")]
if len(parts) != 4:
raise argparse.ArgumentTypeError("bbox must be minlon,minlat,maxlon,maxlat")
minlon, minlat, maxlon, maxlat = parts
if minlon >= maxlon or minlat >= maxlat:
raise argparse.ArgumentTypeError("bbox min values must be lower than max values")
return minlon, minlat, maxlon, maxlat
def intersects(
a: tuple[float, float, float, float],
b: tuple[float, float, float, float],
) -> bool:
return not (a[2] < b[0] or a[0] > b[2] or a[3] < b[1] or a[1] > b[3])
def _walk_coordinates(value: Any) -> Iterable[tuple[float, float]]:
if (
isinstance(value, list)
and len(value) >= 2
and isinstance(value[0], int | float)
and isinstance(value[1], int | float)
):
yield float(value[0]), float(value[1])
return
if isinstance(value, list):
for item in value:
yield from _walk_coordinates(item)
def feature_bbox(feature: dict[str, Any]) -> tuple[float, float, float, float] | None:
bbox = feature.get("bbox")
if isinstance(bbox, list) and len(bbox) >= 4:
return float(bbox[0]), float(bbox[1]), float(bbox[2]), float(bbox[3])
geometry = feature.get("geometry") or {}
coordinates = list(_walk_coordinates(geometry.get("coordinates")))
if not coordinates:
return None
lons = [coord[0] for coord in coordinates]
lats = [coord[1] for coord in coordinates]
return min(lons), min(lats), max(lons), max(lats)
def _asset_url(feature: dict[str, Any], base_url: str = META_CHM_V2_BASE_URL) -> str | None:
properties = feature.get("properties") or {}
for key in ("url", "href", "asset_href", "s3_url", "cog_url"):
value = properties.get(key) or feature.get(key)
if isinstance(value, str) and value:
if value.startswith("s3://dataforgood-fb-data/"):
return value.replace(
"s3://dataforgood-fb-data",
"https://dataforgood-fb-data.s3.amazonaws.com",
1,
)
return value
for key in ("filename", "file", "name", "tile", "quadkey"):
value = properties.get(key) or feature.get(key)
if isinstance(value, str) and value:
filename = value if value.lower().endswith((".tif", ".tiff")) else f"{value}.tif"
return f"{base_url.rstrip('/')}/{filename}"
return None
def validate_urls(urls: list[str], sample_size: int = 2) -> None:
for url in urls[:sample_size]:
request = Request(url, headers={"Range": "bytes=0-0", "User-Agent": USER_AGENT})
try:
with urlopen(request) as response:
if response.status >= 400:
raise SystemExit(f"Canopy URL is not reachable: {url} ({response.status})")
except HTTPError as exc:
raise SystemExit(
f"Canopy URL is not reachable: {url} ({exc.code}). "
"Check --index-url/--urls-file or the dataset URL template."
) from exc
except URLError as exc:
raise SystemExit(f"Canopy URL validation failed: {url}: {exc}") from exc
def load_index(index_url: str) -> dict[str, Any]:
path = Path(index_url)
if path.exists():
return json.loads(path.read_text(encoding="utf-8"))
request = Request(index_url, headers={"User-Agent": USER_AGENT})
with urlopen(request) as response:
return json.loads(response.read().decode("utf-8"))
def select_tile_urls(
index: dict[str, Any],
bbox: tuple[float, float, float, float],
*,
base_url: str = META_CHM_V2_BASE_URL,
) -> list[str]:
urls: list[str] = []
for feature in index.get("features", []):
tile_bbox = feature_bbox(feature)
if tile_bbox is None or not intersects(tile_bbox, bbox):
continue
url = _asset_url(feature, base_url=base_url)
if url is not None:
urls.append(url)
return sorted(set(urls))
def urls_from_file(path: Path) -> list[str]:
return [
line.strip()
for line in path.read_text(encoding="utf-8").splitlines()
if line.strip() and not line.lstrip().startswith("#")
]
def destination_for(url: str, output_dir: Path) -> Path:
name = url.rstrip("/").rsplit("/", 1)[-1]
if not name.lower().endswith((".tif", ".tiff")):
name = f"{name}.tif"
return output_dir / name
def _format_mb(value: int) -> str:
return f"{value / (1024 * 1024):.1f} MB"
def _download_one(url: str, destination: Path, timeout_s: int, chunk_size: int) -> None:
partial = destination.with_suffix(f"{destination.suffix}.part")
request = Request(url, headers={"User-Agent": USER_AGENT})
with urlopen(request, timeout=timeout_s) as response, partial.open("wb") as output:
total_header = response.headers.get("Content-Length")
total = int(total_header) if total_header and total_header.isdigit() else None
read = 0
last_report = time.monotonic()
while True:
chunk = response.read(chunk_size)
if not chunk:
break
output.write(chunk)
read += len(chunk)
now = time.monotonic()
if now - last_report >= 5:
if total:
pct = (read / total) * 100
print(f" {_format_mb(read)} / {_format_mb(total)} ({pct:.1f}%)")
else:
print(f" {_format_mb(read)}")
last_report = now
partial.replace(destination)
def download_urls(
urls: list[str],
output_dir: Path,
overwrite: bool,
*,
timeout_s: int = 60,
chunk_size: int = 1024 * 1024,
) -> tuple[int, int, list[str]]:
downloaded = 0
skipped = 0
failed: list[str] = []
output_dir.mkdir(parents=True, exist_ok=True)
for url in urls:
destination = destination_for(url, output_dir)
if destination.exists() and not overwrite:
skipped += 1
print(f"skip existing {destination}")
continue
partial = destination.with_suffix(f"{destination.suffix}.part")
if partial.exists() and not overwrite:
partial.unlink()
print(f"download {url}")
try:
_download_one(url, destination, timeout_s, chunk_size)
except (HTTPError, URLError) as exc:
failed.append(f"{url}: {exc}")
if partial.exists():
partial.unlink()
continue
downloaded += 1
return downloaded, skipped, failed
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Bootstrap Meta/WRI canopy height COGs")
parser.add_argument("--bbox", type=parse_bbox, help="AOI bbox: minlon,minlat,maxlon,maxlat")
parser.add_argument("--output-dir", default=Path("data/canopy"), type=Path)
parser.add_argument("--source", default="meta-chm-v2", choices=["meta-chm-v2"])
parser.add_argument("--index-url", default=META_CHM_V2_INDEX_URL)
parser.add_argument("--urls-file", type=Path)
parser.add_argument("--skip-url-check", action="store_true")
parser.add_argument("--timeout-s", default=60, type=int)
parser.add_argument(
"--proxy",
default=os.environ.get("CANOPY_PROXY"),
help="SOCKS5 proxy URL, e.g. socks5://user:pass@host:port. "
"Default: CANOPY_PROXY env var",
)
parser.add_argument("--overwrite", action="store_true")
return parser.parse_args()
def main() -> None:
args = parse_args()
if args.urls_file is None and args.bbox is None:
raise SystemExit("--bbox is required unless --urls-file is provided")
if args.proxy:
proxy = configure_socks5_proxy(args.proxy)
print(f"using proxy {proxy.endpoint}")
if args.urls_file is not None:
urls = urls_from_file(args.urls_file)
else:
index = load_index(args.index_url)
urls = select_tile_urls(index, args.bbox)
if not urls:
raise SystemExit("No canopy tiles intersect requested AOI")
if not args.skip_url_check:
validate_urls(urls)
print(f"selected_tiles={len(urls)}")
downloaded, skipped, failed = download_urls(
urls,
args.output_dir,
args.overwrite,
timeout_s=args.timeout_s,
)
print(f"downloaded={downloaded} skipped={skipped} failed={len(failed)}")
if failed:
for item in failed:
print(f"failed {item}")
raise SystemExit(1)
if __name__ == "__main__":
main()