210 lines
8.2 KiB
Python
210 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from uuid import UUID
|
|
|
|
from pyproj import Transformer
|
|
from shapely.geometry import mapping
|
|
from shapely.ops import transform as shapely_transform
|
|
|
|
from app.core.errors import AppError
|
|
from app.models import Dataset
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RasterPartitionSelection:
|
|
datasets: list[Dataset]
|
|
values: Any
|
|
selected_cells: Any
|
|
resolution_x: float
|
|
resolution_y: float
|
|
|
|
|
|
class RasterPartitionAnalysisService:
|
|
MAX_PARTITIONS = 64
|
|
|
|
@staticmethod
|
|
def _bbox_intersects(dataset: Dataset, bbox: tuple[float, float, float, float]) -> bool:
|
|
source_bbox = (dataset.source_metadata or {}).get("bbox_epsg4326")
|
|
if not isinstance(source_bbox, list) or len(source_bbox) != 4:
|
|
return True
|
|
try:
|
|
min_x, min_y, max_x, max_y = (float(value) for value in source_bbox)
|
|
except (TypeError, ValueError):
|
|
return True
|
|
return not (
|
|
max_x <= bbox[0]
|
|
or min_x >= bbox[2]
|
|
or max_y <= bbox[1]
|
|
or min_y >= bbox[3]
|
|
)
|
|
|
|
@staticmethod
|
|
def _candidate_datasets(
|
|
db,
|
|
project_id: UUID,
|
|
*,
|
|
source_name: str,
|
|
product_key: str,
|
|
bbox: tuple[float, float, float, float],
|
|
dataset_ids: list[UUID] | None = None,
|
|
) -> list[Dataset]:
|
|
query = db.query(Dataset).filter(
|
|
Dataset.project_id == project_id,
|
|
Dataset.source_name == source_name,
|
|
Dataset.dataset_type == "raster",
|
|
Dataset.status == "ready",
|
|
)
|
|
if dataset_ids is not None:
|
|
query = query.filter(Dataset.id.in_(dataset_ids))
|
|
rows = query.all()
|
|
candidates = [
|
|
dataset
|
|
for dataset in rows
|
|
if str((dataset.source_metadata or {}).get("product_key") or "") == product_key
|
|
and dataset.storage_path
|
|
and Path(dataset.storage_path).is_file()
|
|
and RasterPartitionAnalysisService._bbox_intersects(dataset, bbox)
|
|
]
|
|
candidates.sort(key=lambda dataset: (str(dataset.area_id or ""), str(dataset.id)))
|
|
if not candidates:
|
|
raise AppError(
|
|
code="RASTER_PARTITIONS_NOT_FOUND",
|
|
message="No persisted raster partitions cover this selection",
|
|
details={"source_name": source_name, "product_key": product_key},
|
|
status_code=404,
|
|
)
|
|
if dataset_ids is not None and {dataset.id for dataset in candidates} != set(dataset_ids):
|
|
raise AppError(
|
|
code="RASTER_PARTITION_SOURCE_MISMATCH",
|
|
message="Every requested raster partition must match the governed source product and selection",
|
|
details={"requested_count": len(dataset_ids), "eligible_count": len(candidates)},
|
|
status_code=409,
|
|
)
|
|
if len(candidates) > RasterPartitionAnalysisService.MAX_PARTITIONS:
|
|
raise AppError(
|
|
code="RASTER_PARTITION_LIMIT_EXCEEDED",
|
|
message="The selection intersects too many raster partitions",
|
|
details={
|
|
"partition_count": len(candidates),
|
|
"max_partitions": RasterPartitionAnalysisService.MAX_PARTITIONS,
|
|
},
|
|
status_code=422,
|
|
)
|
|
return candidates
|
|
|
|
@staticmethod
|
|
def select(
|
|
db,
|
|
project_id: UUID,
|
|
*,
|
|
source_name: str,
|
|
product_key: str,
|
|
selection_geometry_4326,
|
|
nodata: float,
|
|
max_pixels: int,
|
|
dataset_ids: list[UUID] | None = None,
|
|
) -> RasterPartitionSelection:
|
|
try:
|
|
import numpy as np
|
|
import rasterio
|
|
from rasterio.features import geometry_mask
|
|
from rasterio.merge import merge
|
|
except ImportError as exc:
|
|
raise AppError(
|
|
code="RASTER_PROCESSING_UNAVAILABLE",
|
|
message="Rasterio and numpy are required for partitioned raster analysis",
|
|
status_code=503,
|
|
) from exc
|
|
|
|
bbox = tuple(float(value) for value in selection_geometry_4326.bounds)
|
|
datasets = RasterPartitionAnalysisService._candidate_datasets(
|
|
db,
|
|
project_id,
|
|
source_name=source_name,
|
|
product_key=product_key,
|
|
bbox=bbox,
|
|
dataset_ids=dataset_ids,
|
|
)
|
|
transformer = Transformer.from_crs("EPSG:4326", "EPSG:31370", always_xy=True)
|
|
selection_metric = shapely_transform(transformer.transform, selection_geometry_4326)
|
|
min_x, min_y, max_x, max_y = selection_metric.bounds
|
|
|
|
try:
|
|
with ExitStack() as stack:
|
|
sources = [stack.enter_context(rasterio.open(dataset.storage_path)) for dataset in datasets]
|
|
invalid_sources = [
|
|
index
|
|
for index, source in enumerate(sources)
|
|
if source.crs is None or source.crs.to_epsg() != 31370 or source.count != 1
|
|
]
|
|
if invalid_sources:
|
|
raise AppError(
|
|
code="RASTER_PARTITION_MISMATCH",
|
|
message="Raster partitions do not share the governed CRS and band layout",
|
|
details={"invalid_partition_indexes": invalid_sources},
|
|
status_code=409,
|
|
)
|
|
target_resolution = max(abs(float(sources[0].res[0])), abs(float(sources[0].res[1])))
|
|
invalid_resolutions = [
|
|
{
|
|
"partition_index": index,
|
|
"resolution": [abs(float(source.res[0])), abs(float(source.res[1]))],
|
|
}
|
|
for index, source in enumerate(sources)
|
|
if not all(
|
|
math.isclose(abs(float(value)), target_resolution, rel_tol=0.001, abs_tol=0.01)
|
|
for value in source.res
|
|
)
|
|
]
|
|
if invalid_resolutions:
|
|
raise AppError(
|
|
code="RASTER_PARTITION_MISMATCH",
|
|
message="Raster partitions do not share one analysis resolution",
|
|
details={"invalid_resolutions": invalid_resolutions},
|
|
status_code=409,
|
|
)
|
|
width = max(1, math.ceil((max_x - min_x) / target_resolution))
|
|
height = max(1, math.ceil((max_y - min_y) / target_resolution))
|
|
if width * height > max_pixels:
|
|
raise AppError(
|
|
code="RASTER_PARTITION_SELECTION_TOO_LARGE",
|
|
message="Select a smaller rectangle for regional raster analysis",
|
|
details={"pixel_count": width * height, "max_pixels": max_pixels},
|
|
status_code=422,
|
|
)
|
|
mosaic, transform = merge(
|
|
sources,
|
|
bounds=(min_x, min_y, max_x, max_y),
|
|
res=(target_resolution, target_resolution),
|
|
nodata=nodata,
|
|
dtype="float32",
|
|
)
|
|
values = np.asarray(mosaic[0], dtype="float64")
|
|
selected_cells = geometry_mask(
|
|
[mapping(selection_metric)],
|
|
out_shape=values.shape,
|
|
transform=transform,
|
|
invert=True,
|
|
)
|
|
return RasterPartitionSelection(
|
|
datasets=datasets,
|
|
values=values,
|
|
selected_cells=selected_cells,
|
|
resolution_x=target_resolution,
|
|
resolution_y=target_resolution,
|
|
)
|
|
except AppError:
|
|
raise
|
|
except Exception as exc:
|
|
raise AppError(
|
|
code="RASTER_PARTITION_ANALYSIS_FAILED",
|
|
message="Persisted raster partitions could not be assembled for this selection",
|
|
details={"reason": str(exc)},
|
|
status_code=500,
|
|
) from exc
|