# Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
from abc import ABC, abstractmethod
from typing import Iterable, Optional
import geojson
import shapely
from .imaged_feature_property_accessor import ImagedFeaturePropertyAccessor
[docs]
class Feature2DSpatialIndex(ABC):
"""
A query-only spatial index allowing clients to lookup features using 2D geometries
"""
[docs]
@abstractmethod
def find_intersects(self, geometry: shapely.Geometry) -> Iterable[geojson.Feature]:
"""
Return the features intersecting the input geometry.
:param geometry: geometry to query the index
:return: the features
"""
[docs]
@abstractmethod
def find_nearest(self, geometry: shapely.Geometry, max_distance: Optional[float] = None) -> Iterable[geojson.Feature]:
"""
Return the nearest feature for the input geometry based on distance within two-dimensional Cartesian space.
:param geometry: geometry to query the index
:param max_distance: maximum distance
:return: the nearest features
"""
[docs]
class STRFeature2DSpatialIndex(Feature2DSpatialIndex):
"""
Implementation of the 2D spatial index for GeoJSON features using Shapely's Sort-Tile-Recursive (STR)
tree datastructure.
"""
def __init__(
self,
feature_collection: geojson.FeatureCollection,
use_image_geometries: bool = True,
property_accessor: ImagedFeaturePropertyAccessor = ImagedFeaturePropertyAccessor(),
) -> None:
self.use_image_geometries = use_image_geometries
self.features = feature_collection.features
if use_image_geometries and property_accessor is not None:
geometries = [property_accessor.find_image_geometry(feature) for feature in self.features]
else:
geometries = [(shapely.shape(feature.geometry), feature) for feature in self.features]
self.index = shapely.STRtree(geometries)
[docs]
def find_intersects(self, geometry: shapely.Geometry) -> Iterable[geojson.Feature]:
result_indexes = self.index.query(geometry, predicate="intersects")
return [self.features[i] for i in result_indexes]
[docs]
def find_nearest(self, geometry: shapely.Geometry, max_distance: Optional[float] = None) -> Iterable[geojson.Feature]:
if max_distance is None:
if self.use_image_geometries:
max_distance = 50
else:
max_distance = 1.0
result_indexes = self.index.query_nearest(geometry, max_distance=max_distance)
return [self.features[i] for i in result_indexes]