from typing import TYPE_CHECKING, Optional, Sequence, Union
from tqdm.auto import tqdm
from rastervision.core.data.utils.geojson import geoms_to_geojson
if TYPE_CHECKING:
from rastervision.core.box import Box
from rastervision.core.data import ClassConfig, CRSTransformer
PROGRESSBAR_DELAY_SEC = 5
[docs]def boxes_to_geojson(
boxes: Sequence['Box'],
class_ids: Sequence[int],
crs_transformer: 'CRSTransformer',
class_config: 'ClassConfig',
scores: Optional[Sequence[Union[float, Sequence[float]]]] = None,
bbox: Optional['Box'] = None) -> dict:
"""Convert boxes and associated data into a GeoJSON dict.
Args:
boxes (Sequence[Box]): List of Box in pixel row/col format.
class_ids (Sequence[int]): List of int (one for each box)
crs_transformer (CRSTransformer): CRSTransformer used to convert pixel
coords to map coords in the GeoJSON.
class_config (ClassConfig): ClassConfig
scores (Optional[Sequence[Union[float, Sequence[float]]]], optional):
Optional list of score or scores. If floats (one for each box),
property name will be "score". If lists of floats, property name
will be "scores". Defaults to None.
bbox (Optional[Box]): User-specified crop of the extent. Must be
provided if the corresponding RasterSource has bbox != extent.
Returns:
dict: Serialized GeoJSON.
"""
if len(boxes) != len(class_ids):
raise ValueError(f'len(boxes) ({len(boxes)}) != '
f'len(class_ids) ({len(class_ids)})')
if scores is not None and len(boxes) != len(scores):
raise ValueError(f'len(boxes) ({len(boxes)}) != '
f'len(scores) ({len(scores)})')
# boxes to map coords
with tqdm(
boxes,
desc='Transforming boxes to map coords',
delay=PROGRESSBAR_DELAY_SEC) as bar:
geoms = [
crs_transformer.pixel_to_map(box.to_shapely(), bbox=bbox)
for box in bar
]
# add box properties (ID and name of predicted class)
with tqdm(
class_ids,
desc='Attaching class IDs and names to boxes',
delay=PROGRESSBAR_DELAY_SEC) as bar:
properties = [
dict(class_id=id, class_name=class_config.get_name(id))
for id in bar
]
# add box properties (ID and name of predicted class)
if scores is not None:
with tqdm(
zip(properties, scores),
desc='Transforming boxes to map coords',
delay=PROGRESSBAR_DELAY_SEC) as bar:
for prop, box_score in bar:
key = 'score' if isinstance(box_score, float) else 'scores'
prop[key] = box_score
geojson = geoms_to_geojson(geoms, properties)
return geojson