Files
RadioPropagationApi/api/app/core/viewshed.py
T
2026-06-24 09:01:11 +03:00

310 lines
9.6 KiB
Python

from __future__ import annotations
import subprocess
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import numpy as np
import rasterio
from pyproj import Geod, Transformer
from rasterio.crs import CRS
from rasterio.features import shapes
from rasterio.merge import merge
from rasterio.warp import (
Resampling,
calculate_default_transform,
reproject,
transform,
transform_bounds,
)
from shapely.geometry import mapping, shape
from shapely.ops import transform as shapely_transform
from shapely.ops import unary_union
from app.core.dem import _dem_files, _sample_dataset
from app.core.geo import GeoPoint
from app.models.viewshed import ViewshedRequest
@dataclass(frozen=True)
class RasterResult:
uri: str
metadata: dict[str, Any]
_GEOD = Geod(ellps="WGS84")
def _utm_epsg(lon: float, lat: float) -> int:
zone = int((lon + 180) // 6) + 1
return 32600 + zone if lat >= 0 else 32700 + zone
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 _radius_bbox(observer: GeoPoint, radius_m: float) -> tuple[float, float, float, float]:
coords = [
_GEOD.fwd(observer.lon, observer.lat, azimuth, radius_m)[:2]
for azimuth in range(0, 360, 45)
]
lons = [observer.lon, *(coord[0] for coord in coords)]
lats = [observer.lat, *(coord[1] for coord in coords)]
return min(lons), min(lats), max(lons), max(lats)
def _find_dem_tile(lat: float, lon: float, dem_path: Path) -> Path | None:
for path in _dem_files(dem_path):
with rasterio.open(path) as dataset:
if _sample_dataset(dataset, lat, lon) is not None:
return path
return None
def _find_dem_tiles(
dem_path: Path,
bbox_wgs84: tuple[float, float, float, float],
) -> list[Path]:
matches: list[Path] = []
for path in _dem_files(dem_path):
with rasterio.open(path) as dataset:
bounds = dataset.bounds
tile_bbox = (
bounds.left,
bounds.bottom,
bounds.right,
bounds.top,
)
if dataset.crs is not None and dataset.crs != CRS.from_epsg(4326):
tile_bbox = transform_bounds(dataset.crs, CRS.from_epsg(4326), *tile_bbox)
if _intersects(tile_bbox, bbox_wgs84):
matches.append(path)
return matches
def _prepare_metric_dem(
source_paths: list[Path],
observer: GeoPoint,
radius_m: float,
output_path: Path,
) -> tuple[Path, float, float, int]:
if not source_paths:
raise FileNotFoundError("No DEM tiles found for viewshed radius")
dst_crs = CRS.from_epsg(_utm_epsg(observer.lon, observer.lat))
bbox_wgs84 = _radius_bbox(observer, radius_m)
datasets = [rasterio.open(path) for path in source_paths]
try:
src_crs = datasets[0].crs or CRS.from_epsg(4326)
merge_bounds = bbox_wgs84
if src_crs != CRS.from_epsg(4326):
merge_bounds = transform_bounds(CRS.from_epsg(4326), src_crs, *bbox_wgs84)
mosaic, src_transform = merge(datasets, bounds=merge_bounds)
source_data = mosaic[0].astype(np.float32)
height, width = source_data.shape
src_bounds = rasterio.transform.array_bounds(height, width, src_transform)
observer_x, observer_y = transform(
CRS.from_epsg(4326),
dst_crs,
[observer.lon],
[observer.lat],
)
transform_affine, width, height = calculate_default_transform(
src_crs,
dst_crs,
width,
height,
*src_bounds,
)
data = np.empty((height, width), dtype=np.float32)
reproject(
source=source_data,
destination=data,
src_transform=src_transform,
src_crs=src_crs,
dst_transform=transform_affine,
dst_crs=dst_crs,
resampling=Resampling.bilinear,
)
profile = {
"driver": "GTiff",
"crs": dst_crs,
"transform": transform_affine,
"width": width,
"height": height,
"dtype": "float32",
"count": 1,
"nodata": -9999,
}
with rasterio.open(output_path, "w", **profile) as dst:
dst.write(data, 1)
finally:
for dataset in datasets:
dataset.close()
return output_path, observer_x[0], observer_y[0], len(source_paths)
def _run_gdal_viewshed(
dem_path: Path,
observer_x: float,
observer_y: float,
observer_height: float,
target_height: float,
max_distance: float,
output_path: Path,
gdal_bin: str,
) -> None:
command = [
gdal_bin,
"-ox",
str(observer_x),
"-oy",
str(observer_y),
"-oz",
str(observer_height),
"-tz",
str(target_height),
"-md",
str(max_distance),
"-vv",
"255",
"-iv",
"0",
"-ov",
"0",
str(dem_path),
str(output_path),
]
subprocess.run(command, check=True, capture_output=True, text=True)
def _to_wgs84(geom, src_crs: CRS | None):
if src_crs is None:
return geom
crs_text = src_crs.to_string()
if crs_text in {"EPSG:4326", "OGC:CRS84", "WGS84"}:
return geom
transformer = Transformer.from_crs(src_crs, "EPSG:4326", always_xy=True)
return shapely_transform(transformer.transform, geom)
def _viewshed_geojson(viewshed_path: Path, observer: GeoPoint) -> dict[str, Any]:
features: list[dict[str, Any]] = []
with rasterio.open(viewshed_path) as dataset:
data = dataset.read(1, masked=True)
mask = (data == 255).astype(np.uint8)
for geom, value in shapes(mask, mask=mask.astype(bool), transform=dataset.transform):
if int(value) != 1:
continue
polygon = _to_wgs84(shape(geom), dataset.crs)
if polygon.is_empty:
continue
features.append(
{
"type": "Feature",
"properties": {"visible": True},
"geometry": mapping(polygon),
}
)
if not features:
return {
"type": "FeatureCollection",
"features": [],
"properties": {"observer": {"lat": observer.lat, "lon": observer.lon}},
}
merged = unary_union([shape(feature["geometry"]) for feature in features])
return {
"type": "FeatureCollection",
"features": [
{
"type": "Feature",
"properties": {"visible": True},
"geometry": mapping(merged),
}
],
"properties": {"observer": {"lat": observer.lat, "lon": observer.lon}},
}
def compute_viewshed(
request: ViewshedRequest,
*,
dem_path: str | Path,
output_dir: str | Path,
gdal_viewshed_bin: str = "gdal_viewshed",
) -> RasterResult:
dem_root = Path(dem_path)
output_root = Path(output_dir)
output_root.mkdir(parents=True, exist_ok=True)
observer = GeoPoint(lat=request.observer.lat, lon=request.observer.lon)
source_tile = _find_dem_tile(observer.lat, observer.lon, dem_root)
calc_radius_m = request.calc_radius_m or request.radius_m
bbox_wgs84 = _radius_bbox(observer, calc_radius_m)
source_tiles = _find_dem_tiles(dem_root, bbox_wgs84)
if source_tile is not None and source_tile not in source_tiles:
source_tiles.append(source_tile)
if not source_tiles:
raise FileNotFoundError(
f"No DEM tile covers viewshed radius at lat={observer.lat}, lon={observer.lon}"
)
with tempfile.TemporaryDirectory(prefix="viewshed-") as tmp_dir:
tmp = Path(tmp_dir)
metric_dem = tmp / "metric_dem.tif"
viewshed_raster = tmp / "viewshed.tif"
metric_path, observer_x, observer_y, dem_tiles_count = _prepare_metric_dem(
source_tiles,
observer,
calc_radius_m,
metric_dem,
)
_run_gdal_viewshed(
metric_path,
observer_x,
observer_y,
request.observer.height_agl,
request.target_height_agl,
calc_radius_m,
viewshed_raster,
gdal_viewshed_bin,
)
if request.format == "geojson":
result = _viewshed_geojson(viewshed_raster, observer)
output_file = output_root / f"viewshed_{observer.lat:.5f}_{observer.lon:.5f}.json"
output_file.write_text(
__import__("json").dumps(result),
encoding="utf-8",
)
return RasterResult(
uri=str(output_file),
metadata={
"format": "geojson",
"surface": request.surface,
"radius_m": request.radius_m,
"calc_radius_m": calc_radius_m,
"dem_tiles_count": dem_tiles_count,
"feature_count": len(result.get("features", [])),
},
)
output_file = output_root / f"viewshed_{observer.lat:.5f}_{observer.lon:.5f}.tif"
output_file.write_bytes(viewshed_raster.read_bytes())
return RasterResult(
uri=str(output_file),
metadata={
"format": request.format,
"surface": request.surface,
"radius_m": request.radius_m,
"calc_radius_m": calc_radius_m,
"dem_tiles_count": dem_tiles_count,
},
)