wehub-resource-sync
486 lines
16 KiB
Python
486 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import io
|
|
import itertools
|
|
import json
|
|
import re
|
|
import typing
|
|
import zipfile
|
|
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import geopandas as gpd
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import polars
|
|
import pyproj
|
|
import requests
|
|
from pyproj import CRS, Transformer
|
|
from pyproj.aoi import AreaOfInterest
|
|
from pyproj.database import query_utm_crs_info
|
|
from tqdm import tqdm
|
|
|
|
import rerun as rr
|
|
import rerun.blueprint as rrb
|
|
|
|
if typing.TYPE_CHECKING:
|
|
import shapely
|
|
|
|
DATA_DIR = Path(__file__).parent / "dataset"
|
|
MAP_DATA_DIR = DATA_DIR / "map_data"
|
|
if not DATA_DIR.exists():
|
|
DATA_DIR.mkdir()
|
|
|
|
INVOLI_DATASETS = {
|
|
"10min": "https://storage.googleapis.com/rerun-example-datasets/involi/involi_demo_set_1_10min.zip",
|
|
"2h": "https://storage.googleapis.com/rerun-example-datasets/involi/involi_demo_set_2_2h.zip",
|
|
}
|
|
|
|
|
|
def download_with_progress(url: str, what: str) -> io.BytesIO:
|
|
"""Download file with tqdm progress bar."""
|
|
chunk_size = 1024 * 1024
|
|
try:
|
|
resp = requests.get(url, stream=True, timeout=30)
|
|
resp.raise_for_status()
|
|
total_size = int(resp.headers.get("content-length", 0))
|
|
with tqdm(
|
|
desc=f"Downloading {what}…",
|
|
total=total_size,
|
|
unit="iB",
|
|
unit_scale=True,
|
|
unit_divisor=1024,
|
|
) as progress:
|
|
download_file = io.BytesIO()
|
|
for data in resp.iter_content(chunk_size):
|
|
download_file.write(data)
|
|
progress.update(len(data))
|
|
|
|
if download_file.tell() != total_size and total_size > 0:
|
|
raise ValueError("Download incomplete: size mismatch")
|
|
|
|
download_file.seek(0)
|
|
return download_file
|
|
except (requests.RequestException, ValueError) as e:
|
|
raise RuntimeError(f"Failed to download {what} from {url}: {e}") from e
|
|
|
|
|
|
def shapely_geom_to_numpy(geom: shapely.Geometry) -> list[npt.NDArray[np.float64]]:
|
|
"""Convert shapely objects to numpy array suitable for logging as line batches."""
|
|
|
|
if geom.geom_type == "Polygon":
|
|
return [np.array(geom.exterior.coords)] + [np.array(interior.coords) for interior in geom.interiors]
|
|
elif geom.geom_type == "MultiPolygon":
|
|
res = []
|
|
for poly in geom.geoms:
|
|
res.extend(shapely_geom_to_numpy(poly))
|
|
return res
|
|
else:
|
|
print(f"Warning: unknown Shapely object {geom}")
|
|
return []
|
|
|
|
|
|
def log_region_boundaries_for_country(
|
|
country_code: str,
|
|
level: int,
|
|
color: tuple[float, float, float],
|
|
crs: CRS,
|
|
) -> None:
|
|
"""Log some boundaries for the given country and level."""
|
|
|
|
def download_eu_map_data() -> None:
|
|
"""Download some basic EU map data."""
|
|
|
|
if MAP_DATA_DIR.exists():
|
|
return
|
|
|
|
EU_MAP_DATA_URL = "https://gisco-services.ec.europa.eu/distribution/v2/nuts/download/ref-nuts-2021-01m.json.zip"
|
|
zip_data = download_with_progress(EU_MAP_DATA_URL, "map data")
|
|
with zipfile.ZipFile(zip_data) as zip_ref:
|
|
zip_ref.extractall(MAP_DATA_DIR)
|
|
|
|
download_eu_map_data()
|
|
|
|
# cspell:disable-next-line
|
|
map_data = gpd.read_file(MAP_DATA_DIR / f"NUTS_RG_01M_2021_4326_LEVL_{level}.json").set_crs("epsg:4326").to_crs(crs)
|
|
|
|
for _i, row in map_data[map_data.CNTR_CODE == country_code].iterrows():
|
|
entity_path = f"region_boundaries/{country_code}/{level}/{row.NUTS_ID}"
|
|
lines = shapely_geom_to_numpy(row.geometry)
|
|
rr.log(entity_path + "/2D", rr.LineStrips2D(lines, colors=color), static=True)
|
|
rr.log(
|
|
entity_path + "/3D",
|
|
rr.LineStrips3D(
|
|
[np.hstack([line, np.zeros((len(line), 1))]) for line in lines],
|
|
colors=color,
|
|
),
|
|
static=True,
|
|
)
|
|
metadata = row.to_dict()
|
|
metadata.pop("geometry")
|
|
rr.log(entity_path, rr.AnyValues(**metadata), static=True)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Measurement:
|
|
"""One measurement loaded from INVOLI data. Corresponds to an "aircraft" record."""
|
|
|
|
icao_id: str
|
|
latitude: float | None
|
|
longitude: float | None
|
|
barometric_altitude: float | None
|
|
wg84_altitude: float | None
|
|
course: float | None
|
|
ground_speed: float | None
|
|
vertical_speed: float | None
|
|
ground_status: str | None
|
|
timestamp: float
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any]) -> Measurement:
|
|
return cls(
|
|
icao_id=data["ids"]["icao"],
|
|
latitude=data.get("latitude"),
|
|
longitude=data.get("longitude"),
|
|
barometric_altitude=data.get("barometric_altitude"),
|
|
wg84_altitude=data.get("wg84_altitude"),
|
|
course=data.get("course"),
|
|
ground_speed=data.get("ground_speed"),
|
|
vertical_speed=data.get("vertical_speed"),
|
|
ground_status=data.get("ground_status"),
|
|
timestamp=data["timestamp"][0] + data["timestamp"][1] / 1e9,
|
|
)
|
|
|
|
|
|
def find_best_utm_crs(measurements: list[Measurement]) -> CRS:
|
|
"""Returns the best UTM coordinates reference system given a list of measurements."""
|
|
|
|
def get_area_of_interest(measurements: list[Measurement]) -> AreaOfInterest:
|
|
"""Compute the span of coordinates for all provided measurements."""
|
|
|
|
print("Computing area of interest…")
|
|
all_long_lat = [
|
|
(a.longitude, a.latitude) for a in measurements if a.latitude is not None and a.longitude is not None
|
|
]
|
|
return AreaOfInterest(
|
|
west_lon_degree=min(x[0] for x in all_long_lat),
|
|
south_lat_degree=min(x[1] for x in all_long_lat),
|
|
east_lon_degree=max(x[0] for x in all_long_lat),
|
|
north_lat_degree=max(x[1] for x in all_long_lat),
|
|
)
|
|
|
|
area_of_interest = get_area_of_interest(measurements)
|
|
utm_crs_list = query_utm_crs_info(
|
|
datum_name="WGS 84",
|
|
area_of_interest=area_of_interest,
|
|
)
|
|
|
|
return CRS.from_epsg(utm_crs_list[0].code)
|
|
|
|
|
|
def load_measurements(paths: list[Path]) -> list[Measurement]:
|
|
"""Load measurements from a bunch of json files."""
|
|
all_measurements = []
|
|
for path in tqdm(paths, "Loading measurements"):
|
|
data = json.loads(path.read_text())
|
|
for data_rec in data:
|
|
for aircraft in data_rec["records"]:
|
|
all_measurements.append(Measurement.from_dict(aircraft["aircraft"]))
|
|
|
|
return all_measurements
|
|
|
|
|
|
def get_paths_for_directory(directory: Path) -> list[Path]:
|
|
"""
|
|
Get a sorted list of JSON file by recursively walking the provided directory.
|
|
|
|
Note: technically, we don't need the list to be sorted as Rerun accepts out of order data. However, it comes at a
|
|
(small) performance cost and any (cheap) sorting on the logging end is always better.
|
|
"""
|
|
|
|
def atoi(text: str) -> int | str:
|
|
return int(text) if text.isdigit() else text
|
|
|
|
def natural_keys(path: Path) -> list[int | str]:
|
|
"""
|
|
Human sort.
|
|
|
|
alist.sort(key=natural_keys) sorts in human order
|
|
https://nedbatchelder.com/blog/200712/human_sorting.html
|
|
(See Toothy's implementation in the comments)
|
|
"""
|
|
return [atoi(c) for c in re.split(r"(\d+)", str(path))]
|
|
|
|
return sorted(directory.rglob("*.json"), key=natural_keys)
|
|
|
|
|
|
class Logger(typing.Protocol):
|
|
def process_measurement(self, measurement: Measurement) -> None:
|
|
pass
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
|
|
# ================================================================================================
|
|
# Simple logger
|
|
|
|
|
|
class MeasurementLogger:
|
|
"""Logger class that uses regular `rr.log` calls."""
|
|
|
|
def __init__(self, proj: pyproj.Transformer, raw: bool) -> None:
|
|
self._proj = proj
|
|
self._raw = raw
|
|
|
|
self._ignored_fields = [
|
|
"icao_id", # already the entity's path
|
|
"timestamp", # already the clock's value
|
|
]
|
|
|
|
def process_measurement(self, measurement: Measurement) -> None:
|
|
rr.set_time("unix_time", timestamp=measurement.timestamp)
|
|
|
|
if self._raw:
|
|
metadata = dataclasses.asdict(measurement)
|
|
else:
|
|
metadata = dataclasses.asdict(
|
|
measurement,
|
|
dict_factory=lambda x: {k: v for (k, v) in x if k not in self._ignored_fields and v is not None},
|
|
)
|
|
|
|
entity_path = f"aircraft/{measurement.icao_id}"
|
|
color = rr.components.Color.from_string(entity_path)
|
|
|
|
if (
|
|
measurement.latitude is not None
|
|
and measurement.longitude is not None
|
|
and measurement.barometric_altitude is not None
|
|
):
|
|
rr.log(
|
|
entity_path,
|
|
rr.Points3D(
|
|
[
|
|
self._proj.transform(
|
|
measurement.longitude,
|
|
measurement.latitude,
|
|
measurement.barometric_altitude,
|
|
),
|
|
],
|
|
colors=color,
|
|
),
|
|
rr.GeoPoints(lat_lon=[measurement.latitude, measurement.longitude]),
|
|
)
|
|
|
|
if len(metadata) > 0:
|
|
rr.log(entity_path, rr.AnyValues(**metadata))
|
|
|
|
if measurement.barometric_altitude is not None:
|
|
rr.log(
|
|
entity_path + "/barometric_altitude",
|
|
rr.Scalars(measurement.barometric_altitude),
|
|
rr.SeriesLines(colors=color),
|
|
)
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
|
|
# ================================================================================================
|
|
# Batch logger
|
|
|
|
|
|
class MeasurementBatchLogger:
|
|
"""Logger class that batches measurements and uses `rr.send_columns` calls."""
|
|
|
|
def __init__(self, proj: pyproj.Transformer, batch_size: int = 8192) -> None:
|
|
self._proj = proj
|
|
self._batch_size = batch_size
|
|
self._measurements: list[Measurement] = []
|
|
self._position_indicators: set[str] = set()
|
|
|
|
def process_measurement(self, measurement: Measurement) -> None:
|
|
self._measurements.append(measurement)
|
|
|
|
if len(self._measurements) >= 8192:
|
|
self.flush()
|
|
|
|
def flush(self) -> None:
|
|
# !!! the raw data is not sorted by timestamp, so we sort it here
|
|
df = polars.DataFrame(self._measurements).sort("timestamp")
|
|
self._measurements = []
|
|
|
|
for (icao_id,), group in df.group_by("icao_id"):
|
|
icao_id = str(icao_id)
|
|
|
|
# Note: this splitting in 3 different functions is due to the pattern of nulls in the raw data.
|
|
self.log_position_and_altitude(group, icao_id)
|
|
self.log_ground_status(group, icao_id)
|
|
self.log_metadata(group, icao_id)
|
|
|
|
def log_position_and_altitude(self, df: polars.DataFrame, icao_id: str) -> None:
|
|
entity_path = f"aircraft/{icao_id}"
|
|
df = df["timestamp", "latitude", "longitude", "barometric_altitude"].drop_nulls()
|
|
|
|
if df.height == 0:
|
|
return
|
|
|
|
if icao_id not in self._position_indicators:
|
|
color = rr.components.Color.from_string(entity_path)
|
|
rr.log(
|
|
entity_path,
|
|
rr.Points3D.from_fields(colors=color),
|
|
# TODO(cmc): That would be UB right now (and doesn't matter as long as we are on the untagged index).
|
|
# rr.GeoPoints.from_fields(colors=color),
|
|
static=True,
|
|
)
|
|
rr.log(entity_path + "/barometric_altitude", rr.SeriesLines.from_fields(colors=color), static=True)
|
|
self._position_indicators.add(icao_id)
|
|
|
|
timestamps = rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())
|
|
pos = self._proj.transform(df["longitude"], df["latitude"], df["barometric_altitude"])
|
|
|
|
raw_coordinates = rr.AnyValues.columns(
|
|
latitude=df["latitude"].to_numpy(),
|
|
longitude=df["longitude"].to_numpy(),
|
|
barometric_altitude=df["barometric_altitude"].to_numpy(),
|
|
)
|
|
|
|
rr.send_columns(
|
|
entity_path,
|
|
[timestamps],
|
|
[
|
|
*rr.Points3D.columns(positions=np.vstack(pos).T),
|
|
*rr.GeoPoints.columns(positions=np.vstack((df["latitude"], df["longitude"])).T),
|
|
*raw_coordinates,
|
|
],
|
|
)
|
|
|
|
rr.send_columns(
|
|
entity_path + "/barometric_altitude",
|
|
[timestamps],
|
|
rr.Scalars.columns(scalars=df["barometric_altitude"].to_numpy()),
|
|
)
|
|
|
|
def log_ground_status(self, df: polars.DataFrame, icao_id: str) -> None:
|
|
entity_path = f"aircraft/{icao_id}"
|
|
df = df["timestamp", "ground_status"].drop_nulls()
|
|
|
|
if df.height == 0:
|
|
return
|
|
|
|
timestamps = rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())
|
|
columns = rr.AnyValues.columns(ground_status=df["ground_status"].to_numpy())
|
|
|
|
rr.send_columns(entity_path, [timestamps], columns)
|
|
|
|
def log_metadata(self, df: polars.DataFrame, icao_id: str) -> None:
|
|
entity_path = f"aircraft/{icao_id}"
|
|
df = df["timestamp", "course", "ground_speed", "vertical_speed"].drop_nulls()
|
|
|
|
if df.height == 0:
|
|
return
|
|
|
|
metadata = rr.AnyValues.columns(
|
|
course=df["course"].to_numpy(),
|
|
ground_speed=df["ground_speed"].to_numpy(),
|
|
vertical_speed=df["vertical_speed"].to_numpy(),
|
|
)
|
|
|
|
rr.send_columns(
|
|
entity_path,
|
|
[rr.TimeColumn("unix_time", timestamp=df["timestamp"].to_numpy())],
|
|
metadata,
|
|
)
|
|
|
|
|
|
# ================================================================================================
|
|
|
|
|
|
def log_everything(paths: list[Path], raw: bool, batch: bool, batch_size: int) -> None:
|
|
measurements = load_measurements(paths)
|
|
utm_crs = find_best_utm_crs(measurements)
|
|
|
|
proj = Transformer.from_crs("EPSG:4326", utm_crs, always_xy=True)
|
|
|
|
for country_code, (level, color) in itertools.product(["DE", "CH"], [(0, (1, 0.5, 0.5))]):
|
|
log_region_boundaries_for_country(country_code, level, color, utm_crs)
|
|
|
|
# Exaggerate altitudes
|
|
rr.log("aircraft", rr.Transform3D(scale=[1, 1, 10]), static=True)
|
|
|
|
if batch:
|
|
logger: Logger = MeasurementBatchLogger(proj, batch_size)
|
|
else:
|
|
logger = MeasurementLogger(proj, raw)
|
|
|
|
for measurement in tqdm(measurements, "Logging measurements"):
|
|
if measurement.icao_id is None:
|
|
continue
|
|
|
|
logger.process_measurement(measurement)
|
|
logger.flush()
|
|
|
|
|
|
def main() -> None:
|
|
parser = ArgumentParser(description="Visualize INVOLI data")
|
|
parser.add_argument(
|
|
"--dataset",
|
|
choices=INVOLI_DATASETS.keys(),
|
|
default="2h",
|
|
help="Which dataset to automatically download and visualize",
|
|
)
|
|
parser.add_argument(
|
|
"--raw",
|
|
action="store_true",
|
|
help="If true, logs the raw data with all its issues (useful to stress edge cases in the viewer)",
|
|
)
|
|
parser.add_argument(
|
|
"--batch",
|
|
action="store_true",
|
|
default=True,
|
|
help="If true, use the batch logger function (rerun 0.18 required)",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size",
|
|
type=int,
|
|
default=8192,
|
|
help="Batch size for the batch logger",
|
|
)
|
|
parser.add_argument(
|
|
"--dir",
|
|
type=Path,
|
|
help="Use this directory of data instead of downloading a dataset",
|
|
)
|
|
rr.script_add_args(parser)
|
|
args = parser.parse_args()
|
|
|
|
if args.dir:
|
|
dataset_directory = args.dir
|
|
else:
|
|
dataset = args.dataset
|
|
dataset_ulr = INVOLI_DATASETS[dataset]
|
|
dataset_directory = DATA_DIR / dataset
|
|
if not dataset_directory.exists():
|
|
dataset_directory.mkdir()
|
|
zip_data = download_with_progress(dataset_ulr, f"dataset {dataset}")
|
|
with zipfile.ZipFile(zip_data) as zip_ref:
|
|
zip_ref.extractall(dataset_directory)
|
|
|
|
# TODO(ab): this blueprint would be massively improved by setting the 3D view's orbit point to FRA's coordinates.
|
|
blueprint = rrb.Vertical(
|
|
rrb.Horizontal(rrb.Spatial3DView(origin="/"), rrb.MapView(origin="/")),
|
|
rrb.TimeSeriesView(origin="/aircraft"),
|
|
row_shares=[3, 1],
|
|
)
|
|
rr.script_setup(args, "rerun_example_air_traffic_data", default_blueprint=blueprint)
|
|
|
|
paths = get_paths_for_directory(dataset_directory)
|
|
log_everything(paths, args.raw, args.batch, args.batch_size)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|