diff --git a/.flake8 b/.flake8 deleted file mode 100644 index c77f18b6..00000000 --- a/.flake8 +++ /dev/null @@ -1,22 +0,0 @@ -[flake8] -ignore = - # Whitespace before ':' - E203, - # Module level import not at top of file - E402, - # Line break occurred before a binary operator - W503, - # Line break occurred after a binary operator - W504 - # line break before binary operator - E203 - # line too long - E501 - # No lambdas — too strict - E731 - # Local variable name is assigned to but never used - F841 -per-file-ignores = - # imported but unused - __init__.py: F401 -max-line-length = 120 diff --git a/.github/workflows/deploy-docs.yaml b/.github/workflows/deploy-docs.yaml index aff1e7dc..bc90412f 100644 --- a/.github/workflows/deploy-docs.yaml +++ b/.github/workflows/deploy-docs.yaml @@ -3,7 +3,7 @@ name: docs on: push: branches: - - main # Change this to your branch name (e.g., docs, dev, etc.) + - main workflow_dispatch: # Allows manual triggering permissions: diff --git a/.github/workflows/pre-commit.yaml b/.github/workflows/pre-commit.yaml new file mode 100644 index 00000000..2b11178b --- /dev/null +++ b/.github/workflows/pre-commit.yaml @@ -0,0 +1,14 @@ +name: pre-commit + +on: + pull_request: + push: + branches: [main] + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v3 + - uses: pre-commit/action@v3.0.1 diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml new file mode 100644 index 00000000..399d5ee6 --- /dev/null +++ b/.github/workflows/pytest.yaml @@ -0,0 +1,26 @@ +name: pytest + +on: + push: + branches: [main] + +jobs: + build: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"] + + steps: + - uses: actions/checkout@v2 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + pip install uv + uv pip install --system -e ".[dev]" + - name: Test with pytest + run: | + pytest diff --git a/.gitignore b/.gitignore index 971a12d1..c3da1bcc 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,7 @@ *.log *.mp4 exp/ +.coverage # Sphinx documentation docs/_build/ @@ -31,4 +32,6 @@ docs/build/ _build/ .doctrees/ -jbwang_* + +# ruff +.ruff_cache/* diff --git a/.isort.cfg b/.isort.cfg deleted file mode 100644 index c168c36f..00000000 --- a/.isort.cfg +++ /dev/null @@ -1,7 +0,0 @@ -[tool.isort] -include_trailing_comma = true -known_first_party = [] -line_length = 120 -multi_line_output = 3 -profile = "black" -use_parentheses = true diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 18eca0ba..5bdfe331 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -18,31 +18,14 @@ repos: exclude: /README\.rst$|\.pot?$|\.ipynb$ args: ['--no-sort-keys', "--autofix"] -- repo: https://github.com/pycqa/isort - rev: 6.0.1 +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.14.6 hooks: - - id: isort - name: isort (python) - args: ["--profile", "black", "--filter-files", '--line-length', '120'] + - id: ruff-check # Run the linter. + args: [--fix] + exclude: __init__.py$ + - id: ruff-format # Run the formatter. exclude: __init__.py$ -- repo: https://github.com/ambv/black - rev: 25.1.0 - hooks: - - id: black - language_version: python3.12 - args: ['--line-length=120'] - files: "\\.py$" -- repo: https://github.com/myint/autoflake - rev: v2.3.1 - hooks: - - id: autoflake - args: ['--in-place', '--remove-all-unused-imports', '--remove-unused-variable'] - exclude: __init__.py$ - language_version: python3.12 -- repo: https://github.com/pycqa/flake8 - rev: 7.3.0 - hooks: - - id: flake8 - repo: https://github.com/kynan/nbstripout rev: 0.8.1 hooks: diff --git a/README.md b/README.md index 46440e04..e3ab4f2d 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,30 @@

- - - Logo + + + Logo -

123D: One Library for 2D and 3D Driving Datasets

+

123D: A Library for Driving Datasets

+

Video | Documentation

+ + +## Features + +- Unified API for driving data, including sensor data, maps, and labels. +- Support for multiple sensors storage formats. +- Fast dataformat based on [Apache Arrow](https://arrow.apache.org/). +- Visualization tools with [matplotlib](https://matplotlib.org/) and [Viser](https://viser.studio/main/). + + +> **Warning** +> +> This library is under active development and not stable. The API and features may change in future releases. +> Please report issues, feature requests, or other feedback by opening an issue. + + +## Changelog + +- **`[2025-11-21]`** v0.0.8 (silent release) + - Release of package and documentation. + - Demo data for tutorials. diff --git a/assets/logo/123D_favicon.svg b/assets/logo/123D_favicon.svg new file mode 100644 index 00000000..af4a91d3 --- /dev/null +++ b/assets/logo/123D_favicon.svg @@ -0,0 +1,91 @@ + + + + diff --git a/docs/README.md b/docs/README.md new file mode 100644 index 00000000..5a9168d6 --- /dev/null +++ b/docs/README.md @@ -0,0 +1,10 @@ +You can install relevant dependencies for editing the public documentation via: +```sh +pip install -e .[docs] +``` + +It is recommended to uses [sphinx-autobuild](https://github.com/sphinx-doc/sphinx-autobuild) (installed above) to edit and view the documentation. You can run: + +```sh +sphinx-autobuild docs docs/_build/html +``` diff --git a/docs/_static/123D_favicon.svg b/docs/_static/123D_favicon.svg new file mode 100644 index 00000000..af4a91d3 --- /dev/null +++ b/docs/_static/123D_favicon.svg @@ -0,0 +1,91 @@ + + + + diff --git a/docs/_static/123D_logo_transparent_black.svg b/docs/_static/123D_logo_transparent_black.svg new file mode 100644 index 00000000..b9214f93 --- /dev/null +++ b/docs/_static/123D_logo_transparent_black.svg @@ -0,0 +1,97 @@ + + + + diff --git a/docs/_static/123D_logo_transparent_white.svg b/docs/_static/123D_logo_transparent_white.svg new file mode 100644 index 00000000..f98386fa --- /dev/null +++ b/docs/_static/123D_logo_transparent_white.svg @@ -0,0 +1,115 @@ + + + + diff --git a/docs/_static/custom.css b/docs/_static/custom.css new file mode 100644 index 00000000..1c0c529f --- /dev/null +++ b/docs/_static/custom.css @@ -0,0 +1,3 @@ +.small-table { + font-size: 0.9em; +} diff --git a/docs/_static/logo_black.png b/docs/_static/logo_black.png deleted file mode 100644 index 6717f5f8..00000000 Binary files a/docs/_static/logo_black.png and /dev/null differ diff --git a/docs/_static/logo_white.png b/docs/_static/logo_white.png deleted file mode 100644 index 16aec842..00000000 Binary files a/docs/_static/logo_white.png and /dev/null differ diff --git a/docs/api/datatypes/detections/01_box_detections.rst b/docs/api/datatypes/detections/01_box_detections.rst new file mode 100644 index 00000000..3d5e2829 --- /dev/null +++ b/docs/api/datatypes/detections/01_box_detections.rst @@ -0,0 +1,22 @@ +Box Detections +^^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.detections.BoxDetectionWrapper + :members: + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.detections.BoxDetectionMetadata + :members: + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.detections.BoxDetectionSE2 + :members: + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.detections.BoxDetectionSE3 + :members: + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/detections/02_traffic_lights.rst b/docs/api/datatypes/detections/02_traffic_lights.rst new file mode 100644 index 00000000..e0c02778 --- /dev/null +++ b/docs/api/datatypes/detections/02_traffic_lights.rst @@ -0,0 +1,16 @@ +Traffic Lights +^^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.detections.TrafficLightDetectionWrapper + :exclude-members: __init__ + :autoclasstoc: + + +.. autoclass:: py123d.datatypes.detections.TrafficLightDetection + :no-inherited-members: + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.detections.TrafficLightStatus + :members: + :no-inherited-members: diff --git a/docs/api/datatypes/detections/index.rst b/docs/api/datatypes/detections/index.rst new file mode 100644 index 00000000..e05f6ca2 --- /dev/null +++ b/docs/api/datatypes/detections/index.rst @@ -0,0 +1,8 @@ +Detections +---------- + +.. toctree:: + :maxdepth: 1 + + 01_box_detections + 02_traffic_lights diff --git a/docs/api/datatypes/index.rst b/docs/api/datatypes/index.rst new file mode 100644 index 00000000..44585fdd --- /dev/null +++ b/docs/api/datatypes/index.rst @@ -0,0 +1,13 @@ +Datatypes +========= + + +.. toctree:: + :maxdepth: 2 + + sensors/index + detections/index + map_objects/index + metadata/index + vehicle_state/index + time/index diff --git a/docs/api/datatypes/map_objects/01_lane.rst b/docs/api/datatypes/map_objects/01_lane.rst new file mode 100644 index 00000000..f9aacd92 --- /dev/null +++ b/docs/api/datatypes/map_objects/01_lane.rst @@ -0,0 +1,11 @@ +Lane +^^^^ + +.. autoclass:: py123d.datatypes.map_objects.Lane + :exclude-members: __init__ + :autoclasstoc: + + +.. .. autoclass:: py123d.datatypes.map_objects.LaneType +.. :no-inherited-members: +.. :exclude-members: __new__ diff --git a/docs/api/datatypes/map_objects/02_lane_group.rst b/docs/api/datatypes/map_objects/02_lane_group.rst new file mode 100644 index 00000000..07e4301d --- /dev/null +++ b/docs/api/datatypes/map_objects/02_lane_group.rst @@ -0,0 +1,6 @@ +Lane Group +^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.LaneGroup + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/03_intersection.rst b/docs/api/datatypes/map_objects/03_intersection.rst new file mode 100644 index 00000000..ab52d6fe --- /dev/null +++ b/docs/api/datatypes/map_objects/03_intersection.rst @@ -0,0 +1,6 @@ +Intersection +^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.Intersection + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/04_crosswalk.rst b/docs/api/datatypes/map_objects/04_crosswalk.rst new file mode 100644 index 00000000..c025eb3e --- /dev/null +++ b/docs/api/datatypes/map_objects/04_crosswalk.rst @@ -0,0 +1,6 @@ +Crosswalk +^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.Crosswalk + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/05_carpark.rst b/docs/api/datatypes/map_objects/05_carpark.rst new file mode 100644 index 00000000..b269d4aa --- /dev/null +++ b/docs/api/datatypes/map_objects/05_carpark.rst @@ -0,0 +1,6 @@ +Carpark +^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.Carpark + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/06_walkway.rst b/docs/api/datatypes/map_objects/06_walkway.rst new file mode 100644 index 00000000..0b6caf31 --- /dev/null +++ b/docs/api/datatypes/map_objects/06_walkway.rst @@ -0,0 +1,6 @@ +Walkway +^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.Walkway + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/07_generic_drivable.rst b/docs/api/datatypes/map_objects/07_generic_drivable.rst new file mode 100644 index 00000000..4ce7b9c0 --- /dev/null +++ b/docs/api/datatypes/map_objects/07_generic_drivable.rst @@ -0,0 +1,6 @@ +Generic Drivable +^^^^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.GenericDrivable + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/08_stop_zone.rst b/docs/api/datatypes/map_objects/08_stop_zone.rst new file mode 100644 index 00000000..e48379ce --- /dev/null +++ b/docs/api/datatypes/map_objects/08_stop_zone.rst @@ -0,0 +1,6 @@ +Stop Zone +^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.StopZone + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/09_road_edge.rst b/docs/api/datatypes/map_objects/09_road_edge.rst new file mode 100644 index 00000000..b9056a16 --- /dev/null +++ b/docs/api/datatypes/map_objects/09_road_edge.rst @@ -0,0 +1,12 @@ +Road Edge +^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.RoadEdge + :no-inherited-members: + :exclude-members: __init__ + :autoclasstoc: + + +.. autoclass:: py123d.datatypes.map_objects.RoadEdgeType + :no-inherited-members: + :exclude-members: __new__ diff --git a/docs/api/datatypes/map_objects/10_road_line.rst b/docs/api/datatypes/map_objects/10_road_line.rst new file mode 100644 index 00000000..d7f538d3 --- /dev/null +++ b/docs/api/datatypes/map_objects/10_road_line.rst @@ -0,0 +1,10 @@ +Road Line +^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.RoadLine + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.map_objects.RoadLineType + :no-inherited-members: + :exclude-members: __new__ diff --git a/docs/api/datatypes/map_objects/11_base_map_objects.rst b/docs/api/datatypes/map_objects/11_base_map_objects.rst new file mode 100644 index 00000000..675a9557 --- /dev/null +++ b/docs/api/datatypes/map_objects/11_base_map_objects.rst @@ -0,0 +1,14 @@ +Base Map Objects +^^^^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.map_objects.BaseMapObject + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.map_objects.BaseMapSurfaceObject + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.datatypes.map_objects.BaseMapLineObject + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/map_objects/index.rst b/docs/api/datatypes/map_objects/index.rst new file mode 100644 index 00000000..dbd7e215 --- /dev/null +++ b/docs/api/datatypes/map_objects/index.rst @@ -0,0 +1,17 @@ +Map Objects +----------- + +.. toctree:: + :maxdepth: 1 + + 01_lane + 02_lane_group + 03_intersection + 04_crosswalk + 05_carpark + 06_walkway + 07_generic_drivable + 08_stop_zone + 09_road_edge + 10_road_line + 11_base_map_objects diff --git a/docs/api/datatypes/metadata/01_log_metadata.rst b/docs/api/datatypes/metadata/01_log_metadata.rst new file mode 100644 index 00000000..5870408f --- /dev/null +++ b/docs/api/datatypes/metadata/01_log_metadata.rst @@ -0,0 +1,6 @@ +Log Metadata +^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.metadata.LogMetadata + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/metadata/02_map_metadata.rst b/docs/api/datatypes/metadata/02_map_metadata.rst new file mode 100644 index 00000000..017ecac4 --- /dev/null +++ b/docs/api/datatypes/metadata/02_map_metadata.rst @@ -0,0 +1,6 @@ +Map Metadata +^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.metadata.MapMetadata + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/metadata/index.rst b/docs/api/datatypes/metadata/index.rst new file mode 100644 index 00000000..34560993 --- /dev/null +++ b/docs/api/datatypes/metadata/index.rst @@ -0,0 +1,8 @@ +Metadata +-------- + +.. toctree:: + :maxdepth: 1 + + 01_log_metadata + 02_map_metadata diff --git a/docs/api/datatypes/sensors/01_pinhole_camera.rst b/docs/api/datatypes/sensors/01_pinhole_camera.rst new file mode 100644 index 00000000..dd253baf --- /dev/null +++ b/docs/api/datatypes/sensors/01_pinhole_camera.rst @@ -0,0 +1,46 @@ +Pinhole Camera +^^^^^^^^^^^^^^^ + +Pinhole Camera Data +------------------- + +.. autoclass:: py123d.datatypes.sensors.PinholeCamera + :members: + :exclude-members: __init__ + :autoclasstoc: + +Pinhole Metadata +---------------- + +.. autoclass:: py123d.datatypes.sensors.PinholeCameraMetadata + :members: + :exclude-members: __init__ + :autoclasstoc: + + +Pinhole Intrinsics +------------------ + +.. autoclass:: py123d.datatypes.sensors.PinholeIntrinsics + :members: + :exclude-members: __init__ + :no-inherited-members: + :autoclasstoc: + + +Pinhole Distortion +------------------ + +.. autoclass:: py123d.datatypes.sensors.PinholeDistortion + :members: + :exclude-members: __init__ + :no-inherited-members: + :autoclasstoc: + + +Pinhole Camera Types +-------------------- + +.. autoclass:: py123d.datatypes.sensors.PinholeCameraType + :no-inherited-members: + :exclude-members: __init__, __new__ diff --git a/docs/api/datatypes/sensors/02_fisheye_mei_camera.rst b/docs/api/datatypes/sensors/02_fisheye_mei_camera.rst new file mode 100644 index 00000000..b6224164 --- /dev/null +++ b/docs/api/datatypes/sensors/02_fisheye_mei_camera.rst @@ -0,0 +1,45 @@ +Fisheye MEI Camera +^^^^^^^^^^^^^^^^^^ + +Fisheye MEI Camera Data +----------------------- + +.. autoclass:: py123d.datatypes.sensors.FisheyeMEICamera + :members: + :exclude-members: __init__ + :autoclasstoc: + +Fisheye MEI Metadata +-------------------- + +.. autoclass:: py123d.datatypes.sensors.FisheyeMEICameraMetadata + :members: + :exclude-members: __init__ + :autoclasstoc: + + +Fisheye MEI Distortion +---------------------- + +.. autoclass:: py123d.datatypes.sensors.FisheyeMEIDistortion + :members: + :exclude-members: __init__ + :no-inherited-members: + :autoclasstoc: + +Fisheye MEI Projection +---------------------- + +.. autoclass:: py123d.datatypes.sensors.FisheyeMEIProjection + :members: + :exclude-members: __init__ + :no-inherited-members: + :autoclasstoc: + + +Fisheye MEI Camera Types +------------------------ + +.. autoclass:: py123d.datatypes.sensors.FisheyeMEICameraType + :no-inherited-members: + :exclude-members: __init__, __new__ diff --git a/docs/api/datatypes/sensors/03_lidar.rst b/docs/api/datatypes/sensors/03_lidar.rst new file mode 100644 index 00000000..6d4fbee7 --- /dev/null +++ b/docs/api/datatypes/sensors/03_lidar.rst @@ -0,0 +1,27 @@ +LiDAR +^^^^^ + +LiDAR Data +---------- + +.. autoclass:: py123d.datatypes.sensors.LiDAR + :members: + :exclude-members: __init__ + :autoclasstoc: + + +LiDAR Metadata +-------------- + +.. autoclass:: py123d.datatypes.sensors.LiDARMetadata + :members: + :exclude-members: __init__ + :autoclasstoc: + + +LiDAR Types +----------- + +.. autoclass:: py123d.datatypes.sensors.LiDARType + :no-inherited-members: + :exclude-members: __init__, __new__ diff --git a/docs/api/datatypes/sensors/index.rst b/docs/api/datatypes/sensors/index.rst new file mode 100644 index 00000000..f128ad32 --- /dev/null +++ b/docs/api/datatypes/sensors/index.rst @@ -0,0 +1,9 @@ +Sensors +------- + +.. toctree:: + :maxdepth: 1 + + 01_pinhole_camera + 02_fisheye_mei_camera + 03_lidar diff --git a/docs/api/datatypes/time/index.rst b/docs/api/datatypes/time/index.rst new file mode 100644 index 00000000..2b51ca90 --- /dev/null +++ b/docs/api/datatypes/time/index.rst @@ -0,0 +1,6 @@ +Time +---- + +.. autoclass:: py123d.datatypes.time.TimePoint + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/vehicle_state/01_ego_state.rst b/docs/api/datatypes/vehicle_state/01_ego_state.rst new file mode 100644 index 00000000..1cd87405 --- /dev/null +++ b/docs/api/datatypes/vehicle_state/01_ego_state.rst @@ -0,0 +1,19 @@ +Ego Vehicle State +^^^^^^^^^^^^^^^^^ + +Ego State in SE(2) +------------------ + +.. autoclass:: py123d.datatypes.vehicle_state.EgoStateSE2 + :members: + :exclude-members: __init__ + :autoclasstoc: + + +Ego State in SE(3) +------------------ + +.. autoclass:: py123d.datatypes.vehicle_state.EgoStateSE3 + :members: + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/datatypes/vehicle_state/02_dynamic_state.rst b/docs/api/datatypes/vehicle_state/02_dynamic_state.rst new file mode 100644 index 00000000..3e9a06ce --- /dev/null +++ b/docs/api/datatypes/vehicle_state/02_dynamic_state.rst @@ -0,0 +1,10 @@ +Dynamic State +^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.vehicle_state.DynamicStateSE2 + :members: + :autoclasstoc: + +.. autoclass:: py123d.datatypes.vehicle_state.DynamicStateSE3 + :members: + :autoclasstoc: diff --git a/docs/api/datatypes/vehicle_state/03_vehicle_parameters.rst b/docs/api/datatypes/vehicle_state/03_vehicle_parameters.rst new file mode 100644 index 00000000..bd0b1882 --- /dev/null +++ b/docs/api/datatypes/vehicle_state/03_vehicle_parameters.rst @@ -0,0 +1,6 @@ +Vehicle Parameters +^^^^^^^^^^^^^^^^^^ + +.. autoclass:: py123d.datatypes.vehicle_state.VehicleParameters + :members: + :autoclasstoc: diff --git a/docs/api/datatypes/vehicle_state/index.rst b/docs/api/datatypes/vehicle_state/index.rst new file mode 100644 index 00000000..1d4aa450 --- /dev/null +++ b/docs/api/datatypes/vehicle_state/index.rst @@ -0,0 +1,9 @@ +Vehicle State +------------- + +.. toctree:: + :maxdepth: 1 + + 01_ego_state + 02_dynamic_state + 03_vehicle_parameters diff --git a/docs/api/geometry/01_primitives/01_points.rst b/docs/api/geometry/01_primitives/01_points.rst new file mode 100644 index 00000000..9160a839 --- /dev/null +++ b/docs/api/geometry/01_primitives/01_points.rst @@ -0,0 +1,8 @@ +Points +^^^^^^ + +.. autoclass:: py123d.geometry.Point2D + :autoclasstoc: + +.. autoclass:: py123d.geometry.Point3D + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/02_vectors.rst b/docs/api/geometry/01_primitives/02_vectors.rst new file mode 100644 index 00000000..28c78357 --- /dev/null +++ b/docs/api/geometry/01_primitives/02_vectors.rst @@ -0,0 +1,10 @@ +Vectors +^^^^^^^ + +.. autoclass:: py123d.geometry.Vector2D + :members: + :autoclasstoc: + +.. autoclass:: py123d.geometry.Vector3D + :members: + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/03_rotations.rst b/docs/api/geometry/01_primitives/03_rotations.rst new file mode 100644 index 00000000..06d98bb2 --- /dev/null +++ b/docs/api/geometry/01_primitives/03_rotations.rst @@ -0,0 +1,10 @@ +Rotations +^^^^^^^^^ + +.. autoclass:: py123d.geometry.Quaternion + :members: + :autoclasstoc: + +.. autoclass:: py123d.geometry.EulerAngles + :members: + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/04_se.rst b/docs/api/geometry/01_primitives/04_se.rst new file mode 100644 index 00000000..851c7052 --- /dev/null +++ b/docs/api/geometry/01_primitives/04_se.rst @@ -0,0 +1,10 @@ +Special Euclidean Groups +^^^^^^^^^^^^^^^^^^^^^^^^ + +.. autoclass:: py123d.geometry.PoseSE2 + :members: + :autoclasstoc: + +.. autoclass:: py123d.geometry.PoseSE3 + :members: + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/05_bounding_boxes.rst b/docs/api/geometry/01_primitives/05_bounding_boxes.rst new file mode 100644 index 00000000..b7514e9a --- /dev/null +++ b/docs/api/geometry/01_primitives/05_bounding_boxes.rst @@ -0,0 +1,10 @@ +Bounding Boxes +^^^^^^^^^^^^^^ + +.. autoclass:: py123d.geometry.BoundingBoxSE2 + :members: + :autoclasstoc: + +.. autoclass:: py123d.geometry.BoundingBoxSE3 + :members: + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/06_polylines.rst b/docs/api/geometry/01_primitives/06_polylines.rst new file mode 100644 index 00000000..fd572eb4 --- /dev/null +++ b/docs/api/geometry/01_primitives/06_polylines.rst @@ -0,0 +1,16 @@ +Polylines +^^^^^^^^^ + +.. autoclass:: py123d.geometry.Polyline2D + :members: + :autoclasstoc: + +.. autoclass:: py123d.geometry.PolylineSE2 + :members: + :exclude-members: __init__ + :autoclasstoc: + +.. autoclass:: py123d.geometry.Polyline3D + :members: + :exclude-members: __init__ + :autoclasstoc: diff --git a/docs/api/geometry/01_primitives/07_indexing_enums.rst b/docs/api/geometry/01_primitives/07_indexing_enums.rst new file mode 100644 index 00000000..475842dd --- /dev/null +++ b/docs/api/geometry/01_primitives/07_indexing_enums.rst @@ -0,0 +1,66 @@ +Indexing Enums +^^^^^^^^^^^^^^ + +Points +------ + +.. autoclass:: py123d.geometry.Point2DIndex + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.Point3DIndex + :members: + :no-inherited-members: + +Vectors +------- + +.. autoclass:: py123d.geometry.Vector2DIndex + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.Vector3DIndex + :members: + :no-inherited-members: + +Rotations +--------- + +.. autoclass:: py123d.geometry.QuaternionIndex + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.EulerAnglesIndex + :members: + :no-inherited-members: + +Poses +----- + +.. autoclass:: py123d.geometry.PoseSE2Index + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.PoseSE3Index + :members: + :no-inherited-members: + + +Bounding Boxes +-------------- + +.. autoclass:: py123d.geometry.BoundingBoxSE2Index + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.Corners2DIndex + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.BoundingBoxSE3Index + :members: + :no-inherited-members: + +.. autoclass:: py123d.geometry.Corners3DIndex + :members: + :no-inherited-members: diff --git a/docs/api/geometry/01_primitives/index.rst b/docs/api/geometry/01_primitives/index.rst new file mode 100644 index 00000000..fd74f01e --- /dev/null +++ b/docs/api/geometry/01_primitives/index.rst @@ -0,0 +1,14 @@ +Primitives +---------- + +.. toctree:: + :maxdepth: 2 + + + 01_points + 02_vectors + 03_rotations + 04_se + 05_bounding_boxes + 06_polylines + 07_indexing_enums diff --git a/docs/api/geometry/02_transform/01_transform_2d.rst b/docs/api/geometry/02_transform/01_transform_2d.rst new file mode 100644 index 00000000..3bee9a63 --- /dev/null +++ b/docs/api/geometry/02_transform/01_transform_2d.rst @@ -0,0 +1,37 @@ +Transforms in 2D +^^^^^^^^^^^^^^^^ + +Transform 2D points between frames +---------------------------------- + +.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_points_2d_array + +.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_points_2d_array + +.. autofunction:: py123d.geometry.transform.convert_points_2d_array_between_origins + + + +Transform SE2 poses between frames +---------------------------------- + +.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_se2_array + +.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_se2_array + +.. autofunction:: py123d.geometry.transform.convert_se2_array_between_origins + + + +Translation along frame axes +---------------------------- + +.. autofunction:: py123d.geometry.transform.translate_se2_along_body_frame + +.. autofunction:: py123d.geometry.transform.translate_se2_along_x + +.. autofunction:: py123d.geometry.transform.translate_se2_along_y + +.. autofunction:: py123d.geometry.transform.translate_se2_array_along_body_frame + +.. autofunction:: py123d.geometry.transform.translate_2d_along_body_frame diff --git a/docs/api/geometry/02_transform/02_transform_3d.rst b/docs/api/geometry/02_transform/02_transform_3d.rst new file mode 100644 index 00000000..b29c5479 --- /dev/null +++ b/docs/api/geometry/02_transform/02_transform_3d.rst @@ -0,0 +1,37 @@ +Transforms in 3D +^^^^^^^^^^^^^^^^ + +Transform 3D points between frames +---------------------------------- + +.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_points_3d_array + +.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_points_3d_array + +.. autofunction:: py123d.geometry.transform.convert_points_3d_array_between_origins + + + +Transform SE3 poses between frames +---------------------------------- + +.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_se3_array + +.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_se3_array + +.. autofunction:: py123d.geometry.transform.convert_se3_array_between_origins + + + +Translation along frame axes +---------------------------- + +.. autofunction:: py123d.geometry.transform.translate_se3_along_body_frame + +.. autofunction:: py123d.geometry.transform.translate_se3_along_x + +.. autofunction:: py123d.geometry.transform.translate_se3_along_y + +.. autofunction:: py123d.geometry.transform.translate_se3_along_z + +.. autofunction:: py123d.geometry.transform.translate_3d_along_body_frame diff --git a/docs/api/geometry/02_transform/index.rst b/docs/api/geometry/02_transform/index.rst new file mode 100644 index 00000000..60756976 --- /dev/null +++ b/docs/api/geometry/02_transform/index.rst @@ -0,0 +1,10 @@ +Transforms +---------- + + +.. toctree:: + :maxdepth: 2 + + + 01_transform_2d + 02_transform_3d diff --git a/docs/api/geometry/index.rst b/docs/api/geometry/index.rst new file mode 100644 index 00000000..a0be633f --- /dev/null +++ b/docs/api/geometry/index.rst @@ -0,0 +1,8 @@ +Geometry +======== + +.. toctree:: + :maxdepth: 2 + + 01_primitives/index + 02_transform/index diff --git a/docs/api/map/index.rst b/docs/api/map/index.rst new file mode 100644 index 00000000..0c812c93 --- /dev/null +++ b/docs/api/map/index.rst @@ -0,0 +1,9 @@ +Map API +======= + +.. autoclass:: py123d.api.MapAPI + :autoclasstoc: + + +.. autoclass:: py123d.datatypes.map_objects.MapLayer + :autoclasstoc: diff --git a/docs/api/scene/index.rst b/docs/api/scene/index.rst new file mode 100644 index 00000000..f19a179b --- /dev/null +++ b/docs/api/scene/index.rst @@ -0,0 +1,5 @@ +Scene API +========= + +.. autoclass:: py123d.api.SceneAPI + :autoclasstoc: diff --git a/docs/conf.py b/docs/conf.py index 0479a8ec..a1929722 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -8,22 +8,24 @@ project = "py123d" -copyright = "2025, 123D Contributors" -author = "123D Contributors" -release = "v0.0.7" +copyright = "2025" +author = "DanielDauner" +release = "v0.0.8" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration extensions = [ - "sphinx.ext.duration", - "sphinx.ext.doctest", + "autoclasstoc", "sphinx.ext.autodoc", - "sphinx.ext.intersphinx", + "sphinx.ext.viewcode", "sphinx.ext.autosummary", "sphinx.ext.intersphinx", "sphinx.ext.napoleon", "sphinx_copybutton", + "sphinx_autodoc_typehints", + "sphinxcontrib.youtube", + "sphinx_design", "myst_parser", ] @@ -31,6 +33,7 @@ "rtd": ("https://docs.readthedocs.io/en/stable/", None), "python": ("https://docs.python.org/3/", None), "sphinx": ("https://www.sphinx-doc.org/en/master/", None), + "numpy": ("https://numpy.org/doc/stable/", None), } intersphinx_disabled_domains = ["std"] @@ -56,20 +59,34 @@ # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". html_static_path = ["_static"] +html_favicon = "_static/123D_favicon.svg" html_theme_options = {} + autodoc_typehints = "both" autodoc_class_signature = "separated" autodoc_default_options = { "members": True, - "member-order": "bysource", + "special-members": False, + "private-members": False, + "inherited-members": True, "undoc-members": True, - "inherited-members": False, - "exclude-members": "__init__, __post_init__, __new__", + "member-order": "bysource", + "exclude-members": "__post_init__, __new__, __weakref__, __iter__, __hash__, annotations, _array", "imported-members": True, } +autosummary_generate = True + +autoclasstoc_sections = [ + "public-attrs", + "public-methods-without-dunders", + "private-methods", +] +html_css_files = ["css/theme_overrides.css", "css/version_switch.css"] +html_js_files = ["js/version_switch.js"] + # Custom CSS for color theming html_css_files = [ @@ -79,12 +96,35 @@ # Additional theme options for color customization html_theme_options.update( { - "light_logo": "logo_black.png", - "dark_logo": "logo_white.png", + "light_logo": "123D_logo_transparent_black.svg", + "dark_logo": "123D_logo_transparent_white.svg", "sidebar_hide_name": True, + "footer_icons": [ + { + "name": "GitHub", + "url": "https://github.com/autonomousvision/py123d", + "html": """ + + + + """, + "class": "", + }, + ], } ) +html_sidebars = { + "**": [ + "sidebar/brand.html", + "sidebar/search.html", + "sidebar/scroll-start.html", + "sidebar/navigation.html", + "sidebar/scroll-end.html", + "sidebar/variant-selector.html", + ] +} + # This CSS should go in /home/daniel/py123d_workspace/py123d/docs/_static/custom.css # Your conf.py already references it in html_css_files = ["custom.css"] diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 00000000..e26675e3 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,108 @@ +## Contributing + +Contributions to 123D are highly encouraged! This guide will help you get started with the development process. + +### Getting Started + +#### 1. Clone the Repository + +```sh +git clone git@github.com:autonomousvision/py123d.git +cd py123d +``` + +#### 2. Installation + +```sh +conda create -n py123d_dev python=3.12 # Optional +conda activate py123d_dev +pip install -e .[dev] +pre-commit install +``` + +The above installation should also include linting, formatting, type-checking in the pre-commit. +We use [`ruff`](https://docs.astral.sh/ruff/) as linter/formatter, for which you can run: +```sh +ruff check --fix . +ruff format . +``` +Type checking is not strictly enforced, but ideally added with [`pyright`](https://github.com/microsoft/pyright). + + +#### 3. Managing dependencies + +We try to keep dependencies minimal to ensure quick and easy installations. +However, various datasets require dependencies in order to load or preprocess the dataset. +In this case, you can add optional dependencies to the `pyproject.toml` install file. +You can follow examples of nuPlan or nuScenes. These optional dependencies can be install with + +```sh +pip install -e .[dev,nuplan,nuscenes] +``` +where you can combined the different optional dependencies. + +The optional dependencies should only be required for data pre-processing. +When writing a dataset conversion method, you can check if the necessary dependencies are installed by calling with the `check_dependencies` function. + +```python +from py123d.common.utils.dependencies import check_dependencies + +check_dependencies(["optional_package_a", "optional_package_b"], "optional_dataset") +import optional_package_a +import optional_package_b + +def load_camera_from_outdated_dataset(...) -> ...: + optional_package_a.module(...) + optional_package_b.module(...) + pass +``` +This will notify the user if `optional_dataset` is not included in the 123D install. + +Also ensure that functions/modules that require optional installs are only imported when necessary, e.g: + +```python +def load_camera_from_file(file_path: str, dataset: str) -> ...: + ... + if dataset == "optional_dataset": + from py123d.some_module import load_camera_from_outdated_dataset + + return load_camera_from_outdated_dataset(...) + ... +``` + +#### 4. Other useful tools + +If you are using VSCode, it is recommended to install: +- [autodocstring](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) - Creating docstrings (please set `"autoDocstring.docstringFormat": "sphinx-notypes"`). +- [Code Spell Checker](https://marketplace.visualstudio.com/items?itemName=streetsidesoftware.code-spell-checker) - A basic spell checker. + +Or other similar plugins depending on your preference/editor. + +### Documentation Requirements + +#### Docstrings +- **Development:** Docstrings are encouraged but not strictly required during active development +- **Format:** Use [Sphinx-style docstrings](https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html) + + +#### Sphinx documentation + +All datasets should be included in the `/docs/datasets/` documentation. Please follow the documentation format of other datasets. + +You can install relevant dependencies for editing the public documentation via: +```sh +pip install -e .[docs] +``` + +It is recommended to uses [sphinx-autobuild](https://github.com/sphinx-doc/sphinx-autobuild) (installed above) to edit and view the documentation. You can run: +```sh +sphinx-autobuild docs docs/_build/html +``` + +### Adding new datasets +TODO + + +### Questions? + +If you have any questions about contributing, please open an issue or reach out to the maintainers. diff --git a/docs/datasets/av2.rst b/docs/datasets/av2.rst index 3c09f329..f8ff3993 100644 --- a/docs/datasets/av2.rst +++ b/docs/datasets/av2.rst @@ -1,75 +1,178 @@ -Argoverse 2 ------------ +.. _av2_sensor: -.. sidebar:: Dataset Name +Argoverse 2 - Sensor +-------------------- - .. image:: https://www.argoverse.org/assets/images/reference_images/av2_vehicle.jpg - :alt: Dataset sample image +Argoverse 2 (AV2) is a collection of three datasets. +The *Sensor Dataset* includes 1000 logs of ~20 second duration, including multi-view cameras, LiDAR point clouds, maps, ego-vehicle data, and bounding boxes. +This dataset is intended to train 3D perception models for autonomous vehicles. - | **Paper:** `Name of Paper `_ - | **Download:** `Documentation `_ - | **Code:** [Code] - | **Documentation:** [License type] - | **License:** [License type] - | **Duration:** [Duration here] - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] +.. dropdown:: Overview + :open: -Description -~~~~~~~~~~~ + .. list-table:: + :header-rows: 0 + :widths: 20 60 -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + * - + - + * - :octicon:`file` Paper + - + `Argoverse 2: Next Generation Datasets for Self-Driving Perception and Forecasting `_ -Installation -~~~~~~~~~~~~ + * - :octicon:`download` Download + - `argoverse.org `_ + * - :octicon:`mark-github` Code + - `argoverse/av2-api `_ + * - :octicon:`law` License + - + `CC BY-NC-SA 4.0 `_ -[Instructions for installing or accessing the dataset] + `Argoverse Terms of Use `_ -.. code-block:: bash + MIT License + * - :octicon:`database` Available splits + - ``av2-sensor_train``, ``av2-sensor_val``, ``av2-sensor_test`` - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 - :widths: 30 5 70 - + :widths: 30 5 65 * - **Name** - **Available** - **Description** * - Ego Vehicle - - X - - [Description of ego vehicle data] + - ✓ + - State of the ego vehicle, including poses, and vehicle parameters, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - X - - [Description of ego vehicle data] + - (✓) + - The HD-Maps are in 3D, but may have artifacts due to polyline to polygon conversion (see below). For more information, see :class:`~py123d.api.MapAPI`. * - Bounding Boxes - - X - - [Description of ego vehicle data] + - ✓ + - The bounding boxes are available with the :class:`~py123d.conversion.registry.AV2SensorBoxDetectionLabel`. For more information, :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights - X - - [Description of ego vehicle data] - * - Cameras + - n/a + * - Pinhole Cameras + - ✓ + - + Includes 9 cameras, see :class:`~py123d.datatypes.sensors.PinholeCamera`: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_F0` (ring_front_center) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R0` (ring_front_right) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R1` (ring_side_right) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R2` (ring_rear_right) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L0` (ring_front_left) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L1` (ring_side_left) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L2` (ring_rear_left) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_STEREO_R` (stereo_front_right) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_STEREO_L` (stereo_front_left) + + + * - Fisheye Cameras - X - - [Description of ego vehicle data] + - n/a * - LiDARs - - X - - [Description of ego vehicle data] + - ✓ + - + Includes 2 LiDARs, see :class:`~py123d.datatypes.sensors.LiDAR`: + + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP` (top up) + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_DOWN` (top down) + + +.. dropdown:: Dataset Specific + + .. autoclass:: py123d.conversion.registry.AV2SensorBoxDetectionLabel + :members: + :no-index: + :no-inherited-members: + + .. autoclass:: py123d.conversion.registry.AV2SensorLiDARIndex + :members: + :no-index: + :no-inherited-members: + +Download +~~~~~~~~ + +You can download the Argoverse 2 Sensor dataset from the `Argoverse website `_. +You can also use directly the dataset from AWS. For that, you first need to install `s5cmd `_: + +.. code-block:: bash + + pip install s5cmd + + +Next, you can run the following bash script to download the dataset: + +.. code-block:: bash + + DATASET_NAME="sensor" # "sensor" "lidar" "motion-forecasting" "tbv" + AV2_SENSOR_ROOT="/path/to/argoverse/sensor" + + mkdir -p "$AV2_SENSOR_ROOT" + s5cmd --no-sign-request cp "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$AV2_SENSOR_ROOT" + # or: s5cmd --no-sign-request sync "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$AV2_SENSOR_ROOT" + + +The downloaded dataset should have the following structure: + +.. code-block:: none + + $AV2_SENSOR_ROOT + ├── train + │ ├── 00a6ffc1-6ce9-3bc3-a060-6006e9893a1a + │ │ ├── annotations.feather + │ │ ├── calibration + │ │ │ ├── egovehicle_SE3_sensor.feather + │ │ │ └── intrinsics.feather + │ │ ├── city_SE3_egovehicle.feather + │ │ ├── map + │ │ │ ├── 00a6ffc1-6ce9-3bc3-a060-6006e9893a1a_ground_height_surface____PIT.npy + │ │ │ ├── 00a6ffc1-6ce9-3bc3-a060-6006e9893a1a___img_Sim2_city.json + │ │ │ └── log_map_archive_00a6ffc1-6ce9-3bc3-a060-6006e9893a1a____PIT_city_31785.json + │ │ └── sensors + │ │ ├── cameras + │ │ │ └──... + │ │ └── lidar + │ │ └──... + │ └── ... + ├── test + │ └── ... + └── val + └── ... + + +Installation +~~~~~~~~~~~~ + +No additional installation steps are required beyond the standard `py123d`` installation. + + +Conversion +~~~~~~~~~~ + +To run the conversion, you either need to set the environment variable ``$AV2_DATA_ROOT`` or ``$AV2_SENSOR_ROOT``. +You can also override the file path and run: + +.. code-block:: bash + + py123d-conversion datasets=["av2_sensor_dataset"] \ + dataset_paths.av2_data_root=$AV2_DATA_ROOT # optional if env variable is set + + + +Dataset Issues +~~~~~~~~~~~~~~ -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ +- **Ego Vehicle:** The vehicle parameters are partially estimated and may be subject to inaccuracies. -[Document any known issues, limitations, or considerations when using this dataset] -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description Citation ~~~~~~~~ @@ -78,12 +181,9 @@ If you use this dataset in your research, please cite: .. code-block:: bibtex - @article{AuthorYearConference, - title={Dataset Title}, - author={Author, First and Author, Second}, - journal={Journal Name}, - year={2023}, - volume={1}, - pages={1-10}, - doi={10.1000/example} - } + @article{Wilson2023NEURIPS, + author = {Benjamin Wilson and William Qi and Tanmay Agarwal and John Lambert and Jagjeet Singh and Siddhesh Khandelwal and Bowen Pan and Ratnesh Kumar and Andrew Hartnett and Jhony Kaesemodel Pontes and Deva Ramanan and Peter Carr and James Hays}, + title = {Argoverse 2: Next Generation Datasets for Self-Driving Perception and Forecasting}, + booktitle = {Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks (NeurIPS Datasets and Benchmarks 2021)}, + year = {2021} + } diff --git a/docs/datasets/carla.rst b/docs/datasets/carla.rst index ccc921f1..68629072 100644 --- a/docs/datasets/carla.rst +++ b/docs/datasets/carla.rst @@ -1,90 +1,110 @@ +.. _carla: + CARLA ----- -.. sidebar:: Dataset Name - - .. image:: https://carla.org/img/services/getty_center_400_400.jpg - :alt: Dataset sample image - :width: 290px - - | **Paper:** `Name of Paper `_ - | **Download:** `Documentation `_ - | **Code:** [Code] - | **Documentation:** [License type] - | **License:** [License type] - | **Duration:** [Duration here] - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] - -Description -~~~~~~~~~~~ +CARLA is an open-source simulator for autonomous driving research. +As such CARLA data is synthetic and can be generated with varying sensor and environmental conditions. +The following documentation is largely incomplete and merely describes the provided demo data. -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] - -Installation -~~~~~~~~~~~~ +.. dropdown:: Quick Links + :open: -[Instructions for installing or accessing the dataset] + .. list-table:: + :header-rows: 0 + :widths: 40 60 -.. code-block:: bash + * - + - + * - :octicon:`file` Paper + - `CARLA: An Open Urban Driving Simulator `_ + * - :octicon:`globe` Website + - `carla.org/ `_ + * - :octicon:`mark-github` Code + - `github.com/carla-simulator/carla `_ + * - :octicon:`law` License + - MIT License + * - :octicon:`database` Available splits + - n/a - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 - :widths: 30 5 70 - + :widths: 20 5 70 * - **Name** - **Available** - **Description** * - Ego Vehicle - - X - - [Description of ego vehicle data] + - ✓ + - Depending on the collected dataset. For further information, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - X - - [Description of ego vehicle data] + - ✓ + - We included a conversion method of OpenDRIVE maps. For further information, see :class:`~py123d.api.MapAPI`. * - Bounding Boxes - - X - - [Description of ego vehicle data] + - ✓ + - Depending on the collected dataset. For further information, see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights - X - - [Description of ego vehicle data] - * - Cameras + - n/a + * - Pinhole Cameras + - ✓ + - Depending on the collected dataset. For further information, see :class:`~py123d.datatypes.sensors.PinholeCamera`. + * - Fisheye Cameras - X - - [Description of ego vehicle data] + - n/a * - LiDARs - - X - - [Description of ego vehicle data] + - ✓ + - Depending on the collected dataset. For further information, see :class:`~py123d.datatypes.sensors.LiDAR`. + -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ +Download +~~~~~~~~ + +n/a + +Installation +~~~~~~~~~~~~ + +n/a + +Dataset Specific +~~~~~~~~~~~~~~~~ -[Document any known issues, limitations, or considerations when using this dataset] +.. dropdown:: Box Detection Labels + + .. autoclass:: py123d.conversion.registry.DefaultBoxDetectionLabel + :members: + :no-index: + :no-inherited-members: + +.. dropdown:: LiDAR Index + + .. autoclass:: py123d.conversion.registry.DefaultLiDARIndex + :members: + :no-index: + :no-inherited-members: + + + +Dataset Issues +~~~~~~~~~~~~~~ -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description +n/a Citation ~~~~~~~~ -If you use this dataset in your research, please cite: +If you use CARLA in your research, please cite: .. code-block:: bibtex - @article{AuthorYearConference, - title={Dataset Title}, - author={Author, First and Author, Second}, - journal={Journal Name}, - year={2023}, - volume={1}, - pages={1-10}, - doi={10.1000/example} - } + @article{Dosovitskiy2017CORL, + title = {{CARLA}: {An} Open Urban Driving Simulator}, + author = {Alexey Dosovitskiy and German Ros and Felipe Codevilla and Antonio Lopez and Vladlen Koltun}, + booktitle = {Proceedings of the 1st Annual Conference on Robot Learning}, + year = {2017} + } diff --git a/docs/datasets/index.rst b/docs/datasets/index.rst index 851fbf32..7242d523 100644 --- a/docs/datasets/index.rst +++ b/docs/datasets/index.rst @@ -3,15 +3,16 @@ Datasets Brief overview of the datasets section... -This section provides comprehensive documentation for various autonomous driving and computer vision datasets. Each dataset entry includes installation instructions, available data types, known issues, and proper citation formats. +This section provides comprehensive documentation for various autonomous driving and computer vision datasets. +Each dataset entry includes installation instructions, available data types, known issues, and references for further reading. .. toctree:: :maxdepth: 1 av2 - nuplan - nuscenes carla kitti-360 - wopd - template + nuplan + nuscenes + pandaset + wodp diff --git a/docs/datasets/kitti-360.rst b/docs/datasets/kitti-360.rst index d4fc8a9f..e1565c8e 100644 --- a/docs/datasets/kitti-360.rst +++ b/docs/datasets/kitti-360.rst @@ -1,81 +1,191 @@ -KiTTI-360 +KITTI-360 --------- -.. sidebar:: Dataset Name +The KITTI-360 dataset is an extension of the popular KITTI dataset, designed for various perception tasks in autonomous driving. +The dataset includes 9 logs (called "sequences") of varying length with stereo cameras, fisheye cameras, LiDAR data, 3D primitives, and semantic annotations. - .. image:: https://www.cvlibs.net/datasets/kitti-360/images/example/3d/semantic/02400.jpg - :alt: Dataset sample image - :width: 290px +.. dropdown:: Quick Links + :open: - | **Paper:** `KITTI-360: A Novel Dataset and Benchmarks for Urban Scene Understanding in 2D and 3D `_ - | **Download:** `www.cvlibs.net/datasets/kitti-360 `_ - | **Code:** `www.github.com/autonomousvision/kitti360Scripts `_ - | **Documentation:** `kitti-360 Document`_ - | **License:** [License type] - | **Duration:** 320k image - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] + .. list-table:: + :header-rows: 0 + :widths: 20 60 -Description -~~~~~~~~~~~ + * - + - + * - :octicon:`file` Paper + - `KITTI-360: A Novel Dataset and Benchmarks for Urban Scene Understanding in 2D and 3D `_ + * - :octicon:`download` Download + - `cvlibs.net/datasets/kitti-360 `_ + * - :octicon:`mark-github` Code + - `github.com/autonomousvision/kitti360scripts `_ + * - :octicon:`law` License + - + - `CC BY-NC-SA 3.0 `_ + - MIT License + * - :octicon:`database` Available splits + - n/a -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] -Installation -~~~~~~~~~~~~ - -[Instructions for installing or accessing the dataset] - -.. code-block:: bash - - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip - -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 30 5 70 - * - **Name** - **Available** - **Description** * - Ego Vehicle - - X - - [Description of ego vehicle data] + - ✓ + - State of the ego vehicle, including poses, dynamic state, and vehicle parameters, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - X - - [Description of ego vehicle data] + - ✓ + - The maps are in 3D vector format and defined per log, see :class:`~py123d.api.MapAPI`. The map does not include lane-level information. * - Bounding Boxes - - X - - [Description of ego vehicle data] + - ✓ + - The bounding boxes are available and labeled with :class:`~py123d.conversion.registry.KITTI360BoxDetectionLabel`. For further information, see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights - X - - [Description of ego vehicle data] - * - Cameras - - X - - [Description of ego vehicle data] + - n/a + * - Pinhole Cameras + - ✓ + - The dataset has two :class:`~py123d.datatypes.sensors.PinholeCamera` in a stereo setup: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_STEREO_L` (image_00) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_STEREO_R` (image_01) + + * - Fisheye Cameras + - ✓ + - The dataset has two :class:`~py123d.datatypes.sensors.FisheyeMEICamera`: + + - :class:`~py123d.datatypes.sensors.FisheyeMEICameraType.FCAM_L` (image_02) + - :class:`~py123d.datatypes.sensors.FisheyeMEICameraType.FCAM_R` (image_03) * - LiDARs - - X - - [Description of ego vehicle data] + - ✓ + - The dataset has :class:`~py123d.datatypes.sensors.LiDAR` mounted on the roof: + + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP` (velodyne_points) + +.. dropdown:: Dataset Specific + + .. autoclass:: py123d.conversion.registry.KITTI360BoxDetectionLabel + :members: + :no-index: + :no-inherited-members: -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ + .. autoclass:: py123d.conversion.registry.KITTI360LiDARIndex + :members: + :no-index: + :no-inherited-members: + + +Download +~~~~~~~~ + +You can download the KITTI-360 dataset from the `official website `_. Please follow the instructions provided there to obtain the data. +The 123D library supports expect the dataset in the following directory structure: + +.. code-block:: text + + $KITTI360_DATA_ROOT/ + ├── calibration/ + │ ├── calib_cam_to_pose.txt + │ ├── calib_cam_to_velo.txt + │ ├── calib_sick_to_velo.txt + │ ├── image_02.yaml + │ ├── image_03.yaml + │ └── perspective.txt + ├── data_2d_raw/ + │ ├── 2013_05_28_drive_0000_sync/ + │ │ ├── image_00/ + │ │ │ ├── data_rect + │ │ │ │ ├── 0000000000.png + │ │ │ │ ├── ... + │ │ │ │ └── 0000011517.png + │ │ │ └── timestamps.txt + │ │ ├── image_01/ + │ │ │ └── ... + │ │ ├── image_02/ + │ │ │ ├── data_rgb + │ │ │ │ ├── 0000000000.png + │ │ │ │ ├── ... + │ │ │ │ └── 0000011517.png + │ │ │ └── timestamps.txt + │ │ └── image_03/ + │ │ └── ... + │ ├── ... + │ └── 2013_05_28_drive_0018_sync/ + │ └── ... + ├── data_2d_semantics/ (not yet supported) + │ └── ... + ├── data_3d_bboxes/ + │ ├── train + │ │ ├── 2013_05_28_drive_0000_sync.xml + │ │ ├── ... + │ │ └── 2013_05_28_drive_0010_sync.xml + │ └── train_full + │ ├── 2013_05_28_drive_0000_sync.xml + │ ├── ... + │ └── 2013_05_28_drive_0010_sync.xml + ├── data_3d_raw/ + │ ├── 2013_05_28_drive_0000_sync/ + │ │ └── velodyne_points/ + │ │ ├── data + │ │ │ ├── 0000000000.bin + │ │ │ ├── ... + │ │ │ └── 0000011517.bin + │ │ └── timestamps.txt + │ ├── ... + │ └── 2013_05_28_drive_0018_sync/ + │ └── ... + ├── data_3d_semantics/ (not yet supported) + │ └── ... + └── data_poses/ + ├── 2013_05_28_drive_0000_sync/ + │ ├── cam0_to_world.txt + │ ├── oxts/ + │ │ └── ... + │ └── poses.txt + ├── ... + └── 2013_05_28_drive_0018_sync/ + └── ... + +Note that not all data modalities are currently supported in 123D. For example, semantic 2D and 3D data are not yet integrated. + + +Installation +~~~~~~~~~~~~ + +No additional installation steps are required beyond the standard `py123d`` installation. + + +Conversion +~~~~~~~~~~ + +You can convert the KITTI-360 dataset by running: + +.. code-block:: bash + + py123d-conversion datasets=["kitti360_dataset"] + + +Note, that you can assign the logs of KITTI-360 to different splits (e.g., "train", "val", "test") in the ``kitti360_dataset.yaml`` config. + + +Dataset Issues +~~~~~~~~~~~~~~ -[Document any known issues, limitations, or considerations when using this dataset] +* **Ego Vehicle:** The vehicle parameters from the VW station wagon are partially estimated and may be subject to inaccuracies. +* **Map:** The ground primitives in KITTI-360 only cover surfaces, e.g. of the road, but not lane-level information. Drivable areas, road edges, walkways, driveways are included. +* **Bounding Boxes:** Bounding boxes in KITTI-360 annotated globally. We therefore determine which boxes are visible in each frame on the number of LiDAR points contained in the box. -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description Citation ~~~~~~~~ -If you use KiTTI-360 in your research, please cite: +If you use KITTI-360 in your research, please cite: .. code-block:: bibtex diff --git a/docs/datasets/nuplan.rst b/docs/datasets/nuplan.rst index 94b50f08..bc5d9c82 100644 --- a/docs/datasets/nuplan.rst +++ b/docs/datasets/nuplan.rst @@ -1,40 +1,43 @@ +.. _nuplan: + nuPlan ------ -.. sidebar:: nuPlan +nuPlan is a planning simulator that comes with a large-scale dataset for autonomous vehicle research. +This dataset contains ~1282 hours of driving logs, including ego-vehicle data, HD maps, and auto-labeled bounding boxes, spanning 4 cities. +About 120 hours of nuPlan include sensor data from 8 cameras and 5 LiDARs. - .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png - :alt: Dataset sample image - :width: 290px +.. dropdown:: Overview + :open: - | **Paper:** `Towards learning-based planning:The nuPlan benchmark for real-world autonomous driving `_ - | **Download:** `www.nuscenes.org/nuplan `_ - | **Code:** `www.github.com/motional/nuplan-devkit `_ - | **Documentation:** `nuPlan Documentation `_ - | **License:** `CC BY-NC-SA 4.0 `_, `nuPlan Dataset License `_ - | **Duration:** 1282 hours (120 hours of sensor data) - | **Supported Versions:** [TODO] - | **Redistribution:** [TODO] + .. list-table:: + :header-rows: 0 + :widths: 20 60 -Description -~~~~~~~~~~~ + * - + - + * - :octicon:`file` Papers + - + `Towards learning-based planning: The nuplan benchmark for real-world autonomous driving `_ -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + `nuplan: A closed-loop ml-based planning benchmark for autonomous vehicles `_ + * - :octicon:`download` Download + - `nuplan.org `_ + * - :octicon:`mark-github` Code + - `nuplan-devkit `_ + * - :octicon:`law` License + - + `CC BY-NC-SA 4.0 `_ -Installation -~~~~~~~~~~~~ + `nuPlan Terms of Use `_ -[Instructions for installing or accessing the dataset] + Apache License 2.0 + * - :octicon:`database` Available splits + - ``nuplan_train``, ``nuplan_val``, ``nuplan_test``, ``nuplan-mini_train``, ``nuplan-mini_val``, ``nuplan-mini_test`` -.. code-block:: bash - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip - -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 @@ -45,37 +48,240 @@ Available Data - **Description** * - Ego Vehicle - ✓ - - [Description of ego vehicle data] + - State of the ego vehicle, including poses, dynamic state, and vehicle parameters, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - ✓ - - [Description of map data] + - (✓) + - The HD-Maps are in 2D vector format and defined per-location. For more information, see :class:`~py123d.api.MapAPI`. * - Bounding Boxes - - X - - [Description of bounding boxes data] + - ✓ + - The bounding boxes are available, see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights + - ✓ + - Traffic lights include the status and the lane id they are associated with, see :class:`~py123d.datatypes.detections.TrafficLightDetectionWrapper`. + * - Pinhole Cameras + - (✓) + - + Subset of nuPlan includes 8x :class:`~py123d.datatypes.sensors.PinholeCamera`: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_F0`: Front camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R0`: Right front camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R1`: Right middle camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R2`: Right rear camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L0`: Left front camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L1`: Left middle camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L2`: Left rear camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_B0`: Back camera + * - Fisheye Cameras - X - - [Description of traffic lights data] - * - Cameras - - X - - [Description of cameras data] + - * - LiDARs - - X - - [Description of LiDARs data] + - (✓) + - + Subset of nuPlan includes 5x :class:`~py123d.datatypes.sensors.LiDAR`: -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP`: Top + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_FRONT`: Front + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_SIDE_LEFT`: Side left + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_SIDE_RIGHT`: Side right + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_BACK`: Rear -[Document any known issues, limitations, or considerations when using this dataset] +.. dropdown:: Dataset Specific -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description + .. autoclass:: py123d.conversion.registry.NuPlanBoxDetectionLabel + :members: + :no-index: + :no-inherited-members: -Citation + .. autoclass:: py123d.conversion.registry.NuPlanLiDARIndex + :members: + :no-index: + :no-inherited-members: + + + +Download ~~~~~~~~ +You can install the nuPlan dataset either by downloading the files from the `official website `_ or by using the following bash script: + +.. dropdown:: Download Scripts + + **License**: + + .. code-block:: bash + + # NOTE: Please check the LICENSE file when downloading the nuPlan dataset + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/LICENSE + + **Maps** (required for ``nuplan`` and ``nuplan-mini``): + + .. code-block:: bash + + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-maps-v1.1.zip + + + **Logs**: + + .. code-block:: bash + + # 1. nuplan_train + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_boston.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_pittsburgh.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_singapore.zip + for split in {1..6}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_vegas_${split}.zip + done + + # 2. nuplan_val + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_test.zip + + # 3. nuplan_test + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_val.zip + + # 4. nuplan-mini_train, nuplan-mini_val, nuplan-mini_test + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_mini.zip + + + **Sensors**: + + .. code-block:: bash + + # 1. nuplan_train + for split in {0..42}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/train_set/nuplan-v1.1_train_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/train_set/nuplan-v1.1_train_lidar_${split}.zip + done + + # 2. nuplan_val + for split in {0..11}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/val_set/nuplan-v1.1_val_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/val_set/nuplan-v1.1_val_lidar_${split}.zip + done + + # 3. nuplan_test + for split in {0..11}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/test_set/nuplan-v1.1_test_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/test_set/nuplan-v1.1_test_lidar_${split}.zip + done + + # 4. nuplan_mini_train, nuplan_mini_val, nuplan_mini_test + for split in {0..8}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_lidar_${split}.zip + done + +The 123D conversion expects the following directory structure: + +.. code-block:: none + + $NUPLAN_DATA_ROOT + ├── maps (or $NUPLAN_MAPS_ROOT) + │ ├── nuplan-maps-v1.0.json + │ ├── sg-one-north + │ │ └── 9.17.1964 + │ │ └── map.gpkg + │ ├── us-ma-boston + │ │ └── 9.12.1817 + │ │ └── map.gpkg + │ ├── us-nv-las-vegas-strip + │ │ └── 9.15.1915 + │ │ └── map.gpkg + │ └── us-pa-pittsburgh-hazelwood + │ └── 9.17.1937 + │ └── map.gpkg + └── nuplan-v1.1 + ├── splits + │ ├── mini + │ │ ├── 2021.05.12.22.00.38_veh-35_01008_01518.db + │ │ ├── 2021.06.09.17.23.18_veh-38_00773_01140.db + │ │ ├── ... + │ │ └── 2021.10.11.08.31.07_veh-50_01750_01948.db + │ ├── test + │ │ └── ... + │ └── trainval + │ ├── 2021.05.12.22.00.38_veh-35_01008_01518.db + │ ├── 2021.06.09.17.23.18_veh-38_00773_01140.db + │ ├── ... + │ └── 2021.10.11.08.31.07_veh-50_01750_01948.db + └── sensor_blobs (or $NUPLAN_SENSOR_ROOT) + ├── 2021.05.12.22.00.38_veh-35_01008_01518 + │ ├── CAM_F0 + │ │ ├── c082c104b7ac5a71.jpg + │ │ ├── af380db4b4ca5d63.jpg + │ │ ├── ... + │ │ └── 2270fccfb44858b3.jpg + │ ├── CAM_B0 + │ ├── CAM_L0 + │ ├── CAM_L1 + │ ├── CAM_L2 + │ ├── CAM_R0 + │ ├── CAM_R1 + │ ├── CAM_R2 + │ └──MergedPointCloud + │ ├── 03fafcf2c0865668.pcd + │ ├── 5aee37ce29665f1b.pcd + │ ├── ... + │ └── 5fe65ef6a97f5caf.pcd + │ + ├── 2021.06.09.17.23.18_veh-38_00773_01140 + ├── ... + └── 2021.10.11.08.31.07_veh-50_01750_01948 + + +Lastly, you need to add the following environment variables to your ``~/.bashrc`` according to your installation paths: + +.. code-block:: bash + + export NUPLAN_DATA_ROOT=/path/to/nuplan/data/root + export NUPLAN_MAPS_ROOT=/path/to/nuplan/data/root/maps + export NUPLAN_SENSOR_ROOT=/path/to/nuplan/data/root/nuplan-v1.1/sensor_blobs + +Or configure the config ``py123d/script/config/common/default_dataset_paths.yaml`` accordingly. + +Installation +~~~~~~~~~~~~ + +For nuPlan, additional installation that are included as optional dependencies in ``py123d`` are required. You can install them via: + +.. tab-set:: + + .. tab-item:: PyPI + + .. code-block:: bash + + pip install py123d[nuplan] + + .. tab-item:: Source + + .. code-block:: bash + + pip install -e .[nuplan] + +Conversion +~~~~~~~~~~~~ + +You can convert the nuPlan dataset (or mini dataset) by running: + +.. code-block:: bash + + py123d-conversion datasets=["nuplan_dataset"] + # or + py123d-conversion datasets=["nuplan_mini_dataset"] + + + +Dataset Issues +~~~~~~~~~~~~~~ + +* **Map:** The HD-Maps are only available in 2D. +* **Camera & LiDAR:** There are synchronization issues between the sensors and the ego vehicle state. +* **Bounding Boxes:** Due to the auto-labeling process of nuPlan, some bounding boxes may be noisy. +* **Traffic Lights:** The status of the traffic lights are inferred from the vehicle movements. As such, there may be incorrect labels. + +Citation +~~~~~~~~ -If you use this dataset in your research, please cite: +If you use nuPlan in your research, please cite: .. code-block:: bibtex diff --git a/docs/datasets/nuscenes.rst b/docs/datasets/nuscenes.rst index 638c5ad6..a8bcf2f9 100644 --- a/docs/datasets/nuscenes.rst +++ b/docs/datasets/nuscenes.rst @@ -1,90 +1,205 @@ +.. _nuscenes: + nuScenes -------- -.. sidebar:: nuScenes - - .. image:: https://ar5iv.labs.arxiv.org/html/1903.11027/assets/figures/sensors.jpg - :alt: Dataset sample image - :width: 290px +The nuScenes dataset is multi-modal autonomous driving dataset that includes data from cameras, LiDARs, and radars, along with detailed annotations from Boston and Singapore. +In total, the dataset contains 1000 driving logs, each of 20 second duration, resulting in 5.5 hours of data. +All logs include ego-vehicle data, camera images, LiDAR point clouds, bounding boxes, and map data. - | **Paper:** `Name of Paper `_ - | **Download:** `Documentation `_ - | **Code:** [Code] - | **Documentation:** [License type] - | **License:** [License type] - | **Duration:** [Duration here] - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] -Description -~~~~~~~~~~~ +.. dropdown:: Overview + :open: -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + .. list-table:: + :header-rows: 0 + :widths: 20 60 -Installation -~~~~~~~~~~~~ + * - + - + * - :octicon:`file` Papers + - + `nuscenes: A multimodal dataset for autonomous driving `_ + * - :octicon:`download` Download + - `nuscenes.org `_ + * - :octicon:`mark-github` Code + - `nuscenes-devkit `_ + * - :octicon:`law` License + - + `CC BY-NC-SA 4.0 `_ -[Instructions for installing or accessing the dataset] + `nuScenes Terms of Use `_ -.. code-block:: bash + Apache License 2.0 + * - :octicon:`database` Available splits + - ``nuscenes_train``, ``nuscenes_val``, ``nuscenes_test``, ``nuscenes-mini_train``, ``nuscenes-mini_val``, ``nuscenes-mini_test`` - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 30 5 70 - * - **Name** - **Available** - **Description** * - Ego Vehicle - - X - - [Description of ego vehicle data] + - ✓ + - State of the ego vehicle, including poses, dynamic state, and vehicle parameters, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - X - - [Description of ego vehicle data] + - (✓) + - The HD-Maps are in 2D vector format and defined per-location. For more information, see :class:`~py123d.api.MapAPI`. * - Bounding Boxes - - X - - [Description of ego vehicle data] + - ✓ + - The bounding boxes are available with the :class:`~py123d.conversion.registry.NuScenesBoxDetectionLabel`. For more information, see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights - X - - [Description of ego vehicle data] - * - Cameras + - + * - Pinhole Cameras + - ✓ + - + nuScenes includes 6x :class:`~py123d.datatypes.sensors.PinholeCamera`: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_F0`: CAM_FRONT + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R0`: CAM_FRONT_RIGHT + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R1`: CAM_BACK_RIGHT + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L0`: CAM_FRONT_LEFT + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L1`: CAM_BACK_LEFT + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_B0`: CAM_BACK + * - Fisheye Cameras - X - - [Description of ego vehicle data] + - * - LiDARs - - X - - [Description of ego vehicle data] + - ✓ + - nuScenes has one :class:`~py123d.datatypes.sensors.LiDAR` of type :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP`. +.. dropdown:: Dataset Specific + + .. autoclass:: py123d.conversion.registry.NuScenesBoxDetectionLabel + :members: + :no-index: + :no-inherited-members: + + .. autoclass:: py123d.conversion.registry.NuScenesLiDARIndex + :members: + :no-index: + :no-inherited-members: + + +Download +~~~~~~~~ + +You need to download the nuScenes dataset from the `official website `_. +From there, you need the following parts: + +* CAN bus expansion pack +* Map expansion pack (v1.3) +* Full dataset (v1.0) + + * Mini dataset (v1.0-mini) (for quick testing) + * Train/Val split (v1.0-trainval) (for the complete dataset) + * Test split (v1.0-test) (for the complete dataset) + + + +The 123D conversion expects the following directory structure: + +.. code-block:: none + + $NUSCENES_DATA_ROOT + ├── can_bus/ + │ ├── scene-0001_meta.json + │ ├── ... + │ └── scene-1110_zoe_veh_info.json + ├── maps/ + │ ├── 36092f0b03a857c6a3403e25b4b7aab3.png + │ ├── ... + │ ├── 93406b464a165eaba6d9de76ca09f5da.png + │ ├── basemap/ + │ │ └── ... + │ ├── expansion/ + │ │ └── ... + │ └── prediction/ + │ └── ... + ├── samples/ + │ ├── CAM_BACK/ + │ │ └── ... + │ ├── ... + │ └── RADAR_FRONT_RIGHT/ + │ └── ... + ├── sweeps/ + │ └── ... + ├── v1.0-mini/ + │ ├── attribute.json + │ ├── ... + │ └── visibility.json + ├── v1.0-test/ + │ ├── attribute.json + │ ├── ... + │ └── visibility.json + └── v1.0-trainval/ + ├── attribute.json + ├── ... + └── visibility.json + +Lastly, you need to add the following environment variables to your ``~/.bashrc`` according to your installation paths: + +.. code-block:: bash + + export NUSCENES_DATA_ROOT=/path/to/nuplan/data/root + +Or configure the config ``py123d/script/config/common/default_dataset_paths.yaml`` accordingly. -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ +Installation +~~~~~~~~~~~~ + +For nuScenes, additional installation that are included as optional dependencies in ``py123d`` are required. You can install them via: + +.. tab-set:: + + .. tab-item:: PyPI + + .. code-block:: bash + + pip install py123d[nuscenes] + + .. tab-item:: Source + + .. code-block:: bash + + pip install -e .[nuscenes] + +Conversion +~~~~~~~~~~~~ + +You can convert the nuScenes dataset (or mini dataset) by running: + +.. code-block:: bash + + py123d-conversion datasets=["nuscenes_dataset"] + # or + py123d-conversion datasets=["nuscenes_mini_dataset"] + + + +Dataset Issues +~~~~~~~~~~~~~~ -[Document any known issues, limitations, or considerations when using this dataset] +* **Map:** The HD-Maps are only available in 2D. +* ... -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description Citation ~~~~~~~~ -If you use this dataset in your research, please cite: +If you use nuPlan in your research, please cite: .. code-block:: bibtex - @article{AuthorYearConference, - title={Dataset Title}, - author={Author, First and Author, Second}, - journal={Journal Name}, - year={2023}, - volume={1}, - pages={1-10}, - doi={10.1000/example} - } + @article{Caesar2020CVPR, + title={nuscenes: A multimodal dataset for autonomous driving}, + author={Caesar, Holger and Bankiti, Varun and Lang, Alex H and Vora, Sourabh and Liong, Venice Erin and Xu, Qiang and Krishnan, Anush and Pan, Yu and Baldan, Giancarlo and Beijbom, Oscar}, + booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, + year={2020} + } diff --git a/docs/datasets/pandaset.rst b/docs/datasets/pandaset.rst new file mode 100644 index 00000000..8879cf8b --- /dev/null +++ b/docs/datasets/pandaset.rst @@ -0,0 +1,213 @@ +.. _pandaset: + +PandaSet +-------- + +The PandaSet dataset is a multi-modal dataset that includes data from cameras and LiDARs, along with detailed 3D bounding box annotations. +It includes 103 logs of 8 second duration, resulting in about 0.2 hours of data. +PandaSet stands out, due to its no-cost commercial license. + + +.. dropdown:: Overview + :open: + + .. list-table:: + :header-rows: 0 + :widths: 20 60 + + * - + - + * - :octicon:`file` Paper + - `PandaSet: Advanced Sensor Suite Dataset for Autonomous Driving `_ + * - :octicon:`download` Download + - + - `scale.com/open-av-datasets/pandaset `_ (official but discontinued). + - `huggingface.co/datasets/georghess/pandaset `_ (unofficial). + * - :octicon:`mark-github` Code + - `github.com/scaleapi/pandaset-devkit `_ + * - :octicon:`law` License + - + - `CC BY 4.0 `_ + - No-cost commercial license* + - Apache License 2.0 + * - :octicon:`database` Available splits + - n/a + +.. dropdown:: Dataset Terms of Use* + + Dataset Terms of Use + Scale AI, Inc. and Hesai Photonics Technology Co., Ltd and their affiliates (hereinafter "Licensors") strive to enhance public access to and use of data that they collect, annotate, and publish. The data are organized in datasets (the “Datasets”) listed at pandaset.org (the “Website”). The Datasets are collections of data, managed by Licensors and provided in a number of machine-readable formats. Licensors provide any individual or entity (hereinafter You” or “Your”) with access to the Datasets free of charge subject to the terms of this agreement (hereinafter “Dataset Terms”). Use of any data derived from the Datasets, which may appear in any format such as tables, charts, devkit, documentation, or code, is also subject to these Dataset Terms. By downloading any Datasets or using any Datasets, You are agreeing to be bound by the Dataset Terms. If you are downloading any Datasets or using any Datasets for an organization, you are agreeing to these Dataset Terms on behalf of that organization. If you do not have the right to agree to these Dataset Terms, do not download or use the Datasets. + + Licenses + Unless specifically labeled otherwise, these Datasets are provided to You under a Creative Commons Attribution 4.0 International Public License (“CC BY 4.0”), with the additional terms included in these Dataset Terms. The CC BY 4.0 may be accessed at https://creativecommons.org/licenses/by/4.0/. When You download or use the Datasets from the Website or elsewhere, You are agreeing to comply with the terms of CC BY 4.0. Where these Dataset Terms conflict with the terms of CC BY 4.0, these Dataset Terms will control. + + Privacy + Licensors prohibit You from using the Datasets in any manner to identify or invade the privacy of any person whose personally identifiable information or personal data may have been incidentally collected in the creation of this Dataset, even when such use is otherwise legal. An individual with any privacy concerns, including a request to remove your personally identifiable information or personal data from the Dataset, may contact us by sending an e-mail to privacy@scaleapi.com. + + No Publicity Rights + You may not use the name, any trademark, official mark, official emblem, or logo of either Licensor, or any of either Licensor’s other means of promotion or publicity without the applicable Licensor’s prior written consent nor in any event to represent or imply an association or affiliation with a Licensor, except as required to comply with the attribution requirements of the CC BY 4.0 license. + + Termination + Licensors may terminate Your access to all or any part of the Datasets or the Website at any time, with or without cause, with or without notice, effective immediately. All provisions of the Dataset Terms which by their nature should survive termination will survive termination, including, without limitation, warranty disclaimers, indemnity, and limitations of liability. + + Indemnification + You will indemnify and hold Licensors harmless from and against any and all claims, loss, cost, expense, liability, or damage, including, without limitation, all reasonable attorneys’ fees and court costs, arising from (i) Your use or misuse of the Website or the Datasets; (ii) Your access to the Website; (iii) Your violation of the Dataset Terms; or (iv) infringement by You, or any third party using Your account, of any intellectual property or other right of any person or entity. Such losses, costs, expenses, damages, or liabilities will include, without limitation, all actual, general, special, indirect, incidental, and consequential damages. + + Dispute Resolution + These Dataset Terms will be governed by and interpreted in accordance with the laws of California (excluding the conflict of laws rules thereof). All disputes under these Dataset Terms will be resolved in the applicable state or federal courts of San Francisco, California. You consent to the jurisdiction of such courts and waive any jurisdictional or venue defenses otherwise available. + + Miscellaneous + You agree that it is Your responsibility to comply with all applicable laws with respect to Your use and publication of the Datasets or derivatives thereof, including any applicable privacy, data protection, security, and export control laws. These Dataset Terms constitute the entire agreement between You and Licensors with respect to the subject matter of these Dataset Terms and supersedes any prior or contemporaneous agreements whether written or oral. If a court of competent jurisdiction finds any term of these Dataset Terms to be unenforceable, the unenforceable term will be modified to reflect the parties’ intention and only to the extent necessary to make the term enforceable. The remaining provisions of these Dataset Terms will remain in effect. You may not assign these Dataset Terms without the prior written consent of the Licensors. The Licensors may assign, transfer, or delegate any of their rights and obligations under these Dataset Terms without consent. The parties are independent contractors. No failure or delay by either party in exercising a right under these Dataset Terms will constitute a waiver of that right. A waiver of a default is not a waiver of any subsequent default. These Dataset Terms may be amended by the Licensors from time to time in our discretion. If an update affects your use of the Dataset, Licensors will notify you before the updated terms are effective for your use. + + +Available Modalities +~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - ✓ + - The poses and vehicle parameters are provided or inferred from the documentation, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. + * - Map + - X + - n/a + * - Bounding Boxes + - ✓ + - Bounding boxes are available with the :class:`~py123d.conversion.registry.PandasetBoxDetectionLabel`. For more information, see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. + * - Traffic Lights + - X + - n/a + * - Pinhole Cameras + - ✓ + - + Pandaset has 6x :class:`~py123d.datatypes.sensors.PinholeCamera`: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_F0`: front_camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L0`: front_left_camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R0`: front_right_camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L1`: left_camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R1`: right_camera + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_B0`: back_camera + + * - Fisheye Cameras + - X + - n/a + * - LiDARs + - ✓ + - + Pandaset has 2x :class:`~py123d.datatypes.sensors.LiDAR`: + + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP`: main_pandar64 + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_FRONT`: front_gt + + +.. dropdown:: Dataset Specific + + .. autoclass:: py123d.conversion.registry.PandasetBoxDetectionLabel + :members: + :no-index: + :no-inherited-members: + + .. autoclass:: py123d.conversion.registry.PandasetLiDARIndex + :members: + :no-index: + :no-inherited-members: + + +Download +~~~~~~~~ + +Since a few years, the official PandaSet dataset download no longer available (see `Issue #151 `_). +However, unofficial copies of the dataset can be found on `Hugging Face `_ or `Kaggle `_. + +The 123D conversion expects the following directory structure: + +.. code-block:: text + + $PANDASET_DATA_ROOT/ + ├── 001/ + │ ├── annotations/ + │ │ ├── cuboids/ + │ │ │ ├── 00.pkl.gz + │ │ │ ├── ... + │ │ │ └── 79.pkl.gz + │ │ └── semseg/ (currently not used) + │ │ ├── 00.pkl.gz + │ │ ├── ... + │ │ ├── 79.pkl.gz + │ │ └── classes.json + │ ├── camera/ + │ │ ├── back_camera/ + │ │ │ ├── 00.jpg + │ │ │ ├── ... + │ │ │ ├── 79.jpg + │ │ │ ├── intrinsics.json + │ │ │ ├── poses.json + │ │ │ └── timestamps.json + │ │ ├── front_camera/ + │ │ │ └── ... + │ │ ├── front_left_camera/ + │ │ │ └── ... + │ │ ├── front_right_camera/ + │ │ │ └── ... + │ │ ├── left_camera/ + │ │ │ └── ... + │ │ └── right_camera/ + │ │ └── ... + │ ├── LICENSE.txt + │ ├── lidar/ + │ │ ├── 00.pkl.gz + │ │ ├── ... + │ │ ├── 79.pkl.gz + │ │ ├── poses.json + │ │ └── timestamps.json + │ └── meta/ + │ ├── gps.json + │ └── timestamps.json + ├── ... + └── 158/ + └── ... + + + +Installation +~~~~~~~~~~~~ + +No additional installation steps are required beyond the standard ``py123d`` installation. + + +Conversion +~~~~~~~~~~~~ + +You can convert the PandaSet by running: + +.. code-block:: bash + + py123d-conversion datasets=["pandaset_dataset"] + + +Dataset Issues +~~~~~~~~~~~~~~ + +* **Ego Vehicle:** The ego vehicle parameters are estimates from the vehicle model. The exact location of the IMU/GPS sensor and the bounding box dimensions of the ego vehicle may not be accurate. +* **Bounding Boxes:** PandaSet provides bounding boxes that fall in the overlap of the LiDAR region twice (for each point cloud). The current implementation only uses the bounding boxes of the top LiDAR sensor. +* **LiDAR:** PandaSet does not motion compensate the LiDAR sweeps (in contrast to other datasets). Artifacts remain visible. + +Citation +~~~~~~~~ + +If you use PandaSet in your research, please cite: + +.. code-block:: bibtex + + @article{Xiao2021ITSC, + title={Pandaset: Advanced sensor suite dataset for autonomous driving}, + author={Xiao, Pengchuan and Shao, Zhenlei and Hao, Steven and Zhang, Zishuo and Chai, Xiaolin and Jiao, Judy and Li, Zesong and Wu, Jian and Sun, Kai and Jiang, Kun and others}, + booktitle={2021 IEEE international intelligent transportation systems conference (ITSC)}, + year={2021}, + } diff --git a/docs/datasets/template.rst b/docs/datasets/template.rst index 29797269..cac5cde0 100644 --- a/docs/datasets/template.rst +++ b/docs/datasets/template.rst @@ -1,70 +1,103 @@ Template -------- +... -.. sidebar:: Dataset Name +.. dropdown:: Quick Links + :open: - .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png - :alt: Dataset sample image - :width: 290px + .. list-table:: + :header-rows: 0 + :widths: 20 60 - | **Paper:** `Name of Paper `_ - | **Download:** `Documentation `_ - | **Code:** [Code] - | **Documentation:** [License type] - | **License:** [License type] - | **Duration:** [Duration here] - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] + * - + - + * - :octicon:`file` Paper + - ... + * - :octicon:`download` Download + - ... + * - :octicon:`mark-github` Code + - ... + * - :octicon:`law` License + - ... + * - :octicon:`database` Available splits + - ... -Description -~~~~~~~~~~~ -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] - -Installation -~~~~~~~~~~~~ - -[Instructions for installing or accessing the dataset] - -.. code-block:: bash - - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip - -Available Data -~~~~~~~~~~~~~~ +Available Modalities +~~~~~~~~~~~~~~~~~~~~ .. list-table:: :header-rows: 1 :widths: 30 5 70 - * - **Name** - **Available** - **Description** * - Ego Vehicle - - X - - [Description of ego vehicle data] + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. * - Map - - X - - [Description of ego vehicle data] + - ✓ / (✓) / X + - ..., see :class:`~py123d.api.MapAPI`. * - Bounding Boxes - - X - - [Description of ego vehicle data] + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. * - Traffic Lights - - X - - [Description of ego vehicle data] - * - Cameras - - X - - [Description of ego vehicle data] + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.detections.TrafficLightDetectionWrapper`. + * - Pinhole Cameras + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.sensors.PinholeCamera`. + * - Fisheye Cameras + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.sensors.FisheyeCamera`. * - LiDARs - - X - - [Description of ego vehicle data] + - ✓ / (✓) / X + - ..., see :class:`~py123d.datatypes.sensors.LiDAR`. + + +Download +~~~~~~~~ + +... + +The 123D conversion expects the following directory structure: + +Installation +~~~~~~~~~~~~ -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ +For *Template*, additional installation that are included as optional dependencies in ``py123d`` are required. You can install them via: + +.. code-block:: bash + + pip install py123d[template] + +Or if you are installing from source: + +.. code-block:: bash + + pip install -e .[template] + + +Dataset Specific +~~~~~~~~~~~~~~~~ + +.. dropdown:: Box Detection Labels + + .. autoclass:: py123d.conversion.registry.DefaultBoxDetectionLabel + :members: + :no-inherited-members: + +.. dropdown:: LiDAR Index + + .. autoclass:: py123d.conversion.registry.DefaultLiDARIndex + :members: + :no-inherited-members: + + + +Dataset Issues +~~~~~~~~~~~~~~ [Document any known issues, limitations, or considerations when using this dataset] @@ -72,19 +105,17 @@ Dataset Specific Issues * Issue 2: Description * Issue 3: Description + Citation ~~~~~~~~ -If you use this dataset in your research, please cite: +If you use *Template* in your research, please cite: .. code-block:: bibtex - @article{AuthorYearConference, - title={Dataset Title}, - author={Author, First and Author, Second}, - journal={Journal Name}, - year={2023}, - volume={1}, - pages={1-10}, - doi={10.1000/example} - } + @article{AuthorYearConference, + title={Template: Some Dataset for Autonomous Driving}, + author={}, + booktitle={}, + year={} + } diff --git a/docs/datasets/wodp.rst b/docs/datasets/wodp.rst new file mode 100644 index 00000000..72d15a93 --- /dev/null +++ b/docs/datasets/wodp.rst @@ -0,0 +1,178 @@ +Waymo Open Dataset - Perception +------------------------------- + +The Waymo Open Dataset (WOD) is a collective term for publicly available datasets from Waymo. +The *Perception Dataset*, abbreviated as WOD-P, is a high-quality dataset targeted for perceptions tasks, such as +With 1150 logs each spanning 20 seconds, the dataset includes about 6.4 hours + +.. dropdown:: Overview + :open: + + .. list-table:: + :header-rows: 0 + :widths: 20 60 + + * - + - + * - :octicon:`file` Paper + - `Scalability in Perception for Autonomous Driving: Waymo Open Dataset `_ + * - :octicon:`download` Download + - `waymo.com/open `_ + * - :octicon:`mark-github` Code + - `waymo-open-dataset `_ + * - :octicon:`law` License + - + `Waymo Dataset License Agreement for Non-Commercial Use `_ + + Apache License 2.0 + `Code Specific Licenses `_ + + * - :octicon:`database` Available splits + - ``wodp_train``, ``wodp_val``, ``wodp_test`` + + +Available Modalities +~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 20 5 75 + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - ✓ + - State of the ego vehicle, including poses, and vehicle parameters, see :class:`~py123d.datatypes.vehicle_state.EgoStateSE3`. + * - Map + - (✓) + - The HD-Maps are in 3D, but may have artifacts due to polyline to polygon conversion (see below). For more information, see :class:`~py123d.api.MapAPI`. + * - Bounding Boxes + - ✓ + - The bounding boxes are available with the :class:`~py123d.conversion.registry.WOPDBoxDetectionLabel`. For more information, :class:`~py123d.datatypes.detections.BoxDetectionWrapper`. + * - Traffic Lights + - X + - n/a + * - Pinhole Cameras + - ✓ + - + Includes 5 cameras, see :class:`~py123d.datatypes.sensors.PinholeCamera`: + + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_F0` (front_camera) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L0` (front_left_camera) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R0` (front_right_camera) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_L1` (left_camera) + - :class:`~py123d.datatypes.sensors.PinholeCameraType.PCAM_R1` (right_camera) + + * - Fisheye Cameras + - X + - n/a + * - LiDARs + - ✓ + - + Includes 5 LiDARs, see :class:`~py123d.datatypes.sensors.LiDAR`: + + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_TOP` (top) + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_FRONT` (front) + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_SIDE_LEFT` (side_left) + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_SIDE_RIGHT` (side_right) + - :class:`~py123d.datatypes.sensors.LiDARType.LIDAR_BACK` (rear) + +.. dropdown:: Dataset Specific + + + .. autoclass:: py123d.conversion.registry.WOPDBoxDetectionLabel + :members: + :no-inherited-members: + + .. autoclass:: py123d.conversion.registry.WOPDLiDARIndex + :members: + :no-inherited-members: + + + +Download +~~~~~~~~ + +To download the Waymo Open Dataset for Perception, please visit the `official website `_ and follow the instructions provided there. +You will need to register and download the Perception Dataset ``V1.4.3``. +(We currently do not support ``V2.0.1`` due to the missing maps.) +The expected directory structure after downloading and extracting the dataset is as follows: + +.. code-block:: text + + $WODP_DATA_ROOT + ├── testing/ + | ├── segment-10084636266401282188_1120_000_1140_000_with_camera_labels.tfrecord + | ├── ... + | └── segment-9806821842001738961_4460_000_4480_000_with_camera_labels.tfrecord + ├── training/ + | ├── segment-10017090168044687777_6380_000_6400_000_with_camera_labels.tfrecord + | ├── ... + | └── segment-9985243312780923024_3049_720_3069_720_with_camera_labels.tfrecord + └── validation/ + ├── segment-10203656353524179475_7625_000_7645_000_with_camera_labels.tfrecord + ├── ... + └── segment-967082162553397800_5102_900_5122_900_with_camera_labels.tfrecord + +You can add the dataset root directory to the environment variable ``WODP_DATA_ROOT`` for easier access. + +.. code-block:: bash + + export WODP_DATA_ROOT=/path/to/wodp_dataset_root + +Optionally, you can adjust the ``py123d/script/config/common/default_dataset_paths.yaml`` accordingly. + +Installation +~~~~~~~~~~~~ + +The Waymo Open Dataset requires additional dependencies that are included as optional dependencies in ``py123d``. You can install them via: + +.. tab-set:: + + .. tab-item:: PyPI + + .. code-block:: bash + + pip install py123d[waymo] + + .. tab-item:: Source + + .. code-block:: bash + + pip install -e .[waymo] + +These dependencies are notoriously difficult to install due to compatibility issues. +We recommend using a dedicated conda environment for this purpose. Using `uv `_ can significantly speed up the installation. +Here is an example of how to set it up: + +.. code-block:: bash + + conda create -n py123d_waymo python=3.10 + conda activate py123d_waymo + uv pip install -e .[waymo] + # If something goes wrong: conda deactivate; conda remove -n py123d_waymo --all + +You only need the Waymo Open Dataset specific dependencies if you convert the dataset or read from the raw TFRecord files. +After conversion, you may use any other ``py123d`` installation. + + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + + +* **Map:** The HD-Map in Waymo has bugs ... + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @inproceedings{Sun2020CVPR, + title={Scalability in perception for autonomous driving: Waymo open dataset}, + author={Sun, Pei and Kretzschmar, Henrik and Dotiwalla, Xerxes and Chouard, Aurelien and Patnaik, Vijaysai and Tsui, Paul and Guo, James and Zhou, Yin and Chai, Yuning and Caine, Benjamin and others}, + booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, + pages={2446--2454}, + year={2020} + } diff --git a/docs/datasets/wopd.rst b/docs/datasets/wopd.rst deleted file mode 100644 index 3b00d2e6..00000000 --- a/docs/datasets/wopd.rst +++ /dev/null @@ -1,99 +0,0 @@ -Waymo Open Perception Dataset (WOPD) ------------------------------------- - -.. sidebar:: WOPD - - .. image:: https://images.ctfassets.net/e6t5diu0txbw/4LpraC18sHNvS87OFnEGKB/63de105d4ce623d91cfdbc23f77d6a37/Open_Dataset_Download_Hero.jpg?fm=webp&q=90 - :alt: Dataset sample image - :width: 290px - - | **Paper:** `Name of Paper `_ - | **Download:** `Documentation `_ - | **Code:** [Code] - | **Documentation:** [License type] - | **License:** [License type] - | **Duration:** [Duration here] - | **Supported Versions:** [Yes/No/Conditions] - | **Redistribution:** [Yes/No/Conditions] - -Description -~~~~~~~~~~~ - -[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] - -Installation -~~~~~~~~~~~~ - -[Instructions for installing or accessing the dataset] - -.. code-block:: bash - - # Example installation commands - pip install py123d[dataset_name] - # or - wget https://example.com/dataset.zip - - -.. code-block:: bash - - conda create -n py123d_waymo python=3.10 - conda activate py123d_waymo - pip install -e .[waymo] - - # pip install protobuf==6.30.2 - # pip install tensorflow==2.13.0 - # pip install waymo-open-dataset-tf-2-12-0==1.6.6 - -Available Data -~~~~~~~~~~~~~~ - -.. list-table:: - :header-rows: 1 - :widths: 30 5 70 - - - * - **Name** - - **Available** - - **Description** - * - Ego Vehicle - - X - - [Description of ego vehicle data] - * - Map - - X - - [Description of ego vehicle data] - * - Bounding Boxes - - X - - [Description of ego vehicle data] - * - Traffic Lights - - X - - [Description of ego vehicle data] - * - Cameras - - X - - [Description of ego vehicle data] - * - LiDARs - - X - - [Description of ego vehicle data] - -Dataset Specific Issues -~~~~~~~~~~~~~~~~~~~~~~~ - -[Document any known issues, limitations, or considerations when using this dataset] - -* Issue 1: Description -* Issue 2: Description -* Issue 3: Description - -Citation -~~~~~~~~ - -If you use this dataset in your research, please cite: - -.. code-block:: bibtex - - @inproceedings{Sun2020CVPR, - title={Scalability in perception for autonomous driving: Waymo open dataset}, - author={Sun, Pei and Kretzschmar, Henrik and Dotiwalla, Xerxes and Chouard, Aurelien and Patnaik, Vijaysai and Tsui, Paul and Guo, James and Zhou, Yin and Chai, Yuning and Caine, Benjamin and others}, - booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, - pages={2446--2454}, - year={2020} - } diff --git a/docs/development/contributing.md b/docs/development/contributing.md deleted file mode 100644 index e7ac92fe..00000000 --- a/docs/development/contributing.md +++ /dev/null @@ -1,118 +0,0 @@ - -# Contributing to 123D - -Contributions to 123D are highly encouraged! This guide will help you get started with the development process. - -## Getting Started - -### 1. Clone the Repository - -```sh -git clone git@github.com:DanielDauner/py123d.git -cd py123d -``` - -### 2. Install the pip-package - -```sh -conda env create -f environment.yml --name py123d_dev # Optional -conda activate py123d_dev -pip install -e .[dev] -pre-commit install -``` - -.. note:: - We might remove the conda environment in the future, but leave the file in the repo during development. - - -### 3. Managing dependencies - -One principal of 123D is to keep *minimal dependencies*. However, various datasets require dependencies in order to load or preprocess the dataset. In this case, you can add optional dependencies to the `pyproject.toml` install file. You can follow examples of Waymo/nuPlan. These optional dependencies can be install with - -```sh -pip install -e .[dev,waymo,nuplan] -``` -where you can combined the different optional dependencies. - -The optional dependencies should only be required for data pre-processing. If a dataset allows to load sensor data dynamically from the original dataset, please encapsule the import accordingly, e.g. - -```python -import numpy as np -import numpy.typing as npt - -def load_camera_from_outdated_dataset(file_path: str) -> npt.NDArray[np.uint8]: - try: - from optional_dataset import load_camera_image - except ImportError: - raise ImportError( - "Optional dependency 'outdated_dataset' is required to load camera images from this dataset. " - "Please install it using: pip install .[outdated_dataset]" - ) - return load_camera_image(file_path) -``` - - -## Code Style and Formatting - -We maintain consistent code quality using the following tools: -- **[Black](https://black.readthedocs.io/)** - Code formatter -- **[isort](https://pycqa.github.io/isort/)** - Import statement formatter -- **[flake8](https://flake8.pycqa.github.io/)** - Style guide enforcement -- **[pytest](https://docs.pytest.org/)** - Testing framework for unit and integration tests -- **[pre-commit](https://pre-commit.com/)** - Framework for managing and running Git hooks to automate code quality checks - - -.. note:: - If you're using VSCode, it is recommended to install the [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter), [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort), and [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) plugins. - - - -### Editor Setup - -**VS Code Users:** -If you're using VSCode, it is recommended to install the following plguins: -- [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) - see above. -- [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort) - see above. -- [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) - see above. -- [autodocstring](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) - Creating docstrings (please set `"autoDocstring.docstringFormat": "sphinx-notypes"`). -- [Code Spell Checker](https://marketplace.visualstudio.com/items?itemName=streetsidesoftware.code-spell-checker) - A basic spell checker. - - -**Other Editors:** -Similar plugins are available for most popular editors including PyCharm, Vim, Emacs, and Sublime Text. - - -## Documentation Requirements - -### Docstrings -- **Development:** Docstrings are encouraged but not strictly required during active development -- **Release:** All public functions, classes, and modules must have comprehensive docstrings before release -- **Format:** Use [Sphinx-style docstrings](https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html) - -**VS Code Users:** The [autoDocstring extension](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) can help generate properly formatted docstrings. - -### Type Hints -- **Required:** All function parameters and return values must include type hints -- **Style:** Follow [PEP 484](https://peps.python.org/pep-0484/) conventions - -### Sphinx documentation - -All datasets should be included in the `/docs/datasets.md` documentation. Please follow the documentation format of other datasets. - -You can install relevant dependencies for editing the public documentation via: -```sh -pip install -e .[docs] -``` - -It is recommended to uses [sphinx-autobuild](https://github.com/sphinx-doc/sphinx-autobuild) (installed above) to edit and view the documentation. You can run: -```sh -sphinx-autobuild docs docs/_build/html -``` - -## Adding new datasets -TODO - - -## Questions? - -If you have any questions about contributing, please open an issue or reach out to the maintainers. diff --git a/docs/development/index.rst b/docs/development/index.rst deleted file mode 100644 index 579ccee4..00000000 --- a/docs/development/index.rst +++ /dev/null @@ -1,7 +0,0 @@ -Development -=========== - -.. toctree:: - :maxdepth: 0 - - contributing diff --git a/docs/geometry.rst b/docs/geometry.rst deleted file mode 100644 index 0a3215e3..00000000 --- a/docs/geometry.rst +++ /dev/null @@ -1,87 +0,0 @@ - -Geometry -======== - -Geometric Primitives --------------------- - -Points -~~~~~~ -.. autoclass:: py123d.geometry.Point2D() - -.. autoclass:: py123d.geometry.Point3D() - -Vectors -~~~~~~~ -.. autoclass:: py123d.geometry.Vector2D() - -.. autoclass:: py123d.geometry.Vector3D() - -Special Euclidean Group -~~~~~~~~~~~~~~~~~~~~~~~ -.. autoclass:: py123d.geometry.StateSE2() - -.. autoclass:: py123d.geometry.StateSE3() - -Bounding Boxes -~~~~~~~~~~~~~~ -.. autoclass:: py123d.geometry.BoundingBoxSE2() - -.. autoclass:: py123d.geometry.BoundingBoxSE3() - -Indexing Enums -~~~~~~~~~~~~~~ -.. autoclass:: py123d.geometry.Point2DIndex() - -.. autoclass:: py123d.geometry.Point3DIndex() - -.. autoclass:: py123d.geometry.Vector2DIndex() - -.. autoclass:: py123d.geometry.Vector3DIndex() - -.. autoclass:: py123d.geometry.StateSE2Index() - -.. autoclass:: py123d.geometry.StateSE3Index() - -.. autoclass:: py123d.geometry.BoundingBoxSE2Index() - -.. autoclass:: py123d.geometry.BoundingBoxSE3Index() - -.. autoclass:: py123d.geometry.Corners2DIndex() - -.. autoclass:: py123d.geometry.Corners3DIndex() - - -Transformations ---------------- - -Transformations in 2D -~~~~~~~~~~~~~~~~~~~~~ -.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_se2_array - -.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_se2_array - -.. autofunction:: py123d.geometry.transform.translate_se2_along_body_frame - -.. autofunction:: py123d.geometry.transform.translate_se2_along_x - -.. autofunction:: py123d.geometry.transform.translate_se2_along_y - - -Transformations in 3D -~~~~~~~~~~~~~~~~~~~~~ -.. autofunction:: py123d.geometry.transform.convert_absolute_to_relative_se3_array - -.. autofunction:: py123d.geometry.transform.convert_relative_to_absolute_se3_array - -.. autofunction:: py123d.geometry.transform.translate_se3_along_body_frame - -.. autofunction:: py123d.geometry.transform.translate_se3_along_x - -.. autofunction:: py123d.geometry.transform.translate_se3_along_y - -.. autofunction:: py123d.geometry.transform.translate_se3_along_z - -Occupancy Map -------------- -.. autoclass:: py123d.geometry.OccupancyMap2D() diff --git a/docs/index.rst b/docs/index.rst index a9e1d197..c878edf2 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,18 +1,31 @@ -.. 123d documentation master file, created by - sphinx-quickstart on Wed Aug 13 16:57:48 2025. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. - 123D Documentation ================== -Add your content using ``reStructuredText`` syntax. See the -`reStructuredText `_ -documentation for details. +Welcome to the official documentation for 123D, a library for driving datasets in 2D and 3D. + +Features include: + +- Unified API for driving data, including sensor data, maps, and labels. +- Support for multiple sensors storage formats. +- Fast dataformat based on `Apache Arrow `_ +- Visualization tools with `matplotlib `_ and `Viser `_. + + +.. warning:: + + This library is under active development and not stable. The API and features may change in future releases. + Please report issues, feature requests, or other feedback by opening an issue on the project's GitHub repository. + + +.. youtube:: Q4q29fpXnx8 + :width: 800 + :height: 450 + :align: center .. toctree:: :maxdepth: 1 + :hidden: :caption: Overview: installation @@ -21,19 +34,17 @@ documentation for details. .. toctree:: :maxdepth: 2 + :hidden: :caption: API Reference: - geometry - - -.. toctree:: - :maxdepth: 1 - :caption: Visualization: - - visualization + api/scene/index + api/map/index + api/datatypes/index + api/geometry/index .. toctree:: :maxdepth: 1 - :caption: Development: + :hidden: + :caption: Notes - development/index + contributing diff --git a/docs/installation.md b/docs/installation.md index 39846467..3216ad85 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -1,37 +1,79 @@ # Installation -Note, the following installation assumes the following folder structure: TODO UPDATE +## Pip-Install +You can simply install `py123d` for Python versions >=3.9 via PyPI with +```bash +pip install py123d ``` -~/py123d_workspace -├── py123d -├── exp -│ └── ... -└── data - ├── maps - │ ├── carla_town01.gpkg - │ ├── carla_town02.gpkg - │ ├── ... - │ └── nuplan_us-pa-pittsburgh-hazelwood.gpkg - ├── nuplan_mini_test - │ ├── 2021.05.25.14.16.10_veh-35_01690_02183.arrow - │ ├── 2021.06.03.12.02.06_veh-35_00233_00609.arrow - │ ├── ... - │ └── 2021.10.06.07.26.10_veh-52_00006_00398.arrow - ├── nuplan_mini_train - │ └── ... - └── nuplan_mini_test - └── ... +or as editable pip package with +```bash +mkdir -p $HOME/py123d_workspace; cd $HOME/py123d_workspace # Optional +git clone git@github.com:autonomousvision/py123d.git +cd py123d +pip install -e . ``` +## File Structure & Storage +The 123D library converts driving datasets to a unified format. By default, all data is stored in directory of the environment variable `$PY123D_DATA_ROOT`. +For example, you can use. -First you need to create a new conda environment and install `py123d` as editable pip package. ```bash -conda create -n py123d_dev python=3.12 -conda activate py123d_dev -pip install -e . +export PY123D_DATA_ROOT="$HOME/py123d_workspace/data" ``` +which can be added to your `~/.bashrc` or to your bash scripts. Optionally, you can adjust all dataset paths in the hydra config: `py123d/script/config/common/default_dataset_paths.yaml`. + +The 123D conversion includes: +- **Logs:** The logs store continuous driving recordings in a single file, including modalities such as timestamps, ego states, bounding boxes, and sensor references. Logs are stored as `.arrow` files. +- **Maps:** The maps are static and store our unified HD-Map API. Maps can either be defined per-log (e.g. in AV2, Waymo) or globally for a certain location (e.g. nuPlan, nuScenes, CARLA). In the current implementation, we store maps as `.gpkg` files. +- **Sensors:** There are multiple options to store sensor data. Cameras and LiDAR point clouds can either (1) be read from the original dataset or (2) stored within the log file. For cameras, we also support (3) compression with MP4 files, which are written into the `/sensors` directory. + +For example, when converting `nuplan-mini` with MP4 compression and using `PY123D_DATA_ROOT="$HOME/py123d_workspace/data"`, the file structure would look the following way: +``` +~/py123d_workspace/ +├── data/ +│ ├── logs +│ │ ├── nuplan-mini_test +│ │ │ ├── 2021.05.25.14.16.10_veh-35_01690_02183.arrow +│ │ │ ├── ... +│ │ │ └── 2021.10.06.07.26.10_veh-52_00006_00398.arrow +│ │ ├── nuplan-mini_train +│ │ │ └── ... +│ │ ├── nuplan-mini_train +│ │ │ └── ... +│ │ └── ... +│ ├── maps +│ │ ├── nuplan +│ │ │ ├── nuplan_sg-one-north.gpkg +│ │ │ ├── ... +│ │ │ └── nuplan_us-pa-pittsburgh-hazelwood.gpkg +│ │ └── ... +│ └── sensors +│ ├── nuplan-mini_test +│ │ ├── 2021.05.25.14.16.10_veh-35_01690_02183 +│ │ │ ├── pcam_b0.mp4 +│ │ │ ├── ... +│ │ │ └── pcam_r2.mp4 +│ │ └── ... +│ └── ... +└── py123d/ (repository) + └── ... +``` + +## Demo data + + +You can test 123D with demo data from [nuPlan](nuplan), [nuScenes](nuscenes), [PandaSet](pandaset), [Argoverse 2 - Sensor](av2_sensor), and [CARLA](carla). Please be aware of the respective licenses, that are included in the download. You can use the following script: -Next, you need add the following environment variables in your `.bashrc`: ```bash -export PY123D_DATA_ROOT="$HOME/py123d_workspace/data" +# Create the data root and a temporary folder. +mkdir -p $PY123D_DATA_ROOT +mkdir -p ./temp + +# Download the demo data. +wget https://s3.eu-central-1.amazonaws.com/avg-projects-2/123d/demo_v0.0.8/data.zip + +# Unzip, sync, and clean up. +unzip -o data.zip -d ./temp +rsync -av ./temp/data/* $PY123D_DATA_ROOT +rm -r ./temp & rm -r data.zip ``` diff --git a/src/py123d/datatypes/maps/cache/__init__.py b/docs/notes/conventions.rst similarity index 100% rename from src/py123d/datatypes/maps/cache/__init__.py rename to docs/notes/conventions.rst diff --git a/docs/visualization.md b/docs/visualization.md deleted file mode 100644 index 9fe2e3cf..00000000 --- a/docs/visualization.md +++ /dev/null @@ -1,11 +0,0 @@ - -# Visualization - - -## Matplotlib - - -## Viser - - -## Bokeh diff --git a/examples/01_viser.py b/examples/01_viser.py deleted file mode 100644 index a91005cf..00000000 --- a/examples/01_viser.py +++ /dev/null @@ -1,37 +0,0 @@ -from py123d.common.multithreading.worker_sequential import Sequential -from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder -from py123d.datatypes.scene.scene_filter import SceneFilter -from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType -from py123d.visualization.viser.viser_viewer import ViserViewer - -if __name__ == "__main__": - # splits = ["kitti360_train"] - # splits = ["nuscenes-mini_val", "nuscenes-mini_train"] - # splits = ["nuplan-mini_test", "nuplan-mini_train", "nuplan-mini_val"] - # splits = ["nuplan_private_test"] - # splits = ["carla_test"] - splits = ["wopd_val"] - # splits = ["av2-sensor_train"] - # splits = ["pandaset_test", "pandaset_val", "pandaset_train"] - # log_names = ["2021.08.24.13.12.55_veh-45_00386_00472"] - # log_names = ["2013_05_28_drive_0000_sync"] - # log_names = ["2013_05_28_drive_0000_sync"] - log_names = None - # scene_uuids = ["87bf69e4-f2fb-5491-99fa-8b7e89fb697c"] - scene_uuids = None - - scene_filter = SceneFilter( - split_names=splits, - log_names=log_names, - scene_uuids=scene_uuids, - duration_s=None, - history_s=0.0, - timestamp_threshold_s=None, - shuffle=True, - pinhole_camera_types=[PinholeCameraType.PCAM_F0], - ) - scene_builder = ArrowSceneBuilder() - worker = Sequential() - scenes = scene_builder.get_scenes(scene_filter, worker) - print(f"Found {len(scenes)} scenes") - visualization_server = ViserViewer(scenes, scene_index=0) diff --git a/notebooks/bev_matplotlib.ipynb b/notebooks/bev_matplotlib.ipynb deleted file mode 100644 index e7bbc334..00000000 --- a/notebooks/bev_matplotlib.ipynb +++ /dev/null @@ -1,290 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", - "from py123d.datatypes.scene.scene_filter import SceneFilter\n", - "\n", - "\n", - "from py123d.common.multithreading.worker_sequential import Sequential\n", - "# from py123d.datatypes.sensors.pinhole_camera_type import PinholeCameraType" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# splits = [\"kitti360_train\"]\n", - "# splits = [\"nuscenes-mini_val\", \"nuscenes-mini_train\"]\n", - "# splits = [\"nuplan-mini_test\", \"nuplan-mini_train\", \"nuplan-mini_val\"]\n", - "# splits = [\"carla_test\"]\n", - "# splits = [\"wopd_val\"]\n", - "# splits = [\"av2-sensor_train\"]\n", - "splits = [\"pandaset_test\", \"pandaset_val\", \"pandaset_train\"]\n", - "log_names = None\n", - "scene_uuids = None\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=splits,\n", - " log_names=log_names,\n", - " scene_uuids=scene_uuids,\n", - " duration_s=None,\n", - " history_s=0.0,\n", - " timestamp_threshold_s=30.0,\n", - " shuffle=True,\n", - " # camera_types=[PinholeCameraType.CAM_F0],\n", - ")\n", - "scene_builder = ArrowSceneBuilder()\n", - "worker = Sequential()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "print(f\"Found {len(scenes)} scenes\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from typing import List, Optional, Tuple\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "from py123d.geometry import Point2D\n", - "from py123d.visualization.color.color import BLACK, DARK_GREY, DARKER_GREY, LIGHT_GREY, NEW_TAB_10, TAB_10\n", - "from py123d.visualization.color.config import PlotConfig\n", - "from py123d.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", - "from py123d.visualization.matplotlib.observation import (\n", - " add_box_detections_to_ax,\n", - " add_default_map_on_ax,\n", - " add_ego_vehicle_to_ax,\n", - " add_traffic_lights_to_ax,\n", - ")\n", - "from py123d.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", - "from py123d.datatypes.maps.abstract_map import AbstractMap\n", - "from py123d.datatypes.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", - "from py123d.datatypes.maps.gpkg.gpkg_map_objects import GPKGIntersection\n", - "from py123d.datatypes.maps.map_datatypes import MapLayer\n", - "from py123d.datatypes.scene.abstract_scene import AbstractScene\n", - "\n", - "\n", - "import shapely.geometry as geom\n", - "\n", - "LEFT_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=TAB_10[2],\n", - " fill_color_alpha=1.0,\n", - " line_color=TAB_10[2],\n", - " line_color_alpha=0.5,\n", - " line_width=1.0,\n", - " line_style=\"-\",\n", - " zorder=3,\n", - ")\n", - "\n", - "RIGHT_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=TAB_10[3],\n", - " fill_color_alpha=1.0,\n", - " line_color=TAB_10[3],\n", - " line_color_alpha=0.5,\n", - " line_width=1.0,\n", - " line_style=\"-\",\n", - " zorder=22,\n", - ")\n", - "\n", - "\n", - "LANE_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=BLACK,\n", - " fill_color_alpha=1.0,\n", - " line_color=BLACK,\n", - " line_color_alpha=0.0,\n", - " line_width=0.0,\n", - " line_style=\"-\",\n", - " zorder=5,\n", - ")\n", - "\n", - "ROAD_EDGE_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=DARKER_GREY,\n", - " fill_color_alpha=1.0,\n", - " line_color=DARKER_GREY,\n", - " line_color_alpha=1.0,\n", - " line_width=1.0,\n", - " line_style=\"-\",\n", - " zorder=3,\n", - ")\n", - "\n", - "ROAD_LINE_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=NEW_TAB_10[5],\n", - " fill_color_alpha=1.0,\n", - " line_color=NEW_TAB_10[5],\n", - " line_color_alpha=1.0,\n", - " line_width=1.5,\n", - " line_style=\"-\",\n", - " zorder=3,\n", - ")\n", - "\n", - "\n", - "def add_debug_map_on_ax(\n", - " ax: plt.Axes,\n", - " map_api: AbstractMap,\n", - " point_2d: Point2D,\n", - " radius: float,\n", - " route_lane_group_ids: Optional[List[int]] = None,\n", - ") -> None:\n", - " layers: List[MapLayer] = [\n", - " # MapLayer.LANE,\n", - " MapLayer.LANE_GROUP,\n", - " MapLayer.GENERIC_DRIVABLE,\n", - " MapLayer.CARPARK,\n", - " # MapLayer.CROSSWALK,\n", - " # MapLayer.INTERSECTION,\n", - " MapLayer.WALKWAY,\n", - " MapLayer.ROAD_EDGE,\n", - " MapLayer.ROAD_LINE,\n", - " ]\n", - " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", - " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", - " patch = geom.box(x_min, y_min, x_max, y_max)\n", - " map_objects_dict = map_api.query(geometry=patch, layers=layers, predicate=\"intersects\")\n", - " # print(map_objects_dict[MapLayer.ROAD_EDGE])\n", - "\n", - " for layer, map_objects in map_objects_dict.items():\n", - " for map_object in map_objects:\n", - " try:\n", - " if layer in [\n", - " MapLayer.GENERIC_DRIVABLE,\n", - " MapLayer.CARPARK,\n", - " MapLayer.CROSSWALK,\n", - " # MapLayer.INTERSECTION,\n", - " MapLayer.WALKWAY,\n", - " ]:\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - "\n", - " if layer in [MapLayer.LANE_GROUP]:\n", - " map_object: AbstractLaneGroup\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - "\n", - " if map_object.intersection is not None:\n", - " add_shapely_polygon_to_ax(ax, map_object.intersection.shapely_polygon, ROUTE_CONFIG)\n", - "\n", - " for lane in map_object.lanes:\n", - " add_shapely_linestring_to_ax(ax, lane.centerline.linestring, CENTERLINE_CONFIG)\n", - "\n", - " # if layer in [MapLayer.LANE]:\n", - " # add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", - " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - "\n", - " if layer in [MapLayer.ROAD_EDGE]:\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", - "\n", - " if layer in [MapLayer.ROAD_LINE]:\n", - " # line_type = int(map_object.road_line_type)\n", - " # plt_config = PlotConfig(\n", - " # fill_color=NEW_TAB_10[line_type % (len(NEW_TAB_10) - 1)],\n", - " # fill_color_alpha=1.0,\n", - " # line_color=NEW_TAB_10[line_type % (len(NEW_TAB_10) - 1)],\n", - " # line_color_alpha=1.0,\n", - " # line_width=1.5,\n", - " # line_style=\"-\",\n", - " # zorder=10,\n", - " # )\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\n", - "\n", - " except Exception:\n", - " import traceback\n", - "\n", - " print(f\"Error adding map object of type {layer.name} and id {map_object.object_id}\")\n", - " traceback.print_exc()\n", - "\n", - " # ax.set_title(f\"Map: {map_api.map_name}\")\n", - "\n", - "\n", - "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", - "\n", - " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", - " box_detections = scene.get_box_detections_at_iteration(iteration)\n", - " map_api = scene.get_map_api()\n", - "\n", - " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", - " if map_api is not None:\n", - " # add_debug_map_on_ax(ax, scene.get_map_api(), point_2d, radius=radius, route_lane_group_ids=None)\n", - "\n", - "\n", - " add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", - " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.get_map_api())\n", - "\n", - " add_box_detections_to_ax(ax, box_detections)\n", - " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", - "\n", - " zoom = 1.0\n", - " ax.set_xlim(point_2d.x - radius * zoom, point_2d.x + radius * zoom)\n", - " ax.set_ylim(point_2d.y - radius * zoom, point_2d.y + radius * zoom)\n", - "\n", - " ax.set_aspect(\"equal\", adjustable=\"box\")\n", - " return ax\n", - "\n", - "\n", - "def plot_scene_at_iteration(\n", - " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", - ") -> Tuple[plt.Figure, plt.Axes]:\n", - "\n", - " size = 10\n", - "\n", - " fig, ax = plt.subplots(figsize=(size, size))\n", - " _plot_scene_on_ax(ax, scene, iteration, radius)\n", - " return fig, ax\n", - "\n", - "\n", - "# scene_index =\n", - "iteration = 1\n", - "\n", - "scale = 10\n", - "fig, ax = plt.subplots(1, 1, figsize=(scale, scale))\n", - "scene = np.random.choice(scenes)\n", - "_plot_scene_on_ax(ax, scene, iteration, radius=80)\n", - "# _plot_scene_on_ax(ax[1], scene, iteration, radius=50)\n", - "# _plot_scene_on_ax(ax[2], scene, iteration,\n", - "# radius=100)\n", - "\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "py123d", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/bev_render.ipynb b/notebooks/bev_render.ipynb deleted file mode 100644 index 1bc41014..00000000 --- a/notebooks/bev_render.ipynb +++ /dev/null @@ -1,131 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", - "from py123d.datatypes.scene.scene_filter import SceneFilter\n", - "\n", - "from py123d.common.multithreading.worker_sequential import Sequential\n", - "# from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# splits = [\"kitti360\"]\n", - "# splits = [\"nuscenes-mini_val\", \"nuscenes-mini_train\"]\n", - "# splits = [\"nuplan-mini_test\", \"nuplan-mini_train\", \"nuplan-mini_val\"]\n", - "# splits = [\"nuplan_private_test\"]\n", - "# splits = [\"carla_test\"]\n", - "splits = [\"wopd_val\"]\n", - "# splits = [\"av2-sensor_train\"]\n", - "# splits = [\"pandaset_test\", \"pandaset_val\", \"pandaset_train\"]\n", - "# log_names = [\"2021.08.24.13.12.55_veh-45_00386_00472\"]\n", - "# log_names = [\"2013_05_28_drive_0000_sync\"]\n", - "# log_names = [\"2013_05_28_drive_0000_sync\"]\n", - "log_names = None\n", - "scene_uuids = [\"9727e2b3-46b0-51bd-84a9-c516c0993045\"]\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=splits,\n", - " log_names=log_names,\n", - " scene_uuids=scene_uuids,\n", - " duration_s=None,\n", - " history_s=0.0,\n", - " timestamp_threshold_s=None,\n", - " shuffle=True,\n", - " # camera_types=[PinholeCameraType.CAM_F0],\n", - ")\n", - "scene_builder = ArrowSceneBuilder()\n", - "worker = Sequential()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "scenes = [scene for scene in scenes if scene.uuid in scene_uuids]\n", - "print(f\"Found {len(scenes)} scenes\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from py123d.visualization.matplotlib.plots import render_scene_animation\n", - "\n", - "for i in [0]:\n", - " render_scene_animation(scenes[i], output_path=\"test\", format=\"mp4\", fps=20, step=1, radius=50)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "py123d", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/camera_matplotlib.ipynb b/notebooks/camera_matplotlib.ipynb deleted file mode 100644 index b33cfdd8..00000000 --- a/notebooks/camera_matplotlib.ipynb +++ /dev/null @@ -1,162 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", - "from py123d.datatypes.scene.scene_filter import SceneFilter\n", - "\n", - "from py123d.common.multithreading.worker_sequential import Sequential\n", - "from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# splits = [\"wopd_val\"]\n", - "# splits = [\"carla_test\"]\n", - "# splits = [\"nuplan-mini_test\"]\n", - "# splits = [\"av2-sensor-mini_train\"]\n", - "\n", - "\n", - "splits = [\"pandaset_train\"]\n", - "# log_names = None\n", - "\n", - "\n", - "log_names = None\n", - "scene_uuids = None\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=splits,\n", - " log_names=log_names,\n", - " scene_uuids=scene_uuids,\n", - " duration_s=6.0,\n", - " history_s=0.0,\n", - " timestamp_threshold_s=20,\n", - " shuffle=True,\n", - " pinhole_camera_types=[PinholeCameraType.PCAM_F0],\n", - ")\n", - "scene_builder = ArrowSceneBuilder()\n", - "worker = Sequential()\n", - "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "print(f\"Found {len(scenes)} scenes\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from matplotlib import pyplot as plt\n", - "from py123d.datatypes.scene.abstract_scene import AbstractScene\n", - "from py123d.visualization.matplotlib.camera import add_box_detections_to_camera_ax, add_camera_ax\n", - "\n", - "iteration = 0\n", - "scene = scenes[0]\n", - "\n", - "scene: AbstractScene\n", - "print(scene.uuid, scene.available_pinhole_camera_types)\n", - "\n", - "scale = 3.0\n", - "fig, ax = plt.subplots(2, 3, figsize=(scale * 6, scale * 2.5))\n", - "\n", - "\n", - "camera_ax_mapping = {\n", - " PinholeCameraType.PCAM_L0: ax[0, 0],\n", - " PinholeCameraType.PCAM_F0: ax[0, 1],\n", - " PinholeCameraType.PCAM_R0: ax[0, 2],\n", - " PinholeCameraType.PCAM_L1: ax[1, 0],\n", - " PinholeCameraType.PCAM_B0: ax[1, 1],\n", - " PinholeCameraType.PCAM_R1: ax[1, 2],\n", - "}\n", - "\n", - "\n", - "for camera_type, ax_ in camera_ax_mapping.items():\n", - " camera = scene.get_pinhole_camera_at_iteration(iteration, camera_type)\n", - " box_detections = scene.get_box_detections_at_iteration(iteration)\n", - " ego_state = scene.get_ego_state_at_iteration(iteration)\n", - "\n", - " add_box_detections_to_camera_ax(\n", - " ax_,\n", - " camera,\n", - " box_detections,\n", - " ego_state,\n", - " )\n", - " ax_.set_title(f\"Camera: {camera_type.name}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "py123d", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/camera_render.ipynb b/notebooks/camera_render.ipynb deleted file mode 100644 index 4365c424..00000000 --- a/notebooks/camera_render.ipynb +++ /dev/null @@ -1,165 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", - "from py123d.datatypes.scene.scene_filter import SceneFilter\n", - "\n", - "from py123d.common.multithreading.worker_sequential import Sequential\n", - "from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType\n", - "\n", - "KITTI360_DATA_ROOT = \"/home/daniel/kitti_360/KITTI-360\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# splits = [\"kitti360\"]\n", - "# splits = [\"nuscenes-mini_val\", \"nuscenes-mini_train\"]\n", - "# splits = [\"nuplan-mini_test\", \"nuplan-mini_train\", \"nuplan-mini_val\"]\n", - "# splits = [\"nuplan_private_test\"]\n", - "# splits = [\"carla_test\"]\n", - "splits = [\"wopd_val\"]\n", - "# splits = [\"av2-sensor_train\"]\n", - "# splits = [\"pandaset_test\", \"pandaset_val\", \"pandaset_train\"]\n", - "# log_names = [\"2021.08.24.13.12.55_veh-45_00386_00472\"]\n", - "# log_names = [\"2013_05_28_drive_0000_sync\"]\n", - "# log_names = [\"2013_05_28_drive_0000_sync\"]\n", - "log_names = None\n", - "scene_uuids = [\"9727e2b3-46b0-51bd-84a9-c516c0993045\"]\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=splits,\n", - " log_names=log_names,\n", - " scene_uuids=scene_uuids,\n", - " duration_s=None,\n", - " history_s=0.0,\n", - " timestamp_threshold_s=None,\n", - " shuffle=True,\n", - " # camera_types=[PinholeCameraType.CAM_F0],\n", - ")\n", - "scene_builder = ArrowSceneBuilder()\n", - "worker = Sequential()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "scenes = [scene for scene in scenes if scene.uuid in scene_uuids]\n", - "print(f\"Found {len(scenes)} scenes\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from matplotlib import pyplot as plt\n", - "from py123d.datatypes.scene.abstract_scene import AbstractScene\n", - "from py123d.visualization.matplotlib.camera import add_box_detections_to_camera_ax, add_camera_ax\n", - "import imageio\n", - "import numpy as np\n", - "\n", - "iteration = 0\n", - "scene = scenes[0]\n", - "\n", - "scene: AbstractScene\n", - "fps = 15 # frames per second\n", - "output_file = f\"camera_{scene.log_metadata.split}_{scene.uuid}.mp4\"\n", - "\n", - "writer = imageio.get_writer(output_file, fps=fps)\n", - "\n", - "scale = 3.0\n", - "fig, ax = plt.subplots(2, 3, figsize=(scale * 6, scale * 2.5))\n", - "\n", - "\n", - "camera_type = PinholeCameraType.CAM_F0\n", - "\n", - "for i in range(scene.number_of_iterations):\n", - " camera = scene.get_camera_at_iteration(i, camera_type)\n", - " box_detections = scene.get_box_detections_at_iteration(i)\n", - " ego_state = scene.get_ego_state_at_iteration(i)\n", - "\n", - " _, image = add_box_detections_to_camera_ax(\n", - " None,\n", - " camera,\n", - " box_detections,\n", - " ego_state,\n", - " return_image=True,\n", - " )\n", - " writer.append_data(image)\n", - "\n", - "writer.close()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "py123d", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/pyproject.toml b/pyproject.toml index 655a2612..eb0687be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,19 +10,20 @@ classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "License :: OSI Approved :: Apache Software License", ] name = "py123d" -version = "v0.0.7" +version = "v0.0.8" authors = [{ name = "Daniel Dauner", email = "daniel.dauner@gmail.com" }] description = "TODO" readme = "README.md" -requires-python = ">=3.9" +requires-python = ">=3.8" license = {text = "Apache-2.0"} dependencies = [ - "bokeh", "geopandas", - "joblib", "matplotlib", "numpy", "opencv-python", @@ -32,82 +33,60 @@ dependencies = [ "pyarrow", "pyogrio", "pyquaternion", - "pytest", - "rasterio", "ray", - "rtree", "scipy", - "setuptools", "shapely>=2.0.0", "tqdm", "notebook", - "pre-commit", "hydra_colorlog", "hydra-core", - "lxml", "trimesh", "viser", "laspy[lazrs]", "DracoPy", + "omegaconf", + "typing-extensions", + "requests", ] [project.scripts] py123d-viser = "py123d.script.run_viser:main" py123d-conversion = "py123d.script.run_conversion:main" - [project.optional-dependencies] dev = [ - "black", - "isort", - "flake8", + "pyright", + "ruff", "pre-commit", + "pytest", + "pytest-cov", ] docs = [ "Sphinx", "sphinx-rtd-theme", "sphinx-autobuild", - "myst-parser", "sphinx-copybutton", + "myst-parser", "furo", + "autoclasstoc", + "sphinx-autodoc-typehints", + "sphinxcontrib-youtube", + "sphinx-design", ] nuplan = [ "nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2", - "ujson", - "tornado", - "sympy", "SQLAlchemy==1.4.27", - "selenium", - "nest_asyncio", - "cachetools", + "rasterio", "aioboto3", - "aiofiles", - "casadi", - "control", - "pyinstrument", - "Fiona", - "guppy3", "retry", + "cachetools", ] nuscenes = [ - "lanelet2", - "nuscenes-devkit==1.2.0", -] -nuscenes_expanded = [ "nuscenes-devkit==1.2.0", - "pycocotools==2.0.10", - "laspy==2.6.1", - "embreex==2.17.7.post6", - "lanelet2==1.2.2", - "protobuf==4.25.3", - "pycollada==0.9.2", - "vhacdx==0.0.8.post2", - "yourdfpy==0.0.58", + "lanelet2", ] waymo = [ - "protobuf==4.21.0", - "tensorflow==2.13.0", - "waymo-open-dataset-tf-2-12-0==1.6.6", + "waymo-open-dataset-tf-2-12-0==1.6.7", ] ffmpeg = [ "imageio[ffmpeg]", @@ -117,5 +96,40 @@ ffmpeg = [ where = ["src"] [project.urls] -"Homepage" = "https://github.com/DanielDauner/py123d" -"Bug Tracker" = "https://github.com/DanielDauner/py123d/issues" +"Homepage" = "https://github.com/autonomousvision/py123d" +"Bug Tracker" = "https://github.com/autonomousvision/py123d/issues" + + +[tool.ruff] +line-length = 120 +lint.select = [ + "E", # pycodestyle errors. + "F", # Pyflakes rules. + "PLC", # Pylint convention warnings. + "PLE", # Pylint errors. + "PLR", # Pylint refactor recommendations. + "PLW", # Pylint warnings. + "I", # Import sorting. +] +lint.ignore = [ + "E731", # Do not assign a lambda expression, use a def. + "E741", # Ambiguous variable name. (l, O, or I) + "E501", # Line too long. + "E402", # Module level import not at top of file + "PLR2004", # Magic value used in comparison. + "PLR0915", # Too many statements. + "PLR0913", # Too many arguments. + "PLC0414", # Import alias does not rename variable. (this is used for exporting names) + "PLC0415", # Import should be at the top-level of a file. + "PLC1901", # Use falsey strings. + "PLR5501", # Use `elif` instead of `else if`. + "PLR0911", # Too many return statements. + "PLR0912", # Too many branches. + "PLW0603", # Global statement updates are discouraged. + "PLW2901", # For loop variable overwritten. + "PLW0642", # Reassigned self in instance method. +] + +exclude = ["__init__.py"] +fixable = ["ALL"] +unfixable = [] diff --git a/scripts/conversion/av2_sensor_conversion.sh b/scripts/conversion/av2_sensor_conversion.sh deleted file mode 100644 index b386972e..00000000 --- a/scripts/conversion/av2_sensor_conversion.sh +++ /dev/null @@ -1,2 +0,0 @@ -py123d-conversion datasets=["av2_sensor_dataset"] \ -dataset_paths.av2_data_root="/media/nvme1/argoverse" diff --git a/scripts/conversion/kitti360_conversion.sh b/scripts/conversion/kitti360_conversion.sh deleted file mode 100644 index 1e939ad5..00000000 --- a/scripts/conversion/kitti360_conversion.sh +++ /dev/null @@ -1,3 +0,0 @@ -export KITTI360_DATA_ROOT="/home/daniel/kitti_360/KITTI-360" - -py123d-conversion datasets=["kitti360_dataset"] map_writer.remap_ids=true diff --git a/scripts/conversion/nuplan_mini_conversion.sh b/scripts/conversion/nuplan_mini_conversion.sh deleted file mode 100644 index 13ec7a53..00000000 --- a/scripts/conversion/nuplan_mini_conversion.sh +++ /dev/null @@ -1 +0,0 @@ -py123d-conversion datasets=["nuplan_mini_dataset"] diff --git a/scripts/conversion/nuscenes_mini_conversion.sh b/scripts/conversion/nuscenes_mini_conversion.sh deleted file mode 100644 index b9a9a7d1..00000000 --- a/scripts/conversion/nuscenes_mini_conversion.sh +++ /dev/null @@ -1,3 +0,0 @@ -export NUSCENES_DATA_ROOT="/home/daniel/nuscenes_mini/" - -py123d-conversion datasets=["nuscenes_mini_dataset"] map_writer.remap_ids=true diff --git a/scripts/conversion/pandaset_conversion.sh b/scripts/conversion/pandaset_conversion.sh deleted file mode 100644 index 0131e60f..00000000 --- a/scripts/conversion/pandaset_conversion.sh +++ /dev/null @@ -1 +0,0 @@ -py123d-conversion datasets=[pandaset_dataset] diff --git a/scripts/conversion/wopd_conversion.sh b/scripts/conversion/wopd_conversion.sh deleted file mode 100644 index e4f4b3d4..00000000 --- a/scripts/conversion/wopd_conversion.sh +++ /dev/null @@ -1,6 +0,0 @@ -py123d-conversion datasets=[wopd_dataset] - - -# pip install protobuf==6.30.2 -# pip install tensorflow==2.13.0 -# pip install waymo-open-dataset-tf-2-12-0==1.6.6 diff --git a/scripts/download/download_av2.sh b/scripts/download/download_av2.sh deleted file mode 100644 index 039de648..00000000 --- a/scripts/download/download_av2.sh +++ /dev/null @@ -1,19 +0,0 @@ -#!/usr/bin/env bash - -# Dataset URIs -# s3://argoverse/datasets/av2/sensor/ -# s3://argoverse/datasets/av2/lidar/ -# s3://argoverse/datasets/av2/motion-forecasting/ -# s3://argoverse/datasets/av2/tbv/ - -DATASET_NAMES=("sensor" "lidar" "motion-forecasting" "tbv") -TARGET_DIR="/path/to/argoverse" - -for DATASET_NAME in "${DATASET_NAMES[@]}"; do - mkdir -p "$TARGET_DIR/$DATASET_NAME" - s5cmd --no-sign-request cp "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$TARGET_DIR/$DATASET_NAME" -done - - -# wget -r s3://argoverse/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg -# wget http://argoverse.s3.amazonaws.com/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg diff --git a/scripts/download/download_kitti_360.sh b/scripts/download/download_kitti_360.sh deleted file mode 100644 index 1cb3e540..00000000 --- a/scripts/download/download_kitti_360.sh +++ /dev/null @@ -1,86 +0,0 @@ -# 2D data & labels -# ---------------------------------------------------------------------------------------------------------------------- - -# Fisheye Images (355G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/a1d81d9f7fc7195c937f9ad12e2a2c66441ecb4e/download_2d_fisheye.zip - -# Fisheye Calibration Images (11G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/data_2d_raw/data_fisheye_calibration.zip - - -# Perspective Images for Train & Val (128G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/a1d81d9f7fc7195c937f9ad12e2a2c66441ecb4e/download_2d_perspective.zip - -# Test Semantic (1.5G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/data_2d_raw/data_2d_test.zip - -# Test NVS 50% Drop (0.3G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/71f967e900f4e7c2e036a542f150effa31909b53/data_2d_nvs_drop50.zip - -# est NVS 90% Drop (0.2G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/71f967e900f4e7c2e036a542f150effa31909b53/data_2d_nvs_drop90.zip - -# Test SLAM (14G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/data_2d_raw/data_2d_test_slam.zip - - -# Semantics of Left Perspective Camera (1.8G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/ed180d24c0a144f2f1ac71c2c655a3e986517ed8/data_2d_semantics.zip - -# Semantics of Right Perspective Camera (1.8G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/ed180d24c0a144f2f1ac71c2c655a3e986517ed8/data_2d_semantics_image_01.zip - - -# Confidence of Left Perspective Camera (44G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/ed180d24c0a144f2f1ac71c2c655a3e986517ed8/data_2d_confidence.zip - -# Confidence of Right Perspective Camera (44G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/ed180d24c0a144f2f1ac71c2c655a3e986517ed8/data_2d_confidence_image_01.zip - - - -# 3D data & labels -# ---------------------------------------------------------------------------------------------------------------------- - -# Raw Velodyne Scans (119G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/a1d81d9f7fc7195c937f9ad12e2a2c66441ecb4e/download_3d_velodyne.zip - -# Test SLAM (12G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/data_3d_raw/data_3d_test_slam.zip - -# Test Completion (35M) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/6489aabd632d115c4280b978b2dcf72cb0142ad9/data_3d_ssc_test.zip - - -# Raw SICK Scans (0.4G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/a1d81d9f7fc7195c937f9ad12e2a2c66441ecb4e/download_3d_sick.zip - - -# Accumulated Point Clouds for Train & Val (12G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/6489aabd632d115c4280b978b2dcf72cb0142ad9/data_3d_semantics.zip - -# Test Semantic (1.2G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/6489aabd632d115c4280b978b2dcf72cb0142ad9/data_3d_semantics_test.zip - - -# 3D Bounding Boxes (30M) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/ffa164387078f48a20f0188aa31b0384bb19ce60/data_3d_bboxes.zip - - - -# Calibrations & Poses -# ---------------------------------------------------------------------------------------------------------------------- - -# Calibrations (3K) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/384509ed5413ccc81328cf8c55cc6af078b8c444/calibration.zip - - -# Vechicle Poses (8.9M) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/89a6bae3c8a6f789e12de4807fc1e8fdcf182cf4/data_poses.zip - - -# OXTS Sync Measurements (37.3M) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/89a6bae3c8a6f789e12de4807fc1e8fdcf182cf4/data_poses_oxts.zip - -# OXTS Raw Measurements (0.4G) -wget https://s3.eu-central-1.amazonaws.com/avg-projects/KITTI-360/89a6bae3c8a6f789e12de4807fc1e8fdcf182cf4/data_poses_oxts_extract.zip diff --git a/scripts/download/download_lyft.sh b/scripts/download/download_lyft.sh deleted file mode 100644 index f2f1e268..00000000 --- a/scripts/download/download_lyft.sh +++ /dev/null @@ -1,22 +0,0 @@ -# Sample dataset (51 MB) -wget https://woven.toyota/common/assets/data/prediction-sample.tar - - -# Training dataset - Part 1/2 (8.4 GB) -wget https://woven.toyota/common/assets/data/prediction-train.tar - - -# Training dataset - Part 2/2 (70 GB) -wget https://woven.toyota/common/assets/data/prediction-train_full.tar - - -# Validation dataset (8.2 GB) -wget https://woven.toyota/common/assets/data/prediction-validate.tar - - -# Aerial Map (2 GB) -wget https://woven.toyota/common/assets/data/prediction-aerial_map.tar - - -# Semantic Map (3 MB) -wget https://woven.toyota/common/assets/data/prediction-semantic_map.tar diff --git a/scripts/download/download_nuplan_logs.sh b/scripts/download/download_nuplan_logs.sh deleted file mode 100644 index 5eac6ddb..00000000 --- a/scripts/download/download_nuplan_logs.sh +++ /dev/null @@ -1,22 +0,0 @@ -# NOTE: Please check the LICENSE file when downloading the nuPlan dataset -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/LICENSE - -# maps -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-maps-v1.1.zip - -# train: nuplan_train -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_boston.zip -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_pittsburgh.zip -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_singapore.zip -for split in {1..6}; do - wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_train_vegas_${split}.zip -done - -# val: nuplan_val -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_test.zip - -# test: nuplan_test -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_val.zip - -# mini: nuplan_mini_train, nuplan_mini_val, nuplan_mini_test -wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/nuplan-v1.1_mini.zip diff --git a/scripts/download/download_nuplan_sensor.sh b/scripts/download/download_nuplan_sensor.sh deleted file mode 100644 index 6cd3d246..00000000 --- a/scripts/download/download_nuplan_sensor.sh +++ /dev/null @@ -1,14 +0,0 @@ -# NOTE: Please check the LICENSE file when downloading the nuPlan dataset -# wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/LICENSE - -# train: nuplan_train - -# val: nuplan_val - -# test: nuplan_test - -# mini: nuplan_mini_train, nuplan_mini_val, nuplan_mini_test -for split in {0..8}; do - wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_camera_${split}.zip - wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_lidar_${split}.zip -done diff --git a/scripts/viz/run_viser.sh b/scripts/viz/run_viser.sh deleted file mode 100644 index 436e7643..00000000 --- a/scripts/viz/run_viser.sh +++ /dev/null @@ -1,6 +0,0 @@ - - -py123d-viser \ -scene_filter=log_scenes \ -scene_filter.shuffle=True \ -worker=sequential diff --git a/src/py123d/__init__.py b/src/py123d/__init__.py index 6526deb4..a73339bf 100644 --- a/src/py123d/__init__.py +++ b/src/py123d/__init__.py @@ -1 +1 @@ -__version__ = "0.0.7" +__version__ = "0.0.8" diff --git a/src/py123d/api/__init__.py b/src/py123d/api/__init__.py new file mode 100644 index 00000000..89c5f7b3 --- /dev/null +++ b/src/py123d/api/__init__.py @@ -0,0 +1,5 @@ +from py123d.api.map.map_api import MapAPI +from py123d.api.scene.scene_api import SceneAPI +from py123d.api.scene.scene_builder import SceneBuilder +from py123d.api.scene.scene_filter import SceneFilter +from py123d.api.scene.scene_metadata import SceneMetadata diff --git a/src/py123d/datatypes/maps/gpkg/__init__.py b/src/py123d/api/map/__init__.py similarity index 100% rename from src/py123d/datatypes/maps/gpkg/__init__.py rename to src/py123d/api/map/__init__.py diff --git a/src/py123d/api/map/gpkg/gpkg_map_api.py b/src/py123d/api/map/gpkg/gpkg_map_api.py new file mode 100644 index 00000000..e8f6699d --- /dev/null +++ b/src/py123d/api/map/gpkg/gpkg_map_api.py @@ -0,0 +1,497 @@ +from __future__ import annotations + +import ast +import warnings +from collections import defaultdict +from functools import lru_cache +from pathlib import Path +from typing import Callable, Dict, Final, Iterable, List, Optional, Tuple, Union + +import geopandas as gpd +import shapely +import shapely.geometry as geom + +from py123d.api.map.gpkg.gpkg_utils import get_row_with_value, load_gdf_with_geometry_columns +from py123d.api.map.map_api import MapAPI +from py123d.datatypes.map_objects.base_map_objects import BaseMapObject +from py123d.datatypes.map_objects.map_layer_types import MapLayer, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + Walkway, +) +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.geometry import Point2D, Point3D, Polyline3D +from py123d.script.utils.dataset_path_utils import get_dataset_paths + +# TODO: add to some configs +USE_ARROW: bool = True +MAX_LRU_CACHED_TABLES: Final[int] = 128 + + +class GPKGMapAPI(MapAPI): + """The GeoPackage (GPKG) implementation of the :class:`~py123d.api.MapAPI`. + + Notes + ----- + The current implementation is inspired by the nuPlan GPKG map API [1]_. Ideally, the 123D implementation could + be replaced in the future by a more lightweight solution based on GeoParquet or GeoArrow. + + References + ---------- + .. [1] https://github.com/motional/nuplan-devkit/blob/master/nuplan/common/maps/nuplan_map/nuplan_map.py + + """ + + def __init__(self, file_path: Union[Path, str]) -> None: + """Initialize a GPKGMapAPI instance. + + :param file_path: The file path to the GeoPackage. + """ + + self._file_path = Path(file_path) + self._map_object_getter: Dict[MapLayer, Callable[[str], Optional[BaseMapObject]]] = { + MapLayer.LANE: self._get_lane, + MapLayer.LANE_GROUP: self._get_lane_group, + MapLayer.INTERSECTION: self._get_intersection, + MapLayer.CROSSWALK: self._get_crosswalk, + MapLayer.WALKWAY: self._get_walkway, + MapLayer.CARPARK: self._get_carpark, + MapLayer.GENERIC_DRIVABLE: self._get_generic_drivable, + MapLayer.ROAD_EDGE: self._get_road_edge, + MapLayer.ROAD_LINE: self._get_road_line, + } + + # loaded during `.initialize()` + self._gpd_dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + self._map_metadata: Optional[MapMetadata] = None + + def _initialize(self) -> None: + """Loads all available map layers and metadata from the GPKG file into GeoDataFrames.""" + if len(self._gpd_dataframes) == 0: + available_layers = list(gpd.list_layers(self._file_path).name) + for map_layer in list(MapLayer): + map_layer_name = map_layer.serialize() + if map_layer_name in available_layers: + self._gpd_dataframes[map_layer] = gpd.read_file( + self._file_path, layer=map_layer_name, use_arrow=USE_ARROW + ) + load_gdf_with_geometry_columns( + self._gpd_dataframes[map_layer], + geometry_column_names=["centerline", "right_boundary", "left_boundary", "outline"], + ) + # TODO: remove the temporary fix and enforce consistent id types in the GPKG files + if "id" in self._gpd_dataframes[map_layer].columns: + self._gpd_dataframes[map_layer]["id"] = self._gpd_dataframes[map_layer]["id"].astype(str) + else: + warnings.warn(f"GPKGMap: {map_layer_name} not available in {str(self._file_path)}") + self._gpd_dataframes[map_layer] = None + + assert "map_metadata" in list(gpd.list_layers(self._file_path).name) + metadata_gdf = gpd.read_file(self._file_path, layer="map_metadata", use_arrow=USE_ARROW) + self._map_metadata = MapMetadata.from_dict(metadata_gdf.iloc[0].to_dict()) + + def _assert_initialize(self) -> None: + """Checks if `.initialize()` was called, before retrieving data.""" + assert len(self._gpd_dataframes) > 0, "GPKGMap: Call `.initialize()` before retrieving data!" + + def _assert_layer_available(self, layer: MapLayer) -> None: + """Checks if layer is available.""" + assert layer in self.get_available_map_layers(), f"GPKGMap: MapLayer {layer.name} is unavailable." + + def get_map_metadata(self): + """Inherited, see superclass.""" + return self._map_metadata + + def get_available_map_layers(self) -> List[MapLayer]: + """Inherited, see superclass.""" + self._assert_initialize() + return list(self._gpd_dataframes.keys()) + + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[BaseMapObject]: + """Inherited, see superclass.""" + + self._assert_initialize() + self._assert_layer_available(layer) + try: + return self._map_object_getter[layer](object_id) + except KeyError: + raise ValueError(f"Object representation for layer: {layer.name} object: {object_id} is unavailable") + + def get_map_objects_in_radius( + self, + point: Union[Point2D, Point3D], + radius: float, + layers: List[MapLayer], + ) -> Dict[MapLayer, List[BaseMapObject]]: + """Inherited, see superclass.""" + center_point = point.shapely_point + patch = center_point.buffer(radius) + return self.query(geometry=patch, layers=layers, predicate="intersects") + + def query( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]]: + """Inherited, see superclass.""" + supported_layers = self.get_available_map_layers() + unsupported_layers = [layer for layer in layers if layer not in supported_layers] + assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" + object_map: Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]] = defaultdict(list) + for layer in layers: + object_map[layer] = self._query_layer(geometry, layer, predicate, sort, distance) + return object_map + + def query_object_ids( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: + """Inherited, see superclass.""" + supported_layers = self.get_available_map_layers() + unsupported_layers = [layer for layer in layers if layer not in supported_layers] + assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + for layer in layers: + object_map[layer] = self._query_layer_objects_ids(geometry, layer, predicate, sort, distance) + return object_map + + def query_nearest( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + return_all: bool = True, + max_distance: Optional[float] = None, + return_distance: bool = False, + exclusive: bool = False, + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: + """Inherited, see superclass.""" + supported_layers = self.get_available_map_layers() + unsupported_layers = [layer for layer in layers if layer not in supported_layers] + assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + for layer in layers: + object_map[layer] = self._query_layer_nearest( + geometry, layer, return_all, max_distance, return_distance, exclusive + ) + return object_map + + def _query_layer( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layer: MapLayer, + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]: + """Helper method to query a single layer.""" + + queried_indices = self._gpd_dataframes[layer].sindex.query( + geometry, predicate=predicate, sort=sort, distance=distance + ) + + if queried_indices.ndim == 2: + query_dict: Dict[int, List[BaseMapObject]] = defaultdict(list) + for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): + map_object_id = self._gpd_dataframes[layer]["id"].iloc[map_object_idx] + query_dict[int(geometry_idx)].append(self.get_map_object(map_object_id, layer)) + return query_dict + else: + map_object_ids = self._gpd_dataframes[layer]["id"].iloc[queried_indices] + return [self.get_map_object(map_object_id, layer) for map_object_id in map_object_ids] + + def _query_layer_objects_ids( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layer: MapLayer, + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Union[List[str], Dict[int, List[str]]]: + """Helper method to query a single layer.""" + + queried_indices = self._gpd_dataframes[layer].sindex.query( + geometry, predicate=predicate, sort=sort, distance=distance + ) + ids = self._gpd_dataframes[layer]["id"].values # numpy array for fast access + + if queried_indices.ndim == 2: + query_dict: Dict[int, List[str]] = defaultdict(list) + for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): + query_dict[int(geometry_idx)].append(ids[map_object_idx]) + return query_dict + else: + return list(ids[queried_indices]) + + def _query_layer_nearest( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layer: MapLayer, + return_all: bool = True, + max_distance: Optional[float] = None, + return_distance: bool = False, + exclusive: bool = False, + ) -> Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]: + """Helper method to query a single layer.""" + + queried_indices = self._gpd_dataframes[layer].sindex.nearest( + geometry, + return_all=return_all, + max_distance=max_distance, + return_distance=return_distance, + exclusive=exclusive, + ) + ids = self._gpd_dataframes[layer]["id"].values # numpy array for fast access + + if return_distance: + queried_indices, distances = queried_indices + query_dict: Dict[int, List[Tuple[str, float]]] = defaultdict(list) + for geometry_idx, map_object_idx, distance in zip(queried_indices[0], queried_indices[1], distances): + query_dict[int(geometry_idx)].append((ids[map_object_idx], float(distance))) + return query_dict + + elif queried_indices.ndim == 2: + query_dict: Dict[int, List[str]] = defaultdict(list) + for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): + query_dict[int(geometry_idx)].append(ids[map_object_idx]) + return query_dict + else: + return list(ids[queried_indices]) + + def _get_lane(self, id: str) -> Optional[Lane]: + """Helper method for getting a lane by its ID.""" + lane: Optional[Lane] = None + lane_row = get_row_with_value(self._gpd_dataframes[MapLayer.LANE], "id", id) + + if lane_row is not None: + object_id: str = lane_row["id"] + lane_group_id: str = lane_row["lane_group_id"] + left_boundary: Polyline3D = Polyline3D.from_linestring(lane_row["left_boundary"]) + right_boundary: Optional[Polyline3D] = Polyline3D.from_linestring(lane_row["right_boundary"]) + centerline: Polyline3D = Polyline3D.from_linestring(lane_row["centerline"]) + left_lane_id: Optional[str] = lane_row["left_lane_id"] + right_lane_id: Optional[str] = lane_row["right_lane_id"] + predecessor_ids: List[str] = ast.literal_eval(lane_row.predecessor_ids) + successor_ids: List[str] = ast.literal_eval(lane_row.successor_ids) + speed_limit_mps: Optional[float] = lane_row["speed_limit_mps"] + outline: Optional[Polyline3D] = ( + Polyline3D.from_linestring(lane_row["outline"]) if lane_row["outline"] is not None else None + ) + geometry: geom.LineString = lane_row["geometry"] + + lane = Lane( + object_id=object_id, + lane_group_id=lane_group_id, + left_boundary=left_boundary, + right_boundary=right_boundary, + centerline=centerline, + left_lane_id=left_lane_id, + right_lane_id=right_lane_id, + predecessor_ids=predecessor_ids, + successor_ids=successor_ids, + speed_limit_mps=speed_limit_mps, + outline=outline, + shapely_polygon=geometry, + map_api=self, + ) + + return lane + + def _get_lane_group(self, id: str) -> Optional[LaneGroup]: + """Helper method for getting a lane group by its ID.""" + lane_group: Optional[LaneGroup] = None + lane_group_row = get_row_with_value(self._gpd_dataframes[MapLayer.LANE_GROUP], "id", id) + if lane_group_row is not None: + object_id: str = lane_group_row["id"] + lane_ids: List[str] = ast.literal_eval(lane_group_row.lane_ids) + left_boundary: Polyline3D = Polyline3D.from_linestring(lane_group_row["left_boundary"]) + right_boundary: Optional[Polyline3D] = Polyline3D.from_linestring(lane_group_row["right_boundary"]) + intersection_id: Optional[str] = lane_group_row["intersection_id"] + predecessor_ids: Optional[List[str]] = ast.literal_eval(lane_group_row.predecessor_ids) + successor_ids: Optional[List[str]] = ast.literal_eval(lane_group_row.successor_ids) + outline: Optional[Polyline3D] = ( + Polyline3D.from_linestring(lane_group_row["outline"]) if lane_group_row["outline"] is not None else None + ) + geometry: geom.Polygon = lane_group_row["geometry"] + + lane_group = LaneGroup( + object_id=object_id, + lane_ids=lane_ids, + left_boundary=left_boundary, + right_boundary=right_boundary, + intersection_id=intersection_id, + predecessor_ids=predecessor_ids, + successor_ids=successor_ids, + outline=outline, + shapely_polygon=geometry, + map_api=self, + ) + + return lane_group + + def _get_intersection(self, id: str) -> Optional[Intersection]: + """Helper method for getting an intersection by its ID.""" + + intersection: Optional[Intersection] = None + intersection_row = get_row_with_value(self._gpd_dataframes[MapLayer.INTERSECTION], "id", id) + if intersection_row is not None: + object_id: str = intersection_row["id"] + lane_group_ids: List[str] = ast.literal_eval(intersection_row.lane_group_ids) + outline: Optional[Polyline3D] = ( + Polyline3D.from_linestring(intersection_row["outline"]) + if intersection_row["outline"] is not None + else None + ) + geometry: geom.Polygon = intersection_row["geometry"] + + intersection = Intersection( + object_id=object_id, + lane_group_ids=lane_group_ids, + outline=outline, + shapely_polygon=geometry, + map_api=self, + ) + + return intersection + + def _get_crosswalk(self, id: str) -> Optional[Crosswalk]: + """Helper method for getting a crosswalk by its ID.""" + + crosswalk: Optional[Crosswalk] = None + crosswalk_row = get_row_with_value(self._gpd_dataframes[MapLayer.CROSSWALK], "id", id) + if crosswalk_row is not None: + object_id: str = crosswalk_row["id"] + outline: Polyline3D = Polyline3D.from_linestring(crosswalk_row["outline"]) + geometry: geom.Polygon = crosswalk_row["geometry"] + + crosswalk = Crosswalk( + object_id=object_id, + outline=outline, + shapely_polygon=geometry, + ) + + return crosswalk + + def _get_walkway(self, id: str) -> Optional[Walkway]: + """Helper method for getting a walkway by its ID.""" + + walkway: Optional[Walkway] = None + walkway_row = get_row_with_value(self._gpd_dataframes[MapLayer.WALKWAY], "id", id) + if walkway_row is not None: + object_id: str = walkway_row["id"] + outline: Polyline3D = Polyline3D.from_linestring(walkway_row["outline"]) + geometry: geom.Polygon = walkway_row["geometry"] + + walkway = Walkway( + object_id=object_id, + outline=outline, + shapely_polygon=geometry, + ) + + return walkway + + def _get_carpark(self, id: str) -> Optional[Carpark]: + """Helper method for getting a carpark by its ID.""" + + carpark: Optional[Carpark] = None + carpark_row = get_row_with_value(self._gpd_dataframes[MapLayer.CARPARK], "id", id) + if carpark_row is not None: + object_id: str = carpark_row["id"] + outline: Polyline3D = Polyline3D.from_linestring(carpark_row["outline"]) + geometry: geom.Polygon = carpark_row["geometry"] + + carpark = Carpark( + object_id=object_id, + outline=outline, + shapely_polygon=geometry, + ) + + return carpark + + def _get_generic_drivable(self, id: str) -> Optional[GenericDrivable]: + """Helper method for getting a generic drivable area by its ID.""" + + generic_drivable: Optional[GenericDrivable] = None + generic_drivable_row = get_row_with_value(self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE], "id", id) + if generic_drivable_row is not None: + object_id: str = generic_drivable_row["id"] + outline: Polyline3D = Polyline3D.from_linestring(generic_drivable_row["outline"]) + geometry: geom.Polygon = generic_drivable_row["geometry"] + + generic_drivable = GenericDrivable( + object_id=object_id, + outline=outline, + shapely_polygon=geometry, + ) + + return generic_drivable + + def _get_road_edge(self, id: str) -> Optional[RoadEdge]: + """Helper method for getting a road edge by its ID.""" + + road_edge: Optional[RoadEdge] = None + road_edge_row = get_row_with_value(self._gpd_dataframes[MapLayer.ROAD_EDGE], "id", id) + if road_edge_row is not None: + object_id: str = road_edge_row["id"] + polyline: Polyline3D = Polyline3D.from_linestring(road_edge_row["geometry"]) + road_edge_type: RoadEdgeType = RoadEdgeType(road_edge_row["road_edge_type"]) + + road_edge = RoadEdge( + object_id=object_id, + road_edge_type=road_edge_type, + polyline=polyline, + ) + + return road_edge + + def _get_road_line(self, id: str) -> Optional[RoadLine]: + """Helper method for getting a road line by its ID.""" + + road_line: Optional[RoadLine] = None + road_line_row = get_row_with_value(self._gpd_dataframes[MapLayer.ROAD_LINE], "id", id) + if road_line_row is not None: + object_id: str = road_line_row["id"] + polyline: Polyline3D = Polyline3D.from_linestring(road_line_row["geometry"]) + road_line_type: RoadLineType = RoadLineType(road_line_row["road_line_type"]) + + road_line = RoadLine( + object_id=object_id, + road_line_type=road_line_type, + polyline=polyline, + ) + + return road_line + + +@lru_cache(maxsize=MAX_LRU_CACHED_TABLES) +def get_global_map_api(dataset: str, location: str) -> GPKGMapAPI: + """Get the global map API for a given dataset and location.""" + PY123D_MAPS_ROOT: Path = Path(get_dataset_paths().py123d_maps_root) + gpkg_path = PY123D_MAPS_ROOT / dataset / f"{dataset}_{location}.gpkg" + assert gpkg_path.is_file(), f"{dataset}_{location}.gpkg not found in {str(PY123D_MAPS_ROOT)}." + map_api = GPKGMapAPI(gpkg_path) + map_api._initialize() + return map_api + + +def get_local_map_api(split_name: str, log_name: str) -> GPKGMapAPI: + """Get the local map API for a given split name and log name.""" + PY123D_MAPS_ROOT: Path = Path(get_dataset_paths().py123d_maps_root) + gpkg_path = PY123D_MAPS_ROOT / split_name / f"{log_name}.gpkg" + assert gpkg_path.is_file(), f"{log_name}.gpkg not found in {str(PY123D_MAPS_ROOT)}." + map_api = GPKGMapAPI(gpkg_path) + map_api._initialize() + return map_api diff --git a/src/py123d/api/map/gpkg/gpkg_utils.py b/src/py123d/api/map/gpkg/gpkg_utils.py new file mode 100644 index 00000000..ed9e16a7 --- /dev/null +++ b/src/py123d/api/map/gpkg/gpkg_utils.py @@ -0,0 +1,62 @@ +from typing import List, Optional + +import geopandas as gpd +import numpy as np +import pandas as pd +from shapely import wkt + + +def load_gdf_with_geometry_columns(gdf: gpd.GeoDataFrame, geometry_column_names: List[str] = []): + """Convert geometry columns stored as wkt back to shapely geometries. + + :param gdf: input GeoDataFrame. + :param geometry_column_names: List of geometry column names to convert, defaults to [] + """ + + # Convert string geometry columns back to shapely objects + for col in geometry_column_names: + if col in gdf.columns and len(gdf) > 0 and isinstance(gdf[col].iloc[0], str): + try: + gdf[col] = gdf[col].apply(lambda x: wkt.loads(x) if isinstance(x, str) else x) + except Exception as e: + print(f"Warning: Could not convert column {col} to geometry: {str(e)}") + + +def get_all_rows_with_value( + elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str +) -> Optional[gpd.geodataframe.GeoDataFrame]: + """Extract all matching elements. Note, if no matching desired_key is found and empty list is returned. + + :param elements: data frame from MapsDb. + :param column_label: key to extract from a column. + :param desired_value: key which is compared with the values of column_label entry. + :return: a subset of the original GeoDataFrame containing the matching key. + """ + if desired_value is None or pd.isna(desired_value): + return None + + return elements.iloc[np.where(elements[column_label].to_numpy().astype(int) == int(desired_value))] + + +def get_row_with_value( + elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str +) -> Optional[gpd.GeoSeries]: + """Extract a matching element. + + :param elements: data frame from MapsDb. + :param column_label: key to extract from a column. + :param desired_value: key which is compared with the values of column_label entry. + :return row from GeoDataFrame. + """ + if column_label == "fid": + return elements.loc[desired_value] + + geo_series: Optional[gpd.GeoSeries] = None + matching_rows = get_all_rows_with_value(elements, column_label, desired_value) + if matching_rows is not None: + assert len(matching_rows) > 0, f"Could not find the desired key = {desired_value}" + assert len(matching_rows) == 1, ( + f"{len(matching_rows)} matching keys found. Expected to only find one.Try using get_all_rows_with_value" + ) + geo_series = matching_rows.iloc[0] + return geo_series diff --git a/src/py123d/api/map/map_api.py b/src/py123d/api/map/map_api.py new file mode 100644 index 00000000..bfd1d131 --- /dev/null +++ b/src/py123d/api/map/map_api.py @@ -0,0 +1,200 @@ +from __future__ import annotations + +import abc +from typing import Dict, Iterable, List, Literal, Optional, Union + +import shapely + +from py123d.datatypes.map_objects import BaseMapObject, MapLayer +from py123d.datatypes.metadata import MapMetadata +from py123d.geometry import Point2D +from py123d.geometry.point import Point3D + + +class MapAPI(abc.ABC): + """The base class for all map APIs in 123D.""" + + # Abstract Methods, to be implemented by subclasses + # ------------------------------------------------------------------------------------------------------------------ + + @abc.abstractmethod + def get_map_metadata(self) -> MapMetadata: + """Returns the :class:`~p123d.datatypes.metadata.MapMetadata` of the map api. + + :return: The map metadata, e.g. location, dataset, etc. + """ + + @abc.abstractmethod + def get_available_map_layers(self) -> List[MapLayer]: + """Returns the available :class:`~p123d.datatypes.map_objects.map_layer_types.MapLayer`, + e.g. LANE, LANE_GROUP, etc. + + :return: A list of available map layers. + """ + + @abc.abstractmethod + def get_map_object(self, object_id: Union[str, int], layer: MapLayer) -> Optional[BaseMapObject]: + """Returns a :class:`~p123d.datatypes.map_objects.base_map_object.BaseMapObject` by its ID + and :class:`~p123d.datatypes.map_objects.map_layer_types.MapLayer`. + + :param object_id: The ID of the map object. + :param layer: The layer the map object belongs to. + :return: The map object if found, None otherwise. + """ + + @abc.abstractmethod + def get_map_objects_in_radius( + self, + point: Union[Point2D, Point3D], + radius: float, + layers: List[MapLayer], + ) -> Dict[MapLayer, List[BaseMapObject]]: + """Returns a dictionary of :class:`~p123d.datatypes.map_objects.map_layer_types.MapLayer` to a list of + :class:`~p123d.datatypes.map_objects.base_map_object.BaseMapObject` within a given radius + around a center point. + + :param point: The center point to search around. + :param radius: The radius to search within. + :param layers: The map layers to search in. + :return: A dictionary mapping each layer to a list of map objects within the radius. + """ + + @abc.abstractmethod + def query( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[ + Literal[ + "contains", + "contains_properly", + "covered_by", + "covers", + "crosses", + "intersects", + "overlaps", + "touches", + "within", + "dwithin", + ] + ] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]]: + """Queries geometries against the map objects in the specified layers using an optional spatial predicate. + + Notes + ----- + The implementation is aligned with the geopandas spatial index query method [1]_, + used in the :class:`GPKGMapAPI`. It is likely this method will be removed or improved in + future versions of 123D. + + References + ---------- + [1] https://geopandas.org/en/stable/docs/reference/api/geopandas.sindex.SpatialIndex.query.html + + :param geometry: A shapely geometry or an iterable of shapely geometries to query against. + :param layers: The map layers to query against. + :param predicate: An optional spatial predicate to filter the results. + :param sort: Whether to sort the results by distance, defaults to False. + :param distance: An optional maximum distance to filter the results, defaults to None. + :return: A dictionary mapping each layer to a list of map objects or a dictionary of indices to + lists of map objects. + """ + + @abc.abstractmethod + def query_object_ids( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[int], Dict[int, List[int]]]]: + """Queries geometries against the map objects in the specified layers using an optional spatial predicate. + Instead of returning the map objects, it returns their IDs only. + + Notes + ----- + The implementation is aligned with the geopandas spatial index query method [1]_, + used in the :class:`GPKGMapAPI`. It is likely this method will be removed or improved in + future versions of 123D. + + References + ---------- + [1] https://geopandas.org/en/stable/docs/reference/api/geopandas.sindex.SpatialIndex.query.html + + + :param geometry: A shapely geometry or an iterable of shapely geometries to query against. + :param layers: The map layers to query against. + :param predicate: An optional spatial predicate to filter the results. + :param sort: Whether to sort the results by distance, defaults to False. + :param distance: An optional maximum distance to filter the results, defaults to None. + :return: A dictionary mapping each layer to a list of map object IDs or a dictionary of indices to + lists of map object IDs. + """ + + @abc.abstractmethod + def query_nearest( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + return_all: bool = True, + max_distance: Optional[float] = None, + return_distance: bool = False, + exclusive: bool = False, + ) -> Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]]: + """Return the nearest map objects in a spatial tree for each input geometry in + + Notes + ----- + The implementation is aligned with the geopandas spatial index nearest method [1]_, + used in the :class:`GPKGMapAPI`. It is likely this method will be removed or improved in + future versions of 123D. + + References + ---------- + [1] https://geopandas.org/en/stable/docs/reference/api/geopandas.sindex.SpatialIndex.nearest.html + + :param geometry: A shapely geometry or an iterable of shapely geometries to query against. + :param layers: The map layers to query against. + :param return_all: Whether to return all matching objects or just the closest one, defaults to True. + :param max_distance: An optional maximum distance to filter the results, defaults to None. + :param return_distance: Whether to return the distance to the nearest object, defaults to False. + :param exclusive: Whether to exclude the input geometries from the results, defaults to False. + :return: A dictionary mapping each layer to a list of nearest map objects or a dictionary of indices to + lists of nearest map objects. + """ + + # Syntactic Sugar / Properties, for easier access to common attributes + # ------------------------------------------------------------------------------------------------------------------ + + @property + def map_metadata(self) -> MapMetadata: + """The :class:`~py123d.datatypes.metadata.MapMetadata` of the map api.""" + return self.get_map_metadata() + + @property + def dataset(self) -> str: + """The dataset name from the map metadata.""" + return self.map_metadata.dataset + + @property + def location(self) -> str: + """The location from the map metadata.""" + return self.map_metadata.location + + @property + def map_is_local(self) -> bool: + """Indicates if the map is local (map for each log) or global (map for multiple logs in dataset).""" + return self.map_metadata.map_is_local + + @property + def map_has_z(self) -> bool: + """Indicates if the map includes Z (elevation) data.""" + return self.map_metadata.map_has_z + + @property + def version(self) -> str: + """The version of the py123d library used to create this map metadata.""" + return self.map_metadata.version diff --git a/src/py123d/api/scene/__init__.py b/src/py123d/api/scene/__init__.py new file mode 100644 index 00000000..51c4ddc5 --- /dev/null +++ b/src/py123d/api/scene/__init__.py @@ -0,0 +1,3 @@ +from py123d.api.scene.scene_api import SceneAPI +from py123d.api.scene.scene_builder import SceneBuilder +from py123d.api.scene.scene_filter import SceneFilter diff --git a/src/py123d/datatypes/scene/__init__.py b/src/py123d/api/scene/arrow/__init__.py similarity index 100% rename from src/py123d/datatypes/scene/__init__.py rename to src/py123d/api/scene/arrow/__init__.py diff --git a/src/py123d/datatypes/scene/arrow/arrow_scene.py b/src/py123d/api/scene/arrow/arrow_scene.py similarity index 59% rename from src/py123d/datatypes/scene/arrow/arrow_scene.py rename to src/py123d/api/scene/arrow/arrow_scene.py index ee26a5f4..1472c907 100644 --- a/src/py123d/datatypes/scene/arrow/arrow_scene.py +++ b/src/py123d/api/scene/arrow/arrow_scene.py @@ -3,22 +3,25 @@ import pyarrow as pa -from py123d.common.utils.arrow_helper import get_lru_cached_arrow_table -from py123d.datatypes.detections.box_detections import BoxDetectionWrapper -from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper -from py123d.datatypes.maps.abstract_map import AbstractMap -from py123d.datatypes.maps.gpkg.gpkg_map import get_global_map_api, get_local_map_api -from py123d.datatypes.scene.abstract_scene import AbstractScene -from py123d.datatypes.scene.arrow.utils.arrow_getters import ( - get_box_detections_from_arrow_table, +from py123d.api.map.gpkg.gpkg_map_api import get_global_map_api, get_local_map_api +from py123d.api.map.map_api import MapAPI +from py123d.api.scene.arrow.utils.arrow_getters import ( + get_box_detections_se3_from_arrow_table, get_camera_from_arrow_table, - get_ego_vehicle_state_from_arrow_table, + get_ego_state_se3_from_arrow_table, get_lidar_from_arrow_table, + get_route_lane_group_ids_from_arrow_table, get_timepoint_from_arrow_table, get_traffic_light_detections_from_arrow_table, ) -from py123d.datatypes.scene.arrow.utils.arrow_metadata_utils import get_log_metadata_from_arrow -from py123d.datatypes.scene.scene_metadata import LogMetadata, SceneExtractionMetadata +from py123d.api.scene.arrow.utils.arrow_metadata_utils import get_log_metadata_from_arrow_file +from py123d.api.scene.scene_api import SceneAPI +from py123d.api.scene.scene_metadata import SceneMetadata +from py123d.common.utils.arrow_column_names import UUID_COLUMN +from py123d.common.utils.arrow_helper import get_lru_cached_arrow_table +from py123d.datatypes.detections.box_detections import BoxDetectionWrapper +from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper +from py123d.datatypes.metadata.log_metadata import LogMetadata from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICamera, FisheyeMEICameraType from py123d.datatypes.sensors.lidar import LiDAR, LiDARType from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeCameraType @@ -26,40 +29,48 @@ from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -class ArrowScene(AbstractScene): +def _get_complete_log_scene_metadata(arrow_file_path: Union[Path, str], log_metadata: LogMetadata) -> SceneMetadata: + """Helper function to get the scene metadata for a complete log of an Arrow file.""" + table = get_lru_cached_arrow_table(arrow_file_path) + initial_uuid = table[UUID_COLUMN][0].as_py() + num_rows = table.num_rows + return SceneMetadata( + initial_uuid=initial_uuid, + initial_idx=0, + duration_s=log_metadata.timestep_seconds * num_rows, + history_s=0.0, + iteration_duration_s=log_metadata.timestep_seconds, + ) + + +class ArrowSceneAPI(SceneAPI): + """Scene API for Arrow-based scenes. Provides access to all data modalities in an Arrow scene.""" def __init__( self, arrow_file_path: Union[Path, str], - scene_extraction_metadata: Optional[SceneExtractionMetadata] = None, + scene_metadata: Optional[SceneMetadata] = None, ) -> None: + """Initializes the :class:`ArrowSceneAPI`. + + :param arrow_file_path: Path to the Arrow file. + :param scene_metadata: Scene metadata, defaults to None + """ self._arrow_file_path: Path = Path(arrow_file_path) - self._log_metadata: LogMetadata = get_log_metadata_from_arrow(arrow_file_path) - - with pa.memory_map(str(self._arrow_file_path), "r") as source: - reader = pa.ipc.open_file(source) - table = reader.read_all() - num_rows = table.num_rows - initial_uuid = table["uuid"][0].as_py() - - if scene_extraction_metadata is None: - scene_extraction_metadata = SceneExtractionMetadata( - initial_uuid=initial_uuid, - initial_idx=0, - duration_s=self._log_metadata.timestep_seconds * num_rows, - history_s=0.0, - iteration_duration_s=self._log_metadata.timestep_seconds, - ) - self._scene_extraction_metadata: SceneExtractionMetadata = scene_extraction_metadata + self._log_metadata: LogMetadata = get_log_metadata_from_arrow_file(str(arrow_file_path)) + self._scene_metadata: SceneMetadata = ( + scene_metadata + if scene_metadata is not None + else _get_complete_log_scene_metadata(arrow_file_path, self._log_metadata) + ) # NOTE: Lazy load a log-specific map API, and keep reference. # Global maps are LRU cached internally. - self._local_map_api: Optional[AbstractMap] = None + self._local_map_api: Optional[MapAPI] = None - #################################################################################################################### - # Helpers for ArrowScene - #################################################################################################################### + # Helper methods + # ------------------------------------------------------------------------------------------------------------------ def __reduce__(self): """Helper for pickling the object.""" @@ -67,7 +78,7 @@ def __reduce__(self): self.__class__, ( self._arrow_file_path, - self._scene_extraction_metadata, + self._scene_metadata, ), ) @@ -76,22 +87,25 @@ def _get_recording_table(self) -> pa.Table: return get_lru_cached_arrow_table(self._arrow_file_path) def _get_table_index(self, iteration: int) -> int: + """Helper function to get the table index for a given iteration.""" assert -self.number_of_history_iterations <= iteration < self.number_of_iterations, "Iteration out of bounds" - table_index = self._scene_extraction_metadata.initial_idx + iteration + table_index = self._scene_metadata.initial_idx + iteration return table_index - #################################################################################################################### - # Implementation of AbstractScene - #################################################################################################################### + # Implementation of abstract methods + # ------------------------------------------------------------------------------------------------------------------ def get_log_metadata(self) -> LogMetadata: + """Inherited, see superclass.""" return self._log_metadata - def get_scene_extraction_metadata(self) -> SceneExtractionMetadata: - return self._scene_extraction_metadata + def get_scene_metadata(self) -> SceneMetadata: + """Inherited, see superclass.""" + return self._scene_metadata - def get_map_api(self) -> Optional[AbstractMap]: - map_api: Optional[AbstractMap] = None + def get_map_api(self) -> Optional[MapAPI]: + """Inherited, see superclass.""" + map_api: Optional[MapAPI] = None if self.log_metadata.map_metadata is not None: if self.log_metadata.map_metadata.map_is_local: if self._local_map_api is None: @@ -104,37 +118,39 @@ def get_map_api(self) -> Optional[AbstractMap]: return map_api def get_timepoint_at_iteration(self, iteration: int) -> TimePoint: + """Inherited, see superclass.""" return get_timepoint_from_arrow_table(self._get_recording_table(), self._get_table_index(iteration)) def get_ego_state_at_iteration(self, iteration: int) -> Optional[EgoStateSE3]: - return get_ego_vehicle_state_from_arrow_table( + """Inherited, see superclass.""" + return get_ego_state_se3_from_arrow_table( self._get_recording_table(), self._get_table_index(iteration), self.log_metadata.vehicle_parameters, ) def get_box_detections_at_iteration(self, iteration: int) -> Optional[BoxDetectionWrapper]: - return get_box_detections_from_arrow_table( + """Inherited, see superclass.""" + return get_box_detections_se3_from_arrow_table( self._get_recording_table(), self._get_table_index(iteration), self.log_metadata, ) def get_traffic_light_detections_at_iteration(self, iteration: int) -> Optional[TrafficLightDetectionWrapper]: + """Inherited, see superclass.""" return get_traffic_light_detections_from_arrow_table( self._get_recording_table(), self._get_table_index(iteration) ) def get_route_lane_group_ids(self, iteration: int) -> Optional[List[int]]: - route_lane_group_ids: Optional[List[int]] = None - table = self._get_recording_table() - if "route_lane_group_ids" in table.column_names: - route_lane_group_ids = table["route_lane_group_ids"][self._get_table_index(iteration)].as_py() - return route_lane_group_ids + """Inherited, see superclass.""" + return get_route_lane_group_ids_from_arrow_table(self._get_recording_table(), self._get_table_index(iteration)) def get_pinhole_camera_at_iteration( - self, iteration: int, camera_type: Union[PinholeCameraType, FisheyeMEICameraType] - ) -> Optional[Union[PinholeCamera, FisheyeMEICamera]]: + self, iteration: int, camera_type: PinholeCameraType + ) -> Optional[PinholeCamera]: + """Inherited, see superclass.""" pinhole_camera: Optional[PinholeCamera] = None if camera_type in self.available_pinhole_camera_types: pinhole_camera = get_camera_from_arrow_table( @@ -148,6 +164,7 @@ def get_pinhole_camera_at_iteration( def get_fisheye_mei_camera_at_iteration( self, iteration: int, camera_type: FisheyeMEICameraType ) -> Optional[FisheyeMEICamera]: + """Inherited, see superclass.""" fisheye_mei_camera: Optional[FisheyeMEICamera] = None if camera_type in self.available_fisheye_mei_camera_types: fisheye_mei_camera = get_camera_from_arrow_table( @@ -159,6 +176,7 @@ def get_fisheye_mei_camera_at_iteration( return fisheye_mei_camera def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> Optional[LiDAR]: + """Inherited, see superclass.""" lidar: Optional[LiDAR] = None if lidar_type in self.available_lidar_types or lidar_type == LiDARType.LIDAR_MERGED: lidar = get_lidar_from_arrow_table( diff --git a/src/py123d/datatypes/scene/arrow/arrow_scene_builder.py b/src/py123d/api/scene/arrow/arrow_scene_builder.py similarity index 64% rename from src/py123d/datatypes/scene/arrow/arrow_scene_builder.py rename to src/py123d/api/scene/arrow/arrow_scene_builder.py index 2afebbba..2326fb7e 100644 --- a/src/py123d/datatypes/scene/arrow/arrow_scene_builder.py +++ b/src/py123d/api/scene/arrow/arrow_scene_builder.py @@ -1,29 +1,33 @@ import random from functools import partial from pathlib import Path -from typing import Iterator, List, Optional, Set, Union - +from typing import List, Optional, Set, Union + +from py123d.api.scene.arrow.arrow_scene import ArrowSceneAPI +from py123d.api.scene.arrow.utils.arrow_metadata_utils import get_log_metadata_from_arrow_table +from py123d.api.scene.scene_api import SceneAPI +from py123d.api.scene.scene_builder import SceneBuilder +from py123d.api.scene.scene_filter import SceneFilter +from py123d.api.scene.scene_metadata import SceneMetadata from py123d.common.multithreading.worker_utils import WorkerPool, worker_map +from py123d.common.utils.arrow_column_names import FISHEYE_CAMERA_DATA_COLUMN, PINHOLE_CAMERA_DATA_COLUMN, UUID_COLUMN from py123d.common.utils.arrow_helper import get_lru_cached_arrow_table -from py123d.datatypes.scene.abstract_scene import AbstractScene -from py123d.datatypes.scene.abstract_scene_builder import SceneBuilder -from py123d.datatypes.scene.arrow.arrow_scene import ArrowScene -from py123d.datatypes.scene.arrow.utils.arrow_metadata_utils import get_log_metadata_from_arrow -from py123d.datatypes.scene.scene_filter import SceneFilter -from py123d.datatypes.scene.scene_metadata import SceneExtractionMetadata from py123d.script.utils.dataset_path_utils import get_dataset_paths class ArrowSceneBuilder(SceneBuilder): - """ - A class to build a scene from a dataset. - """ + """Class for building scenes from Arrow log files.""" def __init__( self, logs_root: Optional[Union[str, Path]] = None, maps_root: Optional[Union[str, Path]] = None, ): + """Initializes the ArrowSceneBuilder. + + :param logs_root: The root directory fo log files, defaults to None + :param maps_root: The root directory for map files, defaults to None + """ if logs_root is None: logs_root = get_dataset_paths().py123d_logs_root if maps_root is None: @@ -32,8 +36,8 @@ def __init__( self._logs_root = Path(logs_root) self._maps_root = Path(maps_root) - def get_scenes(self, filter: SceneFilter, worker: WorkerPool) -> Iterator[AbstractScene]: - """See superclass.""" + def get_scenes(self, filter: SceneFilter, worker: WorkerPool) -> List[SceneAPI]: + """Inherited, see superclass.""" split_types = set(filter.split_types) if filter.split_types else {"train", "val", "test"} split_names = ( @@ -41,23 +45,24 @@ def get_scenes(self, filter: SceneFilter, worker: WorkerPool) -> Iterator[Abstra ) filter_log_names = set(filter.log_names) if filter.log_names else None log_paths = _discover_log_paths(self._logs_root, split_names, filter_log_names) + if len(log_paths) == 0: return [] - scenes = worker_map(worker, partial(_extract_scenes_from_logs, filter=filter), log_paths) + scenes = worker_map(worker, partial(_extract_scenes_from_logs, filter=filter), log_paths) if filter.shuffle: random.shuffle(scenes) if filter.max_num_scenes is not None: scenes = scenes[: filter.max_num_scenes] - return scenes def _discover_split_names(logs_root: Path, split_types: Set[str]) -> Set[str]: - assert set(split_types).issubset( - {"train", "val", "test"} - ), f"Invalid split types: {split_types}. Valid split types are 'train', 'val', 'test'." + """Discovers split names in the logs root directory based on the specified split types.""" + assert set(split_types).issubset({"train", "val", "test"}), ( + f"Invalid split types: {split_types}. Valid split types are 'train', 'val', 'test'." + ) split_names: List[str] = [] for split in logs_root.iterdir(): split_name = split.name @@ -69,6 +74,7 @@ def _discover_split_names(logs_root: Path, split_types: Set[str]) -> Set[str]: def _discover_log_paths(logs_root: Path, split_names: Set[str], log_names: Optional[List[str]]) -> List[Path]: + """Discovers log file paths in the logs root directory based on the specified split names and log names.""" log_paths: List[Path] = [] for split_name in split_names: for log_path in (logs_root / split_name).iterdir(): @@ -78,8 +84,9 @@ def _discover_log_paths(logs_root: Path, split_names: Set[str], log_names: Optio return log_paths -def _extract_scenes_from_logs(log_paths: List[Path], filter: SceneFilter) -> List[AbstractScene]: - scenes: List[AbstractScene] = [] +def _extract_scenes_from_logs(log_paths: List[Path], filter: SceneFilter) -> List[SceneAPI]: + """Extracts scenes from log files based on the given filter.""" + scenes: List[SceneAPI] = [] for log_path in log_paths: try: scene_extraction_metadatas = _get_scene_extraction_metadatas(log_path, filter) @@ -88,19 +95,19 @@ def _extract_scenes_from_logs(log_paths: List[Path], filter: SceneFilter) -> Lis continue for scene_extraction_metadata in scene_extraction_metadatas: scenes.append( - ArrowScene( + ArrowSceneAPI( arrow_file_path=log_path, - scene_extraction_metadata=scene_extraction_metadata, + scene_metadata=scene_extraction_metadata, ) ) return scenes -def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFilter) -> List[SceneExtractionMetadata]: - scene_extraction_metadatas: List[SceneExtractionMetadata] = [] - - recording_table = get_lru_cached_arrow_table(log_path) - log_metadata = get_log_metadata_from_arrow(log_path) +def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFilter) -> List[SceneMetadata]: + """Gets the scene metadatas from a log file based on the given filter.""" + scene_metadatas: List[SceneMetadata] = [] + recording_table = get_lru_cached_arrow_table(str(log_path)) + log_metadata = get_log_metadata_from_arrow_table(recording_table) start_idx = int(filter.history_s / log_metadata.timestep_seconds) if filter.history_s is not None else 0 end_idx = ( @@ -109,8 +116,10 @@ def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFil else len(recording_table) ) - # 1. Filter location - if ( + # 1. Filter location & whether map API is required + if filter.map_api_required and log_metadata.map_metadata is None: + pass + elif ( filter.locations is not None and log_metadata.map_metadata is not None and log_metadata.map_metadata.location not in filter.locations @@ -118,9 +127,9 @@ def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFil pass elif filter.duration_s is None: - scene_extraction_metadatas.append( - SceneExtractionMetadata( - initial_uuid=str(recording_table["uuid"][start_idx].as_py()), + scene_metadatas.append( + SceneMetadata( + initial_uuid=str(recording_table[UUID_COLUMN][start_idx].as_py()), initial_idx=start_idx, duration_s=(end_idx - start_idx) * log_metadata.timestep_seconds, history_s=filter.history_s if filter.history_s is not None else 0.0, @@ -129,43 +138,47 @@ def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFil ) else: scene_uuid_set = set(filter.scene_uuids) if filter.scene_uuids is not None else None - for idx in range(start_idx, end_idx): - scene_extraction_metadata: Optional[SceneExtractionMetadata] = None + step_idx = int(filter.duration_s / log_metadata.timestep_seconds) + all_row_uuids = recording_table[UUID_COLUMN].to_pylist() + history_s = filter.history_s if filter.history_s is not None else 0.0 + + for idx in range(start_idx, end_idx, step_idx): + scene_extraction_metadata: Optional[SceneMetadata] = None + current_uuid = str(all_row_uuids[idx]) if scene_uuid_set is None: - scene_extraction_metadata = SceneExtractionMetadata( - initial_uuid=str(recording_table["uuid"][idx].as_py()), + scene_extraction_metadata = SceneMetadata( + initial_uuid=current_uuid, initial_idx=idx, duration_s=filter.duration_s, - history_s=filter.history_s, + history_s=history_s, iteration_duration_s=log_metadata.timestep_seconds, ) - elif str(recording_table["uuid"][idx]) in scene_uuid_set: - scene_extraction_metadata = SceneExtractionMetadata( - initial_uuid=str(recording_table["uuid"][idx].as_py()), + elif current_uuid in scene_uuid_set: + scene_extraction_metadata = SceneMetadata( + initial_uuid=current_uuid, initial_idx=idx, duration_s=filter.duration_s, - history_s=filter.history_s, + history_s=history_s, iteration_duration_s=log_metadata.timestep_seconds, ) if scene_extraction_metadata is not None: # Check of timestamp threshold exceeded between previous scene, if specified in filter - if filter.timestamp_threshold_s is not None and len(scene_extraction_metadatas) > 0: - iteration_delta = idx - scene_extraction_metadatas[-1].initial_idx + if filter.timestamp_threshold_s is not None and len(scene_metadatas) > 0: + iteration_delta = idx - scene_metadatas[-1].initial_idx if (iteration_delta * log_metadata.timestep_seconds) < filter.timestamp_threshold_s: continue - scene_extraction_metadatas.append(scene_extraction_metadata) + scene_metadatas.append(scene_extraction_metadata) scene_extraction_metadatas_ = [] - for scene_extraction_metadata in scene_extraction_metadatas: - + for scene_extraction_metadata in scene_metadatas: add_scene = True start_idx = scene_extraction_metadata.initial_idx if filter.pinhole_camera_types is not None: for pinhole_camera_type in filter.pinhole_camera_types: - column_name = f"{pinhole_camera_type.serialize()}_data" + column_name = PINHOLE_CAMERA_DATA_COLUMN(pinhole_camera_type.serialize()) if ( pinhole_camera_type in log_metadata.pinhole_camera_metadata @@ -179,7 +192,7 @@ def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFil if filter.fisheye_mei_camera_types is not None: for fisheye_mei_camera_type in filter.fisheye_mei_camera_types: - column_name = f"{fisheye_mei_camera_type.serialize()}_data" + column_name = FISHEYE_CAMERA_DATA_COLUMN(fisheye_mei_camera_type.serialize()) if ( fisheye_mei_camera_type in log_metadata.fisheye_mei_camera_metadata @@ -194,7 +207,5 @@ def _get_scene_extraction_metadatas(log_path: Union[str, Path], filter: SceneFil if add_scene: scene_extraction_metadatas_.append(scene_extraction_metadata) - scene_extraction_metadatas = scene_extraction_metadatas_ - - del recording_table, log_metadata - return scene_extraction_metadatas + # scene_extraction_metadata = scene_extraction_metadatas_ + return scene_extraction_metadatas_ diff --git a/src/py123d/datatypes/scene/arrow/__init__.py b/src/py123d/api/scene/arrow/utils/__init__.py similarity index 100% rename from src/py123d/datatypes/scene/arrow/__init__.py rename to src/py123d/api/scene/arrow/utils/__init__.py diff --git a/src/py123d/api/scene/arrow/utils/arrow_getters.py b/src/py123d/api/scene/arrow/utils/arrow_getters.py new file mode 100644 index 00000000..c0257917 --- /dev/null +++ b/src/py123d/api/scene/arrow/utils/arrow_getters.py @@ -0,0 +1,409 @@ +from pathlib import Path +from typing import Dict, List, Optional, Type, Union + +import numpy as np +import numpy.typing as npt +import pyarrow as pa +from omegaconf import DictConfig + +from py123d.common.utils.arrow_column_names import ( + BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN, + BOX_DETECTIONS_LABEL_COLUMN, + BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN, + BOX_DETECTIONS_SE3_COLUMNS, + BOX_DETECTIONS_TOKEN_COLUMN, + BOX_DETECTIONS_VELOCITY_3D_COLUMN, + EGO_DYNAMIC_STATE_SE3_COLUMN, + EGO_REAR_AXLE_SE3_COLUMN, + EGO_STATE_SE3_COLUMNS, + FISHEYE_CAMERA_DATA_COLUMN, + FISHEYE_CAMERA_EXTRINSIC_COLUMN, + LIDAR_DATA_COLUMN, + PINHOLE_CAMERA_DATA_COLUMN, + PINHOLE_CAMERA_EXTRINSIC_COLUMN, + ROUTE_LANE_GROUP_IDS_COLUMN, + SCENARIO_TAGS_COLUMN, + TIMESTAMP_US_COLUMN, + TRAFFIC_LIGHTS_COLUMNS, + TRAFFIC_LIGHTS_LANE_ID_COLUMN, + TRAFFIC_LIGHTS_STATUS_COLUMN, +) +from py123d.common.utils.mixin import ArrayMixin +from py123d.conversion.registry import DefaultLiDARIndex +from py123d.conversion.sensor_io.camera.jpeg_camera_io import ( + decode_image_from_jpeg_binary, + is_jpeg_binary, + load_image_from_jpeg_file, +) +from py123d.conversion.sensor_io.camera.mp4_camera_io import get_mp4_reader_from_path +from py123d.conversion.sensor_io.camera.png_camera_io import decode_image_from_png_binary, is_png_binary +from py123d.conversion.sensor_io.lidar.draco_lidar_io import is_draco_binary, load_lidar_from_draco_binary +from py123d.conversion.sensor_io.lidar.file_lidar_io import load_lidar_pcs_from_file +from py123d.conversion.sensor_io.lidar.laz_lidar_io import is_laz_binary, load_lidar_from_laz_binary +from py123d.datatypes.detections import ( + BoxDetectionMetadata, + BoxDetectionSE3, + BoxDetectionWrapper, + TrafficLightDetection, + TrafficLightDetectionWrapper, + TrafficLightStatus, +) +from py123d.datatypes.metadata import LogMetadata +from py123d.datatypes.sensors import ( + FisheyeMEICamera, + FisheyeMEICameraType, + LiDAR, + LiDARMetadata, + LiDARType, + PinholeCamera, + PinholeCameraType, +) +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import DynamicStateSE3, EgoStateSE3, VehicleParameters +from py123d.geometry import BoundingBoxSE3, PoseSE3, Vector3D +from py123d.script.utils.dataset_path_utils import get_dataset_paths + +DATASET_PATHS: DictConfig = get_dataset_paths() +DATASET_SENSOR_ROOT: Dict[str, Path] = { + "av2-sensor": DATASET_PATHS.av2_sensor_data_root, + "nuplan": DATASET_PATHS.nuplan_sensor_root, + "nuscenes": DATASET_PATHS.nuscenes_data_root, + "wopd": DATASET_PATHS.wopd_data_root, + "pandaset": DATASET_PATHS.pandaset_data_root, + "kitti360": DATASET_PATHS.kitti360_data_root, +} + + +def get_timepoint_from_arrow_table(arrow_table: pa.Table, index: int) -> TimePoint: + """Builds a :class:`~py123d.datatypes.time.TimePoint` from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the timepoint data. + :param index: The index to extract the timepoint from. + :return: The TimePoint at the given index. + """ + assert TIMESTAMP_US_COLUMN in arrow_table.schema.names, "Timestamp column not found in Arrow table." + return TimePoint.from_us(arrow_table[TIMESTAMP_US_COLUMN][index].as_py()) + + +def get_ego_state_se3_from_arrow_table( + arrow_table: pa.Table, + index: int, + vehicle_parameters: Optional[VehicleParameters], +) -> Optional[EgoStateSE3]: + """Builds a :class:`~py123d.datatypes.vehicle_state.EgoStateSE3` from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the ego state data. + :param index: The index to extract the ego state from. + :param vehicle_parameters: The vehicle parameters used to build the ego state. + :return: The ego state at the given index, or None if not available. + """ + + ego_state_se3: Optional[EgoStateSE3] = None + if _all_columns_in_schema(arrow_table, EGO_STATE_SE3_COLUMNS) and vehicle_parameters is not None: + timepoint = get_timepoint_from_arrow_table(arrow_table, index) + rear_axle_se3 = PoseSE3.from_list(arrow_table[EGO_REAR_AXLE_SE3_COLUMN][index].as_py()) + dynamic_state_se3 = _get_optional_array_mixin( + arrow_table[EGO_DYNAMIC_STATE_SE3_COLUMN][index].as_py(), + DynamicStateSE3, + ) + ego_state_se3 = EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_se3, + vehicle_parameters=vehicle_parameters, + dynamic_state_se3=dynamic_state_se3, + timepoint=timepoint, + ) + return ego_state_se3 + + +def get_box_detections_se3_from_arrow_table( + arrow_table: pa.Table, + index: int, + log_metadata: LogMetadata, +) -> BoxDetectionWrapper: + """Builds a :class:`~py123d.datatypes.detections.BoxDetectionWrapper` from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the box detections data. + :param index: The index to extract the box detections from. + :param log_metadata: The log metadata, contained the label class information. + :return: The BoxDetectionWrapper at the given index. + """ + + box_detections: Optional[BoxDetectionWrapper] = None + if _all_columns_in_schema(arrow_table, BOX_DETECTIONS_SE3_COLUMNS): + timepoint = get_timepoint_from_arrow_table(arrow_table, index) + box_detections_list: List[BoxDetectionSE3] = [] + box_detection_label_class = log_metadata.box_detection_label_class + assert box_detection_label_class is not None, "Box detection label class mapping not found in log metadata." + for _bounding_box_se3, _token, _label, _velocity, _num_lidar_points in zip( + arrow_table[BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN][index].as_py(), + arrow_table[BOX_DETECTIONS_TOKEN_COLUMN][index].as_py(), + arrow_table[BOX_DETECTIONS_LABEL_COLUMN][index].as_py(), + arrow_table[BOX_DETECTIONS_VELOCITY_3D_COLUMN][index].as_py(), + arrow_table[BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN][index].as_py(), + ): + box_detections_list.append( + BoxDetectionSE3( + metadata=BoxDetectionMetadata( + label=box_detection_label_class(_label), + track_token=_token, + num_lidar_points=_num_lidar_points, + timepoint=timepoint, + ), + bounding_box_se3=BoundingBoxSE3.from_list(_bounding_box_se3), + velocity_3d=_get_optional_array_mixin(_velocity, Vector3D), + ) + ) + box_detections = BoxDetectionWrapper(box_detections=box_detections_list) + + return box_detections + + +def get_traffic_light_detections_from_arrow_table( + arrow_table: pa.Table, + index: int, +) -> Optional[TrafficLightDetectionWrapper]: + """Builds a :class:`~py123d.datatypes.detections.TrafficLightDetectionWrapper` from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the traffic light detections data. + :param index: The index to extract the traffic light detections from. + :return: The TrafficLightDetectionWrapper at the given index, or None if not available. + """ + traffic_lights: Optional[TrafficLightDetectionWrapper] = None + if _all_columns_in_schema(arrow_table, TRAFFIC_LIGHTS_COLUMNS): + timepoint = get_timepoint_from_arrow_table(arrow_table, index) + traffic_light_detections: List[TrafficLightDetection] = [] + for lane_id, status in zip( + arrow_table[TRAFFIC_LIGHTS_LANE_ID_COLUMN][index].as_py(), + arrow_table[TRAFFIC_LIGHTS_STATUS_COLUMN][index].as_py(), + ): + traffic_light_detections.append( + TrafficLightDetection( + timepoint=timepoint, + lane_id=lane_id, + status=TrafficLightStatus(status), + ) + ) + traffic_lights = TrafficLightDetectionWrapper(traffic_light_detections=traffic_light_detections) + return traffic_lights + + +def get_camera_from_arrow_table( + arrow_table: pa.Table, + index: int, + camera_type: Union[PinholeCameraType, FisheyeMEICameraType], + log_metadata: LogMetadata, +) -> Optional[Union[PinholeCamera, FisheyeMEICamera]]: + """Builds a camera object from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the camera data. + :param index: The index to extract the camera data from. + :param camera_type: The type of camera to build (Pinhole or FisheyeMEI). + :param log_metadata: Metadata about the log, including dataset information. + :raises ValueError: If the camera data format is unsupported. + :raises NotImplementedError: If the camera data type is not supported. + :return: The constructed camera object, or None if not available. + """ + + assert isinstance(camera_type, (PinholeCameraType, FisheyeMEICameraType)), ( + f"camera_type must be PinholeCameraType or FisheyeMEICameraType, got {type(camera_type)}" + ) + + camera: Optional[Union[PinholeCamera, FisheyeMEICamera]] = None + + camera_name = camera_type.serialize() + is_pinhole = isinstance(camera_type, PinholeCameraType) + + if is_pinhole: + camera_data_column = PINHOLE_CAMERA_DATA_COLUMN(camera_name) + camera_extrinsic_column = PINHOLE_CAMERA_EXTRINSIC_COLUMN(camera_name) + else: + camera_data_column = FISHEYE_CAMERA_DATA_COLUMN(camera_name) + camera_extrinsic_column = FISHEYE_CAMERA_EXTRINSIC_COLUMN(camera_name) + + if _all_columns_in_schema(arrow_table, [camera_data_column, camera_extrinsic_column]): + table_data = arrow_table[camera_data_column][index].as_py() + extrinsic_data = arrow_table[camera_extrinsic_column][index].as_py() + + if table_data is not None and extrinsic_data is not None: + extrinsic = PoseSE3.from_list(extrinsic_data) + image: Optional[npt.NDArray[np.uint8]] = None + + if isinstance(table_data, str): + sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] + assert sensor_root is not None, ( + f"Dataset path for sensor loading not found for dataset: {log_metadata.dataset}" + ) + full_image_path = Path(sensor_root) / table_data + assert full_image_path.exists(), f"Camera file not found: {full_image_path}" + + image = load_image_from_jpeg_file(full_image_path) + elif isinstance(table_data, bytes): + if is_jpeg_binary(table_data): + image = decode_image_from_jpeg_binary(table_data) + elif is_png_binary(table_data): + image = decode_image_from_png_binary(table_data) + else: + raise ValueError("Camera binary data is neither in JPEG nor PNG format.") + + elif isinstance(table_data, int): + image = _unoptimized_demo_mp4_read(log_metadata, camera_name, table_data) + else: + raise NotImplementedError( + f"Only string file paths, bytes, or int frame indices are supported for camera data, got {type(table_data)}" + ) + + if is_pinhole: + camera_metadata = log_metadata.pinhole_camera_metadata[camera_type] + camera = PinholeCamera( + metadata=camera_metadata, + image=image, + extrinsic=extrinsic, + ) + else: + camera_metadata = log_metadata.fisheye_mei_camera_metadata[camera_type] + camera = FisheyeMEICamera( + metadata=camera_metadata, + image=image, + extrinsic=extrinsic, + ) + + return camera + + +def get_lidar_from_arrow_table( + arrow_table: pa.Table, + index: int, + lidar_type: LiDARType, + log_metadata: LogMetadata, +) -> LiDAR: + """Builds a LiDAR object from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the LiDAR data. + :param index: The index to extract the LiDAR data from. + :param lidar_type: The type of LiDAR to build. + :param log_metadata: Metadata about the log, including the LiDAR metadata. + :raises ValueError: If the LiDAR data format is unsupported. + :raises NotImplementedError: If the LiDAR data type is not supported. + :return: The constructed LiDAR object, or None if not available. + """ + + lidar: Optional[LiDAR] = None + # NOTE @DanielDauner: Some LiDAR are stored together and are separated only during loading. + # In this case, we need to use the merged LiDAR column name. + + lidar_column_name = LIDAR_DATA_COLUMN(lidar_type.serialize()) + lidar_column_name = ( + LIDAR_DATA_COLUMN(LiDARType.LIDAR_MERGED.serialize()) + if lidar_column_name not in arrow_table.schema.names + else lidar_column_name + ) + + if lidar_column_name in arrow_table.schema.names: + lidar_data = arrow_table[lidar_column_name][index].as_py() + if isinstance(lidar_data, str): + lidar_pc_dict = load_lidar_pcs_from_file(relative_path=lidar_data, log_metadata=log_metadata, index=index) + if lidar_type == LiDARType.LIDAR_MERGED: + # Merge all available LiDAR point clouds into one + merged_pc = np.vstack(list(lidar_pc_dict.values())) + lidar = LiDAR( + metadata=LiDARMetadata( + lidar_type=LiDARType.LIDAR_MERGED, + lidar_index=DefaultLiDARIndex, + extrinsic=None, + ), + point_cloud=merged_pc, + ) + elif lidar_type in lidar_pc_dict: + lidar = LiDAR( + metadata=log_metadata.lidar_metadata[lidar_type], + point_cloud=lidar_pc_dict[lidar_type], + ) + elif isinstance(lidar_data, bytes): + lidar_metadata = log_metadata.lidar_metadata[lidar_type] + if is_draco_binary(lidar_data): + # NOTE: DRACO only allows XYZ compression, so we need to override the lidar index here. + lidar_metadata.lidar_index = DefaultLiDARIndex + lidar = load_lidar_from_draco_binary(lidar_data, lidar_metadata) + elif is_laz_binary(lidar_data): + lidar = load_lidar_from_laz_binary(lidar_data, lidar_metadata) + else: + raise ValueError("LiDAR binary data is neither in Draco nor LAZ format.") + elif lidar_data is not None: + raise NotImplementedError( + f"Only string file paths or bytes for LiDAR data are supported, got {type(lidar_data)}" + ) + + return lidar + + +def get_route_lane_group_ids_from_arrow_table(arrow_table: pa.Table, index: int) -> Optional[List[int]]: + """Gets the route lane group IDs from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the route lane group IDs data. + :param index: The index to extract the route lane group IDs from. + :return: The route lane group IDs at the given index, or None if not available + """ + route_lane_group_ids: Optional[List[int]] = None + if _all_columns_in_schema(arrow_table, [ROUTE_LANE_GROUP_IDS_COLUMN]): + route_lane_group_ids = arrow_table[ROUTE_LANE_GROUP_IDS_COLUMN][index].as_py() + return route_lane_group_ids + + +def get_scenario_tags_from_arrow_table(arrow_table: pa.Table, index: int) -> Optional[List[int]]: + """Gets the scenario tags from an Arrow table at a given index. + + :param arrow_table: The Arrow table containing the scenario tags data. + :param index: The index to extract the scenario tags from. + :return: The scenario tags at the given index, or None if not available + """ + scenario_tags: Optional[List[int]] = None + if _all_columns_in_schema(arrow_table, [SCENARIO_TAGS_COLUMN]): + scenario_tags = arrow_table[SCENARIO_TAGS_COLUMN][index].as_py() + return scenario_tags + + +def _unoptimized_demo_mp4_read(log_metadata: LogMetadata, camera_name: str, frame_index: int) -> Optional[np.ndarray]: + """Reads a frame from an MP4 file for demonstration purposes. This features is not optimized for performance. + + :param log_metadata: The metadata of the log containing the MP4 file. + :param camera_name: The name of the camera whose MP4 file is to be read. + :param frame_index: The index of the frame to read from the MP4 file. + :return: The image frame as a numpy array, or None if the file does not exist. + """ + image: Optional[npt.NDArray[np.uint8]] = None + + py123d_sensor_root = Path(DATASET_PATHS.py123d_sensors_root) + mp4_path = py123d_sensor_root / log_metadata.split / log_metadata.log_name / f"{camera_name}.mp4" + if mp4_path.exists(): + reader = get_mp4_reader_from_path(str(mp4_path)) + image = reader.get_frame(frame_index) + + return image + + +def _get_optional_array_mixin(data: Optional[Union[List, npt.NDArray]], cls: Type[ArrayMixin]) -> Optional[ArrayMixin]: + """Builds an optional ArrayMixin if data is provided. + + :param data: The data to convert into an ArrayMixin. + :param cls: The ArrayMixin class to instantiate. + :raises ValueError: If the data type is unsupported. + :return: The instantiated ArrayMixin, or None if data is None. + """ + if data is None: + return None + if isinstance(data, list): + return cls.from_list(data) + elif isinstance(data, np.ndarray): + return cls.from_array(data, copy=False) + else: + raise ValueError(f"Unsupported data type for ArrayMixin conversion: {type(data)}") + + +def _all_columns_in_schema(arrow_table: pa.Table, columns: List[str]) -> bool: + """Checks if all specified columns are present in the Arrow table schema. + + :param arrow_table: The Arrow table to check. + :param columns: The list of column names to check for. + :return: True if all columns are present, False otherwise. + """ + return all(column in arrow_table.schema.names for column in columns) diff --git a/src/py123d/api/scene/arrow/utils/arrow_metadata_utils.py b/src/py123d/api/scene/arrow/utils/arrow_metadata_utils.py new file mode 100644 index 00000000..2dd17daf --- /dev/null +++ b/src/py123d/api/scene/arrow/utils/arrow_metadata_utils.py @@ -0,0 +1,25 @@ +import json +from pathlib import Path +from typing import Union + +import pyarrow as pa + +from py123d.common.utils.arrow_helper import get_lru_cached_arrow_table +from py123d.datatypes.metadata import LogMetadata + + +def get_log_metadata_from_arrow_file(arrow_file_path: Union[Path, str]) -> LogMetadata: + """Gets the log metadata from an Arrow file.""" + table = get_lru_cached_arrow_table(arrow_file_path) + return get_log_metadata_from_arrow_table(table) + + +def get_log_metadata_from_arrow_table(arrow_table: pa.Table) -> LogMetadata: + """Gets the log metadata from an Arrow table.""" + return LogMetadata.from_dict(json.loads(arrow_table.schema.metadata[b"log_metadata"].decode())) + + +def add_log_metadata_to_arrow_schema(schema: pa.schema, log_metadata: LogMetadata) -> pa.schema: + """Adds log metadata to an Arrow schema.""" + schema = schema.with_metadata({"log_metadata": json.dumps(log_metadata.to_dict())}) + return schema diff --git a/src/py123d/api/scene/scene_api.py b/src/py123d/api/scene/scene_api.py new file mode 100644 index 00000000..1277a5c4 --- /dev/null +++ b/src/py123d/api/scene/scene_api.py @@ -0,0 +1,205 @@ +from __future__ import annotations + +import abc +from typing import List, Optional + +from py123d.api.map.map_api import MapAPI +from py123d.api.scene.scene_metadata import SceneMetadata +from py123d.datatypes.detections.box_detections import BoxDetectionWrapper +from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper +from py123d.datatypes.metadata.log_metadata import LogMetadata +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICamera, FisheyeMEICameraType +from py123d.datatypes.sensors.lidar import LiDAR, LiDARType +from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeCameraType +from py123d.datatypes.time.time_point import TimePoint +from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters + + +class SceneAPI(abc.ABC): + """Base class for all scene APIs. The scene API provides access to all data modalities at in a scene.""" + + # Abstract Methods, to be implemented by subclasses + # ------------------------------------------------------------------------------------------------------------------ + + @abc.abstractmethod + def get_log_metadata(self) -> LogMetadata: + """Returns the :class:`~py123d.datatypes.metadata.LogMetadata` of the scene. + + :return: The log metadata. + """ + + @abc.abstractmethod + def get_scene_metadata(self) -> SceneMetadata: + """Returns the :class:`~py123d.api.scene.scene_metadata.SceneMetadata` of the scene. + + :return: The scene metadata. + """ + + @abc.abstractmethod + def get_map_api(self) -> Optional[MapAPI]: + """Returns the :class:`~py123d.api.MapAPI` of the scene, if available. + + :return: The map API, or None if not available. + """ + + @abc.abstractmethod + def get_timepoint_at_iteration(self, iteration: int) -> TimePoint: + """Returns the :class:`~py123d.datatypes.time.TimePoint` at a given iteration. + + :param iteration: The iteration to get the timepoint for. + :return: The timepoint at the given iteration. + """ + + @abc.abstractmethod + def get_ego_state_at_iteration(self, iteration: int) -> Optional[EgoStateSE3]: + """Returns the :class:`~py123d.datatypes.vehicle_state.EgoStateSE3` at a given iteration, if available. + + :param iteration: The iteration to get the ego state for. + :return: The ego state at the given iteration, or None if not available. + """ + + @abc.abstractmethod + def get_box_detections_at_iteration(self, iteration: int) -> Optional[BoxDetectionWrapper]: + """Returns the :class:`~py123d.datatypes.detections.BoxDetectionWrapper` at a given iteration, if available. + + :param iteration: The iteration to get the box detections for. + :return: The box detections at the given iteration, or None if not available. + """ + + @abc.abstractmethod + def get_traffic_light_detections_at_iteration(self, iteration: int) -> Optional[TrafficLightDetectionWrapper]: + """Returns the :class:`~py123d.datatypes.detections.TrafficLightDetectionWrapper` at a given iteration, + if available. + + :param iteration: The iteration to get the traffic light detections for. + :return: The traffic light detections at the given iteration, or None if not available. + """ + + @abc.abstractmethod + def get_route_lane_group_ids(self, iteration: int) -> Optional[List[int]]: + """Returns the list of route lane group IDs at a given iteration, if available. + + :param iteration: The iteration to get the route lane group IDs for. + :return: The list of route lane group IDs at the given iteration, or None if not available. + """ + + @abc.abstractmethod + def get_pinhole_camera_at_iteration( + self, + iteration: int, + camera_type: PinholeCameraType, + ) -> Optional[PinholeCamera]: + """Returns the :class:`~py123d.datatypes.sensors.PinholeCamera` of a given \ + :class:`~py123d.datatypes.sensors.PinholeCameraType` at a given iteration, if available. + + :param iteration: The iteration to get the pinhole camera for. + :param camera_type: The :type:`~py123d.datatypes.sensors.PinholeCameraType` of the pinhole camera. + :return: The pinhole camera, or None if not available. + """ + + @abc.abstractmethod + def get_fisheye_mei_camera_at_iteration( + self, iteration: int, camera_type: FisheyeMEICameraType + ) -> Optional[FisheyeMEICamera]: + """Returns the :class:`~py123d.datatypes.sensors.FisheyeMEICamera` of a given \ + :class:`~py123d.datatypes.sensors.FisheyeMEICameraType` at a given iteration, if available. + + :param iteration: The iteration to get the fisheye MEI camera for. + :param camera_type: The :type:`~py123d.datatypes.sensors.FisheyeMEICameraType` of the fisheye MEI camera. + :return: The fisheye MEI camera, or None if not available. + """ + + @abc.abstractmethod + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> Optional[LiDAR]: + """Returns the :class:`~py123d.datatypes.sensors.LiDAR` of a given :class:`~py123d.datatypes.sensors.LiDARType`\ + at a given iteration, if available. + + :param iteration: The iteration to get the LiDAR for. + :param lidar_type: The :type:`~py123d.datatypes.sensors.LiDARType` of the LiDAR. + :return: The LiDAR, or None if not available. + """ + + # Syntactic Sugar / Properties, for easier access to common attributes + # ------------------------------------------------------------------------------------------------------------------ + + @property + def log_metadata(self) -> LogMetadata: + """The :class:`~py123d.datatypes.metadata.LogMetadata` of the scene.""" + return self.get_log_metadata() + + @property + def scene_metadata(self) -> SceneMetadata: + """The :class:`~py123d.api.scene.SceneMetadata` of the scene.""" + return self.get_scene_metadata() + + @property + def map_metadata(self) -> Optional[MapMetadata]: + """The :class:`~py123d.datatypes.metadata.MapMetadata` of the scene, if available.""" + return self.log_metadata.map_metadata + + @property + def map_api(self) -> Optional[MapAPI]: + """The :class:`~py123d.api.map.MapAPI` of the scene, if available.""" + return self.get_map_api() + + @property + def dataset(self) -> str: + """The dataset name from the log metadata.""" + return self.log_metadata.dataset + + @property + def split(self) -> str: + """The data split name from the log metadata.""" + return self.log_metadata.split + + @property + def location(self) -> str: + """The location from the log metadata.""" + return self.log_metadata.location + + @property + def log_name(self) -> str: + """The log name from the log metadata.""" + return self.log_metadata.log_name + + @property + def version(self) -> str: + """The version of the py123d library used to create this log metadata.""" + return self.log_metadata.version + + @property + def scene_uuid(self) -> str: + """The UUID of the scene.""" + return self.scene_metadata.initial_uuid + + @property + def number_of_iterations(self) -> int: + """The number of iterations in the scene.""" + return self.scene_metadata.number_of_iterations + + @property + def number_of_history_iterations(self) -> int: + """The number of history iterations in the scene.""" + return self.scene_metadata.number_of_history_iterations + + @property + def vehicle_parameters(self) -> Optional[VehicleParameters]: + """The :class:`~py123d.datatypes.vehicle_state.VehicleParameters` of the ego vehicle, if available.""" + return self.log_metadata.vehicle_parameters + + @property + def available_pinhole_camera_types(self) -> List[PinholeCameraType]: + """List of available :class:`~py123d.datatypes.sensors.PinholeCameraType` in the log metadata.""" + return list(self.log_metadata.pinhole_camera_metadata.keys()) + + @property + def available_fisheye_mei_camera_types(self) -> List[FisheyeMEICameraType]: + """List of available :class:`~py123d.datatypes.sensors.FisheyeMEICameraType` in the log metadata.""" + return list(self.log_metadata.fisheye_mei_camera_metadata.keys()) + + @property + def available_lidar_types(self) -> List[LiDARType]: + """List of available :class:`~py123d.datatypes.sensors.LiDARType` in the log metadata.""" + return list(self.log_metadata.lidar_metadata.keys()) diff --git a/src/py123d/datatypes/scene/abstract_scene_builder.py b/src/py123d/api/scene/scene_builder.py similarity index 51% rename from src/py123d/datatypes/scene/abstract_scene_builder.py rename to src/py123d/api/scene/scene_builder.py index 17652549..0e7cedf4 100644 --- a/src/py123d/datatypes/scene/abstract_scene_builder.py +++ b/src/py123d/api/scene/scene_builder.py @@ -1,22 +1,21 @@ import abc -from typing import Iterator +from typing import List +from py123d.api.scene.scene_api import SceneAPI +from py123d.api.scene.scene_filter import SceneFilter from py123d.common.multithreading.worker_utils import WorkerPool -from py123d.datatypes.scene.abstract_scene import AbstractScene -from py123d.datatypes.scene.scene_filter import SceneFilter class SceneBuilder(abc.ABC): - """ - Abstract base class for building scenes from a dataset. + """Base class for all scene builders. The scene builder is responsible for building scene given a \ + :class:`~py123d.api.scene.scene_filter.SceneFilter`. """ @abc.abstractmethod - def get_scenes(self, filter: SceneFilter, worker: WorkerPool) -> Iterator[AbstractScene]: - """ - Returns an iterator over scenes that match the given filter. + def get_scenes(self, filter: SceneFilter, worker: WorkerPool) -> List[SceneAPI]: + """Returns a list of scenes that match the given filter. + :param filter: SceneFilter object to filter the scenes. :param worker: WorkerPool to parallelize the scene extraction. :return: Iterator over AbstractScene objects. """ - raise NotImplementedError diff --git a/src/py123d/datatypes/scene/scene_filter.py b/src/py123d/api/scene/scene_filter.py similarity index 50% rename from src/py123d/datatypes/scene/scene_filter.py rename to src/py123d/api/scene/scene_filter.py index 62ad9301..20b2dc44 100644 --- a/src/py123d/datatypes/scene/scene_filter.py +++ b/src/py123d/api/scene/scene_filter.py @@ -10,34 +10,54 @@ @dataclass class SceneFilter: + """Class to filter scenes when building scenes from logs.""" split_types: Optional[List[str]] = None + """List of split types to filter scenes by (e.g. `train`, `val`, `test`).""" + split_names: Optional[List[str]] = None + """List of split names to filter scenes by (in the form `{dataset_name}_{split_type}`).""" + log_names: Optional[List[str]] = None + """Name of logs to include scenes from.""" - locations: Optional[List[str]] = None # TODO: - scene_uuids: Optional[List[str]] = None # TODO: + locations: Optional[List[str]] = None + """List of locations to filter scenes by.""" - timestamp_threshold_s: Optional[float] = None # TODO: - ego_displacement_minimum_m: Optional[float] = None # TODO: + scene_uuids: Optional[List[str]] = None + """List of scene UUIDs to include.""" + + timestamp_threshold_s: Optional[float] = None + """Minimum time between the start timestamps of two consecutive scenes.""" duration_s: Optional[float] = 10.0 - history_s: Optional[float] = 3.0 + """Duration of each scene in seconds.""" + + history_s: Optional[float] = 0.0 + """History duration of each scene in seconds.""" pinhole_camera_types: Optional[List[PinholeCameraType]] = None + """List of :class:`PinholeCameraType` to include in the scenes.""" + fisheye_mei_camera_types: Optional[List[FisheyeMEICameraType]] = None + """List of :class:`FisheyeMEICameraType` to include in the scenes.""" max_num_scenes: Optional[int] = None + """Maximum number of scenes to return.""" + + map_api_required: bool = False + """Whether to only include scenes with an available map API.""" + shuffle: bool = False + """Whether to shuffle the returned scenes.""" def __post_init__(self): def _resolve_enum_arguments( - serial_enum_cls: SerialIntEnum, input: Optional[List[Union[int, str, SerialIntEnum]]] - ) -> List[SerialIntEnum]: - + serial_enum_cls: SerialIntEnum, + input: Optional[List[Union[int, str, SerialIntEnum]]], + ): if input is None: return None - assert isinstance(input, list), f"input must be a list of {serial_enum_cls.__name__}" return [serial_enum_cls.from_arbitrary(value) for value in input] self.pinhole_camera_types = _resolve_enum_arguments(PinholeCameraType, self.pinhole_camera_types) diff --git a/src/py123d/api/scene/scene_metadata.py b/src/py123d/api/scene/scene_metadata.py new file mode 100644 index 00000000..735ea931 --- /dev/null +++ b/src/py123d/api/scene/scene_metadata.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SceneMetadata: + """Metadata for a scene extracted from a log.""" + + initial_uuid: str + """UUID of the scene, i.e., the UUID of the starting frame of the scene.""" + + initial_idx: int + """Index of the starting frame of the scene in the log.""" + + duration_s: float + """Duration of the scene in seconds.""" + + history_s: float + """History duration of the scene in seconds.""" + + iteration_duration_s: float + """Duration of each iteration in seconds.""" + + @property + def number_of_iterations(self) -> int: + """Number of iterations in the scene.""" + return round(self.duration_s / self.iteration_duration_s) + + @property + def number_of_history_iterations(self) -> int: + """Number of history iterations in the scene.""" + return round(self.history_s / self.iteration_duration_s) + + @property + def end_idx(self) -> int: + """Index of the end frame of the scene.""" + return self.initial_idx + self.number_of_iterations + + def __repr__(self) -> str: + return ( + f"SceneMetadata(initial_uuid={self.initial_uuid}, initial_idx={self.initial_idx}, " + f"duration_s={self.duration_s}, history_s={self.history_s}, " + f"iteration_duration_s={self.iteration_duration_s})" + ) diff --git a/src/py123d/common/multithreading/ray_execution.py b/src/py123d/common/multithreading/ray_execution.py index 2ca1d472..eb769ca9 100644 --- a/src/py123d/common/multithreading/ray_execution.py +++ b/src/py123d/common/multithreading/ray_execution.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + import logging import traceback from functools import partial @@ -14,8 +19,8 @@ def _ray_object_iterator(initial_ids: List[ray.ObjectRef]) -> Iterator[Tuple[ray.ObjectRef, Any]]: - """ - Iterator that waits for each ray object in the input object list to be completed and fetches the result. + """Iterator that waits for each ray object in the input object list to be completed and fetches the result. + :param initial_ids: list of ray object ids :yield: result of worker """ @@ -31,8 +36,8 @@ def _ray_object_iterator(initial_ids: List[ray.ObjectRef]) -> Iterator[Tuple[ray def wrap_function(fn: Callable[..., Any], log_dir: Optional[Path] = None) -> Callable[..., Any]: - """ - Wraps a function to save its logs to a unique file inside the log directory. + """Wraps a function to save its logs to a unique file inside the log directory. + :param fn: function to be wrapped. :param log_dir: directory to store logs (wrapper function does nothing if it's not set). :return: wrapped function which changes logging settings while it runs. @@ -68,8 +73,8 @@ def wrapped_fn(*args: Any, **kwargs: Any) -> Any: def _ray_map_items(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: - """ - Map each item of a list of arguments to a callable and executes in parallel. + """Map each item of a list of arguments to a callable and executes in parallel. + :param fn: callable to be run :param item_list: items to be parallelized :param log_dir: directory to store worker logs @@ -78,7 +83,8 @@ def _ray_map_items(task: Task, *item_lists: Iterable[List[Any]], log_dir: Option assert len(item_lists) > 0, "No map arguments received for mapping" assert all(isinstance(items, list) for items in item_lists), "All map arguments must be lists" assert all( - len(cast(List, items)) == len(item_lists[0]) for items in item_lists # type: ignore + len(cast(List, items)) == len(item_lists[0]) + for items in item_lists # type: ignore ), "All lists must have equal size" fn = task.fn # Wrap function in remote decorator and create ray objects @@ -106,8 +112,8 @@ def _ray_map_items(task: Task, *item_lists: Iterable[List[Any]], log_dir: Option def ray_map(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: - """ - Initialize ray, align item lists and map each item of a list of arguments to a callable and executes in parallel. + """Initialize ray, align item lists and map each item of a list of arguments to a callable and executes in parallel. + :param task: callable to be run :param item_lists: items to be parallelized :param log_dir: directory to store worker logs diff --git a/src/py123d/common/multithreading/worker_parallel.py b/src/py123d/common/multithreading/worker_parallel.py index 9183a0fc..2af320e5 100644 --- a/src/py123d/common/multithreading/worker_parallel.py +++ b/src/py123d/common/multithreading/worker_parallel.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + import concurrent import concurrent.futures import logging diff --git a/src/py123d/common/multithreading/worker_pool.py b/src/py123d/common/multithreading/worker_pool.py index 8ffea8ed..716b37d9 100644 --- a/src/py123d/common/multithreading/worker_pool.py +++ b/src/py123d/common/multithreading/worker_pool.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + import abc import logging from concurrent.futures import Future diff --git a/src/py123d/common/multithreading/worker_ray.py b/src/py123d/common/multithreading/worker_ray.py index 48b06f77..0ac47dc9 100644 --- a/src/py123d/common/multithreading/worker_ray.py +++ b/src/py123d/common/multithreading/worker_ray.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + import logging import os from concurrent.futures import Future diff --git a/src/py123d/common/multithreading/worker_sequential.py b/src/py123d/common/multithreading/worker_sequential.py index c0106e86..792f0d4e 100644 --- a/src/py123d/common/multithreading/worker_sequential.py +++ b/src/py123d/common/multithreading/worker_sequential.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + import logging from concurrent.futures import Future from typing import Any, Iterable, List diff --git a/src/py123d/common/multithreading/worker_utils.py b/src/py123d/common/multithreading/worker_utils.py index ce79d6df..ae077aea 100644 --- a/src/py123d/common/multithreading/worker_utils.py +++ b/src/py123d/common/multithreading/worker_utils.py @@ -1,3 +1,8 @@ +""" +Multi-threading execution code. +Code is adapted from the nuplan-devkit: https://github.com/motional/nuplan-devkit +""" + from typing import Any, Callable, List, Optional import numpy as np diff --git a/src/py123d/common/utils/arrow_column_names.py b/src/py123d/common/utils/arrow_column_names.py new file mode 100644 index 00000000..37a3a6bf --- /dev/null +++ b/src/py123d/common/utils/arrow_column_names.py @@ -0,0 +1,79 @@ +from typing import Callable, Final, List + +# Essential Columns +# ---------------------------------------------------------------------------------------------------------------------- +UUID_COLUMN: Final[str] = "uuid" +TIMESTAMP_US_COLUMN: Final[str] = "timestamp.us" + +# Ego State SE3 +# ---------------------------------------------------------------------------------------------------------------------- +EGO_REAR_AXLE_SE3_COLUMN: Final[str] = "ego.rear_axle_se3" +EGO_DYNAMIC_STATE_SE3_COLUMN: Final[str] = "ego.dynamic_state_se3" + +EGO_STATE_SE3_COLUMNS: Final[List[str]] = [ + EGO_REAR_AXLE_SE3_COLUMN, + EGO_DYNAMIC_STATE_SE3_COLUMN, +] + + +# Box Detections SE3 +# ---------------------------------------------------------------------------------------------------------------------- +BOX_DETECTIONS_PREFIX: Final[str] = "box_detections" +BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN: Final[str] = f"{BOX_DETECTIONS_PREFIX}.bounding_box_se3" +BOX_DETECTIONS_TOKEN_COLUMN: Final[str] = f"{BOX_DETECTIONS_PREFIX}.token" +BOX_DETECTIONS_VELOCITY_3D_COLUMN: Final[str] = f"{BOX_DETECTIONS_PREFIX}.velocity_3d" +BOX_DETECTIONS_LABEL_COLUMN: Final[str] = f"{BOX_DETECTIONS_PREFIX}.label" +BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN: Final[str] = f"{BOX_DETECTIONS_PREFIX}.num_lidar_points" + +BOX_DETECTIONS_SE3_COLUMNS: Final[List[str]] = [ + BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN, + BOX_DETECTIONS_TOKEN_COLUMN, + BOX_DETECTIONS_VELOCITY_3D_COLUMN, + BOX_DETECTIONS_LABEL_COLUMN, + BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN, +] + +# Traffic Lights +# ---------------------------------------------------------------------------------------------------------------------- +TRAFFIC_LIGHTS_PREFIX: Final[str] = "traffic_lights" +TRAFFIC_LIGHTS_LANE_ID_COLUMN: Final[str] = f"{TRAFFIC_LIGHTS_PREFIX}.lane_id" +TRAFFIC_LIGHTS_STATUS_COLUMN: Final[str] = f"{TRAFFIC_LIGHTS_PREFIX}.status" + +TRAFFIC_LIGHTS_COLUMNS: Final[List[str]] = [ + TRAFFIC_LIGHTS_LANE_ID_COLUMN, + TRAFFIC_LIGHTS_STATUS_COLUMN, +] + +# Pinhole Cameras +# ---------------------------------------------------------------------------------------------------------------------- +PINHOLE_PREFIX: Final[str] = "pinhole" +PINHOLE_CAMERA_DATA_COLUMN: Callable[[str], str] = lambda name: f"{PINHOLE_PREFIX}.{name}.data" +PINHOLE_CAMERA_EXTRINSIC_COLUMN: Callable[[str], str] = lambda name: f"{PINHOLE_PREFIX}.{name}.state_se3" + +PINHOLE_CAMERA_COLUMNS: Callable[[str], List[str]] = lambda name: [ + PINHOLE_CAMERA_DATA_COLUMN(name), + PINHOLE_CAMERA_EXTRINSIC_COLUMN(name), +] + + +# Fisheye MEI Cameras +# ---------------------------------------------------------------------------------------------------------------------- +FISHEYE_PREFIX: Final[str] = "fisheye_mei" +FISHEYE_CAMERA_DATA_COLUMN: Callable[[str], str] = lambda name: f"{FISHEYE_PREFIX}.{name}.data" +FISHEYE_CAMERA_EXTRINSIC_COLUMN: Callable[[str], str] = lambda name: f"{FISHEYE_PREFIX}.{name}.state_se3" + +FISHEYE_CAMERA_COLUMNS: Callable[[str], List[str]] = lambda name: [ + FISHEYE_CAMERA_DATA_COLUMN(name), + FISHEYE_CAMERA_EXTRINSIC_COLUMN(name), +] + + +# LiDAR +# ---------------------------------------------------------------------------------------------------------------------- +LIDAR_DATA_COLUMN: Callable[[str], str] = lambda name: f"lidar.{name}.data" + + +# Miscellaneous (Scenario Tags / Route) +# ---------------------------------------------------------------------------------------------------------------------- +SCENARIO_TAGS_COLUMN: str = "scenario_tags" +ROUTE_LANE_GROUP_IDS_COLUMN: str = "route_lane_group_ids" diff --git a/src/py123d/common/utils/arrow_helper.py b/src/py123d/common/utils/arrow_helper.py index 7e1b8e47..75e98f52 100644 --- a/src/py123d/common/utils/arrow_helper.py +++ b/src/py123d/common/utils/arrow_helper.py @@ -5,19 +5,28 @@ import pyarrow as pa # TODO: Tune Parameters and add to config? -MAX_LRU_CACHED_TABLES: Final[int] = 4096 +MAX_LRU_CACHED_TABLES: Final[int] = 50_000 def open_arrow_table(arrow_file_path: Union[str, Path]) -> pa.Table: + """Open an `.arrow` file as memory map. + + :param arrow_file_path: The file path, defined as string or Path. + :return: The memory-mapped arrow table.s + """ + with pa.memory_map(str(arrow_file_path), "rb") as source: table: pa.Table = pa.ipc.open_file(source).read_all() return table def write_arrow_table(table: pa.Table, arrow_file_path: Union[str, Path]) -> None: - # compression: Optional[Literal["lz4", "zstd"]] = "lz4" - # codec = pa.Codec("zstd", compression_level=100) if compression is not None else None - # options = pa.ipc.IpcWriteOptions(compression=codec) + """Writes an arrow table to the file path. + + :param table: The arrow table to write. + :param arrow_file_path: The file path, defined as string or Path. + """ + with pa.OSFile(str(arrow_file_path), "wb") as sink: # with pa.ipc.new_file(sink, table.schema, options=options) as writer: with pa.ipc.new_file(sink, table.schema) as writer: diff --git a/src/py123d/common/utils/enums.py b/src/py123d/common/utils/enums.py index 9f7d233e..453b0945 100644 --- a/src/py123d/common/utils/enums.py +++ b/src/py123d/common/utils/enums.py @@ -1,19 +1,28 @@ from __future__ import annotations -from enum import IntEnum - -from pyparsing import Union +import enum +from typing import Union class classproperty(object): + """Decorator for class-level properties.""" + def __init__(self, f): + """Initialize the classproperty with the given function.""" self.f = f def __get__(self, obj, owner): + """Get the property value.""" return self.f(owner) -class SerialIntEnum(IntEnum): +class SerialIntEnum(enum.Enum): + """Base class for serializable integer enums.""" + + def __int__(self) -> int: + """Get the integer value of the enum.""" + return self.value + def serialize(self, lower: bool = True) -> str: """Serialize the type when saving.""" # Allow for lower/upper case letters during serialize diff --git a/src/py123d/common/utils/mixin.py b/src/py123d/common/utils/mixin.py index 2935f99f..d4452980 100644 --- a/src/py123d/common/utils/mixin.py +++ b/src/py123d/common/utils/mixin.py @@ -1,21 +1,46 @@ from __future__ import annotations +from enum import IntEnum + import numpy as np import numpy.typing as npt - -# import pyarrow as pa +from typing_extensions import Self class ArrayMixin: - """Mixin class for object entities.""" + """Mixin class to provide array-like behavior for classes. + + Example: + >>> import numpy as np + >>> from py123d.common.utils.mixin import ArrayMixin + >>> class MyVector(ArrayMixin): + ... def __init__(self, x: float, y: float): + ... self._array = np.array([x, y], dtype=np.float64) + ... @property + ... def array(self) -> npt.NDArray[np.float64]: + ... return self._array + ... @classmethod + ... def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> MyVector: + ... if copy: + ... array = array.copy() + ... return cls(array[0], array[1]) + >>> vec = MyVector(1.0, 2.0) + >>> print(vec) + MyVector(array=[1. 2.]) + >>> np.array(vec, dtype=np.float32) + array([1., 2.], dtype=float32) + + """ + + __slots__ = () @classmethod - def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> ArrayMixin: + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Self: """Create an instance from a NumPy array.""" raise NotImplementedError @classmethod - def from_list(cls, values: list) -> ArrayMixin: + def from_list(cls, values: list) -> Self: """Create an instance from a list of values.""" return cls.from_array(np.array(values, dtype=np.float64), copy=False) @@ -24,7 +49,7 @@ def array(self) -> npt.NDArray[np.float64]: """The array representation of the geometric entity.""" raise NotImplementedError - def __array__(self, dtype: npt.DtypeLike = None, copy: bool = False) -> npt.NDArray: + def __array__(self, dtype: npt.DTypeLike = None, copy: bool = False) -> npt.NDArray: array = self.array return array if dtype is None else array.astype(dtype=dtype, copy=copy) @@ -51,6 +76,10 @@ def tolist(self) -> list: """Convert the array to a Python list.""" return self.array.tolist() + def to_list(self) -> list: + """Convert the array to a Python list.""" + return self.array.tolist() + def copy(self) -> ArrayMixin: """Return a copy of the object with a copied array.""" return self.__class__.from_array(self.array, copy=True) @@ -58,3 +87,20 @@ def copy(self) -> ArrayMixin: def __repr__(self) -> str: """String representation of the ArrayMixin instance.""" return f"{self.__class__.__name__}(array={self.array})" + + def __hash__(self): + """Hash based on the array values.""" + return hash(self.array.tobytes()) + + +def indexed_array_repr(array_mixin: ArrayMixin, indexing: IntEnum) -> str: + """Generate a string representation of an ArrayMixin instance using an indexing enum. + + :param array_mixin: An instance of ArrayMixin. + :param indexing: An IntEnum used for indexing the array. + :return: A string representation of the ArrayMixin instance with named fields. + """ + args = ", ".join( + f"{index.name.lower()}={array_mixin.array[index.value]}" for index in indexing.__members__.values() + ) + return f"{array_mixin.__class__.__name__}({args})" diff --git a/src/py123d/common/utils/timer.py b/src/py123d/common/utils/timer.py index 17558977..4159f556 100644 --- a/src/py123d/common/utils/timer.py +++ b/src/py123d/common/utils/timer.py @@ -6,18 +6,30 @@ class Timer: - """ - A simple timer class to measure execution time of different parts of the code. + """Simple Timer class to log time taken by code blocks. + + Example + ------- + >>> timer = Timer() + >>> timer.start() + >>> time.sleep(0.1) # Simulate code block + >>> timer.log("block_1") + >>> time.sleep(0.2) # Simulate another code block + >>> timer.log("block_2") + >>> timer.end() + >>> print(timer) # Displays timing statistics (with some variation) + mean min max argmax median + block_1 0.100123 0.100123 0.100123 0 0.100123 + block_2 0.200456 0.200456 0.200456 0 0.200456 + total 0.300579 0.300579 0.300579 0 0.300579 + """ - def __init__(self, name: Optional[str] = None, end_key: str = "total"): - """ - Initializes the Timer instance. - :param name: Name of the Timer, defaults to None - :param end_key: name of the final row, defaults to "total" - """ + def __init__(self, end_key: str = "total"): + """Initializes the :class:`Timer` - self._name = name + :param end_key: The key used to log the total time, defaults to "total" + """ self._end_key: str = end_key self._statistic_functions = { "mean": np.mean, @@ -33,15 +45,17 @@ def __init__(self, name: Optional[str] = None, end_key: str = "total"): self._iteration_time: Optional[float] = None def start(self) -> None: - """Called during the start of the timer .""" + """Called at the start of the timer.""" self._start_time = time.perf_counter() self._iteration_time = time.perf_counter() def log(self, key: str) -> None: """ Called after code block execution. Logs the time taken for the block, given the name (key). - :param key: Name of the code block to log the time for. + :param key: Unique identifier of the code block to log the time for. """ + assert self._iteration_time is not None, "Timer has not been started. Call start() before logging." + if key not in self._time_logs.keys(): self._time_logs[key] = [] @@ -50,14 +64,15 @@ def log(self, key: str) -> None: def end(self) -> None: """Called at the end of the timer.""" + assert self._start_time is not None, "Timer has not been started. Call start() before logging." if self._end_key not in self._time_logs.keys(): self._time_logs[self._end_key] = [] self._time_logs[self._end_key].append(time.perf_counter() - self._start_time) def to_pandas(self) -> Optional[pd.DataFrame]: - """ - Returns a DataFrame with statistics of the logged times. + """Returns a DataFrame with statistics of the logged times. + :return: pandas dataframe. """ @@ -73,13 +88,15 @@ def to_pandas(self) -> Optional[pd.DataFrame]: return dataframe def info(self) -> Dict[str, float]: - """ - Summarized information about the timings. + """Summarized information about the timings. + :return: Dictionary with the mean of each timing. """ info = {} for key, timings in self._time_logs.items(): - info[key] = np.array(timings).mean() + info[key] = {} + for name, function in self._statistic_functions.items(): + info[key][name] = function(np.array(timings)) return info def flush(self) -> None: diff --git a/src/py123d/common/utils/uuid_utils.py b/src/py123d/common/utils/uuid_utils.py index d4d2678a..3c928a08 100644 --- a/src/py123d/common/utils/uuid_utils.py +++ b/src/py123d/common/utils/uuid_utils.py @@ -1,11 +1,11 @@ import uuid -from typing import Final +from typing import Final, Optional # Fixed namespace UUID for all UUIDs generated by 123D, do not change! UUID_NAMESPACE_123D: Final[uuid.UUID] = uuid.UUID("123D123D-123D-123D-123D-123D123D123D") -def create_deterministic_uuid(split: str, log_name: str, timestamp_us: int, misc: str = None) -> uuid.UUID: +def create_deterministic_uuid(split: str, log_name: str, timestamp_us: int, misc: Optional[str] = None) -> uuid.UUID: """Create a universally unique identifier (UUID) based on identifying fields. :param split: The data split (in the format {dataset_name}_{train, val, test}) diff --git a/src/py123d/conversion/dataset_converter_config.py b/src/py123d/conversion/dataset_converter_config.py index f9264cd7..cfb28942 100644 --- a/src/py123d/conversion/dataset_converter_config.py +++ b/src/py123d/conversion/dataset_converter_config.py @@ -6,12 +6,12 @@ @dataclass class DatasetConverterConfig: - force_log_conversion: bool = False force_map_conversion: bool = False # Map include_map: bool = False + remap_map_ids: bool = False # Ego include_ego: bool = False @@ -25,15 +25,15 @@ class DatasetConverterConfig: # Pinhole Cameras include_pinhole_cameras: bool = False - pinhole_camera_store_option: Literal["path", "binary", "mp4"] = "path" + pinhole_camera_store_option: Literal["path", "jpeg_binary", "png_binary", "mp4"] = "path" # Fisheye MEI Cameras include_fisheye_mei_cameras: bool = False - fisheye_mei_camera_store_option: Literal["path", "binary", "mp4"] = "path" + fisheye_mei_camera_store_option: Literal["path", "jpeg_binary", "png_binary", "mp4"] = "path" # LiDARs include_lidars: bool = False - lidar_store_option: Literal["path", "path_merged", "binary"] = "path" + lidar_store_option: Literal["path", "path_merged", "laz_binary", "draco_binary"] = "path" # Scenario tag / Route # NOTE: These are only supported for nuPlan. Consider removing or expanding support. @@ -41,15 +41,23 @@ class DatasetConverterConfig: include_route: bool = False def __post_init__(self): - assert self.pinhole_camera_store_option in [ "path", - "binary", + "jpeg_binary", + "png_binary", + "mp4", + ], f"Invalid Pinhole camera store option, got {self.pinhole_camera_store_option}." + + assert self.fisheye_mei_camera_store_option in [ + "path", + "jpeg_binary", + "png_binary", "mp4", - ], f"Invalid camera store option, got {self.pinhole_camera_store_option}." + ], f"Invalid Fisheye MEI camera store option, got {self.fisheye_mei_camera_store_option}." assert self.lidar_store_option in [ "path", "path_merged", - "binary", + "laz_binary", + "draco_binary", ], f"Invalid LiDAR store option, got {self.lidar_store_option}." diff --git a/src/py123d/datatypes/scene/arrow/utils/__init__.py b/src/py123d/conversion/datasets/av2/__init__.py similarity index 100% rename from src/py123d/datatypes/scene/arrow/utils/__init__.py rename to src/py123d/conversion/datasets/av2/__init__.py diff --git a/src/py123d/conversion/datasets/av2/av2_map_conversion.py b/src/py123d/conversion/datasets/av2/av2_map_conversion.py index a55a9cf4..ecbfd466 100644 --- a/src/py123d/conversion/datasets/av2/av2_map_conversion.py +++ b/src/py123d/conversion/datasets/av2/av2_map_conversion.py @@ -1,4 +1,5 @@ import json +import logging from pathlib import Path from typing import Any, Dict, Final, List @@ -7,6 +8,7 @@ import numpy.typing as npt import shapely import shapely.geometry as geom +from pandas import isna from py123d.conversion.datasets.av2.utils.av2_constants import AV2_ROAD_LINE_TYPE_MAPPING from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter @@ -15,16 +17,16 @@ split_line_geometry_by_max_length, ) from py123d.conversion.utils.map_utils.road_edge.road_edge_3d_utils import lift_road_edges_to_3d -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCrosswalk, - CacheGenericDrivable, - CacheIntersection, - CacheLane, - CacheLaneGroup, - CacheRoadEdge, - CacheRoadLine, +from py123d.datatypes.map_objects.map_layer_types import RoadEdgeType +from py123d.datatypes.map_objects.map_objects import ( + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, ) -from py123d.datatypes.maps.map_datatypes import RoadEdgeType from py123d.geometry import OccupancyMap2D, Point3DIndex, Polyline2D, Polyline3D LANE_GROUP_MARK_TYPES: List[str] = [ @@ -34,16 +36,34 @@ "SOLID_DASH_WHITE", "SOLID_WHITE", ] -MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # TODO: Add to config +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 + + +logger = logging.getLogger(__name__) def convert_av2_map(source_log_path: Path, map_writer: AbstractMapWriter) -> None: + """Converts the AV2 map objects to the 123D objects and writes them using the provided map writer. + + :param source_log_path: Path to the AV2 source log folder. + :param map_writer: An instance of AbstractMapWriter to write the converted map objects. + """ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Polyline3D: + """Helper to instantiate a Polyline3D from AV2 coordinate dicts.""" polyline = np.array([[p["x"], p["y"], p["z"]] for p in data], dtype=np.float64) if close: polyline = np.vstack([polyline, polyline[0]]) + # NOTE @DanielDauner: AV2 map can have NaN values in the Z axis. + # In this case we replace NaNs with zeros with the median (or zeros). + if np.isnan(polyline).any(): + median_xyz = np.nanmedian(polyline[:, 2], axis=-1) + logger.warning(f"Found NaN values {source_log_path} polyline data: {polyline}. Replacing NaNs with zeros.") + for i in range(polyline.shape[0]): + if isna(polyline[i, 2]): + polyline[i, 2] = median_xyz if not isna(median_xyz) else 0.0 + return Polyline3D.from_array(polyline) map_folder = source_log_path / "map" @@ -57,27 +77,31 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly # keys: ["area_boundary", "id"] drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict["area_boundary"], close=True) - for lane_segment_id, lane_segment_dict in log_map_archive["lane_segments"].items(): - # keys = [ - # "id", - # "is_intersection", - # "lane_type", - # "left_lane_boundary", - # "left_lane_mark_type", - # "right_lane_boundary", - # "right_lane_mark_type", - # "successors", - # "predecessors", - # "right_neighbor_id", - # "left_neighbor_id", - # ] + for _, lane_segment_dict in log_map_archive["lane_segments"].items(): + # Available keys: + # - "id", + # - "is_intersection", + # - "lane_type", + # - "left_lane_boundary", + # - "left_lane_mark_type", + # - "right_lane_boundary", + # - "right_lane_mark_type", + # - "successors", + # - "predecessors", + # - "right_neighbor_id", + # - "left_neighbor_id", + + # Convert polyline dicts to Polyline3D objects. lane_segment_dict["left_lane_boundary"] = _extract_polyline(lane_segment_dict["left_lane_boundary"]) lane_segment_dict["right_lane_boundary"] = _extract_polyline(lane_segment_dict["right_lane_boundary"]) - for crosswalk_id, crosswalk_dict in log_map_archive["pedestrian_crossings"].items(): - # keys = ["id", "outline"] - # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/pedestrian_crossing.py + for _, crosswalk_dict in log_map_archive["pedestrian_crossings"].items(): + # Available keys: + # - "id" + # - "edge1" + # - "edge2" + # Convert edge dicts to Polyline3D objects. p1, p2 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge1"]], dtype=np.float64) p3, p4 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge2"]], dtype=np.float64) crosswalk_dict["outline"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64)) @@ -95,16 +119,19 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly def _write_av2_lanes(lanes: Dict[int, Any], map_writer: AbstractMapWriter) -> None: + """Helper to write lanes to map writer.""" def _get_centerline_from_boundaries( - left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 0.1 + left_boundary: Polyline3D, + right_boundary: Polyline3D, + resolution: float = 0.1, ) -> Polyline3D: + """Helper to compute centerline from left and right lane boundaries.""" points_per_meter = 1 / resolution num_points = int(np.ceil(max([right_boundary.length, left_boundary.length]) * points_per_meter)) right_array = right_boundary.interpolate(np.linspace(0, right_boundary.length, num_points, endpoint=True)) left_array = left_boundary.interpolate(np.linspace(0, left_boundary.length, num_points, endpoint=True)) - return Polyline3D.from_array(np.mean([right_array, left_array], axis=0)) for lane_id, lane_dict in lanes.items(): @@ -114,11 +141,11 @@ def _get_centerline_from_boundaries( ) # NOTE @DanielDauner: Some neighbor lane IDs might not be present in the dataset. - left_lane_id = lane_dict["left_neighbor_id"] if lane_dict["left_neighbor_id"] in lanes else None - right_lane_id = lane_dict["right_neighbor_id"] if lane_dict["right_neighbor_id"] in lanes else None + left_lane_id = lane_dict["left_neighbor_id"] if str(lane_dict["left_neighbor_id"]) in lanes else None + right_lane_id = lane_dict["right_neighbor_id"] if str(lane_dict["right_neighbor_id"]) in lanes else None map_writer.write_lane( - CacheLane( + Lane( object_id=lane_id, lane_group_id=lane_dict["lane_group_id"], left_boundary=lane_dict["left_lane_boundary"], @@ -130,17 +157,16 @@ def _get_centerline_from_boundaries( successor_ids=lane_dict["successors"], speed_limit_mps=None, outline=None, # Inferred from boundaries - geometry=None, + shapely_polygon=None, ) ) def _write_av2_lane_group(lane_group_dict: Dict[int, Any], map_writer: AbstractMapWriter) -> None: - + """Helper to write lane groups to map writer.""" for lane_group_id, lane_group_values in lane_group_dict.items(): - map_writer.write_lane_group( - CacheLaneGroup( + LaneGroup( object_id=lane_group_id, lane_ids=lane_group_values["lane_ids"], left_boundary=lane_group_values["left_boundary"], @@ -149,15 +175,16 @@ def _write_av2_lane_group(lane_group_dict: Dict[int, Any], map_writer: AbstractM predecessor_ids=lane_group_values["predecessor_ids"], successor_ids=lane_group_values["successor_ids"], outline=None, - geometry=None, + shapely_polygon=None, ) ) def _write_av2_intersections(intersection_dict: Dict[int, Any], map_writer: AbstractMapWriter) -> None: + """Helper to write intersections to map writer.""" for intersection_id, intersection_values in intersection_dict.items(): map_writer.write_intersection( - CacheIntersection( + Intersection( object_id=intersection_id, lane_group_ids=intersection_values["lane_group_ids"], outline=intersection_values["outline_3d"], @@ -166,9 +193,10 @@ def _write_av2_intersections(intersection_dict: Dict[int, Any], map_writer: Abst def _write_av2_crosswalks(crosswalks: Dict[int, npt.NDArray[np.float64]], map_writer: AbstractMapWriter) -> None: + """Helper to write crosswalks to map writer.""" for cross_walk_id, crosswalk_dict in crosswalks.items(): map_writer.write_crosswalk( - CacheCrosswalk( + Crosswalk( object_id=cross_walk_id, outline=crosswalk_dict["outline"], ) @@ -176,9 +204,10 @@ def _write_av2_crosswalks(crosswalks: Dict[int, npt.NDArray[np.float64]], map_wr def _write_av2_generic_drivable(drivable_areas: Dict[int, Polyline3D], map_writer: AbstractMapWriter) -> None: + """Helper to write generic drivable areas to map writer.""" for drivable_area_id, drivable_area_outline in drivable_areas.items(): map_writer.write_generic_drivable( - CacheGenericDrivable( + GenericDrivable( object_id=drivable_area_id, outline=drivable_area_outline, ) @@ -186,10 +215,10 @@ def _write_av2_generic_drivable(drivable_areas: Dict[int, Polyline3D], map_write def _write_av2_road_edge(drivable_areas: Dict[int, Polyline3D], map_writer: AbstractMapWriter) -> None: + """Helper to write road edges to map writer.""" # NOTE @DanielDauner: We merge all drivable areas in 2D and lift the outlines to 3D. # Currently the method assumes that the drivable areas do not overlap and all road surfaces are included. - drivable_polygons = [geom.Polygon(drivable_area.array[:, :2]) for drivable_area in drivable_areas.values()] road_edges_2d = get_road_edge_linear_rings(drivable_polygons) non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, list(drivable_areas.values())) @@ -197,10 +226,9 @@ def _write_av2_road_edge(drivable_areas: Dict[int, Polyline3D], map_writer: Abst road_edges = split_line_geometry_by_max_length(non_conflicting_road_edges_linestrings, MAX_ROAD_EDGE_LENGTH) for idx, road_edge in enumerate(road_edges): - # TODO @DanielDauner: Figure out if other road edge types should/could be assigned here. map_writer.write_road_edge( - CacheRoadEdge( + RoadEdge( object_id=idx, road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, polyline=Polyline3D.from_linestring(road_edge), @@ -209,30 +237,33 @@ def _write_av2_road_edge(drivable_areas: Dict[int, Polyline3D], map_writer: Abst def _write_av2_road_lines(lanes: Dict[int, Any], map_writer: AbstractMapWriter) -> None: - + """Helper to write road lines to map writer.""" running_road_line_id = 0 for lane in lanes.values(): for side in ["left", "right"]: # NOTE @DanielDauner: We currently ignore lane markings that are NONE in the AV2 dataset. if lane[f"{side}_lane_mark_type"] == "NONE": continue - map_writer.write_road_line( - CacheRoadLine( + RoadLine( object_id=running_road_line_id, road_line_type=AV2_ROAD_LINE_TYPE_MAPPING[lane[f"{side}_lane_mark_type"]], polyline=lane[f"{side}_lane_boundary"], ) ) - running_road_line_id += 1 -def _extract_lane_group_dict(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: +def _extract_lane_group_dict(lanes: Dict[int, Any]) -> Dict[int, Any]: + """Collect lane groups from neighboring lanes. This function first extracts lane groups by traversing + neighboring lanes and then builds a dictionary with lane group information, e.g. boundaries, + predecessors, successors. + :param lanes: Dictionary of lane information, e.g. boundaries, and neighboring lanes. + :return: Dictionary of lane group information. + """ lane_group_sets = _extract_lane_group(lanes) lane_group_set_dict = {i: lane_group for i, lane_group in enumerate(lane_group_sets)} - lane_group_dict: Dict[int, Dict[str, Any]] = {} def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: @@ -244,7 +275,6 @@ def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: return list(set(lane_group_ids_)) for lane_group_id, lane_group_set in lane_group_set_dict.items(): - lane_group_dict[lane_group_id] = {} lane_group_dict[lane_group_id]["id"] = lane_group_id lane_group_dict[lane_group_id]["lane_ids"] = [int(lane_id) for lane_id in lane_group_set] @@ -279,6 +309,11 @@ def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: def _extract_lane_group(lanes) -> List[List[str]]: + """Extract lane groups by traversing neighboring lanes. + + :param lanes: Dictionary of lane information, e.g. boundaries, and neighboring lanes. + :return: List of lane groups, where each lane group is a list of lane IDs + """ visited = set() lane_groups = [] @@ -331,8 +366,17 @@ def _traverse_group(start_lane_id): def _extract_intersection_dict( - lanes: Dict[int, Any], lane_group_dict: Dict[int, Any], max_distance: float = 0.01 + lanes: Dict[int, Any], + lane_group_dict: Dict[int, Any], + max_distance: float = 0.01, ) -> Dict[str, Any]: + """Extract intersection outlines from lane groups. + + :param lanes: Dictionary of lane information, e.g. boundaries, and whether lane is part of intersection. + :param lane_group_dict: Dictionary of lane group information. + :param max_distance: Maximum distance to consider for intersection boundaries, defaults to 0.01 + :return: Dictionary of intersection information. + """ def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float: """Interpolate Z coordinate along a 3D line segment.""" @@ -362,15 +406,27 @@ def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[ lane_group_intersection_dict[lane_group_id] = lane_group # 2. Merge polygons of lane groups that are marked as intersections. - lane_group_intersection_geometry = { - lane_group_id: shapely.Polygon(lane_group["outline"].array[:, Point3DIndex.XY]) - for lane_group_id, lane_group in lane_group_intersection_dict.items() - } + # lane_group_intersection_geometry = { + # lane_group_id: shapely.Polygon(lane_group["outline"].array[:, Point3DIndex.XY]) + # for lane_group_id, lane_group in lane_group_intersection_dict.items() + # } + lane_group_intersection_geometry = {} + for lane_group_id, lane_group in lane_group_intersection_dict.items(): + lane_group_polygon_2d = shapely.Polygon(lane_group["outline"].array[:, Point3DIndex.XY]) + if lane_group_polygon_2d.is_valid: + lane_group_intersection_geometry[lane_group_id] = lane_group_polygon_2d + intersection_polygons = gpd.GeoSeries(lane_group_intersection_geometry).union_all() # 3. Collect all intersection polygons and their lane group IDs. + geometries = [] + if isinstance(intersection_polygons, geom.Polygon): + geometries.append(intersection_polygons) + elif isinstance(intersection_polygons, geom.MultiPolygon): + geometries.extend(intersection_polygons.geoms) + intersection_dict = {} - for intersection_idx, intersection_polygon in enumerate(intersection_polygons.geoms): + for intersection_idx, intersection_polygon in enumerate(geometries): if intersection_polygon.is_empty: continue lane_group_ids = [ @@ -394,23 +450,24 @@ def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[ segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1) boundary_segments.append(segment_coords_boundary) - boundary_segments = np.concatenate(boundary_segments, axis=0) - boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) - occupancy_map = OccupancyMap2D(boundary_segment_linestrings) - - for intersection_id, intersection_data in intersection_dict.items(): - points_2d = intersection_data["outline_2d"].array - points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) - points_3d[:, :2] = points_2d - - query_points = shapely.creation.points(points_2d) - results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) - for query_idx, geometry_idx in zip(*results): - query_point = query_points[query_idx] - segment_coords = boundary_segments[geometry_idx] - best_z = _interpolate_z_on_segment(query_point, segment_coords) - points_3d[query_idx, 2] = best_z - - intersection_dict[intersection_id]["outline_3d"] = Polyline3D.from_array(points_3d) + if len(boundary_segments) >= 1: + boundary_segments = np.concatenate(boundary_segments, axis=0) + boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) + occupancy_map = OccupancyMap2D(boundary_segment_linestrings) + + for intersection_id, intersection_data in intersection_dict.items(): + points_2d = intersection_data["outline_2d"].array + points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) + points_3d[:, :2] = points_2d + + query_points = shapely.creation.points(points_2d) + results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) + for query_idx, geometry_idx in zip(*results): + query_point = query_points[query_idx] + segment_coords = boundary_segments[geometry_idx] + best_z = _interpolate_z_on_segment(query_point, segment_coords) + points_3d[query_idx, 2] = best_z + + intersection_dict[intersection_id]["outline_3d"] = Polyline3D.from_array(points_3d) return intersection_dict diff --git a/src/py123d/conversion/datasets/av2/av2_sensor_converter.py b/src/py123d/conversion/datasets/av2/av2_sensor_converter.py index 1e5a192a..48b444a1 100644 --- a/src/py123d/conversion/datasets/av2/av2_sensor_converter.py +++ b/src/py123d/conversion/datasets/av2/av2_sensor_converter.py @@ -16,62 +16,60 @@ ) from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.registry.box_detection_label_registry import AV2SensorBoxDetectionLabel -from py123d.conversion.registry.lidar_index_registry import AVSensorLiDARIndex -from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import ( +from py123d.conversion.registry import AV2SensorBoxDetectionLabel, AV2SensorLiDARIndex +from py123d.datatypes.detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper +from py123d.datatypes.metadata import LogMetadata, MapMetadata +from py123d.datatypes.sensors import ( + LiDARMetadata, + LiDARType, PinholeCameraMetadata, PinholeCameraType, PinholeDistortion, PinholeIntrinsics, ) -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import ( - get_av2_ford_fusion_hybrid_parameters, - rear_axle_se3_to_center_se3, -) -from py123d.geometry import BoundingBoxSE3Index, StateSE3, Vector3D, Vector3DIndex -from py123d.geometry.bounding_box import BoundingBoxSE3 -from py123d.geometry.transform.transform_se3 import convert_relative_to_absolute_se3_array +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import get_av2_ford_fusion_hybrid_parameters +from py123d.geometry import BoundingBoxSE3, BoundingBoxSE3Index, PoseSE3, Vector3D, Vector3DIndex +from py123d.geometry.transform import convert_relative_to_absolute_se3_array class AV2SensorConverter(AbstractDatasetConverter): + """Dataset converter for the AV2 sensor dataset.""" + def __init__( self, splits: List[str], av2_data_root: Union[Path, str], dataset_converter_config: DatasetConverterConfig, ) -> None: + """Initializes the AV2SensorConverter. + + :param splits: List of dataset splits to convert, e.g. ["av2-sensor_train", "av2-sensor_val", "av2-sensor_test"] + :param av2_data_root: Root directory of the AV2 sensor dataset. + :param dataset_converter_config: Configuration for the dataset converter. + """ super().__init__(dataset_converter_config) assert av2_data_root is not None, "The variable `av2_data_root` must be provided." for split in splits: - assert ( - split in AV2_SENSOR_SPLITS - ), f"Split {split} is not available. Available splits: {self.available_splits}" + assert split in AV2_SENSOR_SPLITS, f"Split {split} is not available. Available splits: {AV2_SENSOR_SPLITS}" self._splits: List[str] = splits self._av2_data_root: Path = Path(av2_data_root) - self._log_paths_and_split: Dict[str, List[Path]] = self._collect_log_paths() + self._log_paths_and_split: List[Tuple[Path, str]] = self._collect_log_paths() - def _collect_log_paths(self) -> Dict[str, List[Path]]: + def _collect_log_paths(self) -> List[Tuple[Path, str]]: + """Collects source log folder paths for the specified splits.""" log_paths_and_split: List[Tuple[Path, str]] = [] - for split in self._splits: dataset_name = split.split("_")[0] split_type = split.split("_")[-1] assert split_type in ["train", "val", "test"] - if "av2-sensor" == dataset_name: log_folder = self._av2_data_root / "sensor" / split_type else: raise ValueError(f"Unknown dataset name {dataset_name} in split {split}.") - log_paths_and_split.extend([(log_path, split) for log_path in log_folder.iterdir()]) - return log_paths_and_split def get_number_of_maps(self) -> int: @@ -125,7 +123,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: # 3. Process source log data if log_needs_writing: - sensor_df = build_sensor_dataframe(source_log_path) synchronization_df = build_synchronization_dataframe(sensor_df) @@ -167,10 +164,13 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: def _get_av2_sensor_map_metadata(split: str, source_log_path: Path) -> MapMetadata: + """Helper to get map metadata for AV2 sensor dataset.""" # NOTE: We need to get the city name from the map folder. # see: https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/av2_sensor_dataloader.py#L163 + map_folder = source_log_path / "map" - log_map_archive_path = next(map_folder.glob("log_map_archive_*.json")) + log_map_archive_path = next(map_folder.glob("log_map_archive_*.json"), None) + assert log_map_archive_path is not None, f"Log map archive file not found in {map_folder}." location = log_map_archive_path.name.split("____")[1].split("_")[0] return MapMetadata( dataset="av2-sensor", @@ -185,7 +185,7 @@ def _get_av2_sensor_map_metadata(split: str, source_log_path: Path) -> MapMetada def _get_av2_pinhole_camera_metadata( source_log_path: Path, dataset_converter_config: DatasetConverterConfig ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: - + """Helper to get pinhole camera metadata for AV2 sensor dataset.""" pinhole_camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = {} if dataset_converter_config.include_pinhole_cameras: intrinsics_file = source_log_path / "calibration" / "intrinsics.feather" @@ -206,11 +206,10 @@ def _get_av2_pinhole_camera_metadata( def _get_av2_lidar_metadata( source_log_path: Path, dataset_converter_config: DatasetConverterConfig ) -> Dict[LiDARType, LiDARMetadata]: + """Helper to get LiDAR metadata for AV2 sensor dataset.""" metadata: Dict[LiDARType, LiDARMetadata] = {} - if dataset_converter_config.include_lidars: - # Load calibration feather file calibration_file = source_log_path / "calibration" / "egovehicle_SE3_sensor.feather" calibration_df = pd.read_feather(calibration_file) @@ -221,16 +220,16 @@ def _get_av2_lidar_metadata( # top lidar: metadata[LiDARType.LIDAR_TOP] = LiDARMetadata( lidar_type=LiDARType.LIDAR_TOP, - lidar_index=AVSensorLiDARIndex, - extrinsic=_row_dict_to_state_se3( + lidar_index=AV2SensorLiDARIndex, + extrinsic=_row_dict_to_pose_se3( calibration_df[calibration_df["sensor_name"] == "up_lidar"].iloc[0].to_dict() ), ) # down lidar: metadata[LiDARType.LIDAR_DOWN] = LiDARMetadata( lidar_type=LiDARType.LIDAR_DOWN, - lidar_index=AVSensorLiDARIndex, - extrinsic=_row_dict_to_state_se3( + lidar_index=AV2SensorLiDARIndex, + extrinsic=_row_dict_to_pose_se3( calibration_df[calibration_df["sensor_name"] == "down_lidar"].iloc[0].to_dict() ), ) @@ -238,10 +237,9 @@ def _get_av2_lidar_metadata( def _extract_av2_sensor_box_detections( - annotations_df: Optional[pd.DataFrame], - lidar_timestamp_ns: int, - ego_state_se3: EgoStateSE3, + annotations_df: Optional[pd.DataFrame], lidar_timestamp_ns: int, ego_state_se3: EgoStateSE3 ) -> BoxDetectionWrapper: + """Extract box detections from AV2 sensor dataset annotations.""" # TODO: Extract velocity from annotations_df if available. @@ -255,20 +253,19 @@ def _extract_av2_sensor_box_detections( detections_velocity = np.zeros((num_detections, len(Vector3DIndex)), dtype=np.float64) detections_token: List[str] = annotations_slice["track_uuid"].tolist() detections_labels: List[AV2SensorBoxDetectionLabel] = [] + detections_num_lidar_points: List[int] = [] for detection_idx, (_, row) in enumerate(annotations_slice.iterrows()): row = row.to_dict() - detections_state[detection_idx, BoundingBoxSE3Index.XYZ] = [row["tx_m"], row["ty_m"], row["tz_m"]] detections_state[detection_idx, BoundingBoxSE3Index.QUATERNION] = [row["qw"], row["qx"], row["qy"], row["qz"]] detections_state[detection_idx, BoundingBoxSE3Index.EXTENT] = [row["length_m"], row["width_m"], row["height_m"]] + detections_labels.append(AV2SensorBoxDetectionLabel.deserialize(row["category"])) + detections_num_lidar_points.append(int(row["num_interior_pts"])) - detections_label = AV2SensorBoxDetectionLabel.deserialize(row["category"]) - detections_labels.append(detections_label) - - detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( + detections_state[:, BoundingBoxSE3Index.SE3] = convert_relative_to_absolute_se3_array( origin=ego_state_se3.rear_axle_se3, - se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3], + se3_array=detections_state[:, BoundingBoxSE3Index.SE3], ) box_detections: List[BoxDetectionSE3] = [] @@ -277,12 +274,11 @@ def _extract_av2_sensor_box_detections( BoxDetectionSE3( metadata=BoxDetectionMetadata( label=detections_labels[detection_idx], - timepoint=None, track_token=detections_token[detection_idx], - confidence=None, + num_lidar_points=detections_num_lidar_points[detection_idx], ), bounding_box_se3=BoundingBoxSE3.from_array(detections_state[detection_idx]), - velocity=Vector3D.from_array(detections_velocity[detection_idx]), + velocity_3d=Vector3D.from_array(detections_velocity[detection_idx]), ) ) @@ -290,28 +286,24 @@ def _extract_av2_sensor_box_detections( def _extract_av2_sensor_ego_state(city_se3_egovehicle_df: pd.DataFrame, lidar_timestamp_ns: int) -> EgoStateSE3: + """Extract ego state from AV2 sensor dataset city_SE3_egovehicle dataframe.""" ego_state_slice = get_slice_with_timestamp_ns(city_se3_egovehicle_df, lidar_timestamp_ns) - assert ( - len(ego_state_slice) == 1 - ), f"Expected exactly one ego state for timestamp {lidar_timestamp_ns}, got {len(ego_state_slice)}." + assert len(ego_state_slice) == 1, ( + f"Expected exactly one ego state for timestamp {lidar_timestamp_ns}, got {len(ego_state_slice)}." + ) ego_pose_dict = ego_state_slice.iloc[0].to_dict() - rear_axle_pose = _row_dict_to_state_se3(ego_pose_dict) + rear_axle_pose = _row_dict_to_pose_se3(ego_pose_dict) vehicle_parameters = get_av2_ford_fusion_hybrid_parameters() - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) # TODO: Add script to calculate the dynamic state from log sequence. - dynamic_state = DynamicStateSE3( - velocity=Vector3D(x=0.0, y=0.0, z=0.0), - acceleration=Vector3D(x=0.0, y=0.0, z=0.0), - angular_velocity=Vector3D(x=0.0, y=0.0, z=0.0), - ) + dynamic_state_se3 = None - return EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters, + dynamic_state_se3=dynamic_state_se3, timepoint=None, ) @@ -323,6 +315,7 @@ def _extract_av2_sensor_pinhole_cameras( source_log_path: Path, dataset_converter_config: DatasetConverterConfig, ) -> List[CameraData]: + """Extract pinhole camera data from AV2 sensor dataset.""" camera_data_list: List[CameraData] = [] split = source_log_path.parent.name @@ -353,7 +346,7 @@ def _extract_av2_sensor_pinhole_cameras( camera_data = CameraData( camera_type=pinhole_camera_type, - extrinsic=_row_dict_to_state_se3(row), + extrinsic=_row_dict_to_pose_se3(row), dataset_root=av2_sensor_data_root, relative_path=relative_image_path, ) @@ -365,6 +358,7 @@ def _extract_av2_sensor_pinhole_cameras( def _extract_av2_sensor_lidars( source_log_path: Path, lidar_timestamp_ns: int, dataset_converter_config: DatasetConverterConfig ) -> List[LiDARData]: + """Extract LiDAR data from AV2 sensor dataset.""" lidars: List[LiDARData] = [] if dataset_converter_config.include_lidars: av2_sensor_data_root = source_log_path.parent.parent @@ -385,9 +379,9 @@ def _extract_av2_sensor_lidars( return lidars -def _row_dict_to_state_se3(row_dict: Dict[str, float]) -> StateSE3: - """Helper function to convert a row dictionary to a StateSE3 object.""" - return StateSE3( +def _row_dict_to_pose_se3(row_dict: Dict[str, float]) -> PoseSE3: + """Helper function to convert a row dictionary to a PoseSE3 object.""" + return PoseSE3( x=row_dict["tx_m"], y=row_dict["ty_m"], z=row_dict["tz_m"], diff --git a/src/py123d/conversion/datasets/av2/av2_sensor_io.py b/src/py123d/conversion/datasets/av2/av2_sensor_io.py index 81a3de3a..348a6b8f 100644 --- a/src/py123d/conversion/datasets/av2/av2_sensor_io.py +++ b/src/py123d/conversion/datasets/av2/av2_sensor_io.py @@ -8,6 +8,8 @@ def load_av2_sensor_lidar_pcs_from_file(feather_path: Union[Path, str]) -> Dict[LiDARType, np.ndarray]: + """Loads AV2 sensor LiDAR point clouds from a feather file.""" + # NOTE: The AV2 dataset stores both top and down LiDAR data in the same feather file. # We need to separate them based on the laser_number field. # See here: https://github.com/argoverse/av2-api/issues/77#issuecomment-1178040867 diff --git a/tests/unit/datatypes/maps/__init__.py b/src/py123d/conversion/datasets/av2/utils/__init__.py similarity index 100% rename from tests/unit/datatypes/maps/__init__.py rename to src/py123d/conversion/datasets/av2/utils/__init__.py diff --git a/src/py123d/conversion/datasets/av2/utils/av2_constants.py b/src/py123d/conversion/datasets/av2/utils/av2_constants.py index 30b59fa2..569691aa 100644 --- a/src/py123d/conversion/datasets/av2/utils/av2_constants.py +++ b/src/py123d/conversion/datasets/av2/utils/av2_constants.py @@ -1,11 +1,11 @@ from typing import Dict, Final, Set -from py123d.datatypes.maps.map_datatypes import RoadLineType -from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType +from py123d.datatypes.map_objects import RoadLineType +from py123d.datatypes.sensors import PinholeCameraType AV2_SENSOR_SPLITS: Set[str] = {"av2-sensor_train", "av2-sensor_val", "av2-sensor_test"} - +# Mapping from AV2 camera names to PinholeCameraType enums. AV2_CAMERA_TYPE_MAPPING: Dict[str, PinholeCameraType] = { "ring_front_center": PinholeCameraType.PCAM_F0, "ring_front_left": PinholeCameraType.PCAM_L0, @@ -18,9 +18,7 @@ "stereo_front_right": PinholeCameraType.PCAM_STEREO_R, } -# AV2_LIDAR_TYPES: Dict[str, str] = { - - +# Mapping from AV2 road line types to RoadLineType enums. AV2_ROAD_LINE_TYPE_MAPPING: Dict[str, RoadLineType] = { "NONE": RoadLineType.NONE, "UNKNOWN": RoadLineType.UNKNOWN, diff --git a/src/py123d/conversion/datasets/av2/utils/av2_helper.py b/src/py123d/conversion/datasets/av2/utils/av2_helper.py index cd0c1f62..0cfb7f15 100644 --- a/src/py123d/conversion/datasets/av2/utils/av2_helper.py +++ b/src/py123d/conversion/datasets/av2/utils/av2_helper.py @@ -11,12 +11,13 @@ def get_dataframe_from_file(file_path: Path) -> pd.DataFrame: + """Get a Pandas DataFrame from parquet or feather files.""" if file_path.suffix == ".parquet": import pyarrow.parquet as pq return pq.read_table(file_path) elif file_path.suffix == ".feather": - import pyarrow.feather as feather + from pyarrow import feather return feather.read_feather(file_path) else: @@ -29,6 +30,7 @@ def get_slice_with_timestamp_ns(dataframe: pd.DataFrame, timestamp_ns: int): def build_sensor_dataframe(source_log_path: Path) -> pd.DataFrame: + """Builds a sensor dataframe from the AV2 source log path.""" # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L209 @@ -64,6 +66,12 @@ def build_synchronization_dataframe( sensor_dataframe: pd.DataFrame, matching_criterion: Literal["nearest", "forward"] = "nearest", ) -> pd.DataFrame: + """Builds a synchronization dataframe, between sensors observations in a log. + + :param sensor_dataframe: DataFrame containing sensor data. + :param matching_criterion: Criterion for matching timestamps, defaults to "nearest" + :return: DataFrame containing synchronized sensor data. + """ # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L382 @@ -113,6 +121,7 @@ def build_synchronization_dataframe( def populate_sensor_records(sensor_path: Path, split: str, log_id: str) -> pd.DataFrame: + """Populate sensor records from a sensor path.""" sensor_name = sensor_path.name sensor_files = list(sensor_path.iterdir()) diff --git a/src/py123d/conversion/datasets/kitti360/kitti360_converter.py b/src/py123d/conversion/datasets/kitti360/kitti360_converter.py index 22d84229..438c9123 100644 --- a/src/py123d/conversion/datasets/kitti360/kitti360_converter.py +++ b/src/py123d/conversion/datasets/kitti360/kitti360_converter.py @@ -20,42 +20,32 @@ ) from py123d.conversion.datasets.kitti360.utils.kitti360_labels import ( BBOX_LABLES_TO_DETECTION_NAME_DICT, - KIITI360_DETECTION_NAME_DICT, + KITTI360_DETECTION_NAME_DICT, kittiId2label, ) from py123d.conversion.datasets.kitti360.utils.preprocess_detection import process_detection from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.registry.box_detection_label_registry import KITTI360BoxDetectionLabel -from py123d.conversion.registry.lidar_index_registry import Kitti360LiDARIndex -from py123d.datatypes.detections.box_detections import ( - BoxDetectionMetadata, - BoxDetectionSE3, - BoxDetectionWrapper, -) -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.fisheye_mei_camera import ( +from py123d.conversion.registry import KITTI360BoxDetectionLabel, KITTI360LiDARIndex +from py123d.datatypes.detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper +from py123d.datatypes.metadata import LogMetadata, MapMetadata +from py123d.datatypes.sensors import ( FisheyeMEICameraMetadata, FisheyeMEICameraType, FisheyeMEIDistortion, FisheyeMEIProjection, -) -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import ( + LiDARMetadata, + LiDARType, PinholeCameraMetadata, PinholeCameraType, PinholeDistortion, PinholeIntrinsics, ) -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import ( - get_kitti360_vw_passat_parameters, - rear_axle_se3_to_center_se3, -) -from py123d.geometry import BoundingBoxSE3, Quaternion, StateSE3, Vector3D -from py123d.geometry.transform.transform_se3 import convert_se3_array_between_origins, translate_se3_along_body_frame +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import DynamicStateSE3, EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import get_kitti360_vw_passat_parameters +from py123d.geometry import BoundingBoxSE3, PoseSE3, Quaternion, Vector3D +from py123d.geometry.transform import convert_se3_array_between_origins, translate_se3_along_body_frame KITTI360_DT: Final[float] = 0.1 @@ -117,6 +107,8 @@ def _get_kitti360_required_modality_roots(kitti360_folders: Dict[str, Path]) -> class Kitti360Converter(AbstractDatasetConverter): + """Converter class for KITTI-360 dataset.""" + def __init__( self, splits: List[str], @@ -128,6 +120,18 @@ def __init__( val_sequences: List[str], test_sequences: List[str], ) -> None: + """Initializes the Kitti360Converter. + + :param splits: List of splits to include in the conversion, e.g. `kitti360_train`, `kitti360_val`, `kitti360_test` + :param kitti360_data_root: Path to the KITTI-360 dataset root directory + :param detection_cache_root: Path to the detection cache directory + :param detection_radius: Radius for the box detections to include. + :param dataset_converter_config: Dataset converter configuration + :param train_sequences: List of sequences to include in the training split + :param val_sequences: List of sequences to include in the validation split + :param test_sequences: List of sequences to include in the test split + """ + assert kitti360_data_root is not None, "The variable `kitti360_data_root` must be provided." super().__init__(dataset_converter_config) for split in splits: @@ -206,19 +210,15 @@ def _has_modality(seq_name: str, modality_name: str, root: Path) -> bool: return log_paths_and_split def get_number_of_maps(self) -> int: - """Returns the number of available raw data maps for conversion.""" + """Inherited, see superclass.""" return self._total_maps def get_number_of_logs(self) -> int: - """Returns the number of available raw data logs for conversion.""" + """Inherited, see superclass.""" return self._total_logs def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: - """ - Convert a single map in raw data format to the uniform 123D format. - :param map_index: The index of the map to convert. - :param map_writer: The map writer to use for writing the converted map. - """ + """Inherited, see superclass.""" log_name, split = self._log_names_and_split[map_index] map_metadata = _get_kitti360_map_metadata(split, log_name) map_needs_writing = map_writer.reset(self.dataset_converter_config, map_metadata) @@ -227,11 +227,7 @@ def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: map_writer.close() def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: - """ - Convert a single log in raw data format to the uniform 123D format. - :param log_index: The index of the log to convert. - :param log_writer: The log writer to use for writing the converted log. - """ + """Inherited, see superclass.""" log_name, split = self._log_names_and_split[log_index] # Create log metadata @@ -263,7 +259,9 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: if log_needs_writing: ts_list: List[TimePoint] = _read_timestamps(log_name, self._kitti360_folders) ego_state_all, valid_timestamp = _extract_ego_state_all(log_name, self._kitti360_folders) - ego_states_xyz = np.array([ego_state.center.array[:3] for ego_state in ego_state_all], dtype=np.float64) + ego_states_xyz = np.array( + [ego_state.center_se3.point_3d.array[:3] for ego_state in ego_state_all], dtype=np.float64 + ) box_detection_wrapper_all = _extract_kitti360_box_detections_all( log_name, len(ts_list), @@ -316,6 +314,7 @@ def _get_kitti360_pinhole_camera_metadata( kitti360_folders: Dict[str, Path], dataset_converter_config: DatasetConverterConfig, ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + """Gets the KITTI-360 pinhole camera metadata from calibration files.""" pinhole_cam_metadatas: Dict[PinholeCameraType, PinholeCameraMetadata] = {} if dataset_converter_config.include_pinhole_cameras: @@ -351,9 +350,10 @@ def _get_kitti360_fisheye_mei_camera_metadata( kitti360_folders: Dict[str, Path], dataset_converter_config: DatasetConverterConfig, ) -> Dict[FisheyeMEICameraType, FisheyeMEICameraMetadata]: + """Gets the KITTI-360 fisheye MEI camera metadata from calibration files.""" + fisheye_cam_metadatas: Dict[FisheyeMEICameraType, FisheyeMEICameraMetadata] = {} if dataset_converter_config.include_fisheye_mei_cameras: - fisheye_camera02_path = kitti360_folders[DIR_CALIB] / "image_02.yaml" fisheye_camera03_path = kitti360_folders[DIR_CALIB] / "image_03.yaml" @@ -363,7 +363,6 @@ def _get_kitti360_fisheye_mei_camera_metadata( fisheye_result = {"image_02": fisheye02, "image_03": fisheye03} for fcam_type, fcam_name in KITTI360_FISHEYE_MEI_CAMERA_TYPES.items(): - distortion_params = fisheye_result[fcam_name]["distortion_parameters"] distortion = FisheyeMEIDistortion( k1=distortion_params["k1"], @@ -393,6 +392,7 @@ def _get_kitti360_fisheye_mei_camera_metadata( def _get_kitti360_map_metadata(split: str, log_name: str) -> MapMetadata: + """Gets the KITTI-360 map metadata.""" return MapMetadata( dataset="kitti360", split=split, @@ -404,6 +404,7 @@ def _get_kitti360_map_metadata(split: str, log_name: str) -> MapMetadata: def _read_projection_matrix(p_line: str) -> np.ndarray: + """Helper function to read projection matrix from calibration file line.""" parts = p_line.split(" ", 1) if len(parts) != 2: raise ValueError(f"Bad projection line: {p_line}") @@ -414,7 +415,7 @@ def _read_projection_matrix(p_line: str) -> np.ndarray: def _readYAMLFile(fileName: Path) -> Dict[str, Any]: - """make OpenCV YAML file compatible with python""" + """Make OpenCV YAML file compatible with python""" ret = {} skip_lines = 1 # Skip the first line which says "%YAML:1.0". Or replace it with "%YAML 1.0" with open(fileName) as fin: @@ -428,26 +429,24 @@ def _readYAMLFile(fileName: Path) -> Dict[str, Any]: def _get_kitti360_lidar_metadata( - kitti360_folders: Dict[str, Path], - dataset_converter_config: DatasetConverterConfig, + kitti360_folders: Dict[str, Path], dataset_converter_config: DatasetConverterConfig ) -> Dict[LiDARType, LiDARMetadata]: + """Gets the KITTI-360 LiDAR metadata from calibration files.""" metadata: Dict[LiDARType, LiDARMetadata] = {} if dataset_converter_config.include_lidars: extrinsic = get_kitti360_lidar_extrinsic(kitti360_folders[DIR_CALIB]) - extrinsic_state_se3 = StateSE3.from_transformation_matrix(extrinsic) - extrinsic_state_se3 = _extrinsic_from_imu_to_rear_axle(extrinsic_state_se3) + extrinsic_pose_se3 = PoseSE3.from_transformation_matrix(extrinsic) + extrinsic_pose_se3 = _extrinsic_from_imu_to_rear_axle(extrinsic_pose_se3) metadata[LiDARType.LIDAR_TOP] = LiDARMetadata( lidar_type=LiDARType.LIDAR_TOP, - lidar_index=Kitti360LiDARIndex, - extrinsic=extrinsic_state_se3, + lidar_index=KITTI360LiDARIndex, + extrinsic=extrinsic_pose_se3, ) return metadata def _read_timestamps(log_name: str, kitti360_folders: Dict[str, Path]) -> Optional[List[TimePoint]]: - """ - Read KITTI-360 timestamps for the given sequence and return Unix epoch timestamps. - """ + """Read KITTI-360 timestamps for the given sequence and return Unix epoch timestamps.""" ts_files = [ kitti360_folders[DIR_3D_RAW] / log_name / "velodyne_points" / "timestamps.txt", kitti360_folders[DIR_2D_RAW] / log_name / "image_00" / "timestamps.txt", @@ -479,8 +478,9 @@ def _read_timestamps(log_name: str, kitti360_folders: Dict[str, Path]) -> Option def _extract_ego_state_all(log_name: str, kitti360_folders: Dict[str, Path]) -> Tuple[List[EgoStateSE3], List[int]]: + """Extracts all ego states for the given sequence.""" - ego_state_all: List[List[float]] = [] + ego_state_all: List[EgoStateSE3] = [] pose_file = kitti360_folders[DIR_POSES] / log_name / "poses.txt" if not pose_file.exists(): raise FileNotFoundError(f"Pose file not found: {pose_file}") @@ -508,7 +508,7 @@ def _extract_ego_state_all(log_name: str, kitti360_folders: Dict[str, Path]) -> R_mat_cali = R_mat @ KITTI3602NUPLAN_IMU_CALIBRATION[:3, :3] ego_quaternion = Quaternion.from_rotation_matrix(R_mat_cali) - imu_pose = StateSE3( + imu_pose = PoseSE3( x=poses[pos, 4], y=poses[pos, 8], z=poses[pos, 12], @@ -518,35 +518,21 @@ def _extract_ego_state_all(log_name: str, kitti360_folders: Dict[str, Path]) -> qz=ego_quaternion.qz, ) - rear_axle_pose = translate_se3_along_body_frame( + rear_axle_se3 = translate_se3_along_body_frame( imu_pose, Vector3D(0.05, -0.32, 0.0), ) - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) - dynamic_state = DynamicStateSE3( - velocity=Vector3D( - x=oxts_data[8], - y=oxts_data[9], - z=oxts_data[10], - ), - acceleration=Vector3D( - x=oxts_data[14], - y=oxts_data[15], - z=oxts_data[16], - ), - angular_velocity=Vector3D( - x=oxts_data[20], - y=oxts_data[21], - z=oxts_data[22], - ), + dynamic_state_se3 = DynamicStateSE3( + velocity=Vector3D(x=oxts_data[8], y=oxts_data[9], z=oxts_data[10]), + acceleration=Vector3D(x=oxts_data[14], y=oxts_data[15], z=oxts_data[16]), + angular_velocity=Vector3D(x=oxts_data[20], y=oxts_data[21], z=oxts_data[22]), ) ego_state_all.append( - EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, + EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_se3, vehicle_parameters=vehicle_parameters, - timepoint=None, + dynamic_state_se3=dynamic_state_se3, ) ) return ego_state_all, valid_timestamp @@ -561,6 +547,7 @@ def _extract_kitti360_box_detections_all( detection_cache_root: Path, detection_radius: float, ) -> List[BoxDetectionWrapper]: + """Extracts all KITTI-360 box detections for the given sequence.""" detections_states: List[List[List[float]]] = [[] for _ in range(ts_len)] detections_velocity: List[List[List[float]]] = [[] for _ in range(ts_len)] @@ -601,7 +588,7 @@ def _extract_kitti360_box_detections_all( else: label = child.find("label").text name = BBOX_LABLES_TO_DETECTION_NAME_DICT.get(label, "unknown") - if child.find("transform") is None or name not in KIITI360_DETECTION_NAME_DICT.keys(): + if child.find("transform") is None or name not in KITTI360_DETECTION_NAME_DICT.keys(): continue obj = KITTI360Bbox3D() obj.parseBbox(child) @@ -617,7 +604,7 @@ def _extract_kitti360_box_detections_all( detections_states[frame].append(obj.get_state_array()) detections_velocity[frame].append(np.array([0.0, 0.0, 0.0])) detections_tokens[frame].append(str(obj.globalID)) - detections_labels[frame].append(KIITI360_DETECTION_NAME_DICT[obj.name]) + detections_labels[frame].append(KITTI360_DETECTION_NAME_DICT[obj.name]) else: global_ID = obj.globalID dynamic_objs[global_ID].append(obj) @@ -654,7 +641,7 @@ def _extract_kitti360_box_detections_all( detections_states[frame].append(obj.get_state_array()) detections_velocity[frame].append(vel) detections_tokens[frame].append(str(obj.globalID)) - detections_labels[frame].append(KIITI360_DETECTION_NAME_DICT[obj.name]) + detections_labels[frame].append(KITTI360_DETECTION_NAME_DICT[obj.name]) box_detection_wrapper_all: List[BoxDetectionWrapper] = [] for frame in range(ts_len): @@ -671,14 +658,13 @@ def _extract_kitti360_box_detections_all( label=detection_label, timepoint=None, track_token=token, - confidence=None, ) bounding_box_se3 = BoundingBoxSE3.from_array(state) velocity_vector = Vector3D.from_array(velocity) box_detection = BoxDetectionSE3( metadata=detection_metadata, bounding_box_se3=bounding_box_se3, - velocity=velocity_vector, + velocity_3d=velocity_vector, ) box_detections.append(box_detection) box_detection_wrapper_all.append(BoxDetectionWrapper(box_detections=box_detections)) @@ -691,6 +677,7 @@ def _extract_kitti360_lidar( kitti360_folders: Dict[str, Path], data_converter_config: DatasetConverterConfig, ) -> List[LiDARData]: + """Extracts KITTI-360 LiDAR data for the given sequence and index.""" lidars: List[LiDARData] = [] if data_converter_config.include_lidars: @@ -718,10 +705,11 @@ def _extract_kitti360_lidar( def _extract_kitti360_pinhole_cameras( log_name: str, idx: int, - camera_calibration: Dict[str, StateSE3], + camera_calibration: Dict[str, PoseSE3], kitti360_folders: Dict[str, Path], data_converter_config: DatasetConverterConfig, ) -> List[CameraData]: + """Extracts KITTI-360 pinhole camera data for the given sequence and index.""" pinhole_camera_data_list: List[CameraData] = [] if data_converter_config.include_pinhole_cameras: @@ -744,11 +732,11 @@ def _extract_kitti360_pinhole_cameras( def _extract_kitti360_fisheye_mei_cameras( log_name: str, idx: int, - camera_calibration: Dict[str, StateSE3], + camera_calibration: Dict[str, PoseSE3], kitti360_folders: Dict[str, Path], data_converter_config: DatasetConverterConfig, ) -> List[CameraData]: - + """Extracts KITTI-360 fisheye MEI camera data for the given sequence and index.""" fisheye_camera_data_list: List[CameraData] = [] if data_converter_config.include_fisheye_mei_cameras: for camera_type, cam_dir_name in KITTI360_FISHEYE_MEI_CAMERA_TYPES.items(): @@ -766,13 +754,14 @@ def _extract_kitti360_fisheye_mei_cameras( return fisheye_camera_data_list -def _load_kitti_360_calibration(kitti_360_data_root: Path) -> Dict[str, StateSE3]: +def _load_kitti_360_calibration(kitti_360_data_root: Path) -> Dict[str, PoseSE3]: + """Helper function to load KITTI-360 camera to IMU calibration.""" calib_file = kitti_360_data_root / DIR_CALIB / "calib_cam_to_pose.txt" if not calib_file.exists(): raise FileNotFoundError(f"Calibration file not found: {calib_file}") lastrow = np.array([0, 0, 0, 1]).reshape(1, 4) - calib_dict: Dict[str, StateSE3] = {} + calib_dict: Dict[str, PoseSE3] = {} with open(calib_file, "r") as f: for line in f: parts = line.strip().split() @@ -781,13 +770,14 @@ def _load_kitti_360_calibration(kitti_360_data_root: Path) -> Dict[str, StateSE3 matrix = np.array(values).reshape(3, 4) cam2pose = np.concatenate((matrix, lastrow)) cam2pose = KITTI3602NUPLAN_IMU_CALIBRATION @ cam2pose - camera_extrinsic = StateSE3.from_transformation_matrix(cam2pose) + camera_extrinsic = PoseSE3.from_transformation_matrix(cam2pose) camera_extrinsic = _extrinsic_from_imu_to_rear_axle(camera_extrinsic) calib_dict[key] = camera_extrinsic return calib_dict -def _extrinsic_from_imu_to_rear_axle(extrinsic: StateSE3) -> StateSE3: - imu_se3 = StateSE3(x=-0.05, y=0.32, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) - rear_axle_se3 = StateSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) - return StateSE3.from_array(convert_se3_array_between_origins(imu_se3, rear_axle_se3, extrinsic.array)) +def _extrinsic_from_imu_to_rear_axle(extrinsic: PoseSE3) -> PoseSE3: + """Convert extrinsic from IMU origin to rear axle origin.""" + imu_se3 = PoseSE3(x=-0.05, y=0.32, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + rear_axle_se3 = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + return PoseSE3.from_array(convert_se3_array_between_origins(imu_se3, rear_axle_se3, extrinsic.array)) diff --git a/src/py123d/conversion/datasets/kitti360/kitti360_map_conversion.py b/src/py123d/conversion/datasets/kitti360/kitti360_map_conversion.py index 847250eb..29f91431 100644 --- a/src/py123d/conversion/datasets/kitti360/kitti360_map_conversion.py +++ b/src/py123d/conversion/datasets/kitti360/kitti360_map_conversion.py @@ -6,30 +6,30 @@ import numpy as np import shapely.geometry as geom -from py123d.conversion.datasets.kitti360.utils.kitti360_helper import KITTI360_MAP_Bbox3D +from py123d.conversion.datasets.kitti360.utils.kitti360_helper import ( + KITTI360_MAP_Bbox3D, +) from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter from py123d.conversion.utils.map_utils.road_edge.road_edge_2d_utils import ( get_road_edge_linear_rings, split_line_geometry_by_max_length, ) -from py123d.conversion.utils.map_utils.road_edge.road_edge_3d_utils import lift_road_edges_to_3d -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCarpark, - CacheGenericDrivable, - CacheRoadEdge, - CacheWalkway, +from py123d.conversion.utils.map_utils.road_edge.road_edge_3d_utils import ( + lift_road_edges_to_3d, +) +from py123d.datatypes.map_objects.map_layer_types import RoadEdgeType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + GenericDrivable, + RoadEdge, + Walkway, ) -from py123d.datatypes.maps.map_datatypes import RoadEdgeType from py123d.geometry.polyline import Polyline3D MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges - KITTI360_DATA_ROOT = Path(os.environ["KITTI360_DATA_ROOT"]) - DIR_3D_BBOX = "data_3d_bboxes" - PATH_3D_BBOX_ROOT: Path = KITTI360_DATA_ROOT / DIR_3D_BBOX - KITTI360_MAP_BBOX = [ "road", "sidewalk", @@ -40,8 +40,7 @@ def convert_kitti360_map_with_writer(log_name: str, map_writer: AbstractMapWriter) -> None: - """ - Convert KITTI-360 map data using the provided map writer. + """Convert KITTI-360 map data using the provided map writer. This function extracts map data from KITTI-360 XML files and writes them using the map writer interface. :param log_name: The name of the log to convert @@ -71,7 +70,7 @@ def convert_kitti360_map_with_writer(log_name: str, map_writer: AbstractMapWrite for obj in objs: if obj.label == "road": map_writer.write_generic_drivable( - CacheGenericDrivable( + GenericDrivable( object_id=obj.id, outline=obj.vertices, geometry=geom.Polygon(obj.vertices.array[:, :3]), @@ -81,7 +80,7 @@ def convert_kitti360_map_with_writer(log_name: str, map_writer: AbstractMapWrite road_outlines_3d.append(Polyline3D.from_array(road_outline_array)) elif obj.label == "sidewalk": map_writer.write_walkway( - CacheWalkway( + Walkway( object_id=obj.id, outline=obj.vertices, geometry=geom.Polygon(obj.vertices.array[:, :3]), @@ -89,7 +88,7 @@ def convert_kitti360_map_with_writer(log_name: str, map_writer: AbstractMapWrite ) elif obj.label == "driveway": map_writer.write_carpark( - CacheCarpark( + Carpark( object_id=obj.id, outline=obj.vertices, geometry=geom.Polygon(obj.vertices.array[:, :3]), @@ -108,7 +107,7 @@ def convert_kitti360_map_with_writer(log_name: str, map_writer: AbstractMapWrite for idx in range(len(road_edges_linestrings_3d)): map_writer.write_road_edge( - CacheRoadEdge( + RoadEdge( object_id=idx, road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, polyline=Polyline3D.from_linestring(road_edges_linestrings_3d[idx]), diff --git a/src/py123d/conversion/datasets/kitti360/kitti360_sensor_io.py b/src/py123d/conversion/datasets/kitti360/kitti360_sensor_io.py index e58b165d..f96dc84a 100644 --- a/src/py123d/conversion/datasets/kitti360/kitti360_sensor_io.py +++ b/src/py123d/conversion/datasets/kitti360/kitti360_sensor_io.py @@ -4,26 +4,28 @@ import numpy as np -from py123d.conversion.registry.lidar_index_registry import Kitti360LiDARIndex -from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.conversion.registry.lidar_index_registry import KITTI360LiDARIndex +from py123d.datatypes.metadata import LogMetadata from py123d.datatypes.sensors.lidar import LiDARType -from py123d.geometry.se import StateSE3 +from py123d.geometry.pose import PoseSE3 from py123d.geometry.transform.transform_se3 import convert_points_3d_array_between_origins def load_kitti360_lidar_pcs_from_file(filepath: Path, log_metadata: LogMetadata) -> Dict[LiDARType, np.ndarray]: + """Loads KITTI-360 LiDAR point clouds the original binary files.""" + if not filepath.exists(): logging.warning(f"LiDAR file does not exist: {filepath}. Returning empty point cloud.") - return {LiDARType.LIDAR_TOP: np.zeros((1, len(Kitti360LiDARIndex)), dtype=np.float32)} + return {LiDARType.LIDAR_TOP: np.zeros((1, len(KITTI360LiDARIndex)), dtype=np.float32)} lidar_extrinsic = log_metadata.lidar_metadata[LiDARType.LIDAR_TOP].extrinsic lidar_pc = np.fromfile(filepath, dtype=np.float32) - lidar_pc = np.reshape(lidar_pc, [-1, len(Kitti360LiDARIndex)]) + lidar_pc = np.reshape(lidar_pc, [-1, len(KITTI360LiDARIndex)]) - lidar_pc[..., Kitti360LiDARIndex.XYZ] = convert_points_3d_array_between_origins( + lidar_pc[..., KITTI360LiDARIndex.XYZ] = convert_points_3d_array_between_origins( from_origin=lidar_extrinsic, - to_origin=StateSE3(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0), - points_3d_array=lidar_pc[..., Kitti360LiDARIndex.XYZ], + to_origin=PoseSE3(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0), + points_3d_array=lidar_pc[..., KITTI360LiDARIndex.XYZ], ) return {LiDARType.LIDAR_TOP: lidar_pc} diff --git a/src/py123d/conversion/datasets/kitti360/utils/kitti360_helper.py b/src/py123d/conversion/datasets/kitti360/utils/kitti360_helper.py index fa3afa77..beb66b1d 100644 --- a/src/py123d/conversion/datasets/kitti360/utils/kitti360_helper.py +++ b/src/py123d/conversion/datasets/kitti360/utils/kitti360_helper.py @@ -6,7 +6,7 @@ from scipy.linalg import polar from py123d.conversion.datasets.kitti360.utils.kitti360_labels import BBOX_LABLES_TO_DETECTION_NAME_DICT, kittiId2label -from py123d.geometry import BoundingBoxSE3, EulerAngles, Polyline3D, StateSE3 +from py123d.geometry import BoundingBoxSE3, EulerAngles, Polyline3D, PoseSE3 # KITTI360_DATA_ROOT = Path(os.environ["KITTI360_DATA_ROOT"]) # DIR_CALIB = "calibration" @@ -42,14 +42,12 @@ def global2local(globalId: int) -> Tuple[int, int]: class KITTI360Bbox3D: - # global id(only used for sequence 0004) dynamic_global_id = 2000000 static_global_id = 1000000 # Constructor def __init__(self): - # the ID of the corresponding object self.semanticId = -1 self.instanceId = -1 @@ -137,7 +135,7 @@ def parse_scale_rotation(self): self.qz = obj_quaternion.qz def get_state_array(self) -> np.ndarray: - center = StateSE3( + center = PoseSE3( x=self.T[0], y=self.T[1], z=self.T[2], diff --git a/src/py123d/conversion/datasets/kitti360/utils/kitti360_labels.py b/src/py123d/conversion/datasets/kitti360/utils/kitti360_labels.py index a40cffca..06aa8524 100644 --- a/src/py123d/conversion/datasets/kitti360/utils/kitti360_labels.py +++ b/src/py123d/conversion/datasets/kitti360/utils/kitti360_labels.py @@ -185,7 +185,7 @@ def assureSingleInstanceName(name): "vendingMachine": "vending machine", } -KIITI360_DETECTION_NAME_DICT = { +KITTI360_DETECTION_NAME_DICT = { "bicycle": KITTI360BoxDetectionLabel.BICYCLE, "box": KITTI360BoxDetectionLabel.BOX, "bus": KITTI360BoxDetectionLabel.BUS, diff --git a/src/py123d/conversion/datasets/kitti360/utils/preprocess_detection.py b/src/py123d/conversion/datasets/kitti360/utils/preprocess_detection.py index 1c2beeca..c6195c5a 100644 --- a/src/py123d/conversion/datasets/kitti360/utils/preprocess_detection.py +++ b/src/py123d/conversion/datasets/kitti360/utils/preprocess_detection.py @@ -28,7 +28,7 @@ ) from py123d.conversion.datasets.kitti360.utils.kitti360_labels import ( BBOX_LABLES_TO_DETECTION_NAME_DICT, - KIITI360_DETECTION_NAME_DICT, + KITTI360_DETECTION_NAME_DICT, kittiId2label, ) @@ -76,7 +76,7 @@ def _collect_static_objects(kitti360_dataset_root: Path, log_name: str) -> List[ label = child.find("label").text name = BBOX_LABLES_TO_DETECTION_NAME_DICT.get(label, "unknown") timestamp = int(child.find("timestamp").text) # -1 for static objects - if child.find("transform") is None or name not in KIITI360_DETECTION_NAME_DICT.keys() or timestamp != -1: + if child.find("transform") is None or name not in KITTI360_DETECTION_NAME_DICT.keys() or timestamp != -1: continue obj = KITTI360Bbox3D() obj.parseBbox(child) diff --git a/src/py123d/conversion/datasets/nuplan/nuplan_converter.py b/src/py123d/conversion/datasets/nuplan/nuplan_converter.py index 36fa3a25..d1e23add 100644 --- a/src/py123d/conversion/datasets/nuplan/nuplan_converter.py +++ b/src/py123d/conversion/datasets/nuplan/nuplan_converter.py @@ -1,6 +1,6 @@ import pickle from pathlib import Path -from typing import Dict, Final, List, Optional, Tuple, Union +from typing import Dict, Final, List, Tuple, Union import numpy as np import yaml @@ -24,28 +24,28 @@ ) from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.registry.box_detection_label_registry import NuPlanBoxDetectionLabel -from py123d.conversion.registry.lidar_index_registry import NuPlanLiDARIndex -from py123d.datatypes.detections.box_detections import BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetection, TrafficLightDetectionWrapper -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import ( +from py123d.conversion.registry import NuPlanBoxDetectionLabel, NuPlanLiDARIndex +from py123d.datatypes.detections import ( + BoxDetectionSE3, + BoxDetectionWrapper, + TrafficLightDetection, + TrafficLightDetectionWrapper, +) +from py123d.datatypes.metadata import LogMetadata, MapMetadata +from py123d.datatypes.sensors import ( + LiDARMetadata, + LiDARType, PinholeCameraMetadata, PinholeCameraType, PinholeDistortion, PinholeIntrinsics, ) -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_chrysler_pacifica_parameters, - rear_axle_se3_to_center_se3, -) -from py123d.geometry import StateSE3, Vector3D +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import DynamicStateSE3, EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuplan_chrysler_pacifica_parameters +from py123d.geometry import PoseSE3, Vector3D -check_dependencies(["nuplan", "sqlalchemy"], "nuplan") +check_dependencies(["nuplan"], "nuplan") from nuplan.database.nuplan_db.nuplan_scenario_queries import get_cameras, get_images_from_lidar_tokens from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB @@ -76,6 +76,8 @@ def create_splits_logs() -> Dict[str, List[str]]: class NuPlanConverter(AbstractDatasetConverter): + """Converter class for the nuPlan dataset.""" + def __init__( self, splits: List[str], @@ -84,27 +86,38 @@ def __init__( nuplan_sensor_root: Union[Path, str], dataset_converter_config: DatasetConverterConfig, ) -> None: + """Initializes the NuPlanConverter. + + :param splits: List of splits to convert, i.e., ["nuplan_train", "nuplan_val", "nuplan_test"] + :param nuplan_data_root: Root directory of the nuPlan data. + :param nuplan_maps_root: Root directory of the nuPlan maps. + :param nuplan_sensor_root: Root directory of the nuPlan sensor data. + :param dataset_converter_config: Configuration for the dataset converter. + """ + super().__init__(dataset_converter_config) assert nuplan_data_root is not None, "The variable `nuplan_data_root` must be provided." assert nuplan_maps_root is not None, "The variable `nuplan_maps_root` must be provided." assert nuplan_sensor_root is not None, "The variable `nuplan_sensor_root` must be provided." for split in splits: - assert ( - split in NUPLAN_DATA_SPLITS - ), f"Split {split} is not available. Available splits: {NUPLAN_DATA_SPLITS}" + assert split in NUPLAN_DATA_SPLITS, ( + f"Split {split} is not available. Available splits: {NUPLAN_DATA_SPLITS}" + ) self._splits: List[str] = splits self._nuplan_data_root: Path = Path(nuplan_data_root) self._nuplan_maps_root: Path = Path(nuplan_maps_root) self._nuplan_sensor_root: Path = Path(nuplan_sensor_root) - self._split_log_path_pairs: List[Tuple[str, List[Path]]] = self._collect_split_log_path_pairs() + self._split_log_path_pairs: List[Tuple[str, Path]] = self._collect_split_log_path_pairs() + + def _collect_split_log_path_pairs(self) -> List[Tuple[str, Path]]: + """Collects the (split, log_path) pairs for the specified splits.""" - def _collect_split_log_path_pairs(self) -> List[Tuple[str, List[Path]]]: # NOTE: the nuplan mini folder has an internal train, val, test structure, all stored in "mini". # The complete dataset is saved in the "trainval" folder (train and val), or in the "test" folder (for test). # Thus, we need filter the logs in a split, based on the internal nuPlan configuration. - split_log_path_pairs: List[Tuple[str, List[Path]]] = [] + split_log_path_pairs: List[Tuple[str, Path]] = [] log_names_per_split = create_splits_logs() for split in self._splits: @@ -117,20 +130,13 @@ def _collect_split_log_path_pairs(self) -> List[Tuple[str, List[Path]]]: nuplan_split_folder = self._nuplan_data_root / "nuplan-v1.1" / "splits" / "test" elif split in ["nuplan-mini_train", "nuplan-mini_val", "nuplan-mini_test"]: nuplan_split_folder = self._nuplan_data_root / "nuplan-v1.1" / "splits" / "mini" - elif split == "nuplan-private_test": - # TODO: Remove private split - nuplan_split_folder = self._nuplan_data_root / "nuplan-v1.1" / "splits" / "private_test" + else: + raise ValueError(f"Unknown nuPlan split: {split}") all_log_files_in_path = [log_file for log_file in nuplan_split_folder.glob("*.db")] - - # TODO: Remove private split - if split == "nuplan-private_test": - valid_log_names = [str(log_file.stem) for log_file in all_log_files_in_path] - else: - all_log_files_in_path = [log_file for log_file in nuplan_split_folder.glob("*.db")] - all_log_names = set([str(log_file.stem) for log_file in all_log_files_in_path]) - log_names_in_split = set(log_names_per_split[split_type]) - valid_log_names = list(all_log_names & log_names_in_split) + all_log_names = set([str(log_file.stem) for log_file in all_log_files_in_path]) + log_names_in_split = set(log_names_per_split[split_type]) + valid_log_names = list(all_log_names & log_names_in_split) for log_name in valid_log_names: log_path = nuplan_split_folder / f"{log_name}.db" @@ -166,9 +172,7 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: """Inherited, see superclass.""" split, source_log_path = self._split_log_path_pairs[log_index] - - nuplan_log_db = NuPlanDB(self._nuplan_data_root, str(source_log_path), None) - + nuplan_log_db = NuPlanDB(str(self._nuplan_data_root), str(source_log_path), None) log_name = nuplan_log_db.log_name # 1. Initialize log metadata @@ -181,10 +185,14 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: vehicle_parameters=get_nuplan_chrysler_pacifica_parameters(), box_detection_label_class=NuPlanBoxDetectionLabel, pinhole_camera_metadata=_get_nuplan_camera_metadata( - source_log_path, self._nuplan_sensor_root, self.dataset_converter_config + source_log_path, + self._nuplan_sensor_root, + self.dataset_converter_config, ), lidar_metadata=_get_nuplan_lidar_metadata( - self._nuplan_sensor_root, log_name, self.dataset_converter_config + self._nuplan_sensor_root, + log_name, + self.dataset_converter_config, ), map_metadata=_get_nuplan_map_metadata(nuplan_log_db.log.map_version), ) @@ -195,7 +203,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: if log_needs_writing: step_interval: float = int(TARGET_DT / NUPLAN_DEFAULT_DT) for nuplan_lidar_pc in nuplan_log_db.lidar_pc[::step_interval]: - lidar_pc_token: str = nuplan_lidar_pc.token log_writer.write( timestamp=TimePoint.from_us(nuplan_lidar_pc.timestamp), @@ -229,6 +236,7 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: def _get_nuplan_map_metadata(location: str) -> MapMetadata: + """Gets the nuPlan map metadata for a given location.""" return MapMetadata( dataset="nuplan", split=None, @@ -244,6 +252,7 @@ def _get_nuplan_camera_metadata( nuplan_sensor_root: Path, dataset_converter_config: DatasetConverterConfig, ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + """Extracts the nuPlan camera metadata for a given log.""" def _get_camera_metadata(camera_type: PinholeCameraType) -> PinholeCameraMetadata: cam = list(get_cameras(source_log_path, [str(NUPLAN_CAMERA_MAPPING[camera_type].value)]))[0] @@ -278,10 +287,9 @@ def _get_nuplan_lidar_metadata( log_name: str, dataset_converter_config: DatasetConverterConfig, ) -> Dict[LiDARType, LiDARMetadata]: - + """Extracts the nuPlan LiDAR metadata for a given log.""" metadata: Dict[LiDARType, LiDARMetadata] = {} log_lidar_folder = nuplan_sensor_root / log_name / "MergedPointCloud" - # NOTE: We first need to check if the LiDAR folder exists, as not all logs have LiDAR data if log_lidar_folder.exists() and log_lidar_folder.is_dir() and dataset_converter_config.include_lidars: for lidar_type in NUPLAN_LIDAR_DICT.values(): @@ -294,9 +302,10 @@ def _get_nuplan_lidar_metadata( def _extract_nuplan_ego_state(nuplan_lidar_pc: LidarPc) -> EgoStateSE3: + """Extracts the nuPlan ego state from a given LidarPc database objects.""" vehicle_parameters = get_nuplan_chrysler_pacifica_parameters() - rear_axle_pose = StateSE3( + rear_axle_pose = PoseSE3( x=nuplan_lidar_pc.ego_pose.x, y=nuplan_lidar_pc.ego_pose.y, z=nuplan_lidar_pc.ego_pose.z, @@ -305,8 +314,7 @@ def _extract_nuplan_ego_state(nuplan_lidar_pc: LidarPc) -> EgoStateSE3: qy=nuplan_lidar_pc.ego_pose.qy, qz=nuplan_lidar_pc.ego_pose.qz, ) - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) - dynamic_state = DynamicStateSE3( + dynamic_state_se3 = DynamicStateSE3( velocity=Vector3D( x=nuplan_lidar_pc.ego_pose.vx, y=nuplan_lidar_pc.ego_pose.vy, @@ -323,22 +331,23 @@ def _extract_nuplan_ego_state(nuplan_lidar_pc: LidarPc) -> EgoStateSE3: z=nuplan_lidar_pc.ego_pose.angular_rate_z, ), ) - return EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters, - timepoint=None, # NOTE: Timepoint is not needed during writing, set to None + dynamic_state_se3=dynamic_state_se3, ) def _extract_nuplan_box_detections(lidar_pc: LidarPc, source_log_path: Path) -> BoxDetectionWrapper: - box_detections: List[BoxDetectionSE3] = list( - get_box_detections_for_lidarpc_token_from_db(source_log_path, lidar_pc.token) + """Extracts the nuPlan box detections from a given LidarPc database objects.""" + box_detections: List[BoxDetectionSE3] = get_box_detections_for_lidarpc_token_from_db( + str(source_log_path), lidar_pc.token ) return BoxDetectionWrapper(box_detections=box_detections) def _extract_nuplan_traffic_lights(log_db: NuPlanDB, lidar_pc_token: str) -> TrafficLightDetectionWrapper: + """Extracts the nuPlan traffic light detections from a given LidarPc database objects.""" traffic_lights_detections: List[TrafficLightDetection] = [ TrafficLightDetection( timepoint=None, # NOTE: Timepoint is not needed during writing, set to None @@ -357,22 +366,19 @@ def _extract_nuplan_cameras( nuplan_sensor_root: Path, dataset_converter_config: DatasetConverterConfig, ) -> List[CameraData]: - + """Extracts the nuPlan camera data from a given LidarPc database objects.""" camera_data_list: List[CameraData] = [] - if dataset_converter_config.include_pinhole_cameras: log_cam_infos = {camera.token: camera for camera in nuplan_log_db.log.cameras} for camera_type, camera_channel in NUPLAN_CAMERA_MAPPING.items(): - camera_data: Optional[Union[str, bytes]] = None image_class = list( - get_images_from_lidar_tokens(source_log_path, [nuplan_lidar_pc.token], [str(camera_channel.value)]) + get_images_from_lidar_tokens(str(source_log_path), [nuplan_lidar_pc.token], [str(camera_channel.value)]) ) if len(image_class) != 0: image = image_class[0] filename_jpg = nuplan_sensor_root / image.filename_jpg if filename_jpg.exists() and filename_jpg.is_file(): - # NOTE: This part of the modified from the MTGS code # In MTGS, a slower method is used to find the nearest ego pose. # The code below uses a direct SQL query to find the nearest ego pose, in a given window. @@ -393,7 +399,7 @@ def _extract_nuplan_cameras( cam_info = log_cam_infos[image.camera_token] c2img_e = cam_info.trans_matrix c2e = img_e2e @ c2img_e - extrinsic = StateSE3.from_transformation_matrix(c2e) + extrinsic = PoseSE3.from_transformation_matrix(c2e) # Store in dictionary camera_data_list.append( @@ -404,7 +410,6 @@ def _extract_nuplan_cameras( relative_path=filename_jpg.relative_to(nuplan_sensor_root), ) ) - return camera_data_list @@ -413,10 +418,9 @@ def _extract_nuplan_lidars( nuplan_sensor_root: Path, dataset_converter_config: DatasetConverterConfig, ) -> List[LiDARData]: - + """Extracts the nuPlan LiDAR data from a given LidarPc database objects.""" lidars: List[LiDARData] = [] if dataset_converter_config.include_lidars: - lidar_full_path = nuplan_sensor_root / nuplan_lidar_pc.filename if lidar_full_path.exists() and lidar_full_path.is_file(): lidars.append( @@ -426,11 +430,11 @@ def _extract_nuplan_lidars( relative_path=nuplan_lidar_pc.filename, ) ) - return lidars def _extract_nuplan_scenario_tag(nuplan_log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: + """Extracts the nuPlan scenario tags from a given LidarPc database objects.""" scenario_tags = [ scenario_tag.type for scenario_tag in nuplan_log_db.scenario_tag.select_many(lidar_pc_token=lidar_pc_token) ] @@ -440,6 +444,7 @@ def _extract_nuplan_scenario_tag(nuplan_log_db: NuPlanDB, lidar_pc_token: str) - def _extract_nuplan_route_lane_group_ids(nuplan_lidar_pc: LidarPc) -> List[int]: + """Extracts the nuPlan route lane group IDs from a given LidarPc database objects.""" return [ int(roadblock_id) for roadblock_id in str(nuplan_lidar_pc.scene.roadblock_ids).split(" ") diff --git a/src/py123d/conversion/datasets/nuplan/nuplan_map_conversion.py b/src/py123d/conversion/datasets/nuplan/nuplan_map_conversion.py index b8b010cb..a27f42dc 100644 --- a/src/py123d/conversion/datasets/nuplan/nuplan_map_conversion.py +++ b/src/py123d/conversion/datasets/nuplan/nuplan_map_conversion.py @@ -7,6 +7,7 @@ import pyogrio from shapely import LineString +from py123d.api.map.gpkg.gpkg_utils import get_all_rows_with_value, get_row_with_value from py123d.conversion.datasets.nuplan.utils.nuplan_constants import ( NUPLAN_MAP_GPKG_LAYERS, NUPLAN_MAP_LOCATION_FILES, @@ -17,25 +18,25 @@ get_road_edge_linear_rings, split_line_geometry_by_max_length, ) -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCarpark, - CacheCrosswalk, - CacheGenericDrivable, - CacheIntersection, - CacheLane, - CacheLaneGroup, - CacheRoadEdge, - CacheRoadLine, - CacheWalkway, +from py123d.datatypes.map_objects.map_layer_types import RoadEdgeType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + Walkway, ) -from py123d.datatypes.maps.gpkg.gpkg_utils import get_all_rows_with_value, get_row_with_value -from py123d.datatypes.maps.map_datatypes import RoadEdgeType -from py123d.geometry.polyline import Polyline2D, Polyline3D +from py123d.geometry import Polyline2D, Polyline3D -MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # meters, used to filter out very long road edges. TODO @add to config? +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # meters, used to filter out very long road edges. TODO add to config? def write_nuplan_map(nuplan_maps_root: Path, location: str, map_writer: AbstractMapWriter) -> None: + """Convert nuPlan map data using the provided map writer.""" assert location in NUPLAN_MAP_LOCATION_FILES.keys(), f"Map name {location} is not supported." source_map_path = nuplan_maps_root / NUPLAN_MAP_LOCATION_FILES[location] assert source_map_path.exists(), f"Map file {source_map_path} does not exist." @@ -55,6 +56,7 @@ def write_nuplan_map(nuplan_maps_root: Path, location: str, map_writer: Abstract def _write_nuplan_lanes(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan lanes to the map writer.""" # NOTE: drops: lane_index (?), creator_id, name (?), road_type_fid (?), lane_type_fid (?), width (?), # left_offset (?), right_offset (?), min_speed (?), max_speed (?), stops, left_has_reflectors (?), @@ -66,7 +68,6 @@ def _write_nuplan_lanes(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: Abs all_geometries = nuplan_gdf["lanes_polygons"].geometry.to_list() for idx, lane_id in enumerate(all_ids): - # 1. predecessor_ids, successor_ids predecessor_ids = get_all_rows_with_value( nuplan_gdf["lane_connectors"], @@ -105,7 +106,7 @@ def _write_nuplan_lanes(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: Abs right_boundary = align_boundary_direction(centerline, right_boundary) map_writer.write_lane( - CacheLane( + Lane( object_id=lane_id, lane_group_id=all_lane_group_ids[idx], left_boundary=Polyline3D.from_linestring(left_boundary), @@ -117,12 +118,13 @@ def _write_nuplan_lanes(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: Abs successor_ids=successor_ids, speed_limit_mps=all_speed_limits_mps[idx], outline=None, - geometry=all_geometries[idx], + shapely_polygon=all_geometries[idx], ) ) def _write_nuplan_lane_connectors(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan lane connectors (lanes on intersections) to the map writer.""" # NOTE: drops: exit_lane_group_fid, entry_lane_group_fid, to_edge_fid, # turn_type_fid (?), bulb_fids (?), traffic_light_stop_line_fids (?), overlap (?), creator_id @@ -132,7 +134,6 @@ def _write_nuplan_lane_connectors(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_w all_speed_limits_mps = nuplan_gdf["lane_connectors"].speed_limit_mps.to_list() for idx, lane_id in enumerate(all_ids): - # 1. predecessor_ids, successor_ids lane_connector_row = get_row_with_value(nuplan_gdf["lane_connectors"], "fid", str(lane_id)) predecessor_ids = [lane_connector_row["entry_lane_fid"]] @@ -158,7 +159,7 @@ def _write_nuplan_lane_connectors(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_w # geometries.append(lane_connector_polygons_row.geometry) map_writer.write_lane( - CacheLane( + Lane( object_id=lane_id, lane_group_id=all_lane_group_ids[idx], left_boundary=Polyline3D.from_linestring(left_boundary), @@ -170,18 +171,19 @@ def _write_nuplan_lane_connectors(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_w successor_ids=successor_ids, speed_limit_mps=all_speed_limits_mps[idx], outline=None, - geometry=lane_connector_polygons_row.geometry, + shapely_polygon=lane_connector_polygons_row.geometry, ) ) def _write_nuplan_lane_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan lane groups to the map writer.""" + # NOTE: drops: creator_id, from_edge_fid, to_edge_fid ids = nuplan_gdf["lane_groups_polygons"].fid.to_list() # all_geometries = nuplan_gdf["lane_groups_polygons"].geometry.to_list() for lane_group_id in ids: - # 1. lane_ids lane_ids = get_all_rows_with_value( nuplan_gdf["lanes_polygons"], @@ -216,7 +218,7 @@ def _write_nuplan_lane_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_write right_boundary = align_boundary_direction(repr_centerline, right_boundary) map_writer.write_lane_group( - CacheLaneGroup( + LaneGroup( object_id=lane_group_id, lane_ids=lane_ids, left_boundary=Polyline3D.from_linestring(left_boundary), @@ -225,19 +227,20 @@ def _write_nuplan_lane_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_write predecessor_ids=predecessor_lane_group_ids, successor_ids=successor_lane_group_ids, outline=None, - geometry=lane_group_row.geometry, + shapely_polygon=lane_group_row.geometry, ) ) def _write_nuplan_lane_connector_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan lane connector groups (lane groups on intersections) to the map writer.""" + # NOTE: drops: creator_id, from_edge_fid, to_edge_fid, intersection_fid ids = nuplan_gdf["lane_group_connectors"].fid.to_list() all_intersection_ids = nuplan_gdf["lane_group_connectors"].intersection_fid.to_list() # all_geometries = nuplan_gdf["lane_group_connectors"].geometry.to_list() for idx, lane_group_connector_id in enumerate(ids): - # 1. lane_ids lane_ids = get_all_rows_with_value( nuplan_gdf["lane_connectors"], "lane_group_connector_fid", lane_group_connector_id @@ -257,7 +260,7 @@ def _write_nuplan_lane_connector_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], right_boundary = get_row_with_value(nuplan_gdf["boundaries"], "fid", str(right_boundary_fid))["geometry"] map_writer.write_lane_group( - CacheLaneGroup( + LaneGroup( object_id=lane_group_connector_id, lane_ids=lane_ids, left_boundary=Polyline3D.from_linestring(left_boundary), @@ -266,12 +269,13 @@ def _write_nuplan_lane_connector_groups(nuplan_gdf: Dict[str, gpd.GeoDataFrame], predecessor_ids=predecessor_lane_group_ids, successor_ids=successor_lane_group_ids, outline=None, - geometry=lane_group_connector_row.geometry, + shapely_polygon=lane_group_connector_row.geometry, ) ) def _write_nuplan_intersections(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan intersections to the map writer.""" # NOTE: drops: creator_id, intersection_type_fid (?), is_mini (?) all_ids = nuplan_gdf["intersections"].fid.to_list() all_geometries = nuplan_gdf["intersections"].geometry.to_list() @@ -281,41 +285,46 @@ def _write_nuplan_intersections(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_wri )["fid"].tolist() map_writer.write_intersection( - CacheIntersection( + Intersection( object_id=intersection_id, lane_group_ids=lane_group_connector_ids, - geometry=all_geometries[idx], + shapely_polygon=all_geometries[idx], ) ) def _write_nuplan_crosswalks(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan crosswalks to the map writer.""" # NOTE: drops: creator_id, intersection_fids, lane_fids, is_marked (?) for id, geometry in zip(nuplan_gdf["crosswalks"].fid.to_list(), nuplan_gdf["crosswalks"].geometry.to_list()): - map_writer.write_crosswalk(CacheCrosswalk(object_id=id, geometry=geometry)) + map_writer.write_crosswalk(Crosswalk(object_id=id, shapely_polygon=geometry)) def _write_nuplan_walkways(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan walkways to the map writer.""" # NOTE: drops: creator_id for id, geometry in zip(nuplan_gdf["walkways"].fid.to_list(), nuplan_gdf["walkways"].geometry.to_list()): - map_writer.write_walkway(CacheWalkway(object_id=id, geometry=geometry)) + map_writer.write_walkway(Walkway(object_id=id, shapely_polygon=geometry)) def _write_nuplan_carparks(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan carparks to the map writer.""" # NOTE: drops: creator_id for id, geometry in zip(nuplan_gdf["carpark_areas"].fid.to_list(), nuplan_gdf["carpark_areas"].geometry.to_list()): - map_writer.write_carpark(CacheCarpark(object_id=id, geometry=geometry)) + map_writer.write_carpark(Carpark(object_id=id, shapely_polygon=geometry)) def _write_nuplan_generic_drivables(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan generic drivable areas to the map writer.""" # NOTE: drops: creator_id for id, geometry in zip( nuplan_gdf["generic_drivable_areas"].fid.to_list(), nuplan_gdf["generic_drivable_areas"].geometry.to_list() ): - map_writer.write_generic_drivable(CacheGenericDrivable(object_id=id, geometry=geometry)) + map_writer.write_generic_drivable(GenericDrivable(object_id=id, shapely_polygon=geometry)) def _write_nuplan_road_edges(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan road edges to the map writer.""" drivable_polygons = ( nuplan_gdf["intersections"].geometry.to_list() + nuplan_gdf["lane_groups_polygons"].geometry.to_list() @@ -327,7 +336,7 @@ def _write_nuplan_road_edges(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer for idx in range(len(road_edges)): map_writer.write_road_edge( - CacheRoadEdge( + RoadEdge( object_id=idx, road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, polyline=Polyline2D.from_linestring(road_edges[idx]), @@ -336,13 +345,14 @@ def _write_nuplan_road_edges(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer def _write_nuplan_road_lines(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer: AbstractMapWriter) -> None: + """Write nuPlan road lines to the map writer.""" boundaries = nuplan_gdf["boundaries"].geometry.to_list() fids = nuplan_gdf["boundaries"].fid.to_list() boundary_types = nuplan_gdf["boundaries"].boundary_type_fid.to_list() for idx in range(len(boundary_types)): map_writer.write_road_line( - CacheRoadLine( + RoadLine( object_id=fids[idx], road_line_type=NUPLAN_ROAD_LINE_CONVERSION[boundary_types[idx]], polyline=Polyline2D.from_linestring(boundaries[idx]), @@ -351,6 +361,7 @@ def _write_nuplan_road_lines(nuplan_gdf: Dict[str, gpd.GeoDataFrame], map_writer def _load_nuplan_gdf(map_file_path: Path) -> Dict[str, gpd.GeoDataFrame]: + """Load nuPlan map data from a GPKG file into a dictionary of GeoDataFrames.""" # The projected coordinate system depends on which UTM zone the mapped location is in. map_meta = gpd.read_file(map_file_path, layer="meta", engine="pyogrio") @@ -377,11 +388,14 @@ def _load_nuplan_gdf(map_file_path: Path) -> Dict[str, gpd.GeoDataFrame]: def _flip_linestring(linestring: LineString) -> LineString: + """Flips the direction of a shapely LineString.""" # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. return LineString(linestring.coords[::-1]) def lines_same_direction(centerline: LineString, boundary: LineString) -> bool: + """Check if the boundary LineString is in the same direction as the centerline LineString.""" + # TODO: refactor helper function. center_start = np.array(centerline.coords[0]) center_end = np.array(centerline.coords[-1]) @@ -396,6 +410,7 @@ def lines_same_direction(centerline: LineString, boundary: LineString) -> bool: def align_boundary_direction(centerline: LineString, boundary: LineString) -> LineString: + """Aligns the boundary LineString direction to be the same as the centerline LineString direction.""" # TODO: refactor helper function. if not lines_same_direction(centerline, boundary): return _flip_linestring(boundary) diff --git a/src/py123d/conversion/datasets/nuplan/nuplan_sensor_io.py b/src/py123d/conversion/datasets/nuplan/nuplan_sensor_io.py index 8c2506f0..5cd1ffbc 100644 --- a/src/py123d/conversion/datasets/nuplan/nuplan_sensor_io.py +++ b/src/py123d/conversion/datasets/nuplan/nuplan_sensor_io.py @@ -6,14 +6,16 @@ from py123d.common.utils.dependencies import check_dependencies from py123d.conversion.datasets.nuplan.utils.nuplan_constants import NUPLAN_LIDAR_DICT -from py123d.conversion.registry.lidar_index_registry import NuPlanLiDARIndex -from py123d.datatypes.sensors.lidar import LiDARType +from py123d.conversion.registry import NuPlanLiDARIndex +from py123d.datatypes.sensors import LiDARType check_dependencies(["nuplan"], "nuplan") from nuplan.database.utils.pointclouds.lidar import LidarPointCloud def load_nuplan_lidar_pcs_from_file(pcd_path: Path) -> Dict[LiDARType, np.ndarray]: + """Loads nuPlan LiDAR point clouds from a ``.pcd`` file.""" + assert pcd_path.exists(), f"LiDAR file not found: {pcd_path}" with open(pcd_path, "rb") as fp: buffer = io.BytesIO(fp.read()) diff --git a/src/py123d/conversion/datasets/nuplan/utils/nuplan_constants.py b/src/py123d/conversion/datasets/nuplan/utils/nuplan_constants.py index d904a698..fd655ff5 100644 --- a/src/py123d/conversion/datasets/nuplan/utils/nuplan_constants.py +++ b/src/py123d/conversion/datasets/nuplan/utils/nuplan_constants.py @@ -1,10 +1,10 @@ from typing import Dict, Final, List, Set -from py123d.conversion.registry.box_detection_label_registry import NuPlanBoxDetectionLabel -from py123d.datatypes.detections.traffic_light_detections import TrafficLightStatus -from py123d.datatypes.maps.map_datatypes import RoadLineType -from py123d.datatypes.sensors.lidar import LiDARType -from py123d.datatypes.time.time_point import TimePoint +from py123d.conversion.registry import NuPlanBoxDetectionLabel +from py123d.datatypes.detections import TrafficLightStatus +from py123d.datatypes.map_objects import RoadLineType +from py123d.datatypes.sensors import LiDARType +from py123d.datatypes.time import TimePoint NUPLAN_DEFAULT_DT: Final[float] = 0.05 diff --git a/src/py123d/conversion/datasets/nuplan/utils/nuplan_sql_helper.py b/src/py123d/conversion/datasets/nuplan/utils/nuplan_sql_helper.py index 866246ba..62c6128f 100644 --- a/src/py123d/conversion/datasets/nuplan/utils/nuplan_sql_helper.py +++ b/src/py123d/conversion/datasets/nuplan/utils/nuplan_sql_helper.py @@ -3,7 +3,7 @@ from py123d.common.utils.dependencies import check_dependencies from py123d.conversion.datasets.nuplan.utils.nuplan_constants import NUPLAN_DETECTION_NAME_DICT from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3 -from py123d.geometry import BoundingBoxSE3, EulerAngles, StateSE3, Vector3D +from py123d.geometry import BoundingBoxSE3, EulerAngles, PoseSE3, Vector3D from py123d.geometry.utils.constants import DEFAULT_PITCH, DEFAULT_ROLL check_dependencies(modules=["nuplan"], optional_name="nuplan") @@ -11,6 +11,7 @@ def get_box_detections_for_lidarpc_token_from_db(log_file: str, token: str) -> List[BoxDetectionSE3]: + """Gets the box detections for a given LiDAR point cloud token from the NuPlan database.""" query = """ SELECT c.name AS category_name, @@ -42,7 +43,7 @@ def get_box_detections_for_lidarpc_token_from_db(log_file: str, token: str) -> L for row in execute_many(query, (bytearray.fromhex(token),), log_file): quaternion = EulerAngles(roll=DEFAULT_ROLL, pitch=DEFAULT_PITCH, yaw=row["yaw"]).quaternion bounding_box = BoundingBoxSE3( - center=StateSE3( + center_se3=PoseSE3( x=row["x"], y=row["y"], z=row["z"], @@ -61,14 +62,15 @@ def get_box_detections_for_lidarpc_token_from_db(log_file: str, token: str) -> L track_token=row["track_token"].hex(), ), bounding_box_se3=bounding_box, - velocity=Vector3D(x=row["vx"], y=row["vy"], z=row["vz"]), + velocity_3d=Vector3D(x=row["vx"], y=row["vy"], z=row["vz"]), ) box_detections.append(box_detection) return box_detections -def get_ego_pose_for_timestamp_from_db(log_file: str, timestamp: int) -> StateSE3: +def get_ego_pose_for_timestamp_from_db(log_file: str, timestamp: int) -> PoseSE3: + """Gets the ego pose for a given timestamp from the NuPlan database.""" query = """ SELECT ep.x, @@ -89,10 +91,8 @@ def get_ego_pose_for_timestamp_from_db(log_file: str, timestamp: int) -> StateSE """ row = execute_one(query, (timestamp,), log_file) - if row is None: - return None - - return StateSE3(x=row["x"], y=row["y"], z=row["z"], qw=row["qw"], qx=row["qx"], qy=row["qy"], qz=row["qz"]) + assert row is not None, f"No ego pose found for timestamp {timestamp} in log file {log_file}" + return PoseSE3(x=row["x"], y=row["y"], z=row["z"], qw=row["qw"], qx=row["qx"], qy=row["qy"], qz=row["qz"]) def get_nearest_ego_pose_for_timestamp_from_db( @@ -101,7 +101,8 @@ def get_nearest_ego_pose_for_timestamp_from_db( tokens: List[str], lookahead_window_us: int = 50000, lookback_window_us: int = 50000, -) -> StateSE3: +) -> PoseSE3: + """Gets the nearest ego pose for a given timestamp from the NuPlan database within a lookahead and lookback window.""" query = f""" SELECT ep.x, @@ -115,7 +116,7 @@ def get_nearest_ego_pose_for_timestamp_from_db( INNER JOIN lidar_pc AS lpc ON ep.timestamp <= lpc.timestamp + ? AND ep.timestamp >= lpc.timestamp - ? - WHERE lpc.token IN ({('?,'*len(tokens))[:-1]}) + WHERE lpc.token IN ({("?," * len(tokens))[:-1]}) ORDER BY ABS(ep.timestamp - ?) LIMIT 1 """ # noqa: E226 @@ -125,4 +126,4 @@ def get_nearest_ego_pose_for_timestamp_from_db( args += [timestamp] for row in execute_many(query, args, log_file): - return StateSE3(x=row["x"], y=row["y"], z=row["z"], qw=row["qw"], qx=row["qx"], qy=row["qy"], qz=row["qz"]) + return PoseSE3(x=row["x"], y=row["y"], z=row["z"], qw=row["qw"], qx=row["qx"], qy=row["qy"], qz=row["qz"]) diff --git a/src/py123d/conversion/datasets/nuscenes/__init__.py b/src/py123d/conversion/datasets/nuscenes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py index 7cc39d34..33706bf6 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py @@ -1,6 +1,6 @@ import gc from pathlib import Path -from typing import Any, Dict, List, Union +from typing import Any, Dict, List, Optional, Union import numpy as np from pyquaternion import Quaternion @@ -8,20 +8,33 @@ from py123d.common.utils.dependencies import check_dependencies from py123d.conversion.abstract_dataset_converter import AbstractDatasetConverter from py123d.conversion.dataset_converter_config import DatasetConverterConfig -from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import NUSCENES_MAPS, write_nuscenes_map +from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import ( + NUSCENES_MAPS, + write_nuscenes_map, +) from py123d.conversion.datasets.nuscenes.utils.nuscenes_constants import ( NUSCENES_CAMERA_TYPES, NUSCENES_DATA_SPLITS, + NUSCENES_DATABASE_VERSION_MAPPING, NUSCENES_DETECTION_NAME_DICT, NUSCENES_DT, ) -from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData +from py123d.conversion.log_writer.abstract_log_writer import ( + AbstractLogWriter, + CameraData, + LiDARData, +) from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.registry.box_detection_label_registry import NuScenesBoxDetectionLabel +from py123d.conversion.registry.box_detection_label_registry import ( + NuScenesBoxDetectionLabel, +) from py123d.conversion.registry.lidar_index_registry import NuScenesLiDARIndex -from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.datatypes.detections.box_detections import ( + BoxDetectionMetadata, + BoxDetectionSE3, + BoxDetectionWrapper, +) +from py123d.datatypes.metadata import LogMetadata, MapMetadata from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType from py123d.datatypes.sensors.pinhole_camera import ( PinholeCameraMetadata, @@ -31,8 +44,10 @@ ) from py123d.datatypes.time.time_point import TimePoint from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuscenes_renault_zoe_parameters -from py123d.geometry import BoundingBoxSE3, StateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import ( + get_nuscenes_renault_zoe_parameters, +) +from py123d.geometry import BoundingBoxSE3, PoseSE3 from py123d.geometry.vector import Vector3D check_dependencies(["nuscenes"], "nuscenes") @@ -43,6 +58,8 @@ class NuScenesConverter(AbstractDatasetConverter): + """Dataset converter for the nuScenes dataset.""" + def __init__( self, splits: List[str], @@ -51,21 +68,24 @@ def __init__( nuscenes_lanelet2_root: Union[Path, str], use_lanelet2: bool, dataset_converter_config: DatasetConverterConfig, - version: str = "v1.0-mini", + nuscenes_dbs: Optional[Dict[str, NuScenes]] = None, ) -> None: + """Initializes the :class:`NuScenesConverter`. + + :param splits: List of splits to include in the conversion, e.g., ["nuscenes_train", "nuscenes_val"] + :param nuscenes_data_root: Path to the root directory of the nuScenes dataset + :param nuscenes_map_root: Path to the root directory of the nuScenes map data + :param nuscenes_lanelet2_root: Path to the root directory of the nuScenes Lanelet2 data + :param use_lanelet2: Whether to use Lanelet2 data for map conversion + :param dataset_converter_config: Configuration for the dataset converter + """ super().__init__(dataset_converter_config) assert nuscenes_data_root is not None, "The variable `nuscenes_data_root` must be provided." assert nuscenes_map_root is not None, "The variable `nuscenes_map_root` must be provided." for split in splits: - assert ( - split in NUSCENES_DATA_SPLITS - ), f"Split {split} is not available. Available splits: {NUSCENES_DATA_SPLITS}" - - if dataset_converter_config.include_lidars: - assert dataset_converter_config.lidar_store_option in ["path", "binary"], ( - f"Invalid lidar_store_option: {dataset_converter_config.lidar_store_option}. " - f"Supported options are 'path' and 'binary'." + assert split in NUSCENES_DATA_SPLITS, ( + f"Split {split} is not available. Available splits: {NUSCENES_DATA_SPLITS}" ) self._splits: List[str] = splits @@ -73,16 +93,30 @@ def __init__( self._nuscenes_data_root: Path = Path(nuscenes_data_root) self._nuscenes_map_root: Path = Path(nuscenes_map_root) self._nuscenes_lanelet2_root: Path = Path(nuscenes_lanelet2_root) - self._use_lanelet2 = use_lanelet2 - self._version = version + + self._nuscenes_dbs: Dict[str, NuScenes] = nuscenes_dbs if nuscenes_dbs is not None else {} self._scene_tokens_per_split: Dict[str, List[str]] = self._collect_scene_tokens() + def __reduce__(self): + return ( + self.__class__, + ( + self._splits, + self._nuscenes_data_root, + self._nuscenes_map_root, + self._nuscenes_lanelet2_root, + self._use_lanelet2, + self.dataset_converter_config, + self._nuscenes_dbs, + ), + ) + def _collect_scene_tokens(self) -> Dict[str, List[str]]: + """Collects scene tokens for the specified splits.""" scene_tokens_per_split: Dict[str, List[str]] = {} - nusc = NuScenes(version=self._version, dataroot=str(self._nuscenes_data_root), verbose=False) - + # Conversion from nuScenes internal split names to our split names nuscenes_split_name_mapping = { "nuscenes_train": "train", "nuscenes_val": "val", @@ -91,17 +125,28 @@ def _collect_scene_tokens(self) -> Dict[str, List[str]]: "nuscenes-mini_val": "mini_val", } + # Loads the mapping from split names to scene names in nuScenes scene_splits = create_splits_scenes() - available_scenes = [scene for scene in nusc.scene] + # Iterate over split names, for split in self._splits: + database_version = NUSCENES_DATABASE_VERSION_MAPPING[split] + nusc = self._nuscenes_dbs.get(database_version) + if nusc is None: + nusc = NuScenes( + version=database_version, + dataroot=str(self._nuscenes_data_root), + verbose=False, + ) + self._nuscenes_dbs[database_version] = nusc + + available_scenes = [scene for scene in nusc.scene] nuscenes_split = nuscenes_split_name_mapping[split] scene_names = scene_splits.get(nuscenes_split, []) # get token scene_tokens = [scene["token"] for scene in available_scenes if scene["name"] in scene_names] scene_tokens_per_split[split] = scene_tokens - return scene_tokens_per_split def get_number_of_maps(self) -> int: @@ -142,7 +187,8 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: split, scene_token = all_scene_tokens[log_index] - nusc = NuScenes(version=self._version, dataroot=str(self._nuscenes_data_root), verbose=False) + database_version = NUSCENES_DATABASE_VERSION_MAPPING[split] + nusc = self._nuscenes_dbs[database_version] scene = nusc.get("scene", scene_token) log_record = nusc.get("log", scene["log_token"]) @@ -206,21 +252,18 @@ def _get_nuscenes_pinhole_camera_metadata( scene: Dict[str, Any], dataset_converter_config: DatasetConverterConfig, ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + """Extracts the pinhole camera metadata from a nuScenes scene.""" camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = {} - if dataset_converter_config.include_pinhole_cameras: first_sample_token = scene["first_sample_token"] first_sample = nusc.get("sample", first_sample_token) - for camera_type, camera_channel in NUSCENES_CAMERA_TYPES.items(): cam_token = first_sample["data"][camera_channel] cam_data = nusc.get("sample_data", cam_token) calib = nusc.get("calibrated_sensor", cam_data["calibrated_sensor_token"]) - intrinsic_matrix = np.array(calib["camera_intrinsic"]) intrinsic = PinholeIntrinsics.from_camera_matrix(intrinsic_matrix) distortion = PinholeDistortion.from_array(np.zeros(5), copy=False) - camera_metadata[camera_type] = PinholeCameraMetadata( camera_type=camera_type, width=cam_data["width"], @@ -228,7 +271,6 @@ def _get_nuscenes_pinhole_camera_metadata( intrinsics=intrinsic, distortion=distortion, ) - return camera_metadata @@ -237,32 +279,30 @@ def _get_nuscenes_lidar_metadata( scene: Dict[str, Any], dataset_converter_config: DatasetConverterConfig, ) -> Dict[LiDARType, LiDARMetadata]: + """Extracts the LiDAR metadata from a nuScenes scene.""" metadata: Dict[LiDARType, LiDARMetadata] = {} - if dataset_converter_config.include_lidars: first_sample_token = scene["first_sample_token"] first_sample = nusc.get("sample", first_sample_token) lidar_token = first_sample["data"]["LIDAR_TOP"] lidar_data = nusc.get("sample_data", lidar_token) calib = nusc.get("calibrated_sensor", lidar_data["calibrated_sensor_token"]) - translation = np.array(calib["translation"]) rotation = Quaternion(calib["rotation"]).rotation_matrix extrinsic = np.eye(4) extrinsic[:3, :3] = rotation extrinsic[:3, 3] = translation - extrinsic = StateSE3.from_transformation_matrix(extrinsic) - + extrinsic = PoseSE3.from_transformation_matrix(extrinsic) metadata[LiDARType.LIDAR_TOP] = LiDARMetadata( lidar_type=LiDARType.LIDAR_TOP, lidar_index=NuScenesLiDARIndex, extrinsic=extrinsic, ) - return metadata def _get_nuscenes_map_metadata(location): + """Creates nuScenes map metadata for a given location.""" return MapMetadata( dataset="nuscenes", split=None, @@ -274,13 +314,12 @@ def _get_nuscenes_map_metadata(location): def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: + """Extracts the ego state from a nuScenes sample.""" lidar_data = nusc.get("sample_data", sample["data"]["LIDAR_TOP"]) ego_pose = nusc.get("ego_pose", lidar_data["ego_pose_token"]) - quat = Quaternion(ego_pose["rotation"]) - vehicle_parameters = get_nuscenes_renault_zoe_parameters() - imu_pose = StateSE3( + imu_pose = PoseSE3( x=ego_pose["translation"][0], y=ego_pose["translation"][1], z=ego_pose["translation"][2], @@ -289,14 +328,11 @@ def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: qy=quat.y, qz=quat.z, ) - scene_name = nusc.get("scene", sample["scene_token"])["name"] - try: pose_msgs = can_bus.get_messages(scene_name, "pose") except Exception: pose_msgs = [] - if pose_msgs: closest_msg = None min_time_diff = float("inf") @@ -320,34 +356,23 @@ def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: acceleration=Vector3D(*acceleration), angular_velocity=Vector3D(*angular_velocity), ) - - # return EgoStateSE3( - # center_se3=pose, - # dynamic_state_se3=dynamic_state, - # vehicle_parameters=vehicle_parameters, - # timepoint=TimePoint.from_us(sample["timestamp"]), - # ) return EgoStateSE3.from_rear_axle( rear_axle_se3=imu_pose, dynamic_state_se3=dynamic_state, vehicle_parameters=vehicle_parameters, - time_point=TimePoint.from_us(sample["timestamp"]), ) def _extract_nuscenes_box_detections(nusc: NuScenes, sample: Dict[str, Any]) -> BoxDetectionWrapper: + """Extracts the box detections from a nuScenes sample.""" box_detections: List[BoxDetectionSE3] = [] - for ann_token in sample["anns"]: ann = nusc.get("sample_annotation", ann_token) box = Box(ann["translation"], ann["size"], Quaternion(ann["rotation"])) - box_quat = box.orientation - euler_angles = box_quat.yaw_pitch_roll # (yaw, pitch, roll) - - # Create StateSE3 for box center and orientation + # Create PoseSE3 for box center and orientation center_quat = box.orientation - center = StateSE3( + center = PoseSE3( box.center[0], box.center[1], box.center[2], @@ -356,7 +381,12 @@ def _extract_nuscenes_box_detections(nusc: NuScenes, sample: Dict[str, Any]) -> center_quat.y, center_quat.z, ) - bounding_box = BoundingBoxSE3(center, box.wlh[1], box.wlh[0], box.wlh[2]) + bounding_box = BoundingBoxSE3( + center_se3=center, + length=box.wlh[1], + width=box.wlh[0], + height=box.wlh[2], + ) # Get detection type category = ann["category_name"] label = NUSCENES_DETECTION_NAME_DICT[category] @@ -369,17 +399,14 @@ def _extract_nuscenes_box_detections(nusc: NuScenes, sample: Dict[str, Any]) -> label=label, track_token=ann["instance_token"], timepoint=TimePoint.from_us(sample["timestamp"]), - confidence=1.0, # nuScenes annotations are ground truth num_lidar_points=ann.get("num_lidar_pts", 0), ) - box_detection = BoxDetectionSE3( metadata=metadata, bounding_box_se3=bounding_box, - velocity=velocity_3d, + velocity_3d=velocity_3d, ) box_detections.append(box_detection) - return BoxDetectionWrapper(box_detections=box_detections) @@ -389,8 +416,8 @@ def _extract_nuscenes_cameras( nuscenes_data_root: Path, dataset_converter_config: DatasetConverterConfig, ) -> List[CameraData]: + """Extracts the pinhole camera metadata from a nuScenes scene.""" camera_data_list: List[CameraData] = [] - if dataset_converter_config.include_pinhole_cameras: for camera_type, camera_channel in NUSCENES_CAMERA_TYPES.items(): cam_token = sample["data"][camera_channel] @@ -407,12 +434,10 @@ def _extract_nuscenes_cameras( extrinsic_matrix = np.eye(4) extrinsic_matrix[:3, :3] = rotation extrinsic_matrix[:3, 3] = translation - extrinsic = StateSE3.from_transformation_matrix(extrinsic_matrix) + extrinsic = PoseSE3.from_transformation_matrix(extrinsic_matrix) cam_path = nuscenes_data_root / str(cam_data["filename"]) - if cam_path.exists() and cam_path.is_file(): - # camera_dict[camera_type] = (camera_data, extrinsic) camera_data_list.append( CameraData( camera_type=camera_type, @@ -431,13 +456,12 @@ def _extract_nuscenes_lidars( nuscenes_data_root: Path, dataset_converter_config: DatasetConverterConfig, ) -> List[LiDARData]: + """Extracts the LiDAR data from a nuScenes sample.""" lidars: List[LiDARData] = [] - if dataset_converter_config.include_lidars: lidar_token = sample["data"]["LIDAR_TOP"] lidar_data = nusc.get("sample_data", lidar_token) absolute_lidar_path = nuscenes_data_root / lidar_data["filename"] - if absolute_lidar_path.exists() and absolute_lidar_path.is_file(): lidar = LiDARData( lidar_type=LiDARType.LIDAR_TOP, diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py index 39ae313a..ce17b87a 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py @@ -18,18 +18,18 @@ split_line_geometry_by_max_length, split_polygon_by_grid, ) -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCarpark, - CacheCrosswalk, - CacheGenericDrivable, - CacheIntersection, - CacheLane, - CacheLaneGroup, - CacheRoadEdge, - CacheRoadLine, - CacheWalkway, +from py123d.datatypes.map_objects.map_layer_types import RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + Walkway, ) -from py123d.datatypes.maps.map_datatypes import RoadEdgeType, RoadLineType from py123d.geometry import OccupancyMap2D, Polyline2D, Polyline3D from py123d.geometry.utils.polyline_utils import offset_points_perpendicular @@ -81,7 +81,7 @@ def write_nuscenes_map( _write_nuscenes_walkways(nuscenes_map, map_writer) _write_nuscenes_carparks(nuscenes_map, map_writer) _write_nuscenes_generic_drivables(nuscenes_map, map_writer) - _write_nuscenes_stop_lines(nuscenes_map, map_writer) + _write_nuscenes_stop_zones(nuscenes_map, map_writer) _write_nuscenes_road_lines(nuscenes_map, map_writer) for lane in lanes + lane_connectors: @@ -94,7 +94,7 @@ def write_nuscenes_map( map_writer.write_lane_group(lane_group) -def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[CacheLane]: +def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[Lane]: """Helper function to extract lanes from a nuScenes map.""" # NOTE: nuScenes does not provide explicitly provide lane groups and does not assign lanes to roadblocks. @@ -114,7 +114,7 @@ def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[CacheLane]: } road_block_map = OccupancyMap2D.from_dict(road_block_dict) - lanes: List[CacheLane] = [] + lanes: List[Lane] = [] for lane_record in nuscenes_map.lane: token = lane_record["token"] @@ -140,7 +140,7 @@ def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[CacheLane]: outgoing = nuscenes_map.get_outgoing_lane_ids(token) lanes.append( - CacheLane( + Lane( object_id=token, lane_group_id=lane_group_id, left_boundary=left_boundary, @@ -152,21 +152,21 @@ def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[CacheLane]: successor_ids=outgoing, speed_limit_mps=None, outline=None, - geometry=None, + shapely_polygon=None, ) ) return lanes -def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: List[CacheRoadEdge]) -> List[CacheLane]: +def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: List[RoadEdge]) -> List[Lane]: """Helper function to extract lane connectors from a nuScenes map.""" # TODO @DanielDauner: consider using connected lanes to estimate the lane width road_edge_map = OccupancyMap2D(geometries=[road_edge.shapely_linestring for road_edge in road_edges]) - lane_connectors: List[CacheLane] = [] + lane_connectors: List[Lane] = [] for lane_record in nuscenes_map.lane_connector: lane_connector_token: str = lane_record["token"] @@ -188,7 +188,7 @@ def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: Lis lane_group_id = lane_connector_token lane_connectors.append( - CacheLane( + Lane( object_id=lane_connector_token, lane_group_id=lane_group_id, left_boundary=Polyline2D.from_array(left_pts), @@ -200,7 +200,7 @@ def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: Lis successor_ids=successor_ids, speed_limit_mps=None, # Default value outline=None, - geometry=None, + shapely_polygon=None, ) ) @@ -208,11 +208,8 @@ def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: Lis def _extract_nuscenes_lane_groups( - nuscenes_map: NuScenesMap, - lanes: List[CacheLane], - lane_connectors: List[CacheLane], - intersection_assignment: Dict[str, int], -) -> List[CacheLaneGroup]: + nuscenes_map: NuScenesMap, lanes: List[Lane], lane_connectors: List[Lane], intersection_assignment: Dict[str, int] +) -> List[LaneGroup]: """Helper function to extract lane groups from a nuScenes map.""" lane_groups = [] @@ -224,7 +221,6 @@ def _extract_nuscenes_lane_groups( lane_group_lane_dict[lane.lane_group_id].append(lane.object_id) for lane_group_id, lane_ids in lane_group_lane_dict.items(): - if len(lane_ids) > 1: lane_centerlines: List[Polyline2D] = [lanes_dict[lane_id].centerline for lane_id in lane_ids] ordered_lane_indices = order_lanes_left_to_right(lane_centerlines) @@ -260,7 +256,7 @@ def _extract_nuscenes_lane_groups( intersection_id = None if len(intersection_ids) == 0 else intersection_ids.pop() lane_groups.append( - CacheLaneGroup( + LaneGroup( object_id=lane_group_id, lane_ids=lane_ids, left_boundary=left_boundary, @@ -269,7 +265,7 @@ def _extract_nuscenes_lane_groups( predecessor_ids=list(predecessor_ids), successor_ids=list(successor_ids), outline=None, - geometry=None, + shapely_polygon=None, ) ) @@ -277,7 +273,7 @@ def _extract_nuscenes_lane_groups( def _write_nuscenes_intersections( - nuscenes_map: NuScenesMap, lane_connectors: List[CacheLane], map_writer: AbstractMapWriter + nuscenes_map: NuScenesMap, lane_connectors: List[Lane], map_writer: AbstractMapWriter ) -> None: """Write intersection data to map_writer and return lane-connector to intersection assignment.""" @@ -303,11 +299,11 @@ def _write_nuscenes_intersections( intersection_assignment[lane_connector_id] = idx map_writer.write_intersection( - CacheIntersection( + Intersection( object_id=idx, lane_group_ids=intersecting_lane_connector_ids, outline=None, - geometry=intersection_polygon, + shapely_polygon=intersection_polygon, ) ) @@ -324,12 +320,7 @@ def _write_nuscenes_crosswalks(nuscenes_map: NuScenesMap, map_writer: AbstractMa crosswalk_polygons.append(polygon) for idx, polygon in enumerate(crosswalk_polygons): - map_writer.write_crosswalk( - CacheCrosswalk( - object_id=idx, - geometry=polygon, - ) - ) + map_writer.write_crosswalk(Crosswalk(object_id=idx, shapely_polygon=polygon)) def _write_nuscenes_walkways(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: @@ -341,12 +332,7 @@ def _write_nuscenes_walkways(nuscenes_map: NuScenesMap, map_writer: AbstractMapW walkway_polygons.append(polygon) for idx, polygon in enumerate(walkway_polygons): - map_writer.write_walkway( - CacheWalkway( - object_id=idx, - geometry=polygon, - ) - ) + map_writer.write_walkway(Walkway(object_id=idx, shapely_polygon=polygon)) def _write_nuscenes_carparks(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: @@ -358,17 +344,12 @@ def _write_nuscenes_carparks(nuscenes_map: NuScenesMap, map_writer: AbstractMapW carpark_polygons.append(polygon) for idx, polygon in enumerate(carpark_polygons): - map_writer.write_carpark( - CacheCarpark( - object_id=idx, - geometry=polygon, - ) - ) + map_writer.write_carpark(Carpark(object_id=idx, shapely_polygon=polygon)) def _write_nuscenes_generic_drivables(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: """Write generic drivable area data to map_writer.""" - cell_size = 10.0 + cell_size = 20.0 drivable_polygons = [] for drivable_area_record in nuscenes_map.drivable_area: drivable_area = nuscenes_map.get("drivable_area", drivable_area_record["token"]) @@ -380,10 +361,10 @@ def _write_nuscenes_generic_drivables(nuscenes_map: NuScenesMap, map_writer: Abs # drivable_polygons.append(polygon) for idx, geometry in enumerate(drivable_polygons): - map_writer.write_generic_drivable(CacheGenericDrivable(object_id=idx, geometry=geometry)) + map_writer.write_generic_drivable(GenericDrivable(object_id=idx, shapely_polygon=geometry)) -def _write_nuscenes_stop_lines(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: +def _write_nuscenes_stop_zones(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: """Write stop line data to map_writer.""" # FIXME: Add stop lines. # stop_lines = nuscenes_map.stop_line @@ -421,10 +402,10 @@ def _write_nuscenes_road_lines(nuscenes_map: NuScenesMap, map_writer: AbstractMa line_type = _get_road_line_type(divider["line_token"], nuscenes_map) map_writer.write_road_line( - CacheRoadLine( + RoadLine( object_id=running_idx, road_line_type=line_type, - polyline=Polyline3D(LineString(line.coords)), + polyline=Polyline3D.from_linestring(LineString(line.coords)), ) ) running_idx += 1 @@ -436,16 +417,16 @@ def _write_nuscenes_road_lines(nuscenes_map: NuScenesMap, map_writer: AbstractMa line_type = _get_road_line_type(divider["line_token"], nuscenes_map) map_writer.write_road_line( - CacheRoadLine( + RoadLine( object_id=running_idx, road_line_type=line_type, - polyline=Polyline3D(LineString(line.coords)), + polyline=Polyline3D.from_linestring(LineString(line.coords)), ) ) running_idx += 1 -def _extract_nuscenes_road_edges(nuscenes_map: NuScenesMap) -> List[CacheRoadEdge]: +def _extract_nuscenes_road_edges(nuscenes_map: NuScenesMap) -> List[RoadEdge]: """Helper function to extract road edges from a nuScenes map.""" drivable_polygons = [] for drivable_area_record in nuscenes_map.drivable_area: @@ -457,10 +438,10 @@ def _extract_nuscenes_road_edges(nuscenes_map: NuScenesMap) -> List[CacheRoadEdg road_edge_linear_rings = get_road_edge_linear_rings(drivable_polygons) road_edges_linestrings = split_line_geometry_by_max_length(road_edge_linear_rings, MAX_ROAD_EDGE_LENGTH) - road_edges_cache: List[CacheRoadEdge] = [] + road_edges_cache: List[RoadEdge] = [] for idx in range(len(road_edges_linestrings)): road_edges_cache.append( - CacheRoadEdge( + RoadEdge( object_id=idx, road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, polyline=Polyline2D.from_linestring(road_edges_linestrings[idx]), diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py index e09caae6..00ff2397 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py @@ -4,20 +4,22 @@ import numpy as np from py123d.conversion.registry.lidar_index_registry import NuScenesLiDARIndex -from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.datatypes.metadata import LogMetadata from py123d.datatypes.sensors.lidar import LiDARType -from py123d.geometry.se import StateSE3 +from py123d.geometry.pose import PoseSE3 from py123d.geometry.transform.transform_se3 import convert_points_3d_array_between_origins def load_nuscenes_lidar_pcs_from_file(pcd_path: Path, log_metadata: LogMetadata) -> Dict[LiDARType, np.ndarray]: + """Loads nuScenes LiDAR point clouds from the original binary files.""" + lidar_pc = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, len(NuScenesLiDARIndex)) # convert lidar to ego frame lidar_extrinsic = log_metadata.lidar_metadata[LiDARType.LIDAR_TOP].extrinsic lidar_pc[..., NuScenesLiDARIndex.XYZ] = convert_points_3d_array_between_origins( from_origin=lidar_extrinsic, - to_origin=StateSE3(0, 0, 0, 1.0, 0, 0, 0), + to_origin=PoseSE3(0, 0, 0, 1.0, 0, 0, 0), points_3d_array=lidar_pc[..., NuScenesLiDARIndex.XYZ], ) return {LiDARType.LIDAR_TOP: lidar_pc} diff --git a/src/py123d/conversion/datasets/nuscenes/utils/__init__.py b/src/py123d/conversion/datasets/nuscenes/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py index 5cafb870..cd89b6ea 100644 --- a/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py +++ b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py @@ -1,4 +1,4 @@ -from typing import Final, List +from typing import Dict, Final, List from py123d.conversion.registry.box_detection_label_registry import NuScenesBoxDetectionLabel from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType @@ -48,6 +48,14 @@ } +NUSCENES_DATABASE_VERSION_MAPPING: Dict[str, str] = { + "nuscenes_train": "v1.0-trainval", + "nuscenes_val": "v1.0-trainval", + "nuscenes_test": "v1.0-test", + "nuscenes-mini_train": "v1.0-mini", + "nuscenes-mini_val": "v1.0-mini", +} + NUSCENES_CAMERA_TYPES = { PinholeCameraType.PCAM_F0: "CAM_FRONT", PinholeCameraType.PCAM_B0: "CAM_BACK", diff --git a/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py index a512dc18..4d229d68 100644 --- a/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py +++ b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py @@ -158,7 +158,6 @@ def order_lanes_left_to_right(polylines: List[Polyline2D]) -> List[int]: # Step 1: Compute the average direction vector across all lanes all_directions = [] for polyline in polylines: - polyline_array = polyline.array if len(polyline_array) < 2: continue diff --git a/src/py123d/conversion/datasets/pandaset/__init__.py b/src/py123d/conversion/datasets/pandaset/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/py123d/conversion/datasets/pandaset/pandaset_converter.py b/src/py123d/conversion/datasets/pandaset/pandaset_converter.py index 008ab4e0..2b0192a6 100644 --- a/src/py123d/conversion/datasets/pandaset/pandaset_converter.py +++ b/src/py123d/conversion/datasets/pandaset/pandaset_converter.py @@ -17,32 +17,35 @@ ) from py123d.conversion.datasets.pandaset.utils.pandaset_utlis import ( main_lidar_to_rear_axle, - pandaset_pose_dict_to_state_se3, + pandaset_pose_dict_to_pose_se3, read_json, read_pkl_gz, rotate_pandaset_pose_to_iso_coordinates, ) from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.registry.box_detection_label_registry import PandasetBoxDetectionLabel -from py123d.conversion.registry.lidar_index_registry import PandasetLiDARIndex -from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCameraMetadata, PinholeCameraType, PinholeIntrinsics -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import ( - get_pandaset_chrysler_pacifica_parameters, - rear_axle_se3_to_center_se3, +from py123d.conversion.registry import PandasetBoxDetectionLabel, PandasetLiDARIndex +from py123d.datatypes.detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper +from py123d.datatypes.metadata import LogMetadata +from py123d.datatypes.sensors import ( + LiDARMetadata, + LiDARType, + PinholeCameraMetadata, + PinholeCameraType, + PinholeIntrinsics, ) -from py123d.geometry import BoundingBoxSE3, BoundingBoxSE3Index, EulerAnglesIndex, StateSE3, Vector3D -from py123d.geometry.transform.transform_se3 import convert_absolute_to_relative_se3_array +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import get_pandaset_chrysler_pacifica_parameters +from py123d.geometry import BoundingBoxSE3, BoundingBoxSE3Index, EulerAnglesIndex, PoseSE3 +from py123d.geometry.transform import convert_absolute_to_relative_se3_array from py123d.geometry.utils.constants import DEFAULT_PITCH, DEFAULT_ROLL from py123d.geometry.utils.rotation_utils import get_quaternion_array_from_euler_array class PandasetConverter(AbstractDatasetConverter): + """Converter for the Pandaset dataset.""" + def __init__( self, splits: List[str], @@ -52,6 +55,16 @@ def __init__( val_log_names: List[str], test_log_names: List[str], ) -> None: + """Initializes the :class:`PandasetConverter`. + + :param splits: List of splits to include in the conversion. \ + Available splits: 'pandaset_train', 'pandaset_val', 'pandaset_test'. + :param pandaset_data_root: Path to the root directory of the Pandaset dataset + :param dataset_converter_config: Configuration for the dataset converter + :param train_log_names: List of log names to include in the training split + :param val_log_names: List of log names to include in the validation split + :param test_log_names: List of log names to include in the test split + """ super().__init__(dataset_converter_config) for split in splits: assert split in PANDASET_SPLITS, f"Split {split} is not available. Available splits: {PANDASET_SPLITS}" @@ -63,9 +76,9 @@ def __init__( self._train_log_names: List[str] = train_log_names self._val_log_names: List[str] = val_log_names self._test_log_names: List[str] = test_log_names - self._log_paths_and_split: Dict[str, List[Path]] = self._collect_log_paths() + self._log_paths_and_split: List[Tuple[Path, str]] = self._collect_log_paths() - def _collect_log_paths(self) -> Dict[str, List[Path]]: + def _collect_log_paths(self) -> List[Tuple[Path, str]]: log_paths_and_split: List[Tuple[Path, str]] = [] for log_folder in self._pandaset_data_root.iterdir(): @@ -85,7 +98,7 @@ def _collect_log_paths(self) -> Dict[str, List[Path]]: def get_number_of_maps(self) -> int: """Inherited, see superclass.""" - return 0 # NOTE: Pandaset does not have maps. + return 0 # NOTE @DanielDauner: Pandaset does not have maps. def get_number_of_logs(self) -> int: """Inherited, see superclass.""" @@ -93,7 +106,7 @@ def get_number_of_logs(self) -> int: def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: """Inherited, see superclass.""" - return None # NOTE: Pandaset does not have maps. + return None # NOTE @DanielDauner: Pandaset does not have maps. def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: """Inherited, see superclass.""" @@ -110,8 +123,8 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: vehicle_parameters=get_pandaset_chrysler_pacifica_parameters(), box_detection_label_class=PandasetBoxDetectionLabel, pinhole_camera_metadata=_get_pandaset_camera_metadata(source_log_path, self.dataset_converter_config), - lidar_metadata=_get_pandaset_lidar_metadata(source_log_path, self.dataset_converter_config), - map_metadata=None, # NOTE: Pandaset does not have maps. + lidar_metadata=_get_pandaset_lidar_metadata(self.dataset_converter_config), + map_metadata=None, # NOTE @DanielDauner: Pandaset does not have maps. ) # 2. Prepare log writer @@ -119,7 +132,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: # 3. Process source log data if log_needs_writing: - # Read files from pandaset timesteps = read_json(source_log_path / "meta" / "timestamps.json") gps: List[Dict[str, float]] = read_json(source_log_path / "meta" / "gps.json") @@ -131,12 +143,11 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: # Write data to log writer for iteration, timestep_s in enumerate(timesteps): - ego_state = _extract_pandaset_sensor_ego_state(gps[iteration], lidar_poses[iteration]) log_writer.write( timestamp=TimePoint.from_s(timestep_s), ego_state=ego_state, - box_detections=_extract_pandaset_box_detections(source_log_path, iteration, ego_state), + box_detections=_extract_pandaset_box_detections(source_log_path, iteration), pinhole_cameras=_extract_pandaset_sensor_camera( source_log_path, iteration, @@ -147,7 +158,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: lidars=_extract_pandaset_lidar( source_log_path, iteration, - ego_state, self.dataset_converter_config, ), ) @@ -159,21 +169,20 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: def _get_pandaset_camera_metadata( source_log_path: Path, dataset_config: DatasetConverterConfig ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + """Extracts the pinhole camera metadata from a Pandaset log folder.""" camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = {} - if dataset_config.include_pinhole_cameras: all_cameras_folder = source_log_path / "camera" for camera_folder in all_cameras_folder.iterdir(): camera_name = camera_folder.name - assert camera_name in PANDASET_CAMERA_MAPPING.keys(), f"Camera name {camera_name} is not recognized." - camera_type = PANDASET_CAMERA_MAPPING[camera_name] + camera_type = PANDASET_CAMERA_MAPPING[camera_name] intrinsics_file = camera_folder / "intrinsics.json" assert intrinsics_file.exists(), f"Camera intrinsics file {intrinsics_file} does not exist." - intrinsics_data = read_json(intrinsics_file) + intrinsics_data = read_json(intrinsics_file) camera_metadata[camera_type] = PinholeCameraMetadata( camera_type=camera_type, width=1920, @@ -190,52 +199,37 @@ def _get_pandaset_camera_metadata( return camera_metadata -def _get_pandaset_lidar_metadata( - log_path: Path, dataset_config: DatasetConverterConfig -) -> Dict[LiDARType, LiDARMetadata]: +def _get_pandaset_lidar_metadata(dataset_config: DatasetConverterConfig) -> Dict[LiDARType, LiDARMetadata]: + """Extracts the LiDAR metadata from a Pandaset log folder.""" lidar_metadata: Dict[LiDARType, LiDARMetadata] = {} - if dataset_config.include_lidars: for lidar_name, lidar_type in PANDASET_LIDAR_MAPPING.items(): lidar_metadata[lidar_type] = LiDARMetadata( lidar_type=lidar_type, lidar_index=PandasetLiDARIndex, - extrinsic=PANDASET_LIDAR_EXTRINSICS[ - lidar_name - ], # TODO: These extrinsics are incorrect, and need to be transformed correctly. + extrinsic=PANDASET_LIDAR_EXTRINSICS[lidar_name], ) return lidar_metadata def _extract_pandaset_sensor_ego_state(gps: Dict[str, float], lidar_pose: Dict[str, Dict[str, float]]) -> EgoStateSE3: - - rear_axle_pose = main_lidar_to_rear_axle(pandaset_pose_dict_to_state_se3(lidar_pose)) - + """Extracts the ego state from Pandaset GPS and LiDAR pose data.""" + rear_axle_se3 = main_lidar_to_rear_axle(pandaset_pose_dict_to_pose_se3(lidar_pose)) vehicle_parameters = get_pandaset_chrysler_pacifica_parameters() - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) - - # TODO: Add script to calculate the dynamic state from log sequence. - dynamic_state = DynamicStateSE3( - # velocity=Vector3D(x=gps["xvel"], y=gps["yvel"], z=0.0), - velocity=Vector3D(x=0.0, y=0.0, z=0.0), - acceleration=Vector3D(x=0.0, y=0.0, z=0.0), - angular_velocity=Vector3D(x=0.0, y=0.0, z=0.0), - ) - - return EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, + dynamic_state_se3 = None + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_se3, vehicle_parameters=vehicle_parameters, + dynamic_state_se3=dynamic_state_se3, timepoint=None, ) -def _extract_pandaset_box_detections( - source_log_path: Path, iteration: int, ego_state_se3: EgoStateSE3 -) -> BoxDetectionWrapper: +def _extract_pandaset_box_detections(source_log_path: Path, iteration: int) -> BoxDetectionWrapper: + """Extracts the box detections from a Pandaset log folder at a given iteration.""" - # NOTE: The following provided quboids annotations are not stored in 123D + # NOTE @DanielDauner: The following provided quboids annotations are not stored in 123D # - stationary # - camera_used # - attributes.object_motion @@ -281,7 +275,7 @@ def _extract_pandaset_box_detections( box_se3_array[:, BoundingBoxSE3Index.QUATERNION] = get_quaternion_array_from_euler_array(box_euler_angles_array) box_se3_array[:, BoundingBoxSE3Index.EXTENT] = np.stack([box_lengths, box_widths, box_heights], axis=-1) - # NOTE: Pandaset annotates moving bounding boxes twice (for synchronization reasons), + # NOTE @DanielDauner: Pandaset annotates moving bounding boxes twice (for synchronization reasons), # if they are in the overlap area between the top 360° lidar and the front-facing lidar (and moving). # The value in `cuboids.sensor_id` is either # - `0` (mechanical 360° LiDAR) @@ -300,7 +294,6 @@ def _extract_pandaset_box_detections( # Fill bounding box detections and return box_detections: List[BoxDetectionSE3] = [] for box_idx in range(num_boxes): - # Skip duplicate box detections from front lidar if sibling exists in top lidar if sensor_ids[box_idx] == 1 and sibling_ids[box_idx] in top_lidar_uuids: continue @@ -309,8 +302,8 @@ def _extract_pandaset_box_detections( # Convert coordinates to ISO 8855 # NOTE: This would be faster over a batch operation. - box_se3_array[box_idx, BoundingBoxSE3Index.STATE_SE3] = rotate_pandaset_pose_to_iso_coordinates( - StateSE3.from_array(box_se3_array[box_idx, BoundingBoxSE3Index.STATE_SE3], copy=False) + box_se3_array[box_idx, BoundingBoxSE3Index.SE3] = rotate_pandaset_pose_to_iso_coordinates( + PoseSE3.from_array(box_se3_array[box_idx, BoundingBoxSE3Index.SE3], copy=False) ).array box_detection_se3 = BoxDetectionSE3( @@ -319,7 +312,7 @@ def _extract_pandaset_box_detections( track_token=box_uuids[box_idx], ), bounding_box_se3=BoundingBoxSE3.from_array(box_se3_array[box_idx]), - velocity=Vector3D(0.0, 0.0, 0.0), # TODO: Add velocity + velocity_3d=None, ) box_detections.append(box_detection_se3) @@ -333,20 +326,18 @@ def _extract_pandaset_sensor_camera( camera_poses: Dict[str, List[Dict[str, Dict[str, float]]]], dataset_converter_config: DatasetConverterConfig, ) -> List[CameraData]: + """Extracts the pinhole camera metadata from a Pandaset scene at a given iteration.""" camera_data_list: List[CameraData] = [] iteration_str = f"{iteration:02d}" if dataset_converter_config.include_pinhole_cameras: - for camera_name, camera_type in PANDASET_CAMERA_MAPPING.items(): - image_abs_path = source_log_path / f"camera/{camera_name}/{iteration_str}.jpg" assert image_abs_path.exists(), f"Camera image file {str(image_abs_path)} does not exist." camera_pose_dict = camera_poses[camera_name][iteration] - camera_extrinsic = pandaset_pose_dict_to_state_se3(camera_pose_dict) - - camera_extrinsic = StateSE3.from_array( + camera_extrinsic = pandaset_pose_dict_to_pose_se3(camera_pose_dict) + camera_extrinsic = PoseSE3.from_array( convert_absolute_to_relative_se3_array(ego_state_se3.rear_axle_se3, camera_extrinsic.array), copy=True ) camera_data_list.append( @@ -362,8 +353,9 @@ def _extract_pandaset_sensor_camera( def _extract_pandaset_lidar( - source_log_path: Path, iteration: int, ego_state_se3: EgoStateSE3, dataset_converter_config: DatasetConverterConfig + source_log_path: Path, iteration: int, dataset_converter_config: DatasetConverterConfig ) -> List[LiDARData]: + """Extracts the LiDAR data from a Pandaset scene at a given iteration.""" lidars: List[LiDARData] = [] if dataset_converter_config.include_lidars: diff --git a/src/py123d/conversion/datasets/pandaset/pandaset_sensor_io.py b/src/py123d/conversion/datasets/pandaset/pandaset_sensor_io.py index 14f1f236..71f7cc6f 100644 --- a/src/py123d/conversion/datasets/pandaset/pandaset_sensor_io.py +++ b/src/py123d/conversion/datasets/pandaset/pandaset_sensor_io.py @@ -6,7 +6,7 @@ from py123d.conversion.datasets.pandaset.utils.pandaset_utlis import ( main_lidar_to_rear_axle, - pandaset_pose_dict_to_state_se3, + pandaset_pose_dict_to_pose_se3, read_json, read_pkl_gz, ) @@ -16,6 +16,8 @@ def load_pandaset_global_lidar_pc_from_path(pkl_gz_path: Union[Path, str]) -> Dict[LiDARType, np.ndarray]: + """Loads Pandaset LiDAR point clouds from a gzip-pickle file (pickled pandas DataFrame).""" + # NOTE: The Pandaset dataset stores both front and top LiDAR data in the same gzip-pickle file. # We need to separate them based on the laser_number field. # See here: https://github.com/scaleapi/pandaset-devkit/blob/master/python/pandaset/sensors.py#L160 @@ -34,17 +36,15 @@ def load_pandaset_global_lidar_pc_from_path(pkl_gz_path: Union[Path, str]) -> Di def load_pandaset_lidars_pcs_from_file( pkl_gz_path: Union[Path, str], iteration: Optional[int], -) -> np.ndarray: +) -> Dict[LiDARType, np.ndarray]: + """Loads Pandaset LiDAR point clouds from a gzip-pickle file and converts them to ego frame.""" pkl_gz_path = Path(pkl_gz_path) assert pkl_gz_path.exists(), f"Pandaset LiDAR file not found: {pkl_gz_path}" - lidar_pc_dict = load_pandaset_global_lidar_pc_from_path(pkl_gz_path) - ego_pose = main_lidar_to_rear_axle( - pandaset_pose_dict_to_state_se3(read_json(pkl_gz_path.parent / "poses.json")[iteration]) + pandaset_pose_dict_to_pose_se3(read_json(pkl_gz_path.parent / "poses.json")[iteration]) ) - for lidar_type in lidar_pc_dict.keys(): lidar_pc_dict[lidar_type][..., PandasetLiDARIndex.XYZ] = convert_absolute_to_relative_points_3d_array( ego_pose, diff --git a/src/py123d/conversion/datasets/pandaset/utils/__init__.py b/src/py123d/conversion/datasets/pandaset/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/py123d/conversion/datasets/pandaset/utils/pandaset_constants.py b/src/py123d/conversion/datasets/pandaset/utils/pandaset_constants.py index 8428aadd..52bb8d16 100644 --- a/src/py123d/conversion/datasets/pandaset/utils/pandaset_constants.py +++ b/src/py123d/conversion/datasets/pandaset/utils/pandaset_constants.py @@ -3,7 +3,7 @@ from py123d.conversion.registry.box_detection_label_registry import PandasetBoxDetectionLabel from py123d.datatypes.sensors.lidar import LiDARType from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType, PinholeDistortion, PinholeIntrinsics -from py123d.geometry import StateSE3 +from py123d.geometry import PoseSE3 PANDASET_SPLITS: List[str] = ["pandaset_train", "pandaset_val", "pandaset_test"] @@ -16,7 +16,10 @@ "right_camera": PinholeCameraType.PCAM_R1, } -PANDASET_LIDAR_MAPPING: Dict[str, LiDARType] = {"main_pandar64": LiDARType.LIDAR_TOP, "front_gt": LiDARType.LIDAR_FRONT} +PANDASET_LIDAR_MAPPING: Dict[str, LiDARType] = { + "main_pandar64": LiDARType.LIDAR_TOP, + "front_gt": LiDARType.LIDAR_FRONT, +} PANDASET_BOX_DETECTION_FROM_STR: Dict[str, PandasetBoxDetectionLabel] = { @@ -51,8 +54,8 @@ # https://github.com/scaleapi/pandaset-devkit/blob/master/docs/static_extrinsic_calibration.yaml -PANDASET_LIDAR_EXTRINSICS: Dict[str, StateSE3] = { - "front_gt": StateSE3( +PANDASET_LIDAR_EXTRINSICS: Dict[str, PoseSE3] = { + "front_gt": PoseSE3( x=-0.000451117754, y=-0.605646431446, z=-0.301525235176, @@ -61,12 +64,12 @@ qy=0.01134678181520767, qz=0.9997028534282365, ), - "main_pandar64": StateSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0), + "main_pandar64": PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0), } # https://github.com/scaleapi/pandaset-devkit/blob/master/docs/static_extrinsic_calibration.yaml -PANDASET_CAMERA_EXTRINSICS: Dict[str, StateSE3] = { - "back_camera": StateSE3( +PANDASET_CAMERA_EXTRINSICS: Dict[str, PoseSE3] = { + "back_camera": PoseSE3( x=-0.0004217634029916384, y=-0.21683144949675118, z=-1.0553445472201475, @@ -75,7 +78,7 @@ qy=-0.001595758695393934, qz=-0.0005330311533742299, ), - "front_camera": StateSE3( + "front_camera": PoseSE3( x=0.0002585796504896516, y=-0.03907777167811011, z=-0.0440125762408362, @@ -84,7 +87,7 @@ qy=0.7114721800418571, qz=-0.7025205466606356, ), - "front_left_camera": StateSE3( + "front_left_camera": PoseSE3( x=-0.25842240863267835, y=-0.3070654284505582, z=-0.9244245686318884, @@ -93,7 +96,7 @@ qy=-0.6283486651480494, qz=0.6206973014480826, ), - "front_right_camera": StateSE3( + "front_right_camera": PoseSE3( x=0.2546935700219631, y=-0.24929449717803095, z=-0.8686597280810242, @@ -102,7 +105,7 @@ qy=0.6120314641083645, qz=-0.6150170047424814, ), - "left_camera": StateSE3( + "left_camera": PoseSE3( x=0.23864835336611942, y=-0.2801448284013492, z=-0.5376795959387791, @@ -111,7 +114,7 @@ qy=-0.4989265501075421, qz=0.503409565706149, ), - "right_camera": StateSE3( + "right_camera": PoseSE3( x=-0.23097163411257893, y=-0.30843497058841024, z=-0.6850441215571058, diff --git a/src/py123d/conversion/datasets/pandaset/utils/pandaset_utlis.py b/src/py123d/conversion/datasets/pandaset/utils/pandaset_utlis.py index 68575e7e..4b34d159 100644 --- a/src/py123d/conversion/datasets/pandaset/utils/pandaset_utlis.py +++ b/src/py123d/conversion/datasets/pandaset/utils/pandaset_utlis.py @@ -5,32 +5,33 @@ from typing import Dict import numpy as np +from pyparsing import Union -from py123d.geometry import StateSE3, Vector3D +from py123d.geometry import PoseSE3, Vector3D from py123d.geometry.transform.transform_se3 import translate_se3_along_body_frame -def read_json(json_file: Path): +def read_json(json_file: Union[Path, str]): """Helper function to read a json file as dict.""" with open(json_file, "r") as f: json_data = json.load(f) return json_data -def read_pkl_gz(pkl_gz_file: Path): +def read_pkl_gz(pkl_gz_file: Union[Path, str]): """Helper function to read a pkl.gz file as dict.""" with gzip.open(pkl_gz_file, "rb") as f: pkl_data = pickle.load(f) return pkl_data -def pandaset_pose_dict_to_state_se3(pose_dict: Dict[str, Dict[str, float]]) -> StateSE3: - """Helper function for pandaset to convert a pose dict to StateSE3. +def pandaset_pose_dict_to_pose_se3(pose_dict: Dict[str, Dict[str, float]]) -> PoseSE3: + """Helper function for pandaset to convert a pose dict to PoseSE3. :param pose_dict: The input pose dict. - :return: The converted StateSE3. + :return: The converted PoseSE3. """ - return StateSE3( + return PoseSE3( x=pose_dict["position"]["x"], y=pose_dict["position"]["y"], z=pose_dict["position"]["z"], @@ -41,7 +42,7 @@ def pandaset_pose_dict_to_state_se3(pose_dict: Dict[str, Dict[str, float]]) -> S ) -def rotate_pandaset_pose_to_iso_coordinates(pose: StateSE3) -> StateSE3: +def rotate_pandaset_pose_to_iso_coordinates(pose: PoseSE3) -> PoseSE3: """Helper function for pandaset to rotate a pose to ISO coordinate system (x: forward, y: left, z: up). NOTE: Pandaset uses a different coordinate system (x: right, y: forward, z: up). @@ -61,11 +62,10 @@ def rotate_pandaset_pose_to_iso_coordinates(pose: StateSE3) -> StateSE3: transformation_matrix = pose.transformation_matrix.copy() transformation_matrix[0:3, 0:3] = transformation_matrix[0:3, 0:3] @ F - return StateSE3.from_transformation_matrix(transformation_matrix) + return PoseSE3.from_transformation_matrix(transformation_matrix) -def main_lidar_to_rear_axle(pose: StateSE3) -> StateSE3: - +def main_lidar_to_rear_axle(pose: PoseSE3) -> PoseSE3: F = np.array( [ [0.0, 1.0, 0.0], # new X = old Y (forward) @@ -77,7 +77,7 @@ def main_lidar_to_rear_axle(pose: StateSE3) -> StateSE3: transformation_matrix = pose.transformation_matrix.copy() transformation_matrix[0:3, 0:3] = transformation_matrix[0:3, 0:3] @ F - rotated_pose = StateSE3.from_transformation_matrix(transformation_matrix) + rotated_pose = PoseSE3.from_transformation_matrix(transformation_matrix) imu_pose = translate_se3_along_body_frame(rotated_pose, vector_3d=Vector3D(x=-0.840, y=0.0, z=0.0)) diff --git a/src/py123d/conversion/datasets/wopd/utils/__init__.py b/src/py123d/conversion/datasets/wopd/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/py123d/conversion/datasets/wopd/waymo_map_utils/womp_boundary_utils.py b/src/py123d/conversion/datasets/wopd/utils/womp_boundary_utils.py similarity index 93% rename from src/py123d/conversion/datasets/wopd/waymo_map_utils/womp_boundary_utils.py rename to src/py123d/conversion/datasets/wopd/utils/womp_boundary_utils.py index e9e6afad..c73591c0 100644 --- a/src/py123d/conversion/datasets/wopd/waymo_map_utils/womp_boundary_utils.py +++ b/src/py123d/conversion/datasets/wopd/utils/womp_boundary_utils.py @@ -4,9 +4,9 @@ import numpy as np import shapely.geometry as geom -from py123d.datatypes.maps.abstract_map_objects import AbstractRoadEdge, AbstractRoadLine -from py123d.datatypes.maps.map_datatypes import LaneType -from py123d.geometry import OccupancyMap2D, Point3D, Polyline3D, PolylineSE2, StateSE2, Vector2D +from py123d.datatypes.map_objects.map_layer_types import LaneType +from py123d.datatypes.map_objects.map_objects import RoadEdge, RoadLine +from py123d.geometry import OccupancyMap2D, Point3D, Polyline3D, PolylineSE2, PoseSE2, Vector2D from py123d.geometry.transform.transform_se2 import translate_se2_along_body_frame from py123d.geometry.utils.rotation_utils import normalize_angle @@ -80,7 +80,7 @@ def hit_polyline_type(self) -> int: def _collect_perpendicular_hits( - lane_query_se2: StateSE2, + lane_query_se2: PoseSE2, lane_token: str, polyline_dict: Dict[str, Dict[int, Polyline3D]], lane_polyline_se2_dict: Dict[int, PolylineSE2], @@ -147,10 +147,8 @@ def _filter_perpendicular_hits( perpendicular_hits: List[PerpendicularHit], lane_point_3d: Point3D, ) -> List[PerpendicularHit]: - filtered_hits = [] for hit in perpendicular_hits: - # 1. filter hits too far in the vertical direction z_distance = np.abs(hit.hit_point_3d.z - lane_point_3d.z) if z_distance > MAX_Z_DISTANCE: @@ -169,9 +167,7 @@ def _filter_perpendicular_hits( def fill_lane_boundaries( - lane_data_dict: Dict[int, WaymoLaneData], - road_lines: List[AbstractRoadLine], - road_edges: List[AbstractRoadEdge], + lane_data_dict: Dict[int, WaymoLaneData], road_lines: List[RoadLine], road_edges: List[RoadEdge] ) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: """Welcome to insanity. @@ -211,24 +207,23 @@ def fill_lane_boundaries( lane_polyline_se2 = lane_polyline_se2_dict[current_lane_token] # 1. sample poses along centerline - distances_se2 = np.linspace( - 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) + num_samples = int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1 + + distances_se2 = np.linspace(0, lane_polyline_se2.length, num_samples, endpoint=True) lane_queries_se2 = [ - StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) + PoseSE2.from_array(pose_se2_array) for pose_se2_array in lane_polyline_se2.interpolate(distances_se2) ] - distances_3d = np.linspace( - 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) + distances_3d = np.linspace(0, lane_polyline.length, num_samples, endpoint=True) lane_queries_3d = [ Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) ] - assert len(lane_queries_se2) == len(lane_queries_3d) + assert len(lane_queries_se2) == len(lane_queries_3d), ( + f"Number of sampled SE2 poses {len(lane_queries_se2)} and 3D points {len(lane_queries_3d)} must be the same" + ) for sign in [1.0, -1.0]: boundary_points_3d: List[Optional[Point3D]] = [] for lane_query_se2, lane_query_3d in zip(lane_queries_se2, lane_queries_3d): - perpendicular_hits = _collect_perpendicular_hits( lane_query_se2=lane_query_se2, lane_token=current_lane_token, @@ -252,12 +247,10 @@ def fill_lane_boundaries( elif first_hit.hit_polyline_type == "road-line": boundary_point_3d = first_hit.hit_point_3d elif first_hit.hit_polyline_type == "lane": - for hit in perpendicular_hits: if hit.hit_polyline_type == "road-edge": continue if hit.hit_polyline_type == "lane": - lane_data_dict[lane_id].predecessor_ids has_same_predecessor = ( @@ -295,9 +288,7 @@ def fill_lane_boundaries( no_boundary_ratio = boundary_points_3d.count(None) / len(boundary_points_3d) final_boundary_points_3d = [] - def _get_default_boundary_point_3d( - lane_query_se2: StateSE2, lane_query_3d: Point3D, sign: float - ) -> Point3D: + def _get_default_boundary_point_3d(lane_query_se2: PoseSE2, lane_query_3d: Point3D, sign: float) -> Point3D: perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 boundary_point_se2 = translate_se2_along_body_frame( lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance) diff --git a/src/py123d/conversion/datasets/wopd/utils/wopd_constants.py b/src/py123d/conversion/datasets/wopd/utils/wopd_constants.py index efa577d6..bda4eced 100644 --- a/src/py123d/conversion/datasets/wopd/utils/wopd_constants.py +++ b/src/py123d/conversion/datasets/wopd/utils/wopd_constants.py @@ -1,7 +1,7 @@ from typing import Dict, List from py123d.conversion.registry.box_detection_label_registry import WOPDBoxDetectionLabel -from py123d.datatypes.maps.map_datatypes import LaneType, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_layer_types import LaneType, RoadEdgeType, RoadLineType from py123d.datatypes.sensors.lidar import LiDARType from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType diff --git a/src/py123d/conversion/datasets/wopd/wopd_converter.py b/src/py123d/conversion/datasets/wopd/wopd_converter.py index c5d7a411..9e486a11 100644 --- a/src/py123d/conversion/datasets/wopd/wopd_converter.py +++ b/src/py123d/conversion/datasets/wopd/wopd_converter.py @@ -16,32 +16,33 @@ WOPD_DETECTION_NAME_DICT, WOPD_LIDAR_TYPES, ) -from py123d.conversion.datasets.wopd.waymo_map_utils.wopd_map_utils import convert_wopd_map +from py123d.conversion.datasets.wopd.wopd_map_conversion import convert_wopd_map from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, CameraData, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter from py123d.conversion.registry.box_detection_label_registry import WOPDBoxDetectionLabel from py123d.conversion.registry.lidar_index_registry import DefaultLiDARIndex, WOPDLiDARIndex from py123d.conversion.utils.sensor_utils.camera_conventions import CameraConvention, convert_camera_convention from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import ( +from py123d.datatypes.metadata.log_metadata import LogMetadata +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.datatypes.sensors import ( + LiDARMetadata, + LiDARType, PinholeCameraMetadata, PinholeCameraType, PinholeDistortion, PinholeIntrinsics, ) from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 +from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 from py123d.datatypes.vehicle_state.vehicle_parameters import get_wopd_chrysler_pacifica_parameters from py123d.geometry import ( BoundingBoxSE3, BoundingBoxSE3Index, EulerAngles, EulerAnglesIndex, - StateSE3, - StateSE3Index, + PoseSE3, + PoseSE3Index, Vector3D, Vector3DIndex, ) @@ -65,6 +66,8 @@ class WOPDConverter(AbstractDatasetConverter): + """Converter for the Waymo Open Perception Dataset (WOPD).""" + def __init__( self, splits: List[str], @@ -74,24 +77,34 @@ def __init__( add_map_pose_offset: bool, dataset_converter_config: DatasetConverterConfig, ) -> None: + """Initializes the :class:`WOPDConverter`. + + :param splits: List of splits to convert, i.e. ``["wopd_train", "wopd_val", "wopd_test"]``. + :param wopd_data_root: Path to the root directory of the WOPD dataset + :param zero_roll_pitch: Whether to zero out roll and pitch angles in the vehicle pose + :param keep_polar_features: Whether to keep polar features in the LiDAR point clouds + :param add_map_pose_offset: Whether to add a pose offset to the map + :param dataset_converter_config: Configuration for the dataset converter + """ + super().__init__(dataset_converter_config) for split in splits: - assert ( - split in WOPD_AVAILABLE_SPLITS - ), f"Split {split} is not available. Available splits: {WOPD_AVAILABLE_SPLITS}" + assert split in WOPD_AVAILABLE_SPLITS, ( + f"Split {split} is not available. Available splits: {WOPD_AVAILABLE_SPLITS}" + ) self._splits: List[str] = splits self._wopd_data_root: Path = Path(wopd_data_root) self._zero_roll_pitch: bool = zero_roll_pitch self._keep_polar_features: bool = keep_polar_features - self._add_map_pose_offset: bool = add_map_pose_offset # TODO: Implement this feature + self._add_map_pose_offset: bool = add_map_pose_offset - self._split_tf_record_pairs: List[Tuple[str, List[Path]]] = self._collect_split_tf_record_pairs() + self._split_tf_record_pairs: List[Tuple[str, Path]] = self._collect_split_tf_record_pairs() - def _collect_split_tf_record_pairs(self) -> Dict[str, List[Path]]: + def _collect_split_tf_record_pairs(self) -> List[Tuple[str, Path]]: """Helper to collect the pairings of the split names and the corresponding tf record file.""" - split_tf_record_pairs: List[Tuple[str, List[Path]]] = [] + split_tf_record_pairs: List[Tuple[str, Path]] = [] split_name_mapping: Dict[str, str] = { "wopd_train": "training", "wopd_val": "validation", @@ -201,6 +214,8 @@ def _get_initial_frame_from_tfrecord( tf_record_path: Path, keep_dataset: bool = False, ) -> Union[dataset_pb2.Frame, Tuple[dataset_pb2.Frame, tf.data.TFRecordDataset]]: + """Helper to get the initial frame from a tf record file.""" + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") for data in dataset: initial_frame = dataset_pb2.Frame() @@ -215,25 +230,23 @@ def _get_initial_frame_from_tfrecord( def _get_wopd_map_metadata(initial_frame: dataset_pb2.Frame, split: str) -> MapMetadata: - + """Gets the WOPD map metadata from the initial frame.""" map_metadata = MapMetadata( dataset="wopd", split=split, log_name=str(initial_frame.context.name), location=None, # TODO: Add location information. map_has_z=True, - map_is_local=True, # True, if map is per log + map_is_local=True, ) - return map_metadata def _get_wopd_camera_metadata( initial_frame: dataset_pb2.Frame, dataset_converter_config: DatasetConverterConfig ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: - + """Get the WOPD camera metadata from the initial frame.""" camera_metadata_dict: Dict[PinholeCameraType, PinholeCameraMetadata] = {} - if dataset_converter_config.pinhole_camera_store_option is not None: for calibration in initial_frame.context.camera_calibrations: camera_type = WOPD_CAMERA_TYPES[calibration.name] @@ -250,7 +263,6 @@ def _get_wopd_camera_metadata( intrinsics=intrinsics, distortion=distortion, ) - return camera_metadata_dict @@ -259,20 +271,17 @@ def _get_wopd_lidar_metadata( keep_polar_features: bool, dataset_converter_config: DatasetConverterConfig, ) -> Dict[LiDARType, LiDARMetadata]: + """Get the WOPD LiDAR metadata from the initial frame.""" laser_metadatas: Dict[LiDARType, LiDARMetadata] = {} - - # NOTE: Using lidar_index = WOPDLiDARIndex if keep_polar_features else DefaultLiDARIndex if dataset_converter_config.lidar_store_option is not None: for laser_calibration in initial_frame.context.laser_calibrations: - lidar_type = WOPD_LIDAR_TYPES[laser_calibration.name] - - extrinsic: Optional[StateSE3] = None + extrinsic: Optional[PoseSE3] = None if laser_calibration.extrinsic: extrinsic_transform = np.array(laser_calibration.extrinsic.transform, dtype=np.float64).reshape(4, 4) - extrinsic = StateSE3.from_transformation_matrix(extrinsic_transform) + extrinsic = PoseSE3.from_transformation_matrix(extrinsic_transform) laser_metadatas[lidar_type] = LiDARMetadata( lidar_type=lidar_type, @@ -283,36 +292,35 @@ def _get_wopd_lidar_metadata( return laser_metadatas -def _get_ego_pose_se3(frame: dataset_pb2.Frame, map_pose_offset: Vector3D) -> StateSE3: +def _get_ego_pose_se3(frame: dataset_pb2.Frame, map_pose_offset: Vector3D) -> PoseSE3: + """Helper to get the ego pose SE3 from a WOPD frame.""" ego_pose_matrix = np.array(frame.pose.transform, dtype=np.float64).reshape(4, 4) - ego_pose_se3 = StateSE3.from_transformation_matrix(ego_pose_matrix) - ego_pose_se3.array[StateSE3Index.XYZ] += map_pose_offset.array[Vector3DIndex.XYZ] + ego_pose_se3 = PoseSE3.from_transformation_matrix(ego_pose_matrix) + ego_pose_se3.array[PoseSE3Index.XYZ] += map_pose_offset.array[Vector3DIndex.XYZ] return ego_pose_se3 def _extract_wopd_ego_state(frame: dataset_pb2.Frame, map_pose_offset: Vector3D) -> List[float]: + """Extracts the ego state from a WOPD frame.""" rear_axle_pose = _get_ego_pose_se3(frame, map_pose_offset) vehicle_parameters = get_wopd_chrysler_pacifica_parameters() # FIXME: Find dynamic state in waymo open perception dataset # https://github.com/waymo-research/waymo-open-dataset/issues/55#issuecomment-546152290 - dynamic_state = DynamicStateSE3( - velocity=Vector3D(*np.zeros(3)), - acceleration=Vector3D(*np.zeros(3)), - angular_velocity=Vector3D(*np.zeros(3)), - ) + dynamic_state_se3 = None return EgoStateSE3.from_rear_axle( rear_axle_se3=rear_axle_pose, - dynamic_state_se3=dynamic_state, + dynamic_state_se3=dynamic_state_se3, vehicle_parameters=vehicle_parameters, - time_point=None, + timepoint=None, ) def _extract_wopd_box_detections( frame: dataset_pb2.Frame, map_pose_offset: Vector3D, zero_roll_pitch: bool = True ) -> BoxDetectionWrapper: + """Extracts the box detections from a WOPD frame.""" ego_pose_se3 = _get_ego_pose_se3(frame, map_pose_offset) @@ -323,7 +331,6 @@ def _extract_wopd_box_detections( detections_token: List[str] = [] for detection_idx, detection in enumerate(frame.laser_labels): - detection_quaternion = EulerAngles( roll=DEFAULT_ROLL, pitch=DEFAULT_PITCH, @@ -350,8 +357,8 @@ def _extract_wopd_box_detections( detections_types.append(WOPD_DETECTION_NAME_DICT[detection.type]) detections_token.append(str(detection.id)) - detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( - origin=ego_pose_se3, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] + detections_state[:, BoundingBoxSE3Index.SE3] = convert_relative_to_absolute_se3_array( + origin=ego_pose_se3, se3_array=detections_state[:, BoundingBoxSE3Index.SE3] ) if zero_roll_pitch: euler_array = get_euler_array_from_quaternion_array(detections_state[:, BoundingBoxSE3Index.QUATERNION]) @@ -367,10 +374,9 @@ def _extract_wopd_box_detections( label=detections_types[detection_idx], timepoint=None, track_token=detections_token[detection_idx], - confidence=None, ), bounding_box_se3=BoundingBoxSE3.from_array(detections_state[detection_idx]), - velocity=Vector3D.from_array(detections_velocity[detection_idx]), + velocity_3d=Vector3D.from_array(detections_velocity[detection_idx]), ) ) @@ -380,19 +386,17 @@ def _extract_wopd_box_detections( def _extract_wopd_cameras( frame: dataset_pb2.Frame, dataset_converter_config: DatasetConverterConfig ) -> List[CameraData]: + """Extracts the camera data from a WOPD frame.""" camera_data_list: List[CameraData] = [] - if dataset_converter_config.include_pinhole_cameras: - # NOTE @DanielDauner: The extrinsic matrix in frame.context.camera_calibration is fixed to model the ego to camera transformation. # The poses in frame.images[idx] are the motion compensated ego poses when the camera triggers. - # TODO: Verify if this is correct. - camera_extrinsic: Dict[str, StateSE3] = {} + camera_extrinsic: Dict[str, PoseSE3] = {} for calibration in frame.context.camera_calibrations: camera_type = WOPD_CAMERA_TYPES[calibration.name] camera_transform = np.array(calibration.extrinsic.transform, dtype=np.float64).reshape(4, 4) - camera_pose = StateSE3.from_transformation_matrix(camera_transform) + camera_pose = PoseSE3.from_transformation_matrix(camera_transform) # NOTE: WOPD uses a different camera convention than py123d # https://arxiv.org/pdf/1912.04838 (Figure 1.) camera_pose = convert_camera_convention( @@ -423,11 +427,10 @@ def _extract_wopd_lidars( absolute_tf_record_path: Path, wopd_data_root: Path, ) -> Dict[LiDARType, npt.NDArray[np.float32]]: + """Extracts the LiDAR data from a WOPD frame.""" lidars: List[LiDARData] = [] - if dataset_converter_config.include_lidars: - relative_path = absolute_tf_record_path.relative_to(wopd_data_root) lidars.append( LiDARData( @@ -437,5 +440,4 @@ def _extract_wopd_lidars( relative_path=relative_path, ) ) - return lidars diff --git a/src/py123d/conversion/datasets/wopd/waymo_map_utils/wopd_map_utils.py b/src/py123d/conversion/datasets/wopd/wopd_map_conversion.py similarity index 86% rename from src/py123d/conversion/datasets/wopd/waymo_map_utils/wopd_map_utils.py rename to src/py123d/conversion/datasets/wopd/wopd_map_conversion.py index 741d9de7..c7301428 100644 --- a/src/py123d/conversion/datasets/wopd/waymo_map_utils/wopd_map_utils.py +++ b/src/py123d/conversion/datasets/wopd/wopd_map_conversion.py @@ -3,23 +3,15 @@ import numpy as np from py123d.common.utils.dependencies import check_dependencies +from py123d.conversion.datasets.wopd.utils.womp_boundary_utils import WaymoLaneData, fill_lane_boundaries from py123d.conversion.datasets.wopd.utils.wopd_constants import ( WAYMO_LANE_TYPE_CONVERSION, WAYMO_ROAD_EDGE_TYPE_CONVERSION, WAYMO_ROAD_LINE_TYPE_CONVERSION, ) -from py123d.conversion.datasets.wopd.waymo_map_utils.womp_boundary_utils import WaymoLaneData, fill_lane_boundaries from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.datatypes.maps.abstract_map_objects import AbstractLane, AbstractRoadEdge, AbstractRoadLine -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCarpark, - CacheCrosswalk, - CacheLane, - CacheLaneGroup, - CacheRoadEdge, - CacheRoadLine, -) -from py123d.datatypes.maps.map_datatypes import LaneType, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_layer_types import LaneType, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_objects import Carpark, Crosswalk, Lane, LaneGroup, RoadEdge, RoadLine from py123d.geometry import Polyline3D from py123d.geometry.utils.units import mph_to_mps @@ -34,7 +26,6 @@ def convert_wopd_map(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> None: - # We first extract all road lines, road edges, and lanes, and write them to the map writer. # NOTE: road lines and edges are used needed to extract lane boundaries. road_lines = _write_and_get_waymo_road_lines(frame, map_writer) @@ -48,17 +39,17 @@ def convert_wopd_map(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> _write_waymo_misc_surfaces(frame, map_writer) -def _write_and_get_waymo_road_lines(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> List[AbstractRoadLine]: +def _write_and_get_waymo_road_lines(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> List[RoadLine]: """Helper function to extract road lines from a Waymo frame proto.""" - road_lines: List[AbstractRoadLine] = [] + road_lines: List[RoadLine] = [] for map_feature in frame.map_features: if map_feature.HasField("road_line"): polyline = _extract_polyline_waymo_proto(map_feature.road_line) if polyline is not None: road_line_type = WAYMO_ROAD_LINE_TYPE_CONVERSION.get(map_feature.road_line.type, RoadLineType.UNKNOWN) road_lines.append( - CacheRoadLine( + RoadLine( object_id=map_feature.id, road_line_type=road_line_type, polyline=polyline, @@ -71,17 +62,17 @@ def _write_and_get_waymo_road_lines(frame: dataset_pb2.Frame, map_writer: Abstra return road_lines -def _write_and_get_waymo_road_edges(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> List[AbstractRoadEdge]: +def _write_and_get_waymo_road_edges(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> List[RoadEdge]: """Helper function to extract road edges from a Waymo frame proto.""" - road_edges: List[AbstractRoadEdge] = [] + road_edges: List[RoadEdge] = [] for map_feature in frame.map_features: if map_feature.HasField("road_edge"): polyline = _extract_polyline_waymo_proto(map_feature.road_edge) if polyline is not None: road_edge_type = WAYMO_ROAD_EDGE_TYPE_CONVERSION.get(map_feature.road_edge.type, RoadEdgeType.UNKNOWN) road_edges.append( - CacheRoadEdge( + RoadEdge( object_id=map_feature.id, road_edge_type=road_edge_type, polyline=polyline, @@ -95,12 +86,8 @@ def _write_and_get_waymo_road_edges(frame: dataset_pb2.Frame, map_writer: Abstra def _write_and_get_waymo_lanes( - frame: dataset_pb2.Frame, - road_lines: List[AbstractRoadLine], - road_edges: List[AbstractRoadEdge], - map_writer: AbstractMapWriter, -) -> List[AbstractLane]: - + frame: dataset_pb2.Frame, road_lines: List[RoadLine], road_edges: List[RoadEdge], map_writer: AbstractMapWriter +) -> List[Lane]: # 1. Load lane data from Waymo frame proto lane_data_dict: Dict[int, WaymoLaneData] = {} for map_feature in frame.map_features: @@ -136,15 +123,14 @@ def _get_majority_neighbor(neighbors: List[Dict[str, int]]) -> Optional[int]: } return str(max(length, key=length.get)) - lanes: List[AbstractLane] = [] + lanes: List[Lane] = [] for lane_data in lane_data_dict.values(): - # Skip lanes without boundaries if lane_data.left_boundary is None or lane_data.right_boundary is None: continue lanes.append( - CacheLane( + Lane( object_id=lane_data.object_id, lane_group_id=lane_data.object_id, left_boundary=lane_data.left_boundary, @@ -164,12 +150,11 @@ def _get_majority_neighbor(neighbors: List[Dict[str, int]]) -> Optional[int]: return lanes -def _write_waymo_lane_groups(lanes: List[AbstractLane], map_writer: AbstractMapWriter) -> None: - +def _write_waymo_lane_groups(lanes: List[Lane], map_writer: AbstractMapWriter) -> None: # NOTE: WOPD does not provide lane groups, so we create a lane group for each lane. for lane in lanes: map_writer.write_lane_group( - CacheLaneGroup( + LaneGroup( object_id=lane.object_id, lane_ids=[lane.object_id], left_boundary=lane.left_boundary, @@ -183,17 +168,16 @@ def _write_waymo_lane_groups(lanes: List[AbstractLane], map_writer: AbstractMapW def _write_waymo_misc_surfaces(frame: dataset_pb2.Frame, map_writer: AbstractMapWriter) -> None: - for map_feature in frame.map_features: if map_feature.HasField("driveway"): # NOTE: We currently only handle classify driveways as carparks. outline = _extract_outline_from_waymo_proto(map_feature.driveway) if outline is not None: - map_writer.write_carpark(CacheCarpark(object_id=map_feature.id, outline=outline)) + map_writer.write_carpark(Carpark(object_id=map_feature.id, outline=outline)) elif map_feature.HasField("crosswalk"): outline = _extract_outline_from_waymo_proto(map_feature.crosswalk) if outline is not None: - map_writer.write_crosswalk(CacheCrosswalk(object_id=map_feature.id, outline=outline)) + map_writer.write_crosswalk(Crosswalk(object_id=map_feature.id, outline=outline)) elif map_feature.HasField("stop_sign"): pass # TODO: Implement stop signs diff --git a/src/py123d/conversion/datasets/wopd/waymo_sensor_io.py b/src/py123d/conversion/datasets/wopd/wopd_sensor_io.py similarity index 92% rename from src/py123d/conversion/datasets/wopd/waymo_sensor_io.py rename to src/py123d/conversion/datasets/wopd/wopd_sensor_io.py index cca4bc53..14dae656 100644 --- a/src/py123d/conversion/datasets/wopd/waymo_sensor_io.py +++ b/src/py123d/conversion/datasets/wopd/wopd_sensor_io.py @@ -33,6 +33,7 @@ def load_jpeg_binary_from_tf_record_file( iteration: int, pinhole_camera_type: PinholeCameraType, ) -> bytes: + """Loads the JPEG binary of a specific pinhole camera from a Waymo TFRecord file at a given iteration.""" frame = _get_frame_at_iteration(tf_record_path, iteration) assert frame is not None, f"Frame at iteration {iteration} not found in Waymo file: {tf_record_path}" @@ -48,12 +49,11 @@ def load_jpeg_binary_from_tf_record_file( def load_wopd_lidar_pcs_from_file( tf_record_path: Path, index: int, keep_polar_features: bool = False ) -> Dict[LiDARType, np.ndarray]: + """Loads Waymo Open Perception Dataset (WOPD) LiDAR point clouds from a TFRecord file at a given iteration.""" frame = _get_frame_at_iteration(tf_record_path, index) assert frame is not None, f"Frame at iteration {index} not found in Waymo file: {tf_record_path}" - (range_images, camera_projections, _, range_image_top_pose) = parse_range_image_and_camera_projection(frame) - points, cp_points = frame_utils.convert_range_image_to_point_cloud( frame=frame, range_images=range_images, @@ -61,7 +61,6 @@ def load_wopd_lidar_pcs_from_file( range_image_top_pose=range_image_top_pose, keep_polar_features=keep_polar_features, ) - lidar_pcs_dict: Dict[LiDARType, np.ndarray] = {} for lidar_idx, frame_lidar in enumerate(frame.lasers): lidar_type = WOPD_LIDAR_TYPES[frame_lidar.name] diff --git a/src/py123d/conversion/log_writer/abstract_log_writer.py b/src/py123d/conversion/log_writer/abstract_log_writer.py index 238919d6..446cd4c1 100644 --- a/src/py123d/conversion/log_writer/abstract_log_writer.py +++ b/src/py123d/conversion/log_writer/abstract_log_writer.py @@ -11,13 +11,13 @@ from py123d.conversion.dataset_converter_config import DatasetConverterConfig from py123d.datatypes.detections.box_detections import BoxDetectionWrapper from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper -from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.datatypes.metadata import LogMetadata from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICameraType from py123d.datatypes.sensors.lidar import LiDARType from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType from py123d.datatypes.time.time_point import TimePoint from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.geometry import StateSE3 +from py123d.geometry import PoseSE3 class AbstractLogWriter(abc.ABC): @@ -32,9 +32,12 @@ def reset( self, dataset_converter_config: DatasetConverterConfig, log_metadata: LogMetadata, - ) -> None: - """ - Reset the log writer for a new log. + ) -> bool: + """Resets the log writer to start writing a new log according to the provided configuration and metadata. + + :param dataset_converter_config: The dataset converter configuration. + :param log_metadata: The metadata for the log. + :return: True if the current logs needs to be written, False otherwise. """ @abc.abstractmethod @@ -51,15 +54,27 @@ def write( route_lane_group_ids: Optional[List[int]] = None, **kwargs, ) -> None: - pass + """Writes a single iteration of data to the log. + + :param timestamp: Required, the timestamp of the iteration. + :param ego_state: Optional, the ego state of the vehicle, defaults to None. + :param box_detections: Optional, the box detections, defaults to None + :param traffic_lights: Optional, the traffic light detections, defaults to None + :param pinhole_cameras: Optional, the pinhole camera data, defaults to None + :param fisheye_mei_cameras: Optional, the fisheye MEI camera data, defaults to None + :param lidars: Optional, the LiDAR data, defaults to None + :param scenario_tags: Optional, the scenario tags, defaults to None + :param route_lane_group_ids: Optional, the route lane group IDs, defaults to None + """ @abc.abstractmethod def close(self) -> None: - pass + """Closes the log writer and finalizes the log io operations.""" @dataclass class LiDARData: + """Helper dataclass to pass LiDAR data to log writers.""" lidar_type: LiDARType @@ -70,9 +85,9 @@ class LiDARData: point_cloud: Optional[npt.NDArray[np.float32]] = None def __post_init__(self): - assert ( - self.has_file_path or self.has_point_cloud - ), "Either file path (dataset_root and relative_path) or point_cloud must be provided for LiDARData." + assert self.has_file_path or self.has_point_cloud, ( + "Either file path (dataset_root and relative_path) or point_cloud must be provided for LiDARData." + ) @property def has_file_path(self) -> bool: @@ -85,9 +100,10 @@ def has_point_cloud(self) -> bool: @dataclass class CameraData: + """Helper dataclass to pass Camera data to log writers.""" camera_type: Union[PinholeCameraType, FisheyeMEICameraType] - extrinsic: StateSE3 + extrinsic: PoseSE3 timestamp: Optional[TimePoint] = None jpeg_binary: Optional[bytes] = None @@ -96,9 +112,9 @@ class CameraData: relative_path: Optional[Union[str, Path]] = None def __post_init__(self): - assert ( - self.has_file_path or self.has_jpeg_binary or self.has_numpy_image - ), "Either file path (dataset_root and relative_path) or jpeg_binary or numpy_image must be provided for CameraData." + assert self.has_file_path or self.has_jpeg_binary or self.has_numpy_image, ( + "Either file path (dataset_root and relative_path) or jpeg_binary or numpy_image must be provided for CameraData." + ) if self.has_file_path: absolute_path = Path(self.dataset_root) / self.relative_path @@ -108,6 +124,14 @@ def __post_init__(self): def has_file_path(self) -> bool: return self.dataset_root is not None and self.relative_path is not None + @property + def has_jpeg_file_path(self) -> bool: + return self.relative_path is not None and str(self.relative_path).lower().endswith((".jpg", ".jpeg")) + + @property + def has_png_file_path(self) -> bool: + return self.relative_path is not None and str(self.relative_path).lower().endswith((".png",)) + @property def has_jpeg_binary(self) -> bool: return self.jpeg_binary is not None diff --git a/src/py123d/conversion/log_writer/arrow_log_writer.py b/src/py123d/conversion/log_writer/arrow_log_writer.py index acbcbb3b..979cfe55 100644 --- a/src/py123d/conversion/log_writer/arrow_log_writer.py +++ b/src/py123d/conversion/log_writer/arrow_log_writer.py @@ -4,6 +4,27 @@ import numpy as np import pyarrow as pa +from py123d.api.scene.arrow.utils.arrow_metadata_utils import add_log_metadata_to_arrow_schema +from py123d.common.utils.arrow_column_names import ( + BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN, + BOX_DETECTIONS_LABEL_COLUMN, + BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN, + BOX_DETECTIONS_TOKEN_COLUMN, + BOX_DETECTIONS_VELOCITY_3D_COLUMN, + EGO_DYNAMIC_STATE_SE3_COLUMN, + EGO_REAR_AXLE_SE3_COLUMN, + FISHEYE_CAMERA_DATA_COLUMN, + FISHEYE_CAMERA_EXTRINSIC_COLUMN, + LIDAR_DATA_COLUMN, + PINHOLE_CAMERA_DATA_COLUMN, + PINHOLE_CAMERA_EXTRINSIC_COLUMN, + ROUTE_LANE_GROUP_IDS_COLUMN, + SCENARIO_TAGS_COLUMN, + TIMESTAMP_US_COLUMN, + TRAFFIC_LIGHTS_LANE_ID_COLUMN, + TRAFFIC_LIGHTS_STATUS_COLUMN, + UUID_COLUMN, +) from py123d.common.utils.uuid_utils import create_deterministic_uuid from py123d.conversion.abstract_dataset_converter import AbstractLogWriter, DatasetConverterConfig from py123d.conversion.log_writer.abstract_log_writer import CameraData, LiDARData @@ -14,35 +35,65 @@ load_jpeg_binary_from_jpeg_file, ) from py123d.conversion.sensor_io.camera.mp4_camera_io import MP4Writer +from py123d.conversion.sensor_io.camera.png_camera_io import ( + encode_image_as_png_binary, + load_image_from_png_file, + load_png_binary_from_png_file, +) from py123d.conversion.sensor_io.lidar.draco_lidar_io import encode_lidar_pc_as_draco_binary from py123d.conversion.sensor_io.lidar.file_lidar_io import load_lidar_pcs_from_file from py123d.conversion.sensor_io.lidar.laz_lidar_io import encode_lidar_pc_as_laz_binary from py123d.datatypes.detections.box_detections import BoxDetectionWrapper from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper -from py123d.datatypes.scene.arrow.utils.arrow_metadata_utils import add_log_metadata_to_arrow_schema -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.lidar import LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType +from py123d.datatypes.metadata import LogMetadata +from py123d.datatypes.sensors import LiDARType, PinholeCameraType from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3, EgoStateSE3Index -from py123d.geometry import BoundingBoxSE3Index, StateSE3, StateSE3Index, Vector3DIndex +from py123d.datatypes.vehicle_state.dynamic_state import DynamicStateSE3Index +from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 +from py123d.geometry import BoundingBoxSE3Index, PoseSE3, PoseSE3Index, Vector3DIndex def _get_logs_root() -> Path: - from py123d.script.utils.dataset_path_utils import get_dataset_paths + from py123d.script.utils.dataset_path_utils import get_dataset_paths # noqa: PLC0415 DATASET_PATHS = get_dataset_paths() return Path(DATASET_PATHS.py123d_logs_root) def _get_sensors_root() -> Path: - from py123d.script.utils.dataset_path_utils import get_dataset_paths + from py123d.script.utils.dataset_path_utils import get_dataset_paths # noqa: PLC0415 DATASET_PATHS = get_dataset_paths() return Path(DATASET_PATHS.py123d_sensors_root) +def _store_option_to_arrow_type( + store_option: Literal["path", "jpeg_binary", "png_binary", "laz_binary", "draco_binary", "mp4"], +) -> pa.DataType: + """Maps the store option literal to the corresponding Arrow data type.""" + data_type_map = { + "path": pa.string(), + "jpeg_binary": pa.binary(), + "png_binary": pa.binary(), + "laz_binary": pa.binary(), + "draco_binary": pa.binary(), + "mp4": pa.int64(), + } + return data_type_map[store_option] + + +def _get_uuid_arrow_type(): + """Gets the appropriate Arrow UUID data type based on pyarrow version.""" + # NOTE @DanielDauner: pyarrow introduced native UUID type in version 18.0.0 + # Easiest option is to require this version or higher, but thanks to the Waymo dataset that's not possible. :( + if pa.__version__ >= "18.0.0": + return pa.uuid() + else: + return pa.binary(16) + + class ArrowLogWriter(AbstractLogWriter): + """Log writer for Arrow-based logs. Writes log data to an Arrow IPC file format.""" def __init__( self, @@ -50,14 +101,19 @@ def __init__( sensors_root: Optional[Union[str, Path]] = None, ipc_compression: Optional[Literal["lz4", "zstd"]] = None, ipc_compression_level: Optional[int] = None, - lidar_compression: Optional[Literal["draco", "laz"]] = "draco", ) -> None: + """Initializes the :class:`ArrowLogWriter`. + + :param logs_root: The root directory for logs, defaults to None + :param sensors_root: The root directory for sensors (i.e. in case of re-writing sensor files), defaults to None + :param ipc_compression: The IPC compression method, defaults to None + :param ipc_compression_level: The IPC compression level, defaults to None + """ self._logs_root = Path(logs_root) if logs_root is not None else _get_logs_root() self._sensors_root = Path(sensors_root) if sensors_root is not None else _get_sensors_root() self._ipc_compression = ipc_compression self._ipc_compression_level = ipc_compression_level - self._lidar_compression = lidar_compression # Loaded during .reset() and cleared during .close() self._dataset_converter_config: Optional[DatasetConverterConfig] = None @@ -69,6 +125,7 @@ def __init__( self._fisheye_mei_mp4_writers: Dict[str, MP4Writer] = {} def reset(self, dataset_converter_config: DatasetConverterConfig, log_metadata: LogMetadata) -> bool: + """Inherited, see superclass.""" log_needs_writing: bool = False sink_log_path: Path = self._logs_root / log_metadata.split / f"{log_metadata.log_name}.arrow" @@ -115,7 +172,9 @@ def write( lidars: Optional[List[LiDARData]] = None, scenario_tags: Optional[List[str]] = None, route_lane_group_ids: Optional[List[int]] = None, + **kwargs, ) -> None: + """Inherited, see superclass.""" assert self._dataset_converter_config is not None, "Log writer is not initialized." assert self._log_metadata is not None, "Log writer is not initialized." @@ -124,14 +183,14 @@ def write( assert self._source is not None, "Log writer is not initialized." record_batch_data = { - "uuid": [ + UUID_COLUMN: [ create_deterministic_uuid( split=self._log_metadata.split, log_name=self._log_metadata.log_name, timestamp_us=timestamp.time_us, ).bytes ], - "timestamp": [timestamp.time_us], + TIMESTAMP_US_COLUMN: [timestamp.time_us], } # -------------------------------------------------------------------------------------------------------------- @@ -139,51 +198,53 @@ def write( # -------------------------------------------------------------------------------------------------------------- if self._dataset_converter_config.include_ego: assert ego_state is not None, "Ego state is required but not provided." - record_batch_data["ego_state"] = [ego_state.array] + record_batch_data[EGO_REAR_AXLE_SE3_COLUMN] = [ego_state.rear_axle_se3] + record_batch_data[EGO_DYNAMIC_STATE_SE3_COLUMN] = [ego_state.dynamic_state_se3] # -------------------------------------------------------------------------------------------------------------- # Box Detections # -------------------------------------------------------------------------------------------------------------- if self._dataset_converter_config.include_box_detections: assert box_detections is not None, "Box detections are required but not provided." - # TODO: Figure out more elegant way without for-loops. # Accumulate box detection data box_detection_state = [] - box_detection_velocity = [] box_detection_token = [] box_detection_label = [] + box_detection_velocity = [] + box_detection_num_lidar_points = [] for box_detection in box_detections: - box_detection_state.append(box_detection.bounding_box.array) - box_detection_velocity.append(box_detection.velocity.array) # TODO: make optional + box_detection_state.append(box_detection.bounding_box_se3) box_detection_token.append(box_detection.metadata.track_token) box_detection_label.append(int(box_detection.metadata.label)) + box_detection_velocity.append(box_detection.velocity_3d) + box_detection_num_lidar_points.append(box_detection.metadata.num_lidar_points) # Add to record batch data - record_batch_data["box_detection_state"] = [box_detection_state] - record_batch_data["box_detection_velocity"] = [box_detection_velocity] - record_batch_data["box_detection_token"] = [box_detection_token] - record_batch_data["box_detection_label"] = [box_detection_label] + record_batch_data[BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN] = [box_detection_state] + record_batch_data[BOX_DETECTIONS_TOKEN_COLUMN] = [box_detection_token] + record_batch_data[BOX_DETECTIONS_LABEL_COLUMN] = [box_detection_label] + record_batch_data[BOX_DETECTIONS_VELOCITY_3D_COLUMN] = [box_detection_velocity] + record_batch_data[BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN] = [box_detection_num_lidar_points] # -------------------------------------------------------------------------------------------------------------- # Traffic Lights # -------------------------------------------------------------------------------------------------------------- if self._dataset_converter_config.include_traffic_lights: assert traffic_lights is not None, "Traffic light detections are required but not provided." - # TODO: Figure out more elegant way without for-loops. # Accumulate traffic light data traffic_light_ids = [] - traffic_light_types = [] + traffic_light_statuses = [] for traffic_light in traffic_lights: traffic_light_ids.append(traffic_light.lane_id) - traffic_light_types.append(int(traffic_light.status)) + traffic_light_statuses.append(int(traffic_light.status)) # Add to record batch data - record_batch_data["traffic_light_ids"] = [traffic_light_ids] - record_batch_data["traffic_light_types"] = [traffic_light_types] + record_batch_data[TRAFFIC_LIGHTS_LANE_ID_COLUMN] = [traffic_light_ids] + record_batch_data[TRAFFIC_LIGHTS_STATUS_COLUMN] = [traffic_light_statuses] # -------------------------------------------------------------------------------------------------------------- # Pinhole Cameras @@ -203,14 +264,16 @@ def write( # NOTE @DanielDauner: Missing cameras are allowed, e.g., for synchronization mismatches. # In this case, we write None/null to the arrow table. + # Theoretically, we could extend the store asynchronous cameras in the future by storing the + # camera data as a dictionary, list or struct-like object in the columns. pinhole_camera_data: Optional[Any] = None - pinhole_camera_pose: Optional[StateSE3] = None + pinhole_camera_pose: Optional[PoseSE3] = None if pinhole_camera_type in provided_pinhole_data: pinhole_camera_data = provided_pinhole_data[pinhole_camera_type] pinhole_camera_pose = provided_pinhole_extrinsics[pinhole_camera_type] - record_batch_data[f"{pinhole_camera_name}_data"] = [pinhole_camera_data] - record_batch_data[f"{pinhole_camera_name}_extrinsic"] = [ + record_batch_data[PINHOLE_CAMERA_DATA_COLUMN(pinhole_camera_name)] = [pinhole_camera_data] + record_batch_data[PINHOLE_CAMERA_EXTRINSIC_COLUMN(pinhole_camera_name)] = [ pinhole_camera_pose.array if pinhole_camera_pose else None ] @@ -233,13 +296,13 @@ def write( # NOTE @DanielDauner: Missing cameras are allowed, e.g., for synchronization mismatches. # In this case, we write None/null to the arrow table. fisheye_mei_camera_data: Optional[Any] = None - fisheye_mei_camera_pose: Optional[StateSE3] = None + fisheye_mei_camera_pose: Optional[PoseSE3] = None if fisheye_mei_camera_type in provided_fisheye_mei_data: fisheye_mei_camera_data = provided_fisheye_mei_data[fisheye_mei_camera_type] fisheye_mei_camera_pose = provided_fisheye_mei_extrinsics[fisheye_mei_camera_type] - record_batch_data[f"{fisheye_mei_camera_name}_data"] = [fisheye_mei_camera_data] - record_batch_data[f"{fisheye_mei_camera_name}_extrinsic"] = [ + record_batch_data[FISHEYE_CAMERA_DATA_COLUMN(fisheye_mei_camera_name)] = [fisheye_mei_camera_data] + record_batch_data[FISHEYE_CAMERA_EXTRINSIC_COLUMN(fisheye_mei_camera_name)] = [ fisheye_mei_camera_pose.array if fisheye_mei_camera_pose else None ] @@ -250,14 +313,15 @@ def write( assert lidars is not None, "LiDAR data is required but not provided." if self._dataset_converter_config.lidar_store_option == "path_merged": - # NOTE @DanielDauner: The path_merged option is necessary for dataset, that natively store multiple + # NOTE @DanielDauner: The path_merged option is necessary for datasets, that natively store multiple # LiDAR point clouds in a single file. In this case, writing the file path several times is wasteful. # Instead, we store the file path once, and divide the point clouds during reading. assert len(lidars) == 1, "Exactly one LiDAR data must be provided for merged LiDAR storage." assert lidars[0].has_file_path, "LiDAR data must provide file path for merged LiDAR storage." merged_lidar_data: Optional[str] = str(lidars[0].relative_path) + lidar_name = LiDARType.LIDAR_MERGED.serialize() - record_batch_data[f"{LiDARType.LIDAR_MERGED.serialize()}_data"] = [merged_lidar_data] + record_batch_data[LIDAR_DATA_COLUMN(lidar_name)] = [merged_lidar_data] else: # NOTE @DanielDauner: for "path" and "binary" options, we write each LiDAR in a separate column. @@ -270,23 +334,24 @@ def write( for lidar_type in expected_lidars: lidar_name = lidar_type.serialize() lidar_data: Optional[Union[str, bytes]] = lidar_data_dict.get(lidar_type, None) - record_batch_data[f"{lidar_name}_data"] = [lidar_data] + record_batch_data[LIDAR_DATA_COLUMN(lidar_name)] = [lidar_data] # -------------------------------------------------------------------------------------------------------------- # Miscellaneous (Scenario Tags / Route) # -------------------------------------------------------------------------------------------------------------- if self._dataset_converter_config.include_scenario_tags: assert scenario_tags is not None, "Scenario tags are required but not provided." - record_batch_data["scenario_tags"] = [scenario_tags] + record_batch_data[SCENARIO_TAGS_COLUMN] = [scenario_tags] if self._dataset_converter_config.include_route: assert route_lane_group_ids is not None, "Route lane group IDs are required but not provided." - record_batch_data["route_lane_group_ids"] = [route_lane_group_ids] + record_batch_data[ROUTE_LANE_GROUP_IDS_COLUMN] = [route_lane_group_ids] record_batch = pa.record_batch(record_batch_data, schema=self._schema) self._record_batch_writer.write_batch(record_batch) def close(self) -> None: + """Inherited, see superclass.""" if self._record_batch_writer is not None: self._record_batch_writer.close() self._record_batch_writer: Optional[pa.ipc.RecordBatchWriter] = None @@ -308,10 +373,16 @@ def close(self) -> None: @staticmethod def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata: LogMetadata) -> pa.Schema: + """Builds the schema for the Arrow table, specifying datatypes and modalities to be stored. + + :param dataset_converter_config: The dataset converter configuration. + :param log_metadata: The metadata for the log. + :return: The Arrow schema object. + """ schema_list: List[Tuple[str, pa.DataType]] = [ - ("uuid", pa.uuid()), - ("timestamp", pa.int64()), + (UUID_COLUMN, _get_uuid_arrow_type()), + (TIMESTAMP_US_COLUMN, pa.int64()), ] # -------------------------------------------------------------------------------------------------------------- @@ -320,7 +391,8 @@ def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata if dataset_converter_config.include_ego: schema_list.extend( [ - ("ego_state", pa.list_(pa.float64(), len(EgoStateSE3Index))), + (EGO_REAR_AXLE_SE3_COLUMN, pa.list_(pa.float64(), len(PoseSE3Index))), + (EGO_DYNAMIC_STATE_SE3_COLUMN, pa.list_(pa.float64(), len(DynamicStateSE3Index))), ] ) @@ -330,10 +402,26 @@ def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata if dataset_converter_config.include_box_detections: schema_list.extend( [ - ("box_detection_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("box_detection_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("box_detection_token", pa.list_(pa.string())), - ("box_detection_label", pa.list_(pa.int16())), + ( + BOX_DETECTIONS_BOUNDING_BOX_SE3_COLUMN, + pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index))), + ), + ( + BOX_DETECTIONS_TOKEN_COLUMN, + pa.list_(pa.string()), + ), + ( + BOX_DETECTIONS_LABEL_COLUMN, + pa.list_(pa.int16()), + ), + ( + BOX_DETECTIONS_VELOCITY_3D_COLUMN, + pa.list_(pa.list_(pa.float64(), len(Vector3DIndex))), + ), + ( + BOX_DETECTIONS_NUM_LIDAR_POINTS_COLUMN, + pa.list_(pa.int64()), + ), ] ) @@ -343,8 +431,8 @@ def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata if dataset_converter_config.include_traffic_lights: schema_list.extend( [ - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), + (TRAFFIC_LIGHTS_LANE_ID_COLUMN, pa.list_(pa.int64())), + (TRAFFIC_LIGHTS_STATUS_COLUMN, pa.list_(pa.int16())), ] ) @@ -354,19 +442,18 @@ def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata if dataset_converter_config.include_pinhole_cameras: for pinhole_camera_type in log_metadata.pinhole_camera_metadata.keys(): pinhole_camera_name = pinhole_camera_type.serialize() - - # Depending on the storage option, define the schema for camera data - if dataset_converter_config.pinhole_camera_store_option == "path": - schema_list.append((f"{pinhole_camera_name}_data", pa.string())) - - elif dataset_converter_config.pinhole_camera_store_option == "binary": - schema_list.append((f"{pinhole_camera_name}_data", pa.binary())) - - elif dataset_converter_config.pinhole_camera_store_option == "mp4": - schema_list.append((f"{pinhole_camera_name}_data", pa.int64())) - - # Add camera pose - schema_list.append((f"{pinhole_camera_name}_extrinsic", pa.list_(pa.float64(), len(StateSE3Index)))) + schema_list.extend( + [ + ( + PINHOLE_CAMERA_DATA_COLUMN(pinhole_camera_name), + _store_option_to_arrow_type(dataset_converter_config.pinhole_camera_store_option), + ), + ( + PINHOLE_CAMERA_EXTRINSIC_COLUMN(pinhole_camera_name), + pa.list_(pa.float64(), len(PoseSE3Index)), + ), + ] + ) # -------------------------------------------------------------------------------------------------------------- # Fisheye MEI Cameras @@ -374,49 +461,54 @@ def _build_schema(dataset_converter_config: DatasetConverterConfig, log_metadata if dataset_converter_config.include_fisheye_mei_cameras: for fisheye_mei_camera_type in log_metadata.fisheye_mei_camera_metadata.keys(): fisheye_mei_camera_name = fisheye_mei_camera_type.serialize() - - # Depending on the storage option, define the schema for camera data - if dataset_converter_config.fisheye_mei_camera_store_option == "path": - schema_list.append((f"{fisheye_mei_camera_name}_data", pa.string())) - - elif dataset_converter_config.fisheye_mei_camera_store_option == "binary": - schema_list.append((f"{fisheye_mei_camera_name}_data", pa.binary())) - - elif dataset_converter_config.fisheye_mei_camera_store_option == "mp4": - schema_list.append((f"{fisheye_mei_camera_name}_data", pa.int64())) - - # Add camera pose - schema_list.append((f"{fisheye_mei_camera_name}_extrinsic", pa.list_(pa.float64(), len(StateSE3Index)))) + schema_list.extend( + [ + ( + FISHEYE_CAMERA_DATA_COLUMN(fisheye_mei_camera_name), + _store_option_to_arrow_type(dataset_converter_config.fisheye_mei_camera_store_option), + ), + ( + FISHEYE_CAMERA_EXTRINSIC_COLUMN(fisheye_mei_camera_name), + pa.list_(pa.float64(), len(PoseSE3Index)), + ), + ] + ) # -------------------------------------------------------------------------------------------------------------- # LiDARs # -------------------------------------------------------------------------------------------------------------- if dataset_converter_config.include_lidars and len(log_metadata.lidar_metadata) > 0: if dataset_converter_config.lidar_store_option == "path_merged": - schema_list.append((f"{LiDARType.LIDAR_MERGED.serialize()}_data", pa.string())) + lidar_name = LiDARType.LIDAR_MERGED.serialize() + schema_list.append((LIDAR_DATA_COLUMN(lidar_name), pa.string())) else: for lidar_type in log_metadata.lidar_metadata.keys(): lidar_name = lidar_type.serialize() - - # Depending on the storage option, define the schema for LiDAR data - if dataset_converter_config.lidar_store_option == "path": - schema_list.append((f"{lidar_name}_data", pa.string())) - - elif dataset_converter_config.lidar_store_option == "binary": - schema_list.append((f"{lidar_name}_data", pa.binary())) + schema_list.append( + ( + LIDAR_DATA_COLUMN(lidar_name), + _store_option_to_arrow_type(dataset_converter_config.lidar_store_option), + ) + ) # -------------------------------------------------------------------------------------------------------------- # Miscellaneous (Scenario Tags / Route) # -------------------------------------------------------------------------------------------------------------- if dataset_converter_config.include_scenario_tags: - schema_list.append(("scenario_tags", pa.list_(pa.string()))) + schema_list.append((SCENARIO_TAGS_COLUMN, pa.list_(pa.string()))) if dataset_converter_config.include_route: - schema_list.append(("route_lane_group_ids", pa.list_(pa.int64()))) + schema_list.append((ROUTE_LANE_GROUP_IDS_COLUMN, pa.list_(pa.int64()))) return add_log_metadata_to_arrow_schema(pa.schema(schema_list), log_metadata) def _prepare_lidar_data_dict(self, lidars: List[LiDARData]) -> Dict[LiDARType, Union[str, bytes]]: + """Helper function to prepare LiDAR data dictionary for the target storage option. + + :param lidars: List of LiDARData objects to be processed. + :return: Dictionary mapping LiDARType to either file path (str) or binary data (bytes) depending on storage option. + """ + lidar_data_dict: Dict[LiDARType, Union[str, bytes]] = {} if self._dataset_converter_config.lidar_store_option == "path": @@ -424,7 +516,7 @@ def _prepare_lidar_data_dict(self, lidars: List[LiDARData]) -> Dict[LiDARType, U assert lidar_data.has_file_path, "LiDAR data must provide file path for path storage." lidar_data_dict[lidar_data.lidar_type] = str(lidar_data.relative_path) - elif self._dataset_converter_config.lidar_store_option == "binary": + elif self._dataset_converter_config.lidar_store_option in ["laz_binary", "draco_binary"]: lidar_pcs_dict: Dict[LiDARType, np.ndarray] = {} # 1. Load point clouds from files @@ -445,10 +537,14 @@ def _prepare_lidar_data_dict(self, lidars: List[LiDARData]) -> Dict[LiDARType, U for lidar_type, point_cloud in lidar_pcs_dict.items(): lidar_metadata = self._log_metadata.lidar_metadata[lidar_type] binary: Optional[bytes] = None - if self._lidar_compression == "draco": + if self._dataset_converter_config.lidar_store_option == "draco_binary": binary = encode_lidar_pc_as_draco_binary(point_cloud, lidar_metadata) - elif self._lidar_compression == "laz": + elif self._dataset_converter_config.lidar_store_option == "laz_binary": binary = encode_lidar_pc_as_laz_binary(point_cloud, lidar_metadata) + else: + raise NotImplementedError( + f"Unsupported LiDAR store option: {self._dataset_converter_config.lidar_store_option}" + ) lidar_data_dict[lidar_type] = binary return lidar_data_dict @@ -456,6 +552,14 @@ def _prepare_lidar_data_dict(self, lidars: List[LiDARData]) -> Dict[LiDARType, U def _prepare_camera_data_dict( self, cameras: List[CameraData], store_option: Literal["path", "binary"] ) -> Dict[PinholeCameraType, Union[str, bytes]]: + """Helper function to prepare camera data dictionary for the target storage option. + + :param cameras: List of CameraData objects to be processed. + :param store_option: The storage option for camera data, either "path" or "binary". + :raises NotImplementedError: If the storage option is not supported. + :raises NotImplementedError: If the camera data does not support the specified storage option. + :return: Dictionary mapping PinholeCameraType to either file path (str) or binary data (bytes) depending on storage option. + """ camera_data_dict: Dict[PinholeCameraType, Union[str, int, bytes]] = {} for camera_data in cameras: @@ -464,8 +568,10 @@ def _prepare_camera_data_dict( camera_data_dict[camera_data.camera_type] = str(camera_data.relative_path) else: raise NotImplementedError("Only file path storage is supported for camera data.") - elif store_option == "binary": + elif store_option == "jpeg_binary": camera_data_dict[camera_data.camera_type] = _get_jpeg_binary_from_camera_data(camera_data) + elif store_option == "png_binary": + camera_data_dict[camera_data.camera_type] = _get_png_binary_from_camera_data(camera_data) elif store_option == "mp4": camera_name = camera_data.camera_type.serialize() if camera_name not in self._pinhole_mp4_writers: @@ -492,15 +598,15 @@ def _get_jpeg_binary_from_camera_data(camera_data: CameraData) -> Optional[bytes if camera_data.has_jpeg_binary: jpeg_binary = camera_data.jpeg_binary + elif camera_data.has_jpeg_file_path: + absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path + jpeg_binary = load_jpeg_binary_from_jpeg_file(absolute_path) + elif camera_data.has_png_file_path: + absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path + numpy_image = load_image_from_png_file(absolute_path) + jpeg_binary = encode_image_as_jpeg_binary(numpy_image) elif camera_data.has_numpy_image: jpeg_binary = encode_image_as_jpeg_binary(camera_data.numpy_image) - elif camera_data.has_file_path: - absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path - - if absolute_path.suffix.lower() in [".jpg", ".jpeg"]: - jpeg_binary = load_jpeg_binary_from_jpeg_file(absolute_path) - else: - raise NotImplementedError(f"Unsupported camera file format: {absolute_path.suffix} for binary storage.") else: raise NotImplementedError("Camera data must provide jpeg_binary, numpy_image, or file path for binary storage.") @@ -508,6 +614,29 @@ def _get_jpeg_binary_from_camera_data(camera_data: CameraData) -> Optional[bytes return jpeg_binary +def _get_png_binary_from_camera_data(camera_data: CameraData) -> Optional[bytes]: + png_binary: Optional[bytes] = None + + if camera_data.has_png_file_path: + absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path + png_binary = load_png_binary_from_png_file(absolute_path) + elif camera_data.has_numpy_image: + png_binary = encode_image_as_png_binary(camera_data.numpy_image) + elif camera_data.has_jpeg_file_path: + absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path + numpy_image = load_image_from_jpeg_file(absolute_path) + png_binary = encode_image_as_png_binary(numpy_image) + + elif camera_data.has_jpeg_binary: + numpy_image = decode_image_from_jpeg_binary(camera_data.jpeg_binary) + png_binary = encode_image_as_png_binary(numpy_image) + else: + raise NotImplementedError("Camera data must provide png_binary, numpy_image, or file path for binary storage.") + + assert png_binary is not None + return png_binary + + def _get_numpy_image_from_camera_data(camera_data: CameraData) -> Optional[np.ndarray]: numpy_image: Optional[np.ndarray] = None @@ -515,9 +644,12 @@ def _get_numpy_image_from_camera_data(camera_data: CameraData) -> Optional[np.nd numpy_image = camera_data.numpy_image elif camera_data.has_jpeg_binary: numpy_image = decode_image_from_jpeg_binary(camera_data.jpeg_binary) - elif camera_data.has_file_path: + elif camera_data.has_jpeg_file_path: absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path numpy_image = load_image_from_jpeg_file(absolute_path) + elif camera_data.has_png_file_path: + absolute_path = Path(camera_data.dataset_root) / camera_data.relative_path + numpy_image = load_image_from_png_file(absolute_path) else: raise NotImplementedError("Camera data must provide numpy_image, jpeg_binary, or file path for numpy image.") diff --git a/src/py123d/conversion/map_writer/abstract_map_writer.py b/src/py123d/conversion/map_writer/abstract_map_writer.py index 99c867e6..cc74e1ad 100644 --- a/src/py123d/conversion/map_writer/abstract_map_writer.py +++ b/src/py123d/conversion/map_writer/abstract_map_writer.py @@ -2,19 +2,19 @@ from abc import abstractmethod from py123d.conversion.dataset_converter_config import DatasetConverterConfig -from py123d.datatypes.maps.abstract_map_objects import ( - AbstractCarpark, - AbstractCrosswalk, - AbstractGenericDrivable, - AbstractIntersection, - AbstractLane, - AbstractLaneGroup, - AbstractRoadEdge, - AbstractRoadLine, - AbstractStopLine, - AbstractWalkway, +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + StopZone, + Walkway, ) -from py123d.datatypes.maps.map_metadata import MapMetadata +from py123d.datatypes.metadata.map_metadata import MapMetadata class AbstractMapWriter(abc.ABC): @@ -25,43 +25,43 @@ def reset(self, dataset_converter_config: DatasetConverterConfig, map_metadata: """Reset the writer to its initial state.""" @abstractmethod - def write_lane(self, lane: AbstractLane) -> None: + def write_lane(self, lane: Lane) -> None: """Write a lane to the map.""" @abstractmethod - def write_lane_group(self, lane: AbstractLaneGroup) -> None: + def write_lane_group(self, lane: LaneGroup) -> None: """Write a group of lanes to the map.""" @abstractmethod - def write_intersection(self, intersection: AbstractIntersection) -> None: + def write_intersection(self, intersection: Intersection) -> None: """Write an intersection to the map.""" @abstractmethod - def write_crosswalk(self, crosswalk: AbstractCrosswalk) -> None: + def write_crosswalk(self, crosswalk: Crosswalk) -> None: """Write a crosswalk to the map.""" @abstractmethod - def write_carpark(self, carpark: AbstractCarpark) -> None: + def write_carpark(self, carpark: Carpark) -> None: """Write a car park to the map.""" @abstractmethod - def write_walkway(self, walkway: AbstractWalkway) -> None: + def write_walkway(self, walkway: Walkway) -> None: """Write a walkway to the map.""" @abstractmethod - def write_generic_drivable(self, obj: AbstractGenericDrivable) -> None: + def write_generic_drivable(self, obj: GenericDrivable) -> None: """Write a generic drivable area to the map.""" @abstractmethod - def write_stop_line(self, stop_line: AbstractStopLine) -> None: - """Write a stop lines to the map.""" + def write_stop_zone(self, stop_zone: StopZone) -> None: + """Write a stop zone to the map.""" @abstractmethod - def write_road_edge(self, road_edge: AbstractRoadEdge) -> None: + def write_road_edge(self, road_edge: RoadEdge) -> None: """Write a road edge to the map.""" @abstractmethod - def write_road_line(self, road_line: AbstractRoadLine) -> None: + def write_road_line(self, road_line: RoadLine) -> None: """Write a road line to the map.""" @abstractmethod diff --git a/src/py123d/conversion/map_writer/gpkg_map_writer.py b/src/py123d/conversion/map_writer/gpkg_map_writer.py index 289e1cc2..09f73636 100644 --- a/src/py123d/conversion/map_writer/gpkg_map_writer.py +++ b/src/py123d/conversion/map_writer/gpkg_map_writer.py @@ -1,3 +1,4 @@ +import logging from collections import defaultdict from pathlib import Path from typing import Dict, List, Optional, Union @@ -9,39 +10,41 @@ from py123d.conversion.dataset_converter_config import DatasetConverterConfig from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter from py123d.conversion.map_writer.utils.gpkg_utils import IntIDMapping -from py123d.datatypes.maps.abstract_map_objects import ( - AbstractCarpark, - AbstractCrosswalk, - AbstractGenericDrivable, - AbstractIntersection, - AbstractLane, - AbstractLaneGroup, - AbstractLineMapObject, - AbstractRoadEdge, - AbstractRoadLine, - AbstractStopLine, - AbstractSurfaceMapObject, - AbstractWalkway, +from py123d.datatypes.map_objects.map_layer_types import MapLayer +from py123d.datatypes.map_objects.map_objects import ( + BaseMapLineObject, + BaseMapSurfaceObject, + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + StopZone, + Walkway, ) -from py123d.datatypes.maps.map_datatypes import MapLayer -from py123d.datatypes.maps.map_metadata import MapMetadata +from py123d.datatypes.metadata.map_metadata import MapMetadata from py123d.geometry.polyline import Polyline3D MAP_OBJECT_DATA = Dict[str, List[Union[str, int, float, bool, geom.base.BaseGeometry]]] +logging.getLogger("pyogrio._io").disabled = True + class GPKGMapWriter(AbstractMapWriter): """Abstract base class for map writers.""" - def __init__(self, maps_root: Union[str, Path], remap_ids: bool = False) -> None: + def __init__(self, maps_root: Union[str, Path]) -> None: self._maps_root = Path(maps_root) self._crs: str = "EPSG:4326" # WGS84 - self._remap_ids = remap_ids # Data to be written to the map for each object type self._map_data: Optional[Dict[MapLayer, MAP_OBJECT_DATA]] = None self._map_file: Optional[Path] = None self._map_metadata: Optional[MapMetadata] = None + self._remap_map_ids: Optional[bool] = None def reset(self, dataset_converter_config: DatasetConverterConfig, map_metadata: MapMetadata) -> bool: """Inherited, see superclass.""" @@ -62,10 +65,11 @@ def reset(self, dataset_converter_config: DatasetConverterConfig, map_metadata: self._map_data = {map_layer: defaultdict(list) for map_layer in MapLayer} self._map_file = map_file self._map_metadata = map_metadata + self._remap_map_ids = dataset_converter_config.remap_map_ids return map_needs_writing - def write_lane(self, lane: AbstractLane) -> None: + def write_lane(self, lane: Lane) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.LANE, lane) self._map_data[MapLayer.LANE]["lane_group_id"].append(lane.lane_group_id) @@ -78,7 +82,7 @@ def write_lane(self, lane: AbstractLane) -> None: self._map_data[MapLayer.LANE]["successor_ids"].append(lane.successor_ids) self._map_data[MapLayer.LANE]["speed_limit_mps"].append(lane.speed_limit_mps) - def write_lane_group(self, lane_group: AbstractLaneGroup) -> None: + def write_lane_group(self, lane_group: LaneGroup) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.LANE_GROUP, lane_group) self._map_data[MapLayer.LANE_GROUP]["lane_ids"].append(lane_group.lane_ids) @@ -88,41 +92,41 @@ def write_lane_group(self, lane_group: AbstractLaneGroup) -> None: self._map_data[MapLayer.LANE_GROUP]["left_boundary"].append(lane_group.left_boundary.linestring) self._map_data[MapLayer.LANE_GROUP]["right_boundary"].append(lane_group.right_boundary.linestring) - def write_intersection(self, intersection: AbstractIntersection) -> None: + def write_intersection(self, intersection: Intersection) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.INTERSECTION, intersection) self._map_data[MapLayer.INTERSECTION]["lane_group_ids"].append(intersection.lane_group_ids) - def write_crosswalk(self, crosswalk: AbstractCrosswalk) -> None: + def write_crosswalk(self, crosswalk: Crosswalk) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.CROSSWALK, crosswalk) - def write_carpark(self, carpark: AbstractCarpark) -> None: + def write_carpark(self, carpark: Carpark) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.CARPARK, carpark) - def write_walkway(self, walkway: AbstractWalkway) -> None: + def write_walkway(self, walkway: Walkway) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.WALKWAY, walkway) - def write_generic_drivable(self, obj: AbstractGenericDrivable) -> None: + def write_generic_drivable(self, obj: GenericDrivable) -> None: """Inherited, see superclass.""" self._write_surface_layer(MapLayer.GENERIC_DRIVABLE, obj) - def write_stop_line(self, stop_line: AbstractStopLine) -> None: + def write_stop_zone(self, stop_zone: StopZone) -> None: """Inherited, see superclass.""" # self._write_line_layer(MapLayer.STOP_LINE, stop_line) - raise NotImplementedError("Stop lines are not yet supported in GPKG maps.") + raise NotImplementedError("Stop zones are not yet supported in GPKG maps.") - def write_road_edge(self, road_edge: AbstractRoadEdge) -> None: + def write_road_edge(self, road_edge: RoadEdge) -> None: """Inherited, see superclass.""" self._write_line_layer(MapLayer.ROAD_EDGE, road_edge) - self._map_data[MapLayer.ROAD_EDGE]["road_edge_type"].append(road_edge.road_edge_type) + self._map_data[MapLayer.ROAD_EDGE]["road_edge_type"].append(int(road_edge.road_edge_type)) - def write_road_line(self, road_line: AbstractRoadLine) -> None: + def write_road_line(self, road_line: RoadLine) -> None: """Inherited, see superclass.""" self._write_line_layer(MapLayer.ROAD_LINE, road_line) - self._map_data[MapLayer.ROAD_LINE]["road_line_type"].append(road_line.road_line_type) + self._map_data[MapLayer.ROAD_LINE]["road_line_type"].append(int(road_line.road_line_type)) def close(self) -> None: """Inherited, see superclass.""" @@ -143,7 +147,7 @@ def close(self) -> None: ) # Optionally remap string IDs to integers - if self._remap_ids: + if self._remap_map_ids: _map_ids_to_integer(map_gdf) # Write each map layer to the GPKG file @@ -163,8 +167,9 @@ def _assert_initialized(self) -> None: assert self._map_data is not None, "Call reset() before writing data." assert self._map_file is not None, "Call reset() before writing data." assert self._map_metadata is not None, "Call reset() before writing data." + assert self._remap_map_ids is not None, "Call reset() before writing data." - def _write_surface_layer(self, layer: MapLayer, surface_object: AbstractSurfaceMapObject) -> None: + def _write_surface_layer(self, layer: MapLayer, surface_object: BaseMapSurfaceObject) -> None: """Helper to write surface map objects. :param layer: map layer of surface object @@ -177,7 +182,7 @@ def _write_surface_layer(self, layer: MapLayer, surface_object: AbstractSurfaceM self._map_data[layer]["outline"].append(surface_object.outline.linestring) self._map_data[layer]["geometry"].append(surface_object.shapely_polygon) - def _write_line_layer(self, layer: MapLayer, line_object: AbstractLineMapObject) -> None: + def _write_line_layer(self, layer: MapLayer, line_object: BaseMapLineObject) -> None: """Helper to write line map objects. :param layer: map layer of line object diff --git a/src/py123d/conversion/map_writer/utils/gpkg_utils.py b/src/py123d/conversion/map_writer/utils/gpkg_utils.py index 2b9ab334..87b99132 100644 --- a/src/py123d/conversion/map_writer/utils/gpkg_utils.py +++ b/src/py123d/conversion/map_writer/utils/gpkg_utils.py @@ -8,6 +8,7 @@ @dataclass class IntIDMapping: + """Class to map string IDs to integer IDs and vice versa.""" str_to_int: Dict[str, int] @@ -16,21 +17,29 @@ def __post_init__(self): @classmethod def from_series(cls, series: pd.Series) -> IntIDMapping: + """Creates an IntIDMapping from a pandas Series of string-like IDs.""" + # Drop NaN values and convert all to strings unique_ids = series.dropna().astype(str).unique() str_to_int = {str_id: idx for idx, str_id in enumerate(unique_ids)} return IntIDMapping(str_to_int) def map(self, str_like: Any) -> Optional[int]: - # Handle NaN and None values + """Maps a string-like ID to its corresponding integer ID.""" + + # NOTE: We need to convert a string-like input to an integer ID if pd.isna(str_like) or str_like is None: return None - # Convert to string for uniform handling - str_key = str(str_like) - return self.str_to_int.get(str_key, None) + if isinstance(str_like, float): + key = str(int(str_like)) # Convert float to int first to avoid decimal point + else: + key = str(str_like) + + return self.str_to_int.get(key, None) def map_list(self, id_list: Optional[List[str]]) -> List[int]: + """Maps a list of string-like IDs to their corresponding integer IDs.""" if id_list is None: return [] list_ = [] @@ -39,18 +48,3 @@ def map_list(self, id_list: Optional[List[str]]) -> List[int]: if mapped_id is not None: list_.append(mapped_id) return list_ - - -class IncrementalIntIDMapping: - - def __init__(self): - self.str_to_int: Dict[str, int] = {} - self.int_to_str: Dict[int, str] = {} - self.next_id: int = 0 - - def get_int_id(self, str_id: str) -> int: - if str_id not in self.str_to_int: - self.str_to_int[str_id] = self.next_id - self.int_to_str[self.next_id] = str_id - self.next_id += 1 - return self.str_to_int[str_id] diff --git a/src/py123d/conversion/registry/__init__.py b/src/py123d/conversion/registry/__init__.py index e69de29b..48dd1b45 100644 --- a/src/py123d/conversion/registry/__init__.py +++ b/src/py123d/conversion/registry/__init__.py @@ -0,0 +1,23 @@ +from py123d.conversion.registry.box_detection_label_registry import ( + BOX_DETECTION_LABEL_REGISTRY, + AV2SensorBoxDetectionLabel, + BoxDetectionLabel, + DefaultBoxDetectionLabel, + KITTI360BoxDetectionLabel, + NuPlanBoxDetectionLabel, + NuScenesBoxDetectionLabel, + PandasetBoxDetectionLabel, + WOPDBoxDetectionLabel, +) +from py123d.conversion.registry.lidar_index_registry import ( + LIDAR_INDEX_REGISTRY, + AV2SensorLiDARIndex, + CARLALiDARIndex, + DefaultLiDARIndex, + KITTI360LiDARIndex, + LiDARIndex, + NuPlanLiDARIndex, + NuScenesLiDARIndex, + PandasetLiDARIndex, + WOPDLiDARIndex, +) diff --git a/src/py123d/conversion/registry/box_detection_label_registry.py b/src/py123d/conversion/registry/box_detection_label_registry.py index 0d39e2d3..5cbfc98a 100644 --- a/src/py123d/conversion/registry/box_detection_label_registry.py +++ b/src/py123d/conversion/registry/box_detection_label_registry.py @@ -8,34 +8,41 @@ def register_box_detection_label(enum_class): + """Decorator to register a BoxDetectionLabel enum class.""" BOX_DETECTION_LABEL_REGISTRY[enum_class.__name__] = enum_class return enum_class class BoxDetectionLabel(SerialIntEnum): + """Base class for all box detection label enums.""" @abc.abstractmethod def to_default(self) -> DefaultBoxDetectionLabel: - raise NotImplementedError("Subclasses must implement this method.") + """Convert to the default box detection label.""" @register_box_detection_label class DefaultBoxDetectionLabel(BoxDetectionLabel): - """ - Enum for agents in py123d. - """ + """Default box detection labels used in 123D. Common labels across datasets.""" - VEHICLE = 0 - BICYCLE = 1 - PEDESTRIAN = 2 + # Vehicles + EGO = 0 + VEHICLE = 1 + TRAIN = 2 - TRAFFIC_CONE = 3 - BARRIER = 4 - CZONE_SIGN = 5 - GENERIC_OBJECT = 6 + # Vulnerable Road Users + BICYCLE = 3 + PERSON = 4 + ANIMAL = 5 - EGO = 7 - SIGN = 8 # TODO: Remove or extent + # Traffic Control + TRAFFIC_SIGN = 6 + TRAFFIC_CONE = 7 + TRAFFIC_LIGHT = 8 + + # Other Obstacles + BARRIER = 9 + GENERIC_OBJECT = 10 def to_default(self) -> DefaultBoxDetectionLabel: """Inherited, see superclass.""" @@ -44,78 +51,79 @@ def to_default(self) -> DefaultBoxDetectionLabel: @register_box_detection_label class AV2SensorBoxDetectionLabel(BoxDetectionLabel): - """Sensor dataset annotation categories.""" + """Argoverse 2 Sensor dataset annotation categories.""" - ANIMAL = 1 - ARTICULATED_BUS = 2 - BICYCLE = 3 - BICYCLIST = 4 - BOLLARD = 5 - BOX_TRUCK = 6 - BUS = 7 - CONSTRUCTION_BARREL = 8 - CONSTRUCTION_CONE = 9 - DOG = 10 - LARGE_VEHICLE = 11 - MESSAGE_BOARD_TRAILER = 12 - MOBILE_PEDESTRIAN_CROSSING_SIGN = 13 - MOTORCYCLE = 14 - MOTORCYCLIST = 15 - OFFICIAL_SIGNALER = 16 - PEDESTRIAN = 17 - RAILED_VEHICLE = 18 - REGULAR_VEHICLE = 19 - SCHOOL_BUS = 20 - SIGN = 21 - STOP_SIGN = 22 - STROLLER = 23 - TRAFFIC_LIGHT_TRAILER = 24 - TRUCK = 25 - TRUCK_CAB = 26 - VEHICULAR_TRAILER = 27 - WHEELCHAIR = 28 - WHEELED_DEVICE = 29 - WHEELED_RIDER = 30 + ANIMAL = 0 + ARTICULATED_BUS = 1 + BICYCLE = 2 + BICYCLIST = 3 + BOLLARD = 4 + BOX_TRUCK = 5 + BUS = 6 + CONSTRUCTION_BARREL = 7 + CONSTRUCTION_CONE = 8 + DOG = 9 + LARGE_VEHICLE = 10 + MESSAGE_BOARD_TRAILER = 11 + MOBILE_PEDESTRIAN_CROSSING_SIGN = 12 + MOTORCYCLE = 13 + MOTORCYCLIST = 14 + OFFICIAL_SIGNALER = 15 + PEDESTRIAN = 16 + RAILED_VEHICLE = 17 + REGULAR_VEHICLE = 18 + SCHOOL_BUS = 19 + SIGN = 20 + STOP_SIGN = 21 + STROLLER = 22 + TRAFFIC_LIGHT_TRAILER = 23 + TRUCK = 24 + TRUCK_CAB = 25 + VEHICULAR_TRAILER = 26 + WHEELCHAIR = 27 + WHEELED_DEVICE = 28 + WHEELED_RIDER = 29 def to_default(self) -> DefaultBoxDetectionLabel: """Inherited, see superclass.""" mapping = { - AV2SensorBoxDetectionLabel.ANIMAL: DefaultBoxDetectionLabel.GENERIC_OBJECT, + AV2SensorBoxDetectionLabel.ANIMAL: DefaultBoxDetectionLabel.ANIMAL, AV2SensorBoxDetectionLabel.ARTICULATED_BUS: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.BICYCLE: DefaultBoxDetectionLabel.BICYCLE, - AV2SensorBoxDetectionLabel.BICYCLIST: DefaultBoxDetectionLabel.PEDESTRIAN, - AV2SensorBoxDetectionLabel.BOLLARD: DefaultBoxDetectionLabel.BARRIER, + AV2SensorBoxDetectionLabel.BICYCLIST: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.BOLLARD: DefaultBoxDetectionLabel.GENERIC_OBJECT, AV2SensorBoxDetectionLabel.BOX_TRUCK: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.BUS: DefaultBoxDetectionLabel.VEHICLE, - AV2SensorBoxDetectionLabel.CONSTRUCTION_BARREL: DefaultBoxDetectionLabel.BARRIER, + AV2SensorBoxDetectionLabel.CONSTRUCTION_BARREL: DefaultBoxDetectionLabel.TRAFFIC_CONE, AV2SensorBoxDetectionLabel.CONSTRUCTION_CONE: DefaultBoxDetectionLabel.TRAFFIC_CONE, - AV2SensorBoxDetectionLabel.DOG: DefaultBoxDetectionLabel.GENERIC_OBJECT, + AV2SensorBoxDetectionLabel.DOG: DefaultBoxDetectionLabel.ANIMAL, AV2SensorBoxDetectionLabel.LARGE_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.MESSAGE_BOARD_TRAILER: DefaultBoxDetectionLabel.VEHICLE, - AV2SensorBoxDetectionLabel.MOBILE_PEDESTRIAN_CROSSING_SIGN: DefaultBoxDetectionLabel.CZONE_SIGN, + AV2SensorBoxDetectionLabel.MOBILE_PEDESTRIAN_CROSSING_SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, AV2SensorBoxDetectionLabel.MOTORCYCLE: DefaultBoxDetectionLabel.BICYCLE, - AV2SensorBoxDetectionLabel.MOTORCYCLIST: DefaultBoxDetectionLabel.BICYCLE, - AV2SensorBoxDetectionLabel.OFFICIAL_SIGNALER: DefaultBoxDetectionLabel.PEDESTRIAN, - AV2SensorBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PEDESTRIAN, - AV2SensorBoxDetectionLabel.RAILED_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, + AV2SensorBoxDetectionLabel.MOTORCYCLIST: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.OFFICIAL_SIGNALER: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.RAILED_VEHICLE: DefaultBoxDetectionLabel.TRAIN, AV2SensorBoxDetectionLabel.REGULAR_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.SCHOOL_BUS: DefaultBoxDetectionLabel.VEHICLE, - AV2SensorBoxDetectionLabel.SIGN: DefaultBoxDetectionLabel.SIGN, - AV2SensorBoxDetectionLabel.STOP_SIGN: DefaultBoxDetectionLabel.SIGN, - AV2SensorBoxDetectionLabel.STROLLER: DefaultBoxDetectionLabel.PEDESTRIAN, + AV2SensorBoxDetectionLabel.SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, + AV2SensorBoxDetectionLabel.STOP_SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, + AV2SensorBoxDetectionLabel.STROLLER: DefaultBoxDetectionLabel.PERSON, AV2SensorBoxDetectionLabel.TRAFFIC_LIGHT_TRAILER: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.TRUCK: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.TRUCK_CAB: DefaultBoxDetectionLabel.VEHICLE, AV2SensorBoxDetectionLabel.VEHICULAR_TRAILER: DefaultBoxDetectionLabel.VEHICLE, - AV2SensorBoxDetectionLabel.WHEELCHAIR: DefaultBoxDetectionLabel.PEDESTRIAN, - AV2SensorBoxDetectionLabel.WHEELED_DEVICE: DefaultBoxDetectionLabel.GENERIC_OBJECT, - AV2SensorBoxDetectionLabel.WHEELED_RIDER: DefaultBoxDetectionLabel.BICYCLE, + AV2SensorBoxDetectionLabel.WHEELCHAIR: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.WHEELED_DEVICE: DefaultBoxDetectionLabel.PERSON, + AV2SensorBoxDetectionLabel.WHEELED_RIDER: DefaultBoxDetectionLabel.PERSON, } return mapping[self] @register_box_detection_label class KITTI360BoxDetectionLabel(BoxDetectionLabel): + """KITTI-360 dataset annotation categories.""" BICYCLE = 0 BOX = 1 @@ -138,6 +146,7 @@ class KITTI360BoxDetectionLabel(BoxDetectionLabel): VENDING_MACHINE = 18 def to_default(self) -> DefaultBoxDetectionLabel: + """Inherited, see superclass.""" mapping = { KITTI360BoxDetectionLabel.BICYCLE: DefaultBoxDetectionLabel.BICYCLE, KITTI360BoxDetectionLabel.BOX: DefaultBoxDetectionLabel.GENERIC_OBJECT, @@ -146,13 +155,13 @@ def to_default(self) -> DefaultBoxDetectionLabel: KITTI360BoxDetectionLabel.CARAVAN: DefaultBoxDetectionLabel.VEHICLE, KITTI360BoxDetectionLabel.LAMP: DefaultBoxDetectionLabel.GENERIC_OBJECT, KITTI360BoxDetectionLabel.MOTORCYCLE: DefaultBoxDetectionLabel.BICYCLE, - KITTI360BoxDetectionLabel.PERSON: DefaultBoxDetectionLabel.PEDESTRIAN, + KITTI360BoxDetectionLabel.PERSON: DefaultBoxDetectionLabel.PERSON, KITTI360BoxDetectionLabel.POLE: DefaultBoxDetectionLabel.GENERIC_OBJECT, - KITTI360BoxDetectionLabel.RIDER: DefaultBoxDetectionLabel.BICYCLE, + KITTI360BoxDetectionLabel.RIDER: DefaultBoxDetectionLabel.PERSON, KITTI360BoxDetectionLabel.SMALLPOLE: DefaultBoxDetectionLabel.GENERIC_OBJECT, - KITTI360BoxDetectionLabel.STOP: DefaultBoxDetectionLabel.SIGN, - KITTI360BoxDetectionLabel.TRAFFIC_LIGHT: DefaultBoxDetectionLabel.SIGN, - KITTI360BoxDetectionLabel.TRAFFIC_SIGN: DefaultBoxDetectionLabel.SIGN, + KITTI360BoxDetectionLabel.STOP: DefaultBoxDetectionLabel.TRAFFIC_SIGN, + KITTI360BoxDetectionLabel.TRAFFIC_LIGHT: DefaultBoxDetectionLabel.TRAFFIC_LIGHT, + KITTI360BoxDetectionLabel.TRAFFIC_SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, KITTI360BoxDetectionLabel.TRAILER: DefaultBoxDetectionLabel.VEHICLE, KITTI360BoxDetectionLabel.TRAIN: DefaultBoxDetectionLabel.VEHICLE, KITTI360BoxDetectionLabel.TRASH_BIN: DefaultBoxDetectionLabel.GENERIC_OBJECT, @@ -164,35 +173,38 @@ def to_default(self) -> DefaultBoxDetectionLabel: @register_box_detection_label class NuPlanBoxDetectionLabel(BoxDetectionLabel): - """ - Semantic labels for nuPlan bounding box detections. - - Descriptions in `.db` files: - - vehicle: Includes all four or more wheeled vehicles, as well as trailers. - - bicycle: Includes bicycles, motorcycles and tricycles. - - pedestrian: All types of pedestrians, incl. strollers and wheelchairs. - - traffic_cone: Cones that are temporarily placed to control the flow of traffic. - - barrier: Solid barriers that can be either temporary or permanent. - - czone_sign: Temporary signs that indicate construction zones. - - generic_object: Animals, debris, pushable/pullable objects, permanent poles. - """ + """Semantic labels for nuPlan bounding box detections.""" VEHICLE = 0 + """Includes all four or more wheeled vehicles, as well as trailers.""" + BICYCLE = 1 + """Includes bicycles, motorcycles and tricycles.""" + PEDESTRIAN = 2 + """All types of pedestrians, incl. strollers and wheelchairs.""" + TRAFFIC_CONE = 3 + """Cones that are temporarily placed to control the flow of traffic.""" + BARRIER = 4 + """Solid barriers that can be either temporary or permanent.""" + CZONE_SIGN = 5 + """Temporary signs that indicate construction zones.""" + GENERIC_OBJECT = 6 + """Animals, debris, pushable/pullable objects, permanent poles.""" def to_default(self) -> DefaultBoxDetectionLabel: + """Inherited, see superclass.""" mapping = { NuPlanBoxDetectionLabel.VEHICLE: DefaultBoxDetectionLabel.VEHICLE, NuPlanBoxDetectionLabel.BICYCLE: DefaultBoxDetectionLabel.BICYCLE, - NuPlanBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PEDESTRIAN, + NuPlanBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PERSON, NuPlanBoxDetectionLabel.TRAFFIC_CONE: DefaultBoxDetectionLabel.TRAFFIC_CONE, NuPlanBoxDetectionLabel.BARRIER: DefaultBoxDetectionLabel.BARRIER, - NuPlanBoxDetectionLabel.CZONE_SIGN: DefaultBoxDetectionLabel.CZONE_SIGN, + NuPlanBoxDetectionLabel.CZONE_SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, NuPlanBoxDetectionLabel.GENERIC_OBJECT: DefaultBoxDetectionLabel.GENERIC_OBJECT, } return mapping[self] @@ -230,6 +242,7 @@ class NuScenesBoxDetectionLabel(BoxDetectionLabel): ANIMAL = 22 def to_default(self): + """Inherited, see superclass.""" mapping = { NuScenesBoxDetectionLabel.VEHICLE_CAR: DefaultBoxDetectionLabel.VEHICLE, NuScenesBoxDetectionLabel.VEHICLE_TRUCK: DefaultBoxDetectionLabel.VEHICLE, @@ -241,19 +254,19 @@ def to_default(self): NuScenesBoxDetectionLabel.VEHICLE_TRAILER: DefaultBoxDetectionLabel.VEHICLE, NuScenesBoxDetectionLabel.VEHICLE_BICYCLE: DefaultBoxDetectionLabel.BICYCLE, NuScenesBoxDetectionLabel.VEHICLE_MOTORCYCLE: DefaultBoxDetectionLabel.BICYCLE, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_ADULT: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_CHILD: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_CONSTRUCTION_WORKER: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_PERSONAL_MOBILITY: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_POLICE_OFFICER: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_STROLLER: DefaultBoxDetectionLabel.PEDESTRIAN, - NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_WHEELCHAIR: DefaultBoxDetectionLabel.PEDESTRIAN, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_ADULT: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_CHILD: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_CONSTRUCTION_WORKER: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_PERSONAL_MOBILITY: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_POLICE_OFFICER: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_STROLLER: DefaultBoxDetectionLabel.PERSON, + NuScenesBoxDetectionLabel.HUMAN_PEDESTRIAN_WHEELCHAIR: DefaultBoxDetectionLabel.PERSON, NuScenesBoxDetectionLabel.MOVABLE_OBJECT_TRAFFICCONE: DefaultBoxDetectionLabel.TRAFFIC_CONE, NuScenesBoxDetectionLabel.MOVABLE_OBJECT_BARRIER: DefaultBoxDetectionLabel.BARRIER, NuScenesBoxDetectionLabel.MOVABLE_OBJECT_PUSHABLE_PULLABLE: DefaultBoxDetectionLabel.GENERIC_OBJECT, NuScenesBoxDetectionLabel.MOVABLE_OBJECT_DEBRIS: DefaultBoxDetectionLabel.GENERIC_OBJECT, NuScenesBoxDetectionLabel.STATIC_OBJECT_BICYCLE_RACK: DefaultBoxDetectionLabel.GENERIC_OBJECT, - NuScenesBoxDetectionLabel.ANIMAL: DefaultBoxDetectionLabel.GENERIC_OBJECT, + NuScenesBoxDetectionLabel.ANIMAL: DefaultBoxDetectionLabel.ANIMAL, } return mapping[self] @@ -294,34 +307,35 @@ class PandasetBoxDetectionLabel(BoxDetectionLabel): TRAM_SUBWAY = 26 def to_default(self) -> DefaultBoxDetectionLabel: + """Inherited, see superclass.""" mapping = { - PandasetBoxDetectionLabel.ANIMALS_BIRD: DefaultBoxDetectionLabel.GENERIC_OBJECT, # TODO: Adjust default types - PandasetBoxDetectionLabel.ANIMALS_OTHER: DefaultBoxDetectionLabel.GENERIC_OBJECT, # TODO: Adjust default types + PandasetBoxDetectionLabel.ANIMALS_BIRD: DefaultBoxDetectionLabel.ANIMAL, + PandasetBoxDetectionLabel.ANIMALS_OTHER: DefaultBoxDetectionLabel.ANIMAL, PandasetBoxDetectionLabel.BICYCLE: DefaultBoxDetectionLabel.BICYCLE, PandasetBoxDetectionLabel.BUS: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.CAR: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.CONES: DefaultBoxDetectionLabel.TRAFFIC_CONE, - PandasetBoxDetectionLabel.CONSTRUCTION_SIGNS: DefaultBoxDetectionLabel.CZONE_SIGN, + PandasetBoxDetectionLabel.CONSTRUCTION_SIGNS: DefaultBoxDetectionLabel.TRAFFIC_SIGN, PandasetBoxDetectionLabel.EMERGENCY_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.MEDIUM_SIZED_TRUCK: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.MOTORCYCLE: DefaultBoxDetectionLabel.BICYCLE, PandasetBoxDetectionLabel.MOTORIZED_SCOOTER: DefaultBoxDetectionLabel.BICYCLE, PandasetBoxDetectionLabel.OTHER_VEHICLE_CONSTRUCTION_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, - PandasetBoxDetectionLabel.OTHER_VEHICLE_PEDICAB: DefaultBoxDetectionLabel.BICYCLE, + PandasetBoxDetectionLabel.OTHER_VEHICLE_PEDICAB: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.OTHER_VEHICLE_UNCOMMON: DefaultBoxDetectionLabel.VEHICLE, - PandasetBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PEDESTRIAN, - PandasetBoxDetectionLabel.PEDESTRIAN_WITH_OBJECT: DefaultBoxDetectionLabel.PEDESTRIAN, + PandasetBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PERSON, + PandasetBoxDetectionLabel.PEDESTRIAN_WITH_OBJECT: DefaultBoxDetectionLabel.PERSON, PandasetBoxDetectionLabel.PERSONAL_MOBILITY_DEVICE: DefaultBoxDetectionLabel.BICYCLE, PandasetBoxDetectionLabel.PICKUP_TRUCK: DefaultBoxDetectionLabel.VEHICLE, PandasetBoxDetectionLabel.PYLONS: DefaultBoxDetectionLabel.TRAFFIC_CONE, PandasetBoxDetectionLabel.ROAD_BARRIERS: DefaultBoxDetectionLabel.BARRIER, PandasetBoxDetectionLabel.ROLLING_CONTAINERS: DefaultBoxDetectionLabel.GENERIC_OBJECT, PandasetBoxDetectionLabel.SEMI_TRUCK: DefaultBoxDetectionLabel.VEHICLE, - PandasetBoxDetectionLabel.SIGNS: DefaultBoxDetectionLabel.SIGN, + PandasetBoxDetectionLabel.SIGNS: DefaultBoxDetectionLabel.TRAFFIC_SIGN, PandasetBoxDetectionLabel.TEMPORARY_CONSTRUCTION_BARRIERS: DefaultBoxDetectionLabel.BARRIER, PandasetBoxDetectionLabel.TOWED_OBJECT: DefaultBoxDetectionLabel.VEHICLE, - PandasetBoxDetectionLabel.TRAIN: DefaultBoxDetectionLabel.GENERIC_OBJECT, # TODO: Adjust default types - PandasetBoxDetectionLabel.TRAM_SUBWAY: DefaultBoxDetectionLabel.GENERIC_OBJECT, # TODO: Adjust default types + PandasetBoxDetectionLabel.TRAIN: DefaultBoxDetectionLabel.TRAIN, # TODO: Adjust default types + PandasetBoxDetectionLabel.TRAM_SUBWAY: DefaultBoxDetectionLabel.TRAIN, # TODO: Adjust default types } return mapping[self] @@ -341,11 +355,12 @@ class WOPDBoxDetectionLabel(BoxDetectionLabel): TYPE_CYCLIST = 4 def to_default(self) -> DefaultBoxDetectionLabel: + """Inherited, see superclass.""" mapping = { WOPDBoxDetectionLabel.TYPE_UNKNOWN: DefaultBoxDetectionLabel.GENERIC_OBJECT, WOPDBoxDetectionLabel.TYPE_VEHICLE: DefaultBoxDetectionLabel.VEHICLE, - WOPDBoxDetectionLabel.TYPE_PEDESTRIAN: DefaultBoxDetectionLabel.PEDESTRIAN, - WOPDBoxDetectionLabel.TYPE_SIGN: DefaultBoxDetectionLabel.SIGN, + WOPDBoxDetectionLabel.TYPE_PEDESTRIAN: DefaultBoxDetectionLabel.PERSON, + WOPDBoxDetectionLabel.TYPE_SIGN: DefaultBoxDetectionLabel.TRAFFIC_SIGN, WOPDBoxDetectionLabel.TYPE_CYCLIST: DefaultBoxDetectionLabel.BICYCLE, } return mapping[self] diff --git a/src/py123d/conversion/registry/lidar_index_registry.py b/src/py123d/conversion/registry/lidar_index_registry.py index a65903b4..593a47c6 100644 --- a/src/py123d/conversion/registry/lidar_index_registry.py +++ b/src/py123d/conversion/registry/lidar_index_registry.py @@ -1,16 +1,21 @@ +from __future__ import annotations + from enum import IntEnum +from typing import Dict from py123d.common.utils.enums import classproperty -LIDAR_INDEX_REGISTRY = {} +LIDAR_INDEX_REGISTRY: Dict[str, LiDARIndex] = {} def register_lidar_index(enum_class): + """Decorator to register a LiDARIndex enum class.""" LIDAR_INDEX_REGISTRY[enum_class.__name__] = enum_class return enum_class class LiDARIndex(IntEnum): + """Base class for all LiDAR Index enums. Defines common indices for LiDAR point clouds.""" @classproperty def XY(self) -> slice: @@ -29,6 +34,8 @@ def XYZ(self) -> slice: @register_lidar_index class DefaultLiDARIndex(LiDARIndex): + """Default LiDAR indices for XYZ point clouds.""" + X = 0 Y = 1 Z = 2 @@ -36,6 +43,8 @@ class DefaultLiDARIndex(LiDARIndex): @register_lidar_index class NuPlanLiDARIndex(LiDARIndex): + """LiDAR Indexing Scheme for the nuPlan dataset.""" + X = 0 Y = 1 Z = 2 @@ -45,6 +54,8 @@ class NuPlanLiDARIndex(LiDARIndex): @register_lidar_index class CARLALiDARIndex(LiDARIndex): + """LiDAR Indexing Scheme for the CARLA.""" + X = 0 Y = 1 Z = 2 @@ -53,6 +64,8 @@ class CARLALiDARIndex(LiDARIndex): @register_lidar_index class WOPDLiDARIndex(LiDARIndex): + """Waymo Open Perception Dataset (WOPD) LiDAR Indexing Scheme, with polar features.""" + RANGE = 0 INTENSITY = 1 ELONGATION = 2 @@ -62,7 +75,9 @@ class WOPDLiDARIndex(LiDARIndex): @register_lidar_index -class Kitti360LiDARIndex(LiDARIndex): +class KITTI360LiDARIndex(LiDARIndex): + """KITTI-360 LiDAR Indexing Scheme.""" + X = 0 Y = 1 Z = 2 @@ -70,11 +85,8 @@ class Kitti360LiDARIndex(LiDARIndex): @register_lidar_index -class AVSensorLiDARIndex(LiDARIndex): - """Argoverse Sensor LiDAR Indexing Scheme. - - NOTE: The LiDAR files also include, 'offset_ns', which we do not currently include. - """ +class AV2SensorLiDARIndex(LiDARIndex): + """Argoverse 2 Sensor LiDAR Indexing Scheme.""" X = 0 Y = 1 @@ -94,6 +106,8 @@ class PandasetLiDARIndex(LiDARIndex): @register_lidar_index class NuScenesLiDARIndex(LiDARIndex): + """NuScenes LiDAR Indexing Scheme.""" + X = 0 Y = 1 Z = 2 diff --git a/src/py123d/conversion/sensor_io/camera/jpeg_camera_io.py b/src/py123d/conversion/sensor_io/camera/jpeg_camera_io.py index d942e729..c2fc4d17 100644 --- a/src/py123d/conversion/sensor_io/camera/jpeg_camera_io.py +++ b/src/py123d/conversion/sensor_io/camera/jpeg_camera_io.py @@ -5,25 +5,41 @@ import numpy.typing as npt +def is_jpeg_binary(jpeg_binary: bytes) -> bool: + """Check if the given binary data represents a JPEG image. + + :param jpeg_binary: The binary data to check. + :return: True if the binary data is a JPEG image, False otherwise. + """ + SOI_MARKER = b"\xff\xd8" # Start Of Image + EOI_MARKER = b"\xff\xd9" # End Of Image + + return jpeg_binary.startswith(SOI_MARKER) and jpeg_binary.endswith(EOI_MARKER) + + def encode_image_as_jpeg_binary(image: npt.NDArray[np.uint8]) -> bytes: + """Encodes a numpy image as JPEG binary.""" _, encoded_img = cv2.imencode(".jpg", image) jpeg_binary = encoded_img.tobytes() return jpeg_binary def decode_image_from_jpeg_binary(jpeg_binary: bytes) -> npt.NDArray[np.uint8]: + """Decodes a numpy image from JPEG binary.""" image = cv2.imdecode(np.frombuffer(jpeg_binary, np.uint8), cv2.IMREAD_UNCHANGED) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) return image def load_jpeg_binary_from_jpeg_file(jpeg_path: Path) -> bytes: + """Loads JPEG binary data from a JPEG file.""" with open(jpeg_path, "rb") as f: jpeg_binary = f.read() return jpeg_binary def load_image_from_jpeg_file(jpeg_path: Path) -> npt.NDArray[np.uint8]: + """Loads a numpy image from a JPEG file.""" image = cv2.imread(str(jpeg_path), cv2.IMREAD_COLOR) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) return image diff --git a/src/py123d/conversion/sensor_io/camera/mp4_camera_io.py b/src/py123d/conversion/sensor_io/camera/mp4_camera_io.py index 9cb91aad..50ce5563 100644 --- a/src/py123d/conversion/sensor_io/camera/mp4_camera_io.py +++ b/src/py123d/conversion/sensor_io/camera/mp4_camera_io.py @@ -1,8 +1,3 @@ -# TODO: add method of handling camera mp4 io -def load_image_from_mp4_file() -> None: - raise NotImplementedError - - from functools import lru_cache from pathlib import Path from typing import Optional, Union @@ -12,16 +7,15 @@ def load_image_from_mp4_file() -> None: class MP4Writer: - """Write images sequentially to an MP4 video file.""" + """Simple implementation of an MP4 video writer, based on OpenCV.""" def __init__(self, output_path: Union[str, Path], fps: float = 30.0, codec: str = "mp4v"): """ Initialize MP4 writer. - Args: - output_path: Path to output MP4 file - fps: Frames per second - codec: Video codec ('mp4v', 'avc1', 'h264') + :param output_path: The output path for the MP4 file. + :param fps: Frames per second, defaults to 30.0 + :param codec: Video codec, defaults to "mp4v" """ self.output_path = Path(output_path) self.fps = fps @@ -31,11 +25,11 @@ def __init__(self, output_path: Union[str, Path], fps: float = 30.0, codec: str self.frame_count = 0 def write_frame(self, frame: np.ndarray) -> int: - """ - Write a single frame to the video. + """Write a single frame to the video. - Args: - frame: Image as numpy array (RGB format) + :param frame: Image as numpy array (RGB format) + :raises ValueError: If frame size does not match video size + :return: Index of the written frame """ frame_idx = int(self.frame_count) if self.writer is None: @@ -47,7 +41,7 @@ def write_frame(self, frame: np.ndarray) -> int: self.writer = cv2.VideoWriter(self.output_path, fourcc, self.fps, self.frame_size) if frame.shape[:2][::-1] != self.frame_size: - raise ValueError(f"Frame size {frame.shape[:2][::-1]} doesn't match " f"video size {self.frame_size}") + raise ValueError(f"Frame size {frame.shape[:2][::-1]} doesn't match video size {self.frame_size}") frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) self.writer.write(frame) @@ -55,21 +49,22 @@ def write_frame(self, frame: np.ndarray) -> int: return frame_idx def close(self): - """Release the video writer.""" + """Close the video writer and finalize the MP4 file.""" if self.writer is not None: self.writer.release() self.writer = None class MP4Reader: - """Read MP4 video with random frame access.""" + """Simple implementation of an MP4 video reader, based on OpenCV.""" def __init__(self, video_path: Union[str, Path], read_all: bool = False): - """ - Initialize MP4 reader. + """Initializes the MP4Reader. - Args: - video_path: Path to MP4 file + :param video_path: Path to the MP4 video file. + :param read_all: Whether to read all frames into memory, defaults to False + :raises FileNotFoundError: If the video file does not exist + :raises ValueError: If the video file cannot be opened """ self.video_path = video_path if not Path(video_path).exists(): @@ -98,18 +93,15 @@ def __init__(self, video_path: Union[str, Path], read_all: bool = False): self.cap = None def get_frame(self, frame_index: int) -> Optional[np.ndarray]: - """ - Get a specific frame by index. + """Get a specific frame, an RBG image as numpy array, by its index. - Args: - frame_index: Zero-based frame index - - Returns: - Frame as numpy array (RGB format) or None if invalid index + :param frame_index: Index of the frame to retrieve. + :raises IndexError: If the frame index is out of range. + :return: The frame as a numpy array (RGB format) or None if the frame could not be read. """ if frame_index < 0 or frame_index >= self.frame_count: - raise IndexError(f"Frame index {frame_index} out of range " f"[0, {len(self.frames)})") + raise IndexError(f"Frame index {frame_index} out of range [0, {len(self.frames)})") if self.read_all: return self.frames[frame_index] @@ -123,9 +115,12 @@ def get_frame(self, frame_index: int) -> Optional[np.ndarray]: return frame if ret else None - def __getitem__(self, index: int) -> np.ndarray: - """Allow indexing like reader[10]""" - return self.get_frame(index) + def __del__(self): + """Destructor to release the video capture.""" + if self.cap is not None: + self.cap.release() + if self.read_all: + self.frames = [] @lru_cache(maxsize=64) diff --git a/src/py123d/conversion/sensor_io/camera/png_camera_io.py b/src/py123d/conversion/sensor_io/camera/png_camera_io.py new file mode 100644 index 00000000..6ab08599 --- /dev/null +++ b/src/py123d/conversion/sensor_io/camera/png_camera_io.py @@ -0,0 +1,44 @@ +from pathlib import Path + +import cv2 +import numpy as np +import numpy.typing as npt + + +def is_png_binary(png_binary: bytes) -> bool: + """Check if the given binary data represents a PNG image. + + :param png_binary: The binary data to check. + :return: True if the binary data is a PNG image, False otherwise. + """ + PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n" # PNG file signature + + return png_binary.startswith(PNG_SIGNATURE) + + +def encode_image_as_png_binary(image: npt.NDArray[np.uint8]) -> bytes: + """Encodes a numpy image as PNG binary.""" + _, encoded_img = cv2.imencode(".png", image) + png_binary = encoded_img.tobytes() + return png_binary + + +def decode_image_from_png_binary(png_binary: bytes) -> npt.NDArray[np.uint8]: + """Decodes a numpy image from PNG binary.""" + image = cv2.imdecode(np.frombuffer(png_binary, np.uint8), cv2.IMREAD_UNCHANGED) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return image + + +def load_png_binary_from_png_file(png_path: Path) -> bytes: + """Loads PNG binary data from a PNG file.""" + with open(png_path, "rb") as f: + png_binary = f.read() + return png_binary + + +def load_image_from_png_file(png_path: Path) -> npt.NDArray[np.uint8]: + """Loads a numpy image from a PNG file.""" + image = cv2.imread(str(png_path), cv2.IMREAD_COLOR) + image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) + return image diff --git a/src/py123d/conversion/sensor_io/lidar/draco_lidar_io.py b/src/py123d/conversion/sensor_io/lidar/draco_lidar_io.py index 61473f08..72715403 100644 --- a/src/py123d/conversion/sensor_io/lidar/draco_lidar_io.py +++ b/src/py123d/conversion/sensor_io/lidar/draco_lidar_io.py @@ -13,6 +13,16 @@ DRACO_PRESERVE_ORDER: Final[bool] = False +def is_draco_binary(draco_binary: bytes) -> bool: + """Check if the given binary data represents a Draco compressed point cloud. + + :param draco_binary: The binary data to check. + :return: True if the binary data is a Draco compressed point cloud, False otherwise. + """ + DRACO_MAGIC_NUMBER = b"DRACO" + return draco_binary.startswith(DRACO_MAGIC_NUMBER) + + def encode_lidar_pc_as_draco_binary(lidar_pc: npt.NDArray[np.float32], lidar_metadata: LiDARMetadata) -> bytes: """Compress LiDAR point cloud data using Draco format. diff --git a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py index 1a9e2583..bfc7ad83 100644 --- a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py +++ b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py @@ -5,7 +5,7 @@ import numpy.typing as npt from omegaconf import DictConfig -from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.datatypes.metadata.log_metadata import LogMetadata from py123d.datatypes.sensors.lidar import LiDARType from py123d.script.utils.dataset_path_utils import get_dataset_paths @@ -26,12 +26,24 @@ def load_lidar_pcs_from_file( index: Optional[int] = None, sensor_root: Optional[Union[str, Path]] = None, ) -> Dict[LiDARType, npt.NDArray[np.float32]]: - assert relative_path is not None, "Relative path to LiDAR file must be provided." + """Loads LiDAR point clouds from a file, based on the dataset specified in the log metadata. + + :param relative_path: Relative path to the LiDAR file. + :param log_metadata: Metadata containing dataset information. + :param index: Optional index for datasets that require it, defaults to None + :param sensor_root: Optional root path for sensor data, defaults to None + :raises NotImplementedError: If the dataset is not supported + :return: Dictionary mapping LiDAR types to their point cloud numpy arrays + """ + # NOTE @DanielDauner: This function is designed s.t. it can load multiple lidar types at the same time. + # Several datasets (e.g., PandaSet, nuScenes) have multiple LiDAR sensors stored in one file. + # Returning this as a dict allows us to handle this case without unnucessary io overhead. + assert relative_path is not None, "Relative path to LiDAR file must be provided." if sensor_root is None: - assert ( - log_metadata.dataset in DATASET_SENSOR_ROOT.keys() - ), f"Dataset path for sensor loading not found for dataset: {log_metadata.dataset}." + assert log_metadata.dataset in DATASET_SENSOR_ROOT.keys(), ( + f"Dataset path for sensor loading not found for dataset: {log_metadata.dataset}." + ) sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] assert sensor_root is not None, f"Dataset path for sensor loading not found for dataset: {log_metadata.dataset}" @@ -50,7 +62,7 @@ def load_lidar_pcs_from_file( lidar_pcs_dict = load_av2_sensor_lidar_pcs_from_file(full_lidar_path) elif log_metadata.dataset == "wopd": - from py123d.conversion.datasets.wopd.waymo_sensor_io import load_wopd_lidar_pcs_from_file + from py123d.conversion.datasets.wopd.wopd_sensor_io import load_wopd_lidar_pcs_from_file lidar_pcs_dict = load_wopd_lidar_pcs_from_file(full_lidar_path, index, keep_polar_features=False) diff --git a/src/py123d/conversion/sensor_io/lidar/laz_lidar_io.py b/src/py123d/conversion/sensor_io/lidar/laz_lidar_io.py index b109c7ca..2ad33714 100644 --- a/src/py123d/conversion/sensor_io/lidar/laz_lidar_io.py +++ b/src/py123d/conversion/sensor_io/lidar/laz_lidar_io.py @@ -7,6 +7,16 @@ from py123d.datatypes.sensors.lidar import LiDAR, LiDARMetadata +def is_laz_binary(laz_binary: bytes) -> bool: + """Check if the given binary data represents a LAZ compressed point cloud. + + :param laz_binary: The binary data to check. + :return: True if the binary data is a LAZ compressed point cloud, False otherwise. + """ + LAS_MAGIC_NUMBER = b"LASF" + return laz_binary[0:4] == LAS_MAGIC_NUMBER + + def encode_lidar_pc_as_laz_binary(point_cloud: npt.NDArray[np.float32], lidar_metadata: LiDARMetadata) -> bytes: """Compress LiDAR point cloud data using LAZ format. diff --git a/src/py123d/conversion/utils/map_utils/opendrive/opendrive_map_conversion.py b/src/py123d/conversion/utils/map_utils/opendrive/opendrive_map_conversion.py index 2bb831c8..0d1ccd78 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/opendrive_map_conversion.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/opendrive_map_conversion.py @@ -6,7 +6,7 @@ import shapely from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.conversion.utils.map_utils.opendrive.parser.opendrive import Junction, OpenDrive +from py123d.conversion.utils.map_utils.opendrive.parser.opendrive import XODR, Junction from py123d.conversion.utils.map_utils.opendrive.utils.collection import collect_element_helpers from py123d.conversion.utils.map_utils.opendrive.utils.lane_helper import ( OpenDriveLaneGroupHelper, @@ -21,18 +21,19 @@ get_road_edges_3d_from_drivable_surfaces, lift_outlines_to_3d, ) -from py123d.datatypes.maps.cache.cache_map_objects import ( - CacheCarpark, - CacheCrosswalk, - CacheGenericDrivable, - CacheIntersection, - CacheLane, - CacheLaneGroup, - CacheRoadEdge, - CacheRoadLine, - CacheWalkway, +from py123d.datatypes.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadEdgeType, + RoadLine, + RoadLineType, + Walkway, ) -from py123d.datatypes.maps.map_datatypes import RoadEdgeType, RoadLineType from py123d.geometry.geometry_index import Point3DIndex from py123d.geometry.polyline import Polyline3D @@ -47,8 +48,15 @@ def convert_xodr_map( interpolation_step_size: float = 1.0, connection_distance_threshold: float = 0.1, ) -> None: + """Converts an OpenDRIVE map file and the map objects to an 123D map using a map writer. - opendrive = OpenDrive.parse_from_file(xordr_file) + :param xordr_file: Path to the OpenDRIVE (.xodr) file. + :param map_writer: Map writer to write the extracted map objects. + :param interpolation_step_size: Step size for interpolating polylines, defaults to 1.0 + :param connection_distance_threshold: Distance threshold for connecting road elements, defaults to 0.1 + """ + + opendrive = XODR.parse_from_file(xordr_file) _, junction_dict, lane_helper_dict, lane_group_helper_dict, object_helper_dict = collect_element_helpers( opendrive, interpolation_step_size, connection_distance_threshold @@ -79,9 +87,10 @@ def convert_xodr_map( def _extract_and_write_lanes( lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper], map_writer: AbstractMapWriter, -) -> List[CacheLane]: +) -> List[Lane]: + """Extracts lanes from lane group helpers and writes them using the map writer.""" - lanes: List[CacheLane] = [] + lanes: List[Lane] = [] for lane_group_helper in lane_group_helper_dict.values(): lane_group_id = lane_group_helper.lane_group_id lane_helpers = lane_group_helper.lane_helpers @@ -90,7 +99,7 @@ def _extract_and_write_lanes( for lane_idx, lane_helper in enumerate(lane_helpers): left_lane_id = lane_helpers[lane_idx - 1].lane_id if lane_idx > 0 else None right_lane_id = lane_helpers[lane_idx + 1].lane_id if lane_idx < num_lanes - 1 else None - lane = CacheLane( + lane = Lane( object_id=lane_helper.lane_id, lane_group_id=lane_group_id, left_boundary=lane_helper.inner_polyline_3d, @@ -102,7 +111,6 @@ def _extract_and_write_lanes( successor_ids=lane_helper.successor_lane_ids, speed_limit_mps=lane_helper.speed_limit_mps, outline=lane_helper.outline_polyline_3d, - geometry=None, ) lanes.append(lane) map_writer.write_lane(lane) @@ -112,12 +120,13 @@ def _extract_and_write_lanes( def _extract_and_write_lane_groups( lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper], map_writer: AbstractMapWriter -) -> List[CacheLaneGroup]: +) -> List[LaneGroup]: + """Extracts lane groups from lane group helpers and writes them using the map writer.""" - lane_groups: List[CacheLaneGroup] = [] + lane_groups: List[LaneGroup] = [] for lane_group_helper in lane_group_helper_dict.values(): lane_group_helper: OpenDriveLaneGroupHelper - lane_group = CacheLaneGroup( + lane_group = LaneGroup( object_id=lane_group_helper.lane_group_id, lane_ids=[lane_helper.lane_id for lane_helper in lane_group_helper.lane_helpers], left_boundary=lane_group_helper.inner_polyline_3d, @@ -126,7 +135,6 @@ def _extract_and_write_lane_groups( predecessor_ids=lane_group_helper.predecessor_lane_group_ids, successor_ids=lane_group_helper.successor_lane_group_ids, outline=lane_group_helper.outline_polyline_3d, - geometry=None, ) lane_groups.append(lane_group) map_writer.write_lane_group(lane_group) @@ -135,46 +143,44 @@ def _extract_and_write_lane_groups( def _write_walkways(lane_helper_dict: Dict[str, OpenDriveLaneHelper], map_writer: AbstractMapWriter) -> None: + """Writes walkways from lane helpers using the map writer.""" for lane_helper in lane_helper_dict.values(): if lane_helper.type == "sidewalk": map_writer.write_walkway( - CacheWalkway( + Walkway( object_id=lane_helper.lane_id, outline=lane_helper.outline_polyline_3d, - geometry=None, ) ) def _extract_and_write_carparks( lane_helper_dict: Dict[str, OpenDriveLaneHelper], map_writer: AbstractMapWriter -) -> List[CacheCarpark]: - - carparks: List[CacheCarpark] = [] +) -> List[Carpark]: + """Extracts carparks from lane helpers and writes them using the map writer.""" + carparks: List[Carpark] = [] for lane_helper in lane_helper_dict.values(): if lane_helper.type == "parking": - carpark = CacheCarpark( + carpark = Carpark( object_id=lane_helper.lane_id, outline=lane_helper.outline_polyline_3d, - geometry=None, ) carparks.append(carpark) map_writer.write_carpark(carpark) - return carparks def _extract_and_write_generic_drivables( lane_helper_dict: Dict[str, OpenDriveLaneHelper], map_writer: AbstractMapWriter -) -> List[CacheGenericDrivable]: +) -> List[GenericDrivable]: + """Extracts generic drivables from lane helpers and writes them using the map writer.""" - generic_drivables: List[CacheGenericDrivable] = [] + generic_drivables: List[GenericDrivable] = [] for lane_helper in lane_helper_dict.values(): if lane_helper.type in ["none", "border", "bidirectional"]: - generic_drivable = CacheGenericDrivable( + generic_drivable = GenericDrivable( object_id=lane_helper.lane_id, outline=lane_helper.outline_polyline_3d, - geometry=None, ) generic_drivables.append(generic_drivable) map_writer.write_generic_drivable(generic_drivable) @@ -186,7 +192,6 @@ def _write_intersections( lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper], map_writer: AbstractMapWriter, ) -> None: - def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriveLaneGroupHelper]: return [ lane_group_helper @@ -204,11 +209,10 @@ def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriv # TODO @DanielDauner: Create a method that extracts 3D outlines of intersections. outline = _extract_intersection_outline(lane_group_helpers, junction.id) map_writer.write_intersection( - CacheIntersection( + Intersection( object_id=junction.id, lane_group_ids=lane_group_ids_, outline=outline, - geometry=None, ) ) @@ -216,16 +220,14 @@ def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriv def _write_crosswalks(object_helper_dict: Dict[int, OpenDriveObjectHelper], map_writer: AbstractMapWriter) -> None: for object_helper in object_helper_dict.values(): map_writer.write_crosswalk( - CacheCrosswalk( + Crosswalk( object_id=object_helper.object_id, outline=object_helper.outline_polyline_3d, - geometry=None, ) ) -def _write_road_lines(lanes: List[CacheLane], lane_groups: List[CacheLaneGroup], map_writer: AbstractMapWriter) -> None: - +def _write_road_lines(lanes: List[Lane], lane_groups: List[LaneGroup], map_writer: AbstractMapWriter) -> None: # NOTE @DanielDauner: This method of extracting road lines is very simplistic and needs improvement. # The OpenDRIVE format provides lane boundary types that could be used here. # Additionally, the logic of inferring road lines is somewhat flawed, e.g, assuming constant types/colors of lines. @@ -240,7 +242,6 @@ def _write_road_lines(lanes: List[CacheLane], lane_groups: List[CacheLaneGroup], running_id = 0 for lane in lanes: - on_intersection = lane_group_on_intersection.get(lane.lane_group_id, False) if on_intersection: # Skip road lines on intersections @@ -266,23 +267,16 @@ def _write_road_lines(lanes: List[CacheLane], lane_groups: List[CacheLaneGroup], running_id += 1 for object_id, road_line_type, polyline in zip(ids, road_line_types, polylines): - map_writer.write_road_line( - CacheRoadLine( - object_id=object_id, - road_line_type=road_line_type, - polyline=polyline, - ) - ) + map_writer.write_road_line(RoadLine(object_id=object_id, road_line_type=road_line_type, polyline=polyline)) def _write_road_edges( - lanes: List[CacheLane], - lane_groups: List[CacheLaneGroup], - car_parks: List[CacheCarpark], - generic_drivables: List[CacheGenericDrivable], + lanes: List[Lane], + lane_groups: List[LaneGroup], + car_parks: List[Carpark], + generic_drivables: List[GenericDrivable], map_writer: AbstractMapWriter, ) -> None: - road_edges_ = get_road_edges_3d_from_drivable_surfaces( lanes=lanes, lane_groups=lane_groups, @@ -297,7 +291,7 @@ def _write_road_edges( for road_edge_linestring in road_edge_linestrings: # TODO @DanielDauner: Figure out if other types should/could be assigned here. map_writer.write_road_edge( - CacheRoadEdge( + RoadEdge( object_id=running_id, road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, polyline=Polyline3D.from_linestring(road_edge_linestring), diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/elevation.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/elevation.py index 2a1261ca..599fde63 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/elevation.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/elevation.py @@ -4,35 +4,35 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import Polynomial +from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import XODRPolynomial @dataclass -class ElevationProfile: +class XORDElevationProfile: """ Models elevation along s-axis of reference line. https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-1d876c00-d69e-46d9-bbcd-709ab48f14b1 """ - elevations: List[Elevation] + elevations: List[XODRElevation] def __post_init__(self): self.elevations.sort(key=lambda x: x.s, reverse=False) @classmethod - def parse(cls, elevation_profile_element: Optional[Element]) -> ElevationProfile: + def parse(cls, elevation_profile_element: Optional[Element]) -> XORDElevationProfile: args = {} - elevations: List[Elevation] = [] + elevations: List[XODRElevation] = [] if elevation_profile_element is not None: for elevation_element in elevation_profile_element.findall("elevation"): - elevations.append(Elevation.parse(elevation_element)) + elevations.append(XODRElevation.parse(elevation_element)) args["elevations"] = elevations - return ElevationProfile(**args) + return XORDElevationProfile(**args) @dataclass -class Elevation(Polynomial): +class XODRElevation(XODRPolynomial): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 @@ -41,41 +41,41 @@ class Elevation(Polynomial): @dataclass -class LateralProfile: +class XODRLateralProfile: """ Models elevation along t-axis of reference line. https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 """ - super_elevations: List[SuperElevation] - shapes: List[Shape] + super_elevations: List[XODRSuperElevation] + shapes: List[XODRShape] def __post_init__(self): self.super_elevations.sort(key=lambda x: x.s, reverse=False) self.shapes.sort(key=lambda x: x.s, reverse=False) @classmethod - def parse(cls, lateral_profile_element: Optional[Element]) -> LateralProfile: + def parse(cls, lateral_profile_element: Optional[Element]) -> XODRLateralProfile: args = {} - super_elevations: List[SuperElevation] = [] - shapes: List[Shape] = [] + super_elevations: List[XODRSuperElevation] = [] + shapes: List[XODRShape] = [] if lateral_profile_element is not None: for super_elevation_element in lateral_profile_element.findall("superelevation"): - super_elevations.append(SuperElevation.parse(super_elevation_element)) + super_elevations.append(XODRSuperElevation.parse(super_elevation_element)) for shape_element in lateral_profile_element.findall("shape"): - shapes.append(Shape.parse(shape_element)) + shapes.append(XODRShape.parse(shape_element)) args["super_elevations"] = super_elevations args["shapes"] = shapes - return LateralProfile(**args) + return XODRLateralProfile(**args) @dataclass -class SuperElevation(Polynomial): +class XODRSuperElevation(XODRPolynomial): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-4abf7baf-fb2f-4263-8133-ad0f64f0feac @@ -84,7 +84,7 @@ class SuperElevation(Polynomial): @dataclass -class Shape: +class XODRShape: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/1.8.0/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 @@ -99,9 +99,9 @@ class Shape: d: float @classmethod - def parse(cls, shape_element: Element) -> Shape: + def parse(cls, shape_element: Element) -> XODRShape: args = {key: float(shape_element.get(key)) for key in ["s", "t", "a", "b", "c", "d"]} - return Shape(**args) + return XODRShape(**args) def get_value(self, dt: float) -> float: """ diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/geometry.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/geometry.py index 3ddb24a8..bc8ed734 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/geometry.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/geometry.py @@ -8,11 +8,11 @@ import numpy.typing as npt from scipy.special import fresnel -from py123d.geometry import StateSE2Index +from py123d.geometry import PoseSE2Index @dataclass -class Geometry: +class XODRGeometry: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_02_road_reference_line.html """ @@ -25,10 +25,10 @@ class Geometry: @property def start_se2(self) -> npt.NDArray[np.float64]: - start_se2 = np.zeros(len(StateSE2Index), dtype=np.float64) - start_se2[StateSE2Index.X] = self.x - start_se2[StateSE2Index.Y] = self.y - start_se2[StateSE2Index.YAW] = self.hdg + start_se2 = np.zeros(len(PoseSE2Index), dtype=np.float64) + start_se2[PoseSE2Index.X] = self.x + start_se2[PoseSE2Index.Y] = self.y + start_se2[PoseSE2Index.YAW] = self.hdg return start_se2 def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @@ -36,31 +36,30 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @dataclass -class Line(Geometry): +class XODRLine(XODRGeometry): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_03_straight_line.html """ @classmethod - def parse(cls, geometry_element: Element) -> Geometry: + def parse(cls, geometry_element: Element) -> XODRGeometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} return cls(**args) def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - interpolated_se2 = self.start_se2.copy() - interpolated_se2[StateSE2Index.X] += s * np.cos(self.hdg) - interpolated_se2[StateSE2Index.Y] += s * np.sin(self.hdg) + interpolated_se2[PoseSE2Index.X] += s * np.cos(self.hdg) + interpolated_se2[PoseSE2Index.Y] += s * np.sin(self.hdg) if t != 0.0: - interpolated_se2[StateSE2Index.X] += t * np.cos(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) - interpolated_se2[StateSE2Index.Y] += t * np.sin(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.X] += t * np.cos(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.Y] += t * np.sin(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) return interpolated_se2 @dataclass -class Arc(Geometry): +class XODRArc(XODRGeometry): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_05_arc.html """ @@ -72,13 +71,12 @@ def __post_init__(self): raise ValueError("Curvature cannot be zero for Arc geometry.") @classmethod - def parse(cls, geometry_element: Element) -> Geometry: + def parse(cls, geometry_element: Element) -> XODRGeometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} args["curvature"] = float(geometry_element.find("arc").get("curvature")) return cls(**args) def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - kappa = self.curvature radius = 1.0 / kappa if kappa != 0 else float("inf") @@ -90,19 +88,19 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: dy = -radius * (np.cos(final_heading) - np.cos(initial_heading)) interpolated_se2 = self.start_se2.copy() - interpolated_se2[StateSE2Index.X] += dx - interpolated_se2[StateSE2Index.Y] += dy - interpolated_se2[StateSE2Index.YAW] = final_heading + interpolated_se2[PoseSE2Index.X] += dx + interpolated_se2[PoseSE2Index.Y] += dy + interpolated_se2[PoseSE2Index.YAW] = final_heading if t != 0.0: - interpolated_se2[StateSE2Index.X] += t * np.cos(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) - interpolated_se2[StateSE2Index.Y] += t * np.sin(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.X] += t * np.cos(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.Y] += t * np.sin(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) return interpolated_se2 @dataclass -class Spiral(Geometry): +class XODRSpiral(XODRGeometry): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_04_spiral.html https://en.wikipedia.org/wiki/Euler_spiral @@ -117,7 +115,7 @@ def __post_init__(self): raise ValueError("Curvature change is too small, cannot define a valid spiral.") @classmethod - def parse(cls, geometry_element: Element) -> Geometry: + def parse(cls, geometry_element: Element) -> XODRGeometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} spiral_element = geometry_element.find("spiral") args["curvature_start"] = float(spiral_element.get("curvStart")) @@ -131,18 +129,17 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: dx, dy = self._compute_spiral_position(s, gamma) - interpolated_se2[StateSE2Index.X] += dx - interpolated_se2[StateSE2Index.Y] += dy - interpolated_se2[StateSE2Index.YAW] += gamma * s**2 / 2 + self.curvature_start * s + interpolated_se2[PoseSE2Index.X] += dx + interpolated_se2[PoseSE2Index.Y] += dy + interpolated_se2[PoseSE2Index.YAW] += gamma * s**2 / 2 + self.curvature_start * s if t != 0.0: - interpolated_se2[StateSE2Index.X] += t * np.cos(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) - interpolated_se2[StateSE2Index.Y] += t * np.sin(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.X] += t * np.cos(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) + interpolated_se2[PoseSE2Index.Y] += t * np.sin(interpolated_se2[PoseSE2Index.YAW] + np.pi / 2) return interpolated_se2 def _compute_spiral_position(self, s: float, gamma: float) -> Tuple[float, float]: - # Transform to normalized Fresnel spiral parameter # Standard Fresnel spiral has κ(u) = u, so we need to scale # Our spiral: κ(s) = κ₀ + γs diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/lane.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/lane.py index bb6df9f5..85867670 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/lane.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/lane.py @@ -4,35 +4,35 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import Polynomial +from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import XODRPolynomial @dataclass -class Lanes: +class XODRLanes: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_01_introduction.html """ - lane_offsets: List[LaneOffset] - lane_sections: List[LaneSection] + lane_offsets: List[XODRLaneOffset] + lane_sections: List[XODRLaneSection] def __post_init__(self): self.lane_offsets.sort(key=lambda x: x.s, reverse=False) self.lane_sections.sort(key=lambda x: x.s, reverse=False) @classmethod - def parse(cls, lanes_element: Optional[Element]) -> Lanes: + def parse(cls, lanes_element: Optional[Element]) -> XODRLanes: args = {} - lane_offsets: List[LaneOffset] = [] + lane_offsets: List[XODRLaneOffset] = [] for lane_offset_element in lanes_element.findall("laneOffset"): - lane_offsets.append(LaneOffset.parse(lane_offset_element)) + lane_offsets.append(XODRLaneOffset.parse(lane_offset_element)) args["lane_offsets"] = lane_offsets - lane_sections: List[LaneSection] = [] + lane_sections: List[XODRLaneSection] = [] for lane_section_element in lanes_element.findall("laneSection"): - lane_sections.append(LaneSection.parse(lane_section_element)) + lane_sections.append(XODRLaneSection.parse(lane_section_element)) args["lane_sections"] = lane_sections - return Lanes(**args) + return XODRLanes(**args) @property def num_lane_sections(self) -> int: @@ -44,7 +44,7 @@ def last_lane_section_idx(self) -> int: @dataclass -class LaneOffset(Polynomial): +class XODRLaneOffset(XODRPolynomial): """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_04_lane_offset.html @@ -53,15 +53,15 @@ class LaneOffset(Polynomial): @dataclass -class LaneSection: +class XODRLaneSection: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_03_lane_sections.html """ s: float - left_lanes: List[Lane] - center_lanes: List[Lane] - right_lanes: List[Lane] + left_lanes: List[XODRLane] + center_lanes: List[XODRLane] + right_lanes: List[XODRLane] def __post_init__(self): self.left_lanes.sort(key=lambda x: x.id, reverse=False) @@ -69,33 +69,33 @@ def __post_init__(self): # NOTE: added assertion/filtering to check for element type or consistency @classmethod - def parse(cls, lane_section_element: Optional[Element]) -> LaneSection: + def parse(cls, lane_section_element: Optional[Element]) -> XODRLaneSection: args = {} args["s"] = float(lane_section_element.get("s")) - left_lanes: List[Lane] = [] + left_lanes: List[XODRLane] = [] if lane_section_element.find("left") is not None: for lane_element in lane_section_element.find("left").findall("lane"): - left_lanes.append(Lane.parse(lane_element)) + left_lanes.append(XODRLane.parse(lane_element)) args["left_lanes"] = left_lanes - center_lanes: List[Lane] = [] + center_lanes: List[XODRLane] = [] if lane_section_element.find("center") is not None: for lane_element in lane_section_element.find("center").findall("lane"): - center_lanes.append(Lane.parse(lane_element)) + center_lanes.append(XODRLane.parse(lane_element)) args["center_lanes"] = center_lanes - right_lanes: List[Lane] = [] + right_lanes: List[XODRLane] = [] if lane_section_element.find("right") is not None: for lane_element in lane_section_element.find("right").findall("lane"): - right_lanes.append(Lane.parse(lane_element)) + right_lanes.append(XODRLane.parse(lane_element)) args["right_lanes"] = right_lanes - return LaneSection(**args) + return XODRLaneSection(**args) @dataclass -class Lane: +class XODRLane: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_05_lane_link.html """ @@ -104,8 +104,8 @@ class Lane: type: str level: bool - widths: List[Width] - road_marks: List[RoadMark] + widths: List[XODRWidth] + road_marks: List[XODRRoadMark] predecessor: Optional[int] = None successor: Optional[int] = None @@ -115,7 +115,7 @@ def __post_init__(self): # NOTE: added assertion/filtering to check for element type or consistency @classmethod - def parse(cls, lane_element: Optional[Element]) -> Lane: + def parse(cls, lane_element: Optional[Element]) -> XODRLane: args = {} args["id"] = int(lane_element.get("id")) args["type"] = lane_element.get("type") @@ -127,21 +127,21 @@ def parse(cls, lane_element: Optional[Element]) -> Lane: if lane_element.find("link").find("successor") is not None: args["successor"] = int(lane_element.find("link").find("successor").get("id")) - widths: List[Width] = [] + widths: List[XODRWidth] = [] for width_element in lane_element.findall("width"): - widths.append(Width.parse(width_element)) + widths.append(XODRWidth.parse(width_element)) args["widths"] = widths - road_marks: List[Width] = [] + road_marks: List[XODRWidth] = [] for road_mark_element in lane_element.findall("roadMark"): - road_marks.append(RoadMark.parse(road_mark_element)) + road_marks.append(XODRRoadMark.parse(road_mark_element)) args["road_marks"] = road_marks - return Lane(**args) + return XODRLane(**args) @dataclass -class Width: +class XODRWidth: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_06_lane_geometry.html#sec-8d8ac2e0-b3d6-4048-a9ed-d5191af5c74b """ @@ -153,20 +153,20 @@ class Width: d: Optional[float] = None @classmethod - def parse(cls, width_element: Optional[Element]) -> Width: + def parse(cls, width_element: Optional[Element]) -> XODRWidth: args = {} args["s_offset"] = float(width_element.get("sOffset")) args["a"] = float(width_element.get("a")) args["b"] = float(width_element.get("b")) args["c"] = float(width_element.get("c")) args["d"] = float(width_element.get("d")) - return Width(**args) + return XODRWidth(**args) - def get_polynomial(self, t_sign: float = 1.0) -> Polynomial: + def get_polynomial(self, t_sign: float = 1.0) -> XODRPolynomial: """ Returns the polynomial representation of the width. """ - return Polynomial( + return XODRPolynomial( s=self.s_offset, a=self.a * t_sign, b=self.b * t_sign, @@ -176,7 +176,7 @@ def get_polynomial(self, t_sign: float = 1.0) -> Polynomial: @dataclass -class RoadMark: +class XODRRoadMark: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_08_road_markings.html """ @@ -193,7 +193,7 @@ def __post_init__(self): pass @classmethod - def parse(cls, road_mark_element: Optional[Element]) -> RoadMark: + def parse(cls, road_mark_element: Optional[Element]) -> XODRRoadMark: args = {} args["s_offset"] = float(road_mark_element.get("sOffset")) args["type"] = road_mark_element.get("type") @@ -202,4 +202,4 @@ def parse(cls, road_mark_element: Optional[Element]) -> RoadMark: if road_mark_element.get("width") is not None: args["width"] = float(road_mark_element.get("width")) args["lane_change"] = road_mark_element.get("lane_change") - return RoadMark(**args) + return XODRRoadMark(**args) diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/objects.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/objects.py index 4c409c54..bd844895 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/objects.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/objects.py @@ -6,7 +6,7 @@ @dataclass -class Object: +class XODRObject: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/13_objects/13_01_introduction.html """ @@ -31,8 +31,7 @@ def __post_init__(self): pass @classmethod - def parse(cls, object_element: Optional[Element]) -> Object: - + def parse(cls, object_element: Optional[Element]) -> XODRObject: args = {} args["id"] = int(object_element.get("id")) args["name"] = object_element.get("name") @@ -53,7 +52,7 @@ def parse(cls, object_element: Optional[Element]) -> Object: outline.append(CornerLocal.parse(corner_element)) args["outline"] = outline - return Object(**args) + return XODRObject(**args) @dataclass diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/opendrive.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/opendrive.py index 719f7bf3..7ee2fce9 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/opendrive.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/opendrive.py @@ -6,31 +6,30 @@ from typing import List, Literal, Optional from xml.etree.ElementTree import Element, parse -from py123d.conversion.utils.map_utils.opendrive.parser.road import Road +from py123d.conversion.utils.map_utils.opendrive.parser.road import XODRRoad @dataclass -class OpenDrive: +class XODR: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/06_general_architecture/06_03_root_element.html """ header: Header - roads: List[Road] + roads: List[XODRRoad] controllers: List[Controller] junctions: List[Junction] @classmethod - def parse(cls, root_element: Element) -> OpenDrive: - + def parse(cls, root_element: Element) -> XODR: args = {} args["header"] = Header.parse(root_element.find("header")) - roads: List[Road] = [] + roads: List[XODRRoad] = [] for road_element in root_element.findall("road"): try: - roads.append(Road.parse(road_element)) + roads.append(XODRRoad.parse(road_element)) except Exception as e: print( f"Error parsing road element with id/name {road_element.get('id')}/{road_element.get('name')}: {e}" @@ -48,12 +47,12 @@ def parse(cls, root_element: Element) -> OpenDrive: junctions.append(Junction.parse(junction_element)) args["junctions"] = junctions - return OpenDrive(**args) + return XODR(**args) @classmethod - def parse_from_file(cls, file_path: Path) -> OpenDrive: + def parse_from_file(cls, file_path: Path) -> XODR: tree = parse(file_path) - return OpenDrive.parse(tree.getroot()) + return XODR.parse(tree.getroot()) @dataclass @@ -111,7 +110,6 @@ class Controller: @classmethod def parse(cls, controller_element: Optional[Element]) -> Junction: - args = {} args["name"] = controller_element.get("name") args["id"] = float(controller_element.get("id")) diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/polynomial.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/polynomial.py index 497ffed5..5a716c52 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/polynomial.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/polynomial.py @@ -3,7 +3,7 @@ @dataclass -class Polynomial: +class XODRPolynomial: """ Multiple OpenDRIVE elements use polynomial coefficients, e.g. Elevation, LaneOffset, etc. This class provides a common interface to parse and access polynomial coefficients. @@ -19,7 +19,7 @@ class Polynomial: d: float @classmethod - def parse(cls: type["Polynomial"], element: Element) -> "Polynomial": + def parse(cls: type["XODRPolynomial"], element: Element) -> "XODRPolynomial": args = {key: float(element.get(key)) for key in ["s", "a", "b", "c", "d"]} return cls(**args) diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/reference.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/reference.py index 9ea3f0f0..236ba5a0 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/reference.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/reference.py @@ -9,39 +9,38 @@ import numpy as np import numpy.typing as npt -from py123d.conversion.utils.map_utils.opendrive.parser.elevation import Elevation -from py123d.conversion.utils.map_utils.opendrive.parser.geometry import Arc, Geometry, Line, Spiral -from py123d.conversion.utils.map_utils.opendrive.parser.lane import LaneOffset, Width -from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import Polynomial -from py123d.geometry import Point3DIndex, StateSE2Index +from py123d.conversion.utils.map_utils.opendrive.parser.elevation import XODRElevation +from py123d.conversion.utils.map_utils.opendrive.parser.geometry import XODRArc, XODRGeometry, XODRLine, XODRSpiral +from py123d.conversion.utils.map_utils.opendrive.parser.lane import XODRLaneOffset, XODRWidth +from py123d.conversion.utils.map_utils.opendrive.parser.polynomial import XODRPolynomial +from py123d.geometry import Point3DIndex, PoseSE2Index TOLERANCE: Final[float] = 1e-3 @dataclass -class PlanView: - - geometries: List[Geometry] +class XODRPlanView: + geometries: List[XODRGeometry] def __post_init__(self): # Ensure geometries are sorted by their starting position 's' self.geometries.sort(key=lambda x: x.s) @classmethod - def parse(cls, plan_view_element: Optional[Element]) -> PlanView: - geometries: List[Geometry] = [] + def parse(cls, plan_view_element: Optional[Element]) -> XODRPlanView: + geometries: List[XODRGeometry] = [] for geometry_element in plan_view_element.findall("geometry"): if geometry_element.find("line") is not None: - geometry = Line.parse(geometry_element) + geometry = XODRLine.parse(geometry_element) elif geometry_element.find("arc") is not None: - geometry = Arc.parse(geometry_element) + geometry = XODRArc.parse(geometry_element) elif geometry_element.find("spiral") is not None: - geometry = Spiral.parse(geometry_element) + geometry = XODRSpiral.parse(geometry_element) else: geometry_str = ET.tostring(geometry_element, encoding="unicode") raise NotImplementedError(f"Geometry not implemented: {geometry_str}") geometries.append(geometry) - return PlanView(geometries=geometries) + return XODRPlanView(geometries=geometries) @cached_property def geometry_lengths(self) -> npt.NDArray[np.float64]: @@ -71,11 +70,10 @@ def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = Fal @dataclass -class ReferenceLine: - - reference_line: Union[ReferenceLine, PlanView] - width_polynomials: List[Polynomial] - elevations: List[Elevation] +class XODRReferenceLine: + reference_line: Union[XODRReferenceLine, XODRPlanView] + width_polynomials: List[XODRPolynomial] + elevations: List[XODRElevation] s_offset: float @property @@ -85,51 +83,49 @@ def length(self) -> float: @classmethod def from_plan_view( cls, - plan_view: PlanView, - lane_offsets: List[LaneOffset], - elevations: List[Elevation], - ) -> ReferenceLine: + plan_view: XODRPlanView, + lane_offsets: List[XODRLaneOffset], + elevations: List[XODRElevation], + ) -> XODRReferenceLine: args = {} args["reference_line"] = plan_view args["width_polynomials"] = lane_offsets args["elevations"] = elevations args["s_offset"] = 0.0 - return ReferenceLine(**args) + return XODRReferenceLine(**args) @classmethod def from_reference_line( cls, - reference_line: ReferenceLine, - widths: List[Width], + reference_line: XODRReferenceLine, + widths: List[XODRWidth], s_offset: float = 0.0, t_sign: float = 1.0, - ) -> ReferenceLine: + ) -> XODRReferenceLine: assert t_sign in [1.0, -1.0], "t_sign must be either 1.0 or -1.0" args = {} args["reference_line"] = reference_line - width_polynomials: List[Polynomial] = [] + width_polynomials: List[XODRPolynomial] = [] for width in widths: width_polynomials.append(width.get_polynomial(t_sign=t_sign)) args["width_polynomials"] = width_polynomials args["s_offset"] = s_offset args["elevations"] = reference_line.elevations - return ReferenceLine(**args) + return XODRReferenceLine(**args) @staticmethod - def _find_polynomial(s: float, polynomials: List[Polynomial], lane_section_end: bool = False) -> Polynomial: - + def _find_polynomial(s: float, polynomials: List[XODRPolynomial], lane_section_end: bool = False) -> XODRPolynomial: out_polynomial = polynomials[-1] for polynomial in polynomials[::-1]: if lane_section_end: if polynomial.s < s: out_polynomial = polynomial break - else: - if polynomial.s <= s: - out_polynomial = polynomial - break + elif polynomial.s <= s: + out_polynomial = polynomial + break # s_values = np.array([poly.s for poly in polynomials]) # side = "left" if lane_section_end else "right" @@ -139,7 +135,6 @@ def _find_polynomial(s: float, polynomials: List[Polynomial], lane_section_end: return out_polynomial def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: - width_polynomial = self._find_polynomial(s, self.width_polynomials, lane_section_end=lane_section_end) t_offset = width_polynomial.get_value(s - width_polynomial.s) se2 = self.reference_line.interpolate_se2(self.s_offset + s, t=t_offset + t, lane_section_end=lane_section_end) @@ -147,12 +142,11 @@ def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = Fal return se2 def interpolate_3d(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: - se2 = self.interpolate_se2(s, t, lane_section_end=lane_section_end) elevation_polynomial = self._find_polynomial(s, self.elevations, lane_section_end=lane_section_end) point_3d = np.zeros(len(Point3DIndex), dtype=np.float64) - point_3d[Point3DIndex.XY] = se2[StateSE2Index.XY] + point_3d[Point3DIndex.XY] = se2[PoseSE2Index.XY] point_3d[Point3DIndex.Z] = elevation_polynomial.get_value(s - elevation_polynomial.s) return point_3d diff --git a/src/py123d/conversion/utils/map_utils/opendrive/parser/road.py b/src/py123d/conversion/utils/map_utils/opendrive/parser/road.py index 7db82b69..ebb9c6a8 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/parser/road.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/parser/road.py @@ -4,14 +4,14 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from py123d.conversion.utils.map_utils.opendrive.parser.elevation import ElevationProfile, LateralProfile -from py123d.conversion.utils.map_utils.opendrive.parser.lane import Lanes -from py123d.conversion.utils.map_utils.opendrive.parser.objects import Object -from py123d.conversion.utils.map_utils.opendrive.parser.reference import PlanView +from py123d.conversion.utils.map_utils.opendrive.parser.elevation import XODRLateralProfile, XORDElevationProfile +from py123d.conversion.utils.map_utils.opendrive.parser.lane import XODRLanes +from py123d.conversion.utils.map_utils.opendrive.parser.objects import XODRObject +from py123d.conversion.utils.map_utils.opendrive.parser.reference import XODRPlanView @dataclass -class Road: +class XODRRoad: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_01_introduction.html """ @@ -21,13 +21,13 @@ class Road: length: float name: Optional[str] - link: Link - road_types: List[RoadType] - plan_view: PlanView - elevation_profile: ElevationProfile - lateral_profile: LateralProfile - lanes: Lanes - objects: List[Object] + link: XODRLink + road_types: List[XODRRoadType] + plan_view: XODRPlanView + elevation_profile: XORDElevationProfile + lateral_profile: XODRLateralProfile + lanes: XODRLanes + objects: List[XODRObject] rule: Optional[str] = None # NOTE: ignored @@ -37,7 +37,7 @@ def __post_init__(self): ) # FIXME: Find out the purpose RHT=right-hand traffic, LHT=left-hand traffic @classmethod - def parse(cls, road_element: Element) -> Road: + def parse(cls, road_element: Element) -> XODRRoad: args = {} args["id"] = int(road_element.get("id")) @@ -45,51 +45,51 @@ def parse(cls, road_element: Element) -> Road: args["length"] = float(road_element.get("length")) args["name"] = road_element.get("name") - args["link"] = Link.parse(road_element.find("link")) + args["link"] = XODRLink.parse(road_element.find("link")) - road_types: List[RoadType] = [] + road_types: List[XODRRoadType] = [] for road_type_element in road_element.findall("type"): - road_types.append(RoadType.parse(road_type_element)) + road_types.append(XODRRoadType.parse(road_type_element)) args["road_types"] = road_types - args["plan_view"] = PlanView.parse(road_element.find("planView")) - args["elevation_profile"] = ElevationProfile.parse(road_element.find("elevationProfile")) - args["lateral_profile"] = LateralProfile.parse(road_element.find("lateralProfile")) + args["plan_view"] = XODRPlanView.parse(road_element.find("planView")) + args["elevation_profile"] = XORDElevationProfile.parse(road_element.find("elevationProfile")) + args["lateral_profile"] = XODRLateralProfile.parse(road_element.find("lateralProfile")) - args["lanes"] = Lanes.parse(road_element.find("lanes")) + args["lanes"] = XODRLanes.parse(road_element.find("lanes")) - objects: List[Object] = [] + objects: List[XODRObject] = [] if road_element.find("objects") is not None: for object_element in road_element.find("objects").findall("object"): - objects.append(Object.parse(object_element)) + objects.append(XODRObject.parse(object_element)) args["objects"] = objects - return Road(**args) + return XODRRoad(**args) @dataclass -class Link: +class XODRLink: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html """ - predecessor: Optional[PredecessorSuccessor] = None - successor: Optional[PredecessorSuccessor] = None + predecessor: Optional[XODRPredecessorSuccessor] = None + successor: Optional[XODRPredecessorSuccessor] = None @classmethod - def parse(cls, link_element: Optional[Element]) -> PlanView: + def parse(cls, link_element: Optional[Element]) -> XODRPlanView: args = {} if link_element is not None: if link_element.find("predecessor") is not None: - args["predecessor"] = PredecessorSuccessor.parse(link_element.find("predecessor")) + args["predecessor"] = XODRPredecessorSuccessor.parse(link_element.find("predecessor")) if link_element.find("successor") is not None: - args["successor"] = PredecessorSuccessor.parse(link_element.find("successor")) - return Link(**args) + args["successor"] = XODRPredecessorSuccessor.parse(link_element.find("successor")) + return XODRLink(**args) @dataclass -class PredecessorSuccessor: +class XODRPredecessorSuccessor: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html """ @@ -103,36 +103,36 @@ def __post_init__(self): assert self.contact_point is None or self.contact_point in ["start", "end"] @classmethod - def parse(cls, element: Element) -> PredecessorSuccessor: + def parse(cls, element: Element) -> XODRPredecessorSuccessor: args = {} args["element_type"] = element.get("elementType") args["element_id"] = int(element.get("elementId")) args["contact_point"] = element.get("contactPoint") - return PredecessorSuccessor(**args) + return XODRPredecessorSuccessor(**args) @dataclass -class RoadType: +class XODRRoadType: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html """ s: Optional[float] = None type: Optional[str] = None - speed: Optional[Speed] = None + speed: Optional[XODRSpeed] = None @classmethod - def parse(cls, road_type_element: Optional[Element]) -> RoadType: + def parse(cls, road_type_element: Optional[Element]) -> XODRRoadType: args = {} if road_type_element is not None: args["s"] = float(road_type_element.get("s")) args["type"] = road_type_element.get("type") - args["speed"] = Speed.parse(road_type_element.find("speed")) - return RoadType(**args) + args["speed"] = XODRSpeed.parse(road_type_element.find("speed")) + return XODRRoadType(**args) @dataclass -class Speed: +class XODRSpeed: """ https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html#sec-33dc6899-854e-4533-a3d9-76e9e1518ee7 """ @@ -141,9 +141,9 @@ class Speed: unit: Optional[str] = None @classmethod - def parse(cls, speed_element: Optional[Element]) -> Speed: + def parse(cls, speed_element: Optional[Element]) -> XODRSpeed: args = {} if speed_element is not None: args["max"] = float(speed_element.get("max")) args["unit"] = speed_element.get("unit") - return Speed(**args) + return XODRSpeed(**args) diff --git a/src/py123d/conversion/utils/map_utils/opendrive/utils/collection.py b/src/py123d/conversion/utils/map_utils/opendrive/utils/collection.py index fc6f96f3..efeb5a29 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/utils/collection.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/utils/collection.py @@ -3,9 +3,9 @@ import numpy as np -from py123d.conversion.utils.map_utils.opendrive.parser.opendrive import Junction, OpenDrive -from py123d.conversion.utils.map_utils.opendrive.parser.reference import ReferenceLine -from py123d.conversion.utils.map_utils.opendrive.parser.road import Road +from py123d.conversion.utils.map_utils.opendrive.parser.opendrive import XODR, Junction +from py123d.conversion.utils.map_utils.opendrive.parser.reference import XODRReferenceLine +from py123d.conversion.utils.map_utils.opendrive.parser.road import XODRRoad from py123d.conversion.utils.map_utils.opendrive.utils.id_system import ( build_lane_id, derive_lane_section_id, @@ -23,25 +23,24 @@ def collect_element_helpers( - opendrive: OpenDrive, + opendrive: XODR, interpolation_step_size: float, connection_distance_threshold: float, ) -> Tuple[ - Dict[int, Road], + Dict[int, XODRRoad], Dict[int, Junction], Dict[str, OpenDriveLaneHelper], Dict[str, OpenDriveLaneGroupHelper], Dict[int, OpenDriveObjectHelper], ]: - # 1. Fill the road and junction dictionaries - road_dict: Dict[int, Road] = {road.id: road for road in opendrive.roads} + road_dict: Dict[int, XODRRoad] = {road.id: road for road in opendrive.roads} junction_dict: Dict[int, Junction] = {junction.id: junction for junction in opendrive.junctions} # 2. Create lane helpers from the roads lane_helper_dict: Dict[str, OpenDriveLaneHelper] = {} for road in opendrive.roads: - reference_line = ReferenceLine.from_plan_view( + reference_line = XODRReferenceLine.from_plan_view( road.plan_view, road.lanes.lane_offsets, road.elevation_profile.elevations, @@ -81,14 +80,16 @@ def collect_element_helpers( return (road_dict, junction_dict, lane_helper_dict, lane_group_helper_dict, crosswalk_dict) -def _update_connection_from_links(lane_helper_dict: Dict[str, OpenDriveLaneHelper], road_dict: Dict[int, Road]) -> None: +def _update_connection_from_links( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], road_dict: Dict[int, XODRRoad] +) -> None: """ Uses the links of the roads to update the connections between lane helpers. :param lane_helper_dict: Dictionary of lane helpers indexed by lane id. :param road_dict: Dictionary of roads indexed by road id. """ - for lane_id in lane_helper_dict.keys(): + for lane_id in lane_helper_dict.keys(): # noqa: PLC0206 road_idx, lane_section_idx, _, lane_idx = lane_id.split("_") road_idx, lane_section_idx, lane_idx = int(road_idx), int(lane_section_idx), int(lane_idx) @@ -164,7 +165,7 @@ def _update_connection_from_links(lane_helper_dict: Dict[str, OpenDriveLaneHelpe def _update_connection_from_junctions( lane_helper_dict: Dict[str, OpenDriveLaneHelper], junction_dict: Dict[int, Junction], - road_dict: Dict[int, Road], + road_dict: Dict[int, XODRRoad], ) -> None: """ Helper function to update the lane connections based on junctions. @@ -180,7 +181,6 @@ def _update_connection_from_junctions( connecting_road = road_dict[connection.connecting_road] for lane_link in connection.lane_links: - incoming_lane_id: Optional[str] = None connecting_lane_id: Optional[str] = None @@ -216,7 +216,7 @@ def _flip_and_set_connections(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. """ - for lane_id in lane_helper_dict.keys(): + for lane_id in lane_helper_dict.keys(): # noqa: PLC0206 if lane_helper_dict[lane_id].id > 0: successors_temp = lane_helper_dict[lane_id].successor_lane_ids lane_helper_dict[lane_id].successor_lane_ids = lane_helper_dict[lane_id].predecessor_lane_ids @@ -235,7 +235,7 @@ def _post_process_connections( :param connection_distance_threshold: Threshold distance for valid connections. """ - for lane_id in lane_helper_dict.keys(): + for lane_id in lane_helper_dict.keys(): # noqa: PLC0206 lane_helper_dict[lane_id] centerline = lane_helper_dict[lane_id].center_polyline_se2 @@ -245,7 +245,6 @@ def _post_process_connections( successor_centerline = lane_helper_dict[successor_lane_id].center_polyline_se2 distance = np.linalg.norm(centerline[-1, :2] - successor_centerline[0, :2]) if distance > connection_distance_threshold: - logger.debug( f"OpenDRIVE: Removing connection {lane_id} -> {successor_lane_id} with distance {distance}" ) @@ -269,9 +268,8 @@ def _post_process_connections( def _collect_lane_groups( lane_helper_dict: Dict[str, OpenDriveLaneHelper], junction_dict: Dict[int, Junction], - road_dict: Dict[int, Road], + road_dict: Dict[int, XODRRoad], ) -> None: - lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = {} def _collect_lane_helper_of_id(lane_group_id: str) -> List[OpenDriveLaneHelper]: @@ -306,13 +304,12 @@ def _collect_lane_group_ids_of_road(road_id: int) -> List[str]: return lane_group_helper_dict -def _collect_crosswalks(opendrive: OpenDrive) -> Dict[int, OpenDriveObjectHelper]: - +def _collect_crosswalks(opendrive: XODR) -> Dict[int, OpenDriveObjectHelper]: object_helper_dict: Dict[int, OpenDriveObjectHelper] = {} for road in opendrive.roads: if len(road.objects) == 0: continue - reference_line = ReferenceLine.from_plan_view( + reference_line = XODRReferenceLine.from_plan_view( road.plan_view, road.lanes.lane_offsets, road.elevation_profile.elevations, diff --git a/src/py123d/conversion/utils/map_utils/opendrive/utils/lane_helper.py b/src/py123d/conversion/utils/map_utils/opendrive/utils/lane_helper.py index 5b8045d2..e7083a98 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/utils/lane_helper.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/utils/lane_helper.py @@ -6,28 +6,27 @@ import numpy.typing as npt import shapely -from py123d.conversion.utils.map_utils.opendrive.parser.lane import Lane, LaneSection -from py123d.conversion.utils.map_utils.opendrive.parser.reference import ReferenceLine -from py123d.conversion.utils.map_utils.opendrive.parser.road import RoadType +from py123d.conversion.utils.map_utils.opendrive.parser.lane import XODRLane, XODRLaneSection +from py123d.conversion.utils.map_utils.opendrive.parser.reference import XODRReferenceLine +from py123d.conversion.utils.map_utils.opendrive.parser.road import XODRRoadType from py123d.conversion.utils.map_utils.opendrive.utils.id_system import ( derive_lane_group_id, derive_lane_id, lane_group_id_from_lane_id, ) -from py123d.geometry import StateSE2Index +from py123d.geometry import PoseSE2Index from py123d.geometry.polyline import Polyline3D, PolylineSE2 from py123d.geometry.utils.units import kmph_to_mps, mph_to_mps @dataclass class OpenDriveLaneHelper: - lane_id: str - open_drive_lane: Lane + open_drive_lane: XODRLane s_inner_offset: float s_range: Tuple[float, float] - inner_boundary: ReferenceLine - outer_boundary: ReferenceLine + inner_boundary: XODRReferenceLine + outer_boundary: XODRReferenceLine speed_limit_mps: Optional[float] interpolation_step_size: float @@ -156,8 +155,8 @@ def outline_polyline_3d(self) -> Polyline3D: @property def shapely_polygon(self) -> shapely.Polygon: - inner_polyline = self.inner_polyline_se2[..., StateSE2Index.XY] - outer_polyline = self.outer_polyline_se2[..., StateSE2Index.XY][::-1] + inner_polyline = self.inner_polyline_se2[..., PoseSE2Index.XY] + outer_polyline = self.outer_polyline_se2[..., PoseSE2Index.XY][::-1] polygon_exterior = np.concatenate( [ inner_polyline, @@ -172,7 +171,6 @@ def shapely_polygon(self) -> shapely.Polygon: @dataclass class OpenDriveLaneGroupHelper: - lane_group_id: str lane_helpers: List[OpenDriveLaneHelper] @@ -182,7 +180,6 @@ class OpenDriveLaneGroupHelper: junction_id: Optional[int] = None def __post_init__(self): - predecessor_lane_group_ids = [] successor_lane_group__ids = [] for lane_helper in self.lane_helpers: @@ -239,8 +236,8 @@ def outline_polyline_3d(self) -> Polyline3D: @property def shapely_polygon(self) -> shapely.Polygon: - inner_polyline = self.inner_polyline_se2[..., StateSE2Index.XY] - outer_polyline = self.outer_polyline_se2[..., StateSE2Index.XY][::-1] + inner_polyline = self.inner_polyline_se2[..., PoseSE2Index.XY] + outer_polyline = self.outer_polyline_se2[..., PoseSE2Index.XY][::-1] polygon_exterior = np.concatenate( [ inner_polyline, @@ -255,14 +252,13 @@ def shapely_polygon(self) -> shapely.Polygon: def lane_section_to_lane_helpers( lane_section_id: str, - lane_section: LaneSection, - reference_line: ReferenceLine, + lane_section: XODRLaneSection, + reference_line: XODRReferenceLine, s_min: float, s_max: float, - road_types: List[RoadType], + road_types: List[XODRRoadType], interpolation_step_size: float, ) -> Dict[str, OpenDriveLaneHelper]: - lane_helpers: Dict[str, OpenDriveLaneHelper] = {} for lanes, t_sign, side in zip([lane_section.left_lanes, lane_section.right_lanes], [1.0, -1.0], ["left", "right"]): @@ -272,7 +268,7 @@ def lane_section_to_lane_helpers( lane_id = derive_lane_id(lane_group_id, lane.id) s_inner_offset = lane_section.s if len(lane_boundaries) == 1 else 0.0 lane_boundaries.append( - ReferenceLine.from_reference_line( + XODRReferenceLine.from_reference_line( reference_line=lane_boundaries[-1], widths=lane.widths, s_offset=s_inner_offset, @@ -294,8 +290,7 @@ def lane_section_to_lane_helpers( return lane_helpers -def _get_speed_limit_mps(s: float, road_types: List[RoadType]) -> Optional[float]: - +def _get_speed_limit_mps(s: float, road_types: List[XODRRoadType]) -> Optional[float]: # NOTE: Likely not correct way to extract speed limit from CARLA maps, but serves as a placeholder speed_limit_mps: Optional[float] = None s_road_types = [road_type.s for road_type in road_types] + [float("inf")] diff --git a/src/py123d/conversion/utils/map_utils/opendrive/utils/objects_helper.py b/src/py123d/conversion/utils/map_utils/opendrive/utils/objects_helper.py index 2e0117a4..899a7491 100644 --- a/src/py123d/conversion/utils/map_utils/opendrive/utils/objects_helper.py +++ b/src/py123d/conversion/utils/map_utils/opendrive/utils/objects_helper.py @@ -5,10 +5,10 @@ import numpy.typing as npt import shapely -from py123d.conversion.utils.map_utils.opendrive.parser.objects import Object -from py123d.conversion.utils.map_utils.opendrive.parser.reference import ReferenceLine -from py123d.geometry import Point3D, Point3DIndex, StateSE2, Vector2D -from py123d.geometry.geometry_index import StateSE2Index +from py123d.conversion.utils.map_utils.opendrive.parser.objects import XODRObject +from py123d.conversion.utils.map_utils.opendrive.parser.reference import XODRReferenceLine +from py123d.geometry import Point3D, Point3DIndex, PoseSE2, Vector2D +from py123d.geometry.geometry_index import PoseSE2Index from py123d.geometry.polyline import Polyline3D from py123d.geometry.transform.transform_se2 import translate_se2_along_body_frame from py123d.geometry.utils.rotation_utils import normalize_angle @@ -16,7 +16,6 @@ @dataclass class OpenDriveObjectHelper: - object_id: int outline_3d: npt.NDArray[np.float64] @@ -33,18 +32,17 @@ def shapely_polygon(self) -> shapely.Polygon: return shapely.geometry.Polygon(self.outline_3d[:, Point3DIndex.XY]) -def get_object_helper(object: Object, reference_line: ReferenceLine) -> OpenDriveObjectHelper: - +def get_object_helper(object: XODRObject, reference_line: XODRReferenceLine) -> OpenDriveObjectHelper: object_helper: Optional[OpenDriveObjectHelper] = None # 1. Extract object position in frenet frame of the reference line - object_se2: StateSE2 = StateSE2.from_array(reference_line.interpolate_se2(s=object.s, t=object.t)) + object_se2: PoseSE2 = PoseSE2.from_array(reference_line.interpolate_se2(s=object.s, t=object.t)) object_3d: Point3D = Point3D.from_array(reference_line.interpolate_3d(s=object.s, t=object.t)) # Adjust yaw angle from object data # TODO: Consider adding setters to StateSE2 to make this cleaner - object_se2._array[StateSE2Index.YAW] = normalize_angle(object_se2.yaw + object.hdg) + object_se2._array[PoseSE2Index.YAW] = normalize_angle(object_se2.yaw + object.hdg) if len(object.outline) == 0: outline_3d = np.zeros((4, len(Point3DIndex)), dtype=np.float64) diff --git a/src/py123d/conversion/utils/map_utils/road_edge/road_edge_3d_utils.py b/src/py123d/conversion/utils/map_utils/road_edge/road_edge_3d_utils.py index 42d01faf..c5591635 100644 --- a/src/py123d/conversion/utils/map_utils/road_edge/road_edge_3d_utils.py +++ b/src/py123d/conversion/utils/map_utils/road_edge/road_edge_3d_utils.py @@ -9,12 +9,12 @@ import shapely.geometry as geom from py123d.conversion.utils.map_utils.road_edge.road_edge_2d_utils import get_road_edge_linear_rings -from py123d.datatypes.maps.abstract_map_objects import ( - AbstractCarpark, - AbstractGenericDrivable, - AbstractLane, - AbstractLaneGroup, - AbstractSurfaceMapObject, +from py123d.datatypes.map_objects.map_objects import ( + BaseMapSurfaceObject, + Carpark, + GenericDrivable, + Lane, + LaneGroup, MapObjectIDType, ) from py123d.geometry import Point3DIndex @@ -25,10 +25,10 @@ def get_road_edges_3d_from_drivable_surfaces( - lanes: List[AbstractLane], - lane_groups: List[AbstractLaneGroup], - car_parks: List[AbstractCarpark], - generic_drivables: List[AbstractGenericDrivable], + lanes: List[Lane], + lane_groups: List[LaneGroup], + car_parks: List[Carpark], + generic_drivables: List[GenericDrivable], ) -> List[Polyline3D]: """Generates 3D road edges from drivable surfaces, i.e., lane groups, car parks, and generic drivables. This method merges polygons in 2D and lifts them to 3D using the boundaries/outlines of elements. @@ -47,7 +47,7 @@ def get_road_edges_3d_from_drivable_surfaces( # 2. Extract road edges in 2D (including conflicting lane groups) drivable_polygons: List[shapely.Polygon] = [] for map_surface in lane_groups + car_parks + generic_drivables: - map_surface: AbstractSurfaceMapObject + map_surface: BaseMapSurfaceObject drivable_polygons.append(map_surface.shapely_polygon) road_edges_2d = get_road_edge_linear_rings(drivable_polygons) @@ -73,7 +73,7 @@ def get_road_edges_3d_from_drivable_surfaces( def _get_conflicting_lane_groups( - lane_groups: List[AbstractLaneGroup], lanes: List[AbstractLane], z_threshold: float = 5.0 + lane_groups: List[LaneGroup], lanes: List[Lane], z_threshold: float = 5.0 ) -> Dict[int, List[int]]: """Identifies conflicting lane groups based on their 2D footprints and Z-values. The z-values are inferred from the centerlines of the lanes within each lane group. @@ -85,9 +85,7 @@ def _get_conflicting_lane_groups( """ # Convert to regular dictionaries for simpler access - lane_group_dict: Dict[MapObjectIDType, AbstractLaneGroup] = { - lane_group.object_id: lane_group for lane_group in lane_groups - } + lane_group_dict: Dict[MapObjectIDType, LaneGroup] = {lane_group.object_id: lane_group for lane_group in lane_groups} lane_centerline_dict: Dict[MapObjectIDType, Polyline3D] = {lane.object_id: lane.centerline for lane in lanes} # Pre-compute all centerlines @@ -164,7 +162,6 @@ def lift_road_edges_to_3d( road_edges_3d: List[Polyline3D] = [] if len(road_edges_2d) >= 1 and len(boundaries) >= 1: - # 1. Build comprehensive spatial index with all boundary segments # NOTE @DanielDauner: We split each boundary polyline into small segments. # The spatial indexing uses axis-aligned bounding boxes, where small geometries lead to better performance. @@ -255,7 +252,7 @@ def lift_outlines_to_3d( def _resolve_conflicting_lane_groups( conflicting_lane_groups: Dict[MapObjectIDType, List[MapObjectIDType]], - lane_groups: List[AbstractLaneGroup], + lane_groups: List[LaneGroup], ) -> List[Polyline3D]: """Resolve conflicting lane groups by merging their geometries. @@ -265,9 +262,7 @@ def _resolve_conflicting_lane_groups( """ # Helper dictionary for easy access to lane group data - lane_group_dict: Dict[MapObjectIDType, AbstractLaneGroup] = { - lane_group.object_id: lane_group for lane_group in lane_groups - } + lane_group_dict: Dict[MapObjectIDType, LaneGroup] = {lane_group.object_id: lane_group for lane_group in lane_groups} # NOTE @DanielDauner: A non-conflicting set has overlapping lane groups separated into different layers (e.g., bridges). # For each non-conflicting set, we can repeat the process of merging polygons in 2D and lifting to 3D. @@ -279,7 +274,6 @@ def _resolve_conflicting_lane_groups( road_edges_3d: List[Polyline3D] = [] for non_conflicting_set in non_conflicting_sets: - # Collect 2D polygons of non-conflicting lane group set and their neighbors merge_lane_group_data: Dict[MapObjectIDType, geom.Polygon] = {} for lane_group_id in non_conflicting_set: @@ -307,9 +301,9 @@ def _resolve_conflicting_lane_groups( def _get_nearest_z_from_points_3d(points_3d: npt.NDArray[np.float64], query_point: npt.NDArray[np.float64]) -> float: """Helpers function to get the Z-value of the nearest 3D point to a query point.""" - assert points_3d.ndim == 2 and points_3d.shape[1] == len( - Point3DIndex - ), "points_3d must be a 2D array with shape (N, 3)" + assert points_3d.ndim == 2 and points_3d.shape[1] == len(Point3DIndex), ( + "points_3d must be a 2D array with shape (N, 3)" + ) distances = np.linalg.norm(points_3d[..., Point3DIndex.XY] - query_point[..., Point3DIndex.XY], axis=1) closest_point = points_3d[np.argmin(distances)] return closest_point[2] diff --git a/src/py123d/conversion/utils/sensor_utils/camera_conventions.py b/src/py123d/conversion/utils/sensor_utils/camera_conventions.py index 75837932..6de9860c 100644 --- a/src/py123d/conversion/utils/sensor_utils/camera_conventions.py +++ b/src/py123d/conversion/utils/sensor_utils/camera_conventions.py @@ -28,7 +28,7 @@ import numpy as np -from py123d.geometry import StateSE3 +from py123d.geometry import PoseSE3 class CameraConvention(Enum): @@ -47,17 +47,17 @@ class CameraConvention(Enum): def convert_camera_convention( - from_pose: StateSE3, + from_pose: PoseSE3, from_convention: Union[CameraConvention, str], to_convention: Union[CameraConvention, str], -) -> StateSE3: +) -> PoseSE3: """Convert camera pose between different conventions. 123D default is pZmYpX (+Z forward, -Y up, +X right). - :param from_pose: StateSE3 representing the camera pose to convert + :param from_pose: PoseSE3 representing the camera pose to convert :param from_convention: CameraConvention representing the current convention of the pose :param to_convention: CameraConvention representing the target convention to convert to - :return: StateSE3 representing the converted camera pose + :return: PoseSE3 representing the converted camera pose """ # TODO: Write tests for this function # TODO: Create function over batch/array of poses @@ -92,4 +92,4 @@ def convert_camera_convention( pose_transformation = from_pose.transformation_matrix.copy() F = flip_matrices[(from_convention, to_convention)] pose_transformation[:3, :3] = pose_transformation[:3, :3] @ F - return StateSE3.from_transformation_matrix(pose_transformation) + return PoseSE3.from_transformation_matrix(pose_transformation) diff --git a/src/py123d/datatypes/detections/__init__.py b/src/py123d/datatypes/detections/__init__.py index e69de29b..be1c4bb0 100644 --- a/src/py123d/datatypes/detections/__init__.py +++ b/src/py123d/datatypes/detections/__init__.py @@ -0,0 +1,12 @@ +from py123d.datatypes.detections.box_detections import ( + BoxDetectionMetadata, + BoxDetectionSE2, + BoxDetectionSE3, + BoxDetection, + BoxDetectionWrapper, +) +from py123d.datatypes.detections.traffic_light_detections import ( + TrafficLightDetection, + TrafficLightDetectionWrapper, + TrafficLightStatus, +) diff --git a/src/py123d/datatypes/detections/box_detections.py b/src/py123d/datatypes/detections/box_detections.py index 64ebb851..dd3d9c6f 100644 --- a/src/py123d/datatypes/detections/box_detections.py +++ b/src/py123d/datatypes/detections/box_detections.py @@ -1,102 +1,235 @@ -from dataclasses import dataclass +from __future__ import annotations + from functools import cached_property from typing import List, Optional, Union import shapely -from py123d.conversion.registry.box_detection_label_registry import BoxDetectionLabel, DefaultBoxDetectionLabel -from py123d.datatypes.time.time_point import TimePoint -from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, OccupancyMap2D, StateSE2, StateSE3, Vector2D, Vector3D +from py123d.conversion.registry import BoxDetectionLabel, DefaultBoxDetectionLabel +from py123d.datatypes.time import TimePoint +from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, OccupancyMap2D, PoseSE2, PoseSE3, Vector2D, Vector3D -@dataclass class BoxDetectionMetadata: + """Stores data about the box detection, including its label, track token, number of LiDAR points, and timepoint.""" + + __slots__ = ("_label", "_track_token", "_num_lidar_points", "_timepoint") + + def __init__( + self, + label: BoxDetectionLabel, + track_token: str, + num_lidar_points: Optional[int] = None, + timepoint: Optional[TimePoint] = None, + ) -> None: + """Initialize a BoxDetectionMetadata instance. + + :param label: The label of the detection. + :param track_token: The track token of the detection. + :param num_lidar_points: The number of LiDAR points, defaults to None. + :param timepoint: The timepoint of the detection, defaults to None. + """ + self._label = label + self._track_token = track_token + self._num_lidar_points = num_lidar_points + self._timepoint = timepoint + + @property + def label(self) -> BoxDetectionLabel: + """The :class:`~py123d.datatypes.detections.BoxDetectionLabel`, from the original dataset's label set.""" + return self._label - label: BoxDetectionLabel - track_token: str - confidence: Optional[float] = None - num_lidar_points: Optional[int] = None - timepoint: Optional[TimePoint] = None + @property + def track_token(self) -> str: + """The unique track token of the detection, consistent across frames.""" + return self._track_token + + @property + def num_lidar_points(self) -> Optional[int]: + """Optionally, the number of LiDAR points associated with the detection.""" + return self._num_lidar_points + + @property + def timepoint(self) -> Optional[TimePoint]: + """Optionally, the :class:`~py123d.datatypes.time.TimePoint` of the detection.""" + return self._timepoint @property def default_label(self) -> DefaultBoxDetectionLabel: + """The unified :class:`~py123d.conversion.registry.DefaultBoxDetectionLabel` + corresponding to the detection's label. + """ return self.label.to_default() -@dataclass class BoxDetectionSE2: + """Detected, tracked, and oriented bounding box 2D space.""" - metadata: BoxDetectionMetadata - bounding_box_se2: BoundingBoxSE2 - velocity: Optional[Vector2D] = None + __slots__ = ("_metadata", "_bounding_box_se2", "_velocity_2d") - @property - def shapely_polygon(self) -> shapely.geometry.Polygon: - return self.bounding_box_se2.shapely_polygon + def __init__( + self, + metadata: BoxDetectionMetadata, + bounding_box_se2: BoundingBoxSE2, + velocity_2d: Optional[Vector2D] = None, + ) -> None: + """Initialize a BoxDetectionSE2 instance. - @property - def center(self) -> StateSE2: - return self.bounding_box_se2.center + :param metadata: The :class:`BoxDetectionMetadata` of the detection. + :param bounding_box_se2: The :class:`~py123d.datatypes.geometry.BoundingBoxSE2` of the detection. + :param velocity: Optionally, a :class:`~py123d.geometry.Vector2D` representing the velocity. + """ + + self._metadata = metadata + self._bounding_box_se2 = bounding_box_se2 + self._velocity_2d = velocity_2d @property - def bounding_box(self) -> BoundingBoxSE2: - return self.bounding_box_se2 + def metadata(self) -> BoxDetectionMetadata: + """The :class:`BoxDetectionMetadata` of the detection.""" + return self._metadata + @property + def bounding_box_se2(self) -> BoundingBoxSE2: + """The :class:`~py123d.geometry.BoundingBoxSE2` of the detection.""" + return self._bounding_box_se2 -@dataclass -class BoxDetectionSE3: + @property + def velocity_2d(self) -> Optional[Vector2D]: + """The :class:`~py123d.geometry.Vector2D` representing the velocity.""" + return self._velocity_2d - metadata: BoxDetectionMetadata - bounding_box_se3: BoundingBoxSE3 - velocity: Optional[Vector3D] = None + @property + def center_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` representing the center of the bounding box.""" + return self.bounding_box_se2.center_se2 @property def shapely_polygon(self) -> shapely.geometry.Polygon: - return self.bounding_box_se3.shapely_polygon + """The shapely polygon of the bounding box in 2D space.""" + return self.bounding_box_se2.shapely_polygon @property - def center(self) -> StateSE3: - return self.bounding_box_se3.center + def box_detection_se2(self) -> BoxDetectionSE2: + """Returns self to maintain interface consistency.""" + return self + + +class BoxDetectionSE3: + """Detected, tracked, and oriented bounding box 3D space.""" + + __slots__ = ("_metadata", "_bounding_box_se3", "_velocity") + + def __init__( + self, + metadata: BoxDetectionMetadata, + bounding_box_se3: BoundingBoxSE3, + velocity_3d: Optional[Vector3D] = None, + ) -> None: + """Initialize a BoxDetectionSE3 instance. + + :param metadata: The :class:`BoxDetectionMetadata` of the detection. + :param bounding_box_se3: The :class:`~py123d.datatypes.geometry.BoundingBoxSE3` of the detection. + :param velocity_3d: Optionally, a :class:`~py123d.geometry.Vector3D` representing the velocity. + """ + self._metadata = metadata + self._bounding_box_se3 = bounding_box_se3 + self._velocity = velocity_3d @property - def center_se3(self) -> StateSE3: - return self.bounding_box_se3.center_se3 + def metadata(self) -> BoxDetectionMetadata: + """The :class:`BoxDetectionMetadata` of the detection.""" + return self._metadata @property - def bounding_box(self) -> BoundingBoxSE3: - return self.bounding_box_se3 + def bounding_box_se3(self) -> BoundingBoxSE3: + """The :class:`~py123d.geometry.BoundingBoxSE3` of the detection.""" + return self._bounding_box_se3 @property def bounding_box_se2(self) -> BoundingBoxSE2: + """The SE2 projection :class:`~py123d.geometry.BoundingBoxSE2` of the SE3 bounding box.""" return self.bounding_box_se3.bounding_box_se2 + @property + def velocity_3d(self) -> Optional[Vector3D]: + """The :class:`~py123d.geometry.Vector3D` representing the velocity.""" + return self._velocity + + @property + def velocity_2d(self) -> Optional[Vector2D]: + """The 2D projection :class:`~py123d.geometry.Vector2D` of the 3D velocity.""" + return Vector2D(self._velocity.x, self._velocity.y) if self._velocity else None + + @property + def center_se3(self) -> PoseSE3: + """The :class:`~py123d.geometry.PoseSE3` representing the center of the bounding box.""" + return self.bounding_box_se3.center_se3 + + @property + def center_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` representing the center of the SE2 bounding box.""" + return self.bounding_box_se2.center_se2 + @property def box_detection_se2(self) -> BoxDetectionSE2: + """The :class:`~py123d.datatypes.detections.BoxDetectionSE2` projection of this SE3 box detection.""" return BoxDetectionSE2( metadata=self.metadata, bounding_box_se2=self.bounding_box_se2, - velocity=Vector2D(self.velocity.x, self.velocity.y) if self.velocity else None, + velocity_2d=Vector2D(self.velocity_3d.x, self.velocity_3d.y) if self.velocity_3d else None, ) + @property + def shapely_polygon(self) -> shapely.geometry.Polygon: + """The shapely polygon of the bounding box in 2D space.""" + return self.bounding_box_se3.shapely_polygon + BoxDetection = Union[BoxDetectionSE2, BoxDetectionSE3] -@dataclass class BoxDetectionWrapper: + """The BoxDetectionWrapper is a container for multiple box detections. + It provides methods to access individual detections as well as to retrieve a detection by track token. + The wrapper is used to read and write box detections from/to logs. + """ + + def __init__(self, box_detections: List[BoxDetection]) -> None: + """Initialize a BoxDetectionWrapper instance. - box_detections: List[BoxDetection] + :param box_detections: A list of :class:`BoxDetection` instances. + """ + self._box_detections = box_detections + + @property + def box_detections(self) -> List[BoxDetection]: + """List of individual :class:`BoxDetectionSE2` or :class:`BoxDetectionSE3`.""" + return self._box_detections def __getitem__(self, index: int) -> BoxDetection: - return self.box_detections[index] + """Retrieve a box detection by its index. + + :param index: The index of the box detection. + :return: The box detection at the given index. + """ + return self._box_detections[index] def __len__(self) -> int: - return len(self.box_detections) + """Number of box detections.""" + return len(self._box_detections) def __iter__(self): - return iter(self.box_detections) + """Iterator over box detections.""" + return iter(self._box_detections) + + def get_detection_by_track_token(self, track_token: str) -> Optional[Union[BoxDetectionSE2, BoxDetectionSE3]]: + """Retrieve a box detection by its track token. + + :param track_token: The track token of the box detection. + :return: The box detection with the given track token, or None if not found. + """ - def get_detection_by_track_token(self, track_token: str) -> Optional[BoxDetection]: box_detection: Optional[BoxDetection] = None for detection in self.box_detections: if detection.metadata.track_token == track_token: @@ -105,7 +238,8 @@ def get_detection_by_track_token(self, track_token: str) -> Optional[BoxDetectio return box_detection @cached_property - def occupancy_map(self) -> OccupancyMap2D: + def occupancy_map_2d(self) -> OccupancyMap2D: + """The :class:`~py123d.geometry.OccupancyMap2D` representing the 2D occupancy of all box detections.""" ids = [detection.metadata.track_token for detection in self.box_detections] - geometries = [detection.bounding_box.shapely_polygon for detection in self.box_detections] + geometries = [detection.shapely_polygon for detection in self.box_detections] return OccupancyMap2D(geometries=geometries, ids=ids) diff --git a/src/py123d/datatypes/detections/traffic_light_detections.py b/src/py123d/datatypes/detections/traffic_light_detections.py index c4f1c57a..30601ae9 100644 --- a/src/py123d/datatypes/detections/traffic_light_detections.py +++ b/src/py123d/datatypes/detections/traffic_light_detections.py @@ -1,45 +1,108 @@ -from dataclasses import dataclass from typing import List, Optional from py123d.common.utils.enums import SerialIntEnum -from py123d.datatypes.time.time_point import TimePoint +from py123d.datatypes.time import TimePoint class TrafficLightStatus(SerialIntEnum): """ - Enum for TrafficLightStatus. + Enum for that represents the status of a traffic light. """ GREEN = 0 + """Green light is on.""" + YELLOW = 1 + """Yellow light is on.""" + RED = 2 + """Red light is on.""" + OFF = 3 + """Traffic light is off.""" + UNKNOWN = 4 + """Traffic light status is unknown.""" -@dataclass class TrafficLightDetection: + """ + Single traffic light detection if a lane, that includes the lane id, status (green, yellow, red, off, unknown), + and optional timepoint of the detection. + """ + + __slots__ = ("_lane_id", "_status", "_timepoint") + + def __init__(self, lane_id: int, status: TrafficLightStatus, timepoint: Optional[TimePoint] = None) -> None: + """Initialize a TrafficLightDetection instance. + + :param lane_id: The lane id associated with the traffic light detection. + :param status: The status of the traffic light (green, yellow, red, off, unknown). + :param timepoint: The optional timepoint of the detection. + """ - lane_id: int - status: TrafficLightStatus - timepoint: Optional[TimePoint] = None + self._lane_id = lane_id + self._status = status + self._timepoint = timepoint + + @property + def lane_id(self) -> int: + """The lane id associated with the traffic light detection.""" + return self._lane_id + + @property + def status(self) -> TrafficLightStatus: + """The :class:`TrafficLightStatus` of the traffic light detection.""" + return self._status + + @property + def timepoint(self) -> Optional[TimePoint]: + """The optional :class:`~py123d.datatypes.time.TimePoint` of the traffic light detection.""" + return self._timepoint -@dataclass class TrafficLightDetectionWrapper: + """The TrafficLightDetectionWrapper is a container for multiple traffic light detections. + It provides methods to access individual detections as well as to retrieve a detection by lane id. + The wrapper is is used in to read and write traffic light detections from/to logs. + """ + + __slots__ = ("_traffic_light_detections",) - traffic_light_detections: List[TrafficLightDetection] + def __init__(self, traffic_light_detections: List[TrafficLightDetection]) -> None: + """Initialize a TrafficLightDetectionWrapper instance. + + :param traffic_light_detections: List of :class:`TrafficLightDetection`. + """ + self._traffic_light_detections = traffic_light_detections + + @property + def traffic_light_detections(self) -> List[TrafficLightDetection]: + """List of individual :class:`TrafficLightDetection`.""" + return self._traffic_light_detections def __getitem__(self, index: int) -> TrafficLightDetection: + """Retrieve a traffic light detection by its index. + + :param index: The index of the traffic light detection. + :return: :class:`TrafficLightDetection` at the given index. + """ return self.traffic_light_detections[index] def __len__(self) -> int: + """The number of traffic light detections in the wrapper.""" return len(self.traffic_light_detections) def __iter__(self): + """Iterator over the traffic light detections in the wrapper.""" return iter(self.traffic_light_detections) def get_detection_by_lane_id(self, lane_id: int) -> Optional[TrafficLightDetection]: + """Retrieve a traffic light detection by its lane id. + + :param lane_id: The lane id to search for. + :return: The traffic light detection for the given lane id, or None if not found. + """ traffic_light_detection: Optional[TrafficLightDetection] = None for detection in self.traffic_light_detections: if int(detection.lane_id) == int(lane_id): diff --git a/src/py123d/datatypes/map_objects/__init__.py b/src/py123d/datatypes/map_objects/__init__.py new file mode 100644 index 00000000..7a944629 --- /dev/null +++ b/src/py123d/datatypes/map_objects/__init__.py @@ -0,0 +1,14 @@ +from py123d.datatypes.map_objects.base_map_objects import BaseMapLineObject, BaseMapObject, BaseMapSurfaceObject +from py123d.datatypes.map_objects.map_layer_types import LaneType, MapLayer, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + StopZone, + Walkway, +) diff --git a/src/py123d/datatypes/map_objects/base_map_objects.py b/src/py123d/datatypes/map_objects/base_map_objects.py new file mode 100644 index 00000000..f1d70d7b --- /dev/null +++ b/src/py123d/datatypes/map_objects/base_map_objects.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import abc +from typing import Optional, Union + +import numpy as np +import shapely.geometry as geom +import trimesh +from typing_extensions import TypeAlias + +from py123d.datatypes.map_objects.map_layer_types import MapLayer +from py123d.geometry import Point3DIndex, Polyline2D, Polyline3D + +MapObjectIDType: TypeAlias = Union[str, int] + + +class BaseMapObject(abc.ABC): + """Base interface representation of all map objects.""" + + __slots__ = ("_object_id",) + + def __init__(self, object_id: MapObjectIDType): + """Constructor of the base map object type. + + :param object_id: unique identifier of the map object. + """ + self._object_id: MapObjectIDType = object_id + + @property + def object_id(self) -> MapObjectIDType: + """The unique identifier of the map object (unique within a map layer).""" + return self._object_id + + @property + @abc.abstractmethod + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + + +class BaseMapSurfaceObject(BaseMapObject): + """Base interface representation of all map objects that represent surfaces.""" + + __slots__ = ("_outline", "_shapely_polygon") + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ) -> None: + """Initialize a BaseMapSurfaceObject instance. Either outline or shapely_polygon must be provided. + + :param object_id: Unique identifier for the map object. + :param outline: Outline of the surface, either 2D or 3D, defaults to None. + :param shapely_polygon: A shapely Polygon representing the surface geometry, defaults to None. + :raises ValueError: If both outline and shapely_polygon are not provided. + """ + super().__init__(object_id) + + if outline is None and shapely_polygon is None: + raise ValueError("Either outline or shapely_polygon must be provided.") + if outline is None: + outline = Polyline3D.from_linestring(shapely_polygon.exterior) # type: ignore + if shapely_polygon is None: + shapely_polygon = geom.Polygon(outline.array[:, :2]) + + self._object_id = object_id + self._outline = outline + self._shapely_polygon = shapely_polygon + + @property + def outline(self) -> Union[Polyline2D, Polyline3D]: + """The outline of the surface as either :class:`~py123d.geometry.Polyline2D` + or :class:`~py123d.geometry.Polyline3D`.""" + return self._outline + + @property + def outline_2d(self) -> Polyline2D: + """The outline of the surface as :class:`~py123d.geometry.Polyline2D`.""" + if isinstance(self._outline, Polyline2D): + return self._outline + # Converts 3D polyline to 2D by dropping the z-coordinate + return Polyline2D.from_linestring(self._outline.linestring) + + @property + def outline_3d(self) -> Polyline3D: + """The outline of the surface as :class:`~py123d.geometry.Polyline3D` (zero-padded to 3D if necessary).""" + if isinstance(self._outline, Polyline3D): + return self._outline + # Converts 2D polyline to 3D by adding a default (zero) z-coordinate + return Polyline3D.from_linestring(self._outline.linestring) + + @property + def shapely_polygon(self) -> geom.Polygon: + """The shapely polygon of the surface.""" + return self._shapely_polygon + + @property + def trimesh_mesh(self) -> trimesh.Trimesh: + """The trimesh mesh representation of the surface.""" + # Fallback to geometry if no boundaries are available + outline_3d_array = self.outline_3d.array + vertices_2d, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) + if len(vertices_2d) == len(outline_3d_array): + # Regular case, where vertices match outline_3d_array + vertices_3d = outline_3d_array + elif len(vertices_2d) == len(outline_3d_array) + 1: + # outline array was not closed, so we need to add the first vertex again + vertices_3d = np.vstack((outline_3d_array, outline_3d_array[0])) + else: + raise ValueError("No vertices found for triangulation.") + trimesh_mesh = trimesh.Trimesh(vertices=vertices_3d, faces=faces) + return trimesh_mesh + + +class BaseMapLineObject(BaseMapObject): + """Base interface representation of all line map objects.""" + + __slots__ = ("_polyline",) + + def __init__(self, object_id: MapObjectIDType, polyline: Union[Polyline2D, Polyline3D]) -> None: + """Initialize a BaseMapLineObject instance. + + :param object_id: Unique identifier for the map object. + :param polyline: The polyline representation of the line object. + """ + super().__init__(object_id) + self._polyline = polyline + + @property + def polyline(self) -> Union[Polyline2D, Polyline3D]: + """The polyline representation, either :class:`~py123d.geometry.Polyline2D` or + :class:`~py123d.geometry.Polyline3D`.""" + return self._polyline + + @property + def polyline_2d(self) -> Polyline2D: + """The polyline representation as :class:`~py123d.geometry.Polyline2D`.""" + if isinstance(self._polyline, Polyline2D): + return self._polyline + # Converts 3D polyline to 2D by dropping the z-coordinate + return Polyline2D.from_linestring(self._polyline.linestring) + + @property + def polyline_3d(self) -> Polyline3D: + """The polyline representation as :class:`~py123d.geometry.Polyline3D` (zero-padded to 3D if necessary).""" + if isinstance(self._polyline, Polyline3D): + return self._polyline + # Converts 2D polyline to 3D by adding a default (zero) z-coordinate + return Polyline3D.from_linestring(self._polyline.linestring) + + @property + def shapely_linestring(self) -> geom.LineString: + """The shapely LineString representation of the polyline.""" + return self._polyline.linestring diff --git a/src/py123d/datatypes/map_objects/map_layer_types.py b/src/py123d/datatypes/map_objects/map_layer_types.py new file mode 100644 index 00000000..abf44382 --- /dev/null +++ b/src/py123d/datatypes/map_objects/map_layer_types.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from py123d.common.utils.enums import SerialIntEnum + +# TODO @DanielDauner: +# - Implement stop zone types. +# - Consider adding types for intersections or other layers. + + +class MapLayer(SerialIntEnum): + """Enum for different map layers (i.e. object types) in a map.""" + + LANE = 0 + """Lanes (surface).""" + + LANE_GROUP = 1 + """Lane groups (surface).""" + + INTERSECTION = 2 + """Intersections (surface).""" + + CROSSWALK = 3 + """Crosswalks (surface).""" + + WALKWAY = 4 + """Walkways (surface).""" + + CARPARK = 5 + """Carparks (surface).""" + + GENERIC_DRIVABLE = 6 + """Generic drivable (surface).""" + + STOP_ZONE = 7 + """Stop zones (surface).""" + + ROAD_EDGE = 8 + """Road edges (lines).""" + + ROAD_LINE = 9 + """Road lines (lines).""" + + +class LaneType(SerialIntEnum): + """Enum for different lane types.""" + + # NOTE @DanielDauner: We currently do not include the lane types, but should add them in the future. + # Some maps (e.g. nuPlan, Waymo) have bike lanes, which need to be distinguished from regular lanes. + + UNDEFINED = 0 + FREEWAY = 1 + SURFACE_STREET = 2 + BIKE_LANE = 3 + + +class RoadEdgeType(SerialIntEnum): + """Enum for different road edge types. + + Notes + ----- + The road edge types follow the Waymo specification [1]_. + + References + ---------- + .. [1] https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L188 + """ + + UNKNOWN = 0 + """Unknown road edge type.""" + + ROAD_EDGE_BOUNDARY = 1 + """Physical road boundary that doesn't have traffic on the other side.""" + + ROAD_EDGE_MEDIAN = 2 + """Physical road boundary that separates the car from other traffic.""" + + +class RoadLineType(SerialIntEnum): + """Enum for different road line types. + + Notes + ----- + The road line types follow the Argoverse 2 specification [1]_. + + References + ---------- + .. [1] https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/lane_segment.py#L33 + """ + + NONE = 0 + """No painted line is present.""" + + UNKNOWN = 1 + """Unknown or unclassified painted line type.""" + + DASH_SOLID_YELLOW = 2 + """Yellow line with dashed marking on one side and solid on the other.""" + + DASH_SOLID_WHITE = 3 + """White line with dashed marking on one side and solid on the other.""" + + DASHED_WHITE = 4 + """White dashed line marking.""" + + DASHED_YELLOW = 5 + """Yellow dashed line marking.""" + + DOUBLE_SOLID_YELLOW = 6 + """Double yellow solid line marking.""" + + DOUBLE_SOLID_WHITE = 7 + """Double white solid line marking.""" + + DOUBLE_DASH_YELLOW = 8 + """Double yellow dashed line marking.""" + + DOUBLE_DASH_WHITE = 9 + """Double white dashed line marking.""" + + SOLID_YELLOW = 10 + """Single solid yellow line marking.""" + + SOLID_WHITE = 11 + """Single solid white line marking.""" + + SOLID_DASH_WHITE = 12 + """Single solid white line with dashed marking on one side.""" + + SOLID_DASH_YELLOW = 13 + """Single solid yellow line with dashed marking on one side.""" + + SOLID_BLUE = 14 + """Single solid blue line marking.""" diff --git a/src/py123d/datatypes/map_objects/map_objects.py b/src/py123d/datatypes/map_objects/map_objects.py new file mode 100644 index 00000000..60298905 --- /dev/null +++ b/src/py123d/datatypes/map_objects/map_objects.py @@ -0,0 +1,582 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List, Optional, Union + +import numpy as np +import shapely.geometry as geom +from trimesh import Trimesh + +from py123d.datatypes.map_objects.base_map_objects import BaseMapLineObject, BaseMapSurfaceObject, MapObjectIDType +from py123d.datatypes.map_objects.map_layer_types import MapLayer, RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.utils import get_trimesh_from_boundaries +from py123d.geometry import Polyline2D, Polyline3D + +if TYPE_CHECKING: + from py123d.api import MapAPI + + +class Lane(BaseMapSurfaceObject): + """Class representing a lane in a map.""" + + __slots__ = ( + "_lane_group_id", + "_left_boundary", + "_right_boundary", + "_centerline", + "_left_lane_id", + "_right_lane_id", + "_predecessor_ids", + "_successor_ids", + "_speed_limit_mps", + "_map_api", + ) + + def __init__( + self, + object_id: MapObjectIDType, + lane_group_id: MapObjectIDType, + left_boundary: Polyline3D, + right_boundary: Polyline3D, + centerline: Polyline3D, + left_lane_id: Optional[MapObjectIDType] = None, + right_lane_id: Optional[MapObjectIDType] = None, + predecessor_ids: List[MapObjectIDType] = [], + successor_ids: List[MapObjectIDType] = [], + speed_limit_mps: Optional[float] = None, + outline: Optional[Polyline3D] = None, + shapely_polygon: Optional[geom.Polygon] = None, + map_api: Optional["MapAPI"] = None, + ) -> None: + """Initialize a :class:`Lane` instance. + + Notes + ----- + If the map_api is provided, neighboring lanes and lane group can be accessed through the properties. + If the outline is not provided, it will be constructed from the left and right boundaries. + If the shapely_polygon is not provided, it will be constructed from the outline. + + :param object_id: The unique identifier for the lane. + :param lane_group_id: The unique identifier for the lane group this lane belongs to. + :param left_boundary: Polyline of left boundary of the lane. + :param right_boundary: Polyline of right boundary of the lane. + :param centerline: Polyline of centerline of the lane. + :param left_lane_id: The unique identifier for the left neighboring lane, defaults to None. + :param right_lane_id: The unique identifier for the right neighboring lane, defaults to None. + :param predecessor_ids: The unique identifiers for the predecessor lanes, defaults to []. + :param successor_ids: The unique identifiers for the successor lanes, defaults to []. + :param speed_limit_mps: The speed limit for the lane in meters per second, defaults to None. + :param outline: The outline of the lane, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the lane, defaults to None. + :param map_api: The MapAPI instance for accessing map objects, defaults to None. + """ + if outline is None: + outline_array = np.vstack( + ( + left_boundary.array, + right_boundary.array[::-1], + left_boundary.array[0], + ) + ) + outline = Polyline3D.from_array(outline_array) + super().__init__(object_id, outline, shapely_polygon) + + self._lane_group_id = lane_group_id + self._left_boundary = left_boundary + self._right_boundary = right_boundary + self._centerline = centerline + self._left_lane_id = left_lane_id + self._right_lane_id = right_lane_id + self._predecessor_ids = predecessor_ids + self._successor_ids = successor_ids + self._speed_limit_mps = speed_limit_mps + self._map_api = map_api + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.LANE + + @property + def lane_group_id(self) -> MapObjectIDType: + """ID of the lane group this lane belongs to.""" + return self._lane_group_id + + @property + def lane_group(self) -> Optional[LaneGroup]: + """The :class:`LaneGroup` this lane belongs to.""" + if self._map_api is not None: + return self._map_api.get_map_object(self.lane_group_id, MapLayer.LANE_GROUP) # type: ignore + return None + + @property + def left_boundary(self) -> Polyline3D: + """The left boundary of the lane as a :class:`~py123d.geometry.Polyline3D`.""" + return self._left_boundary + + @property + def right_boundary(self) -> Polyline3D: + """The right boundary of the lane as a :class:`~py123d.geometry.Polyline3D`.""" + return self._right_boundary + + @property + def centerline(self) -> Polyline3D: + """The centerline of the lane as a :class:`~py123d.geometry.Polyline3D`.""" + return self._centerline + + @property + def left_lane_id(self) -> Optional[MapObjectIDType]: + """ID of the left neighboring lane.""" + return self._left_lane_id + + @property + def left_lane(self) -> Optional[Lane]: + """The left neighboring :class:`Lane`, if available.""" + if self._map_api is not None and self.left_lane_id is not None: + return self._map_api.get_map_object(self.left_lane_id, self.layer) # type: ignore + return None + + @property + def right_lane_id(self) -> Optional[MapObjectIDType]: + """ID of the right neighboring lane.""" + return self._right_lane_id + + @property + def right_lane(self) -> Optional[Lane]: + """The right neighboring :class:`Lane`, if available.""" + if self._map_api is not None and self.right_lane_id is not None: + return self._map_api.get_map_object(self.right_lane_id, self.layer) # type: ignore + return None + + @property + def predecessor_ids(self) -> List[MapObjectIDType]: + """List of IDs of the predecessor lanes.""" + return self._predecessor_ids + + @property + def predecessors(self) -> Optional[List[Lane]]: + """List of predecessor :class:`Lane` instances.""" + predecessors: Optional[List[Lane]] = None + if self._map_api is not None: + predecessors = [self._map_api.get_map_object(lane_id, self.layer) for lane_id in self.predecessor_ids] # type: ignore + return predecessors + + @property + def successor_ids(self) -> List[MapObjectIDType]: + """List of IDs of the successor lanes.""" + return self._successor_ids + + @property + def successors(self) -> Optional[List[Lane]]: + """List of successor :class:`Lane` instances.""" + successors: Optional[List[Lane]] = None + if self._map_api is not None: + successors = [self._map_api.get_map_object(lane_id, self.layer) for lane_id in self.successor_ids] # type: ignore + return successors + + @property + def speed_limit_mps(self) -> Optional[float]: + """The speed limit of the lane in meters per second.""" + return self._speed_limit_mps + + @property + def trimesh_mesh(self) -> Trimesh: + """The trimesh mesh representation of the lane.""" + return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) + + +class LaneGroup(BaseMapSurfaceObject): + """Class representing a group of lanes going in the same direction.""" + + __slots__ = ( + "_lane_ids", + "_left_boundary", + "_right_boundary", + "_intersection_id", + "_predecessor_ids", + "_successor_ids", + "_map_api", + ) + + def __init__( + self, + object_id: MapObjectIDType, + lane_ids: List[MapObjectIDType], + left_boundary: Polyline3D, + right_boundary: Polyline3D, + intersection_id: Optional[MapObjectIDType] = None, + predecessor_ids: List[MapObjectIDType] = [], + successor_ids: List[MapObjectIDType] = [], + outline: Optional[Polyline3D] = None, + shapely_polygon: Optional[geom.Polygon] = None, + map_api: Optional["MapAPI"] = None, + ): + """Initialize a :class:`LaneGroup` instance. + + Notes + ----- + If the map_api is provided, neighboring lane groups and intersection can be accessed through the properties. + If the outline is not provided, it will be constructed from the left and right boundaries. + If the shapely_polygon is not provided, it will be constructed from the outline. + + :param object_id: The ID of the lane group. + :param lane_ids: The IDs of the lanes in the group. + :param left_boundary: The left boundary of the lane group. + :param right_boundary: The right boundary of the lane group. + :param intersection_id: The ID of the intersection, defaults to None + :param predecessor_ids: The IDs of the predecessor lanes, defaults to [] + :param successor_ids: The IDs of the successor lanes, defaults to [] + :param outline: The outline of the lane group, defaults to None + :param shapely_polygon: The shapely polygon representation of the lane group, defaults to None + :param map_api: The map API instance, defaults to None + """ + if outline is None: + outline_array = np.vstack( + ( + left_boundary.array, + right_boundary.array[::-1], + left_boundary.array[0], + ) + ) + outline = Polyline3D.from_array(outline_array) + super().__init__(object_id, outline, shapely_polygon) + + self._lane_ids = lane_ids + self._left_boundary = left_boundary + self._right_boundary = right_boundary + self._intersection_id = intersection_id + self._predecessor_ids = predecessor_ids + self._successor_ids = successor_ids + self._map_api = map_api + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.LANE_GROUP + + @property + def lane_ids(self) -> List[MapObjectIDType]: + """List of IDs of the lanes in the group.""" + return self._lane_ids + + @property + def lanes(self) -> List[Lane]: + """List of :class:`Lane` instances in the group.""" + lanes: Optional[List[Lane]] = None + if self._map_api is not None: + lanes = [self._map_api.get_map_object(lane_id, MapLayer.LANE) for lane_id in self.lane_ids] # type: ignore + return lanes # type: ignore + + @property + def left_boundary(self) -> Polyline3D: + """The left boundary of the lane group.""" + return self._left_boundary + + @property + def right_boundary(self) -> Polyline3D: + """The right boundary of the lane group.""" + return self._right_boundary + + @property + def intersection_id(self) -> Optional[MapObjectIDType]: + """ID of the intersection the lane group belongs to, if available.""" + return self._intersection_id + + @property + def intersection(self) -> Optional[Intersection]: + """The :class:`Intersection` the lane group belongs to, if available.""" + intersection: Optional[Intersection] = None + if self._map_api is not None and self.intersection_id is not None: + intersection = self._map_api.get_map_object(self.intersection_id, MapLayer.INTERSECTION) # type: ignore + return intersection + + @property + def predecessor_ids(self) -> List[MapObjectIDType]: + """List of IDs of the predecessor lane groups.""" + return self._predecessor_ids + + @property + def predecessors(self) -> List[LaneGroup]: + """List of predecessor :class:`LaneGroup` instances.""" + predecessors: Optional[List[LaneGroup]] = None + if self._map_api is not None: + predecessors = [ + self._map_api.get_map_object(lane_group_id, self.layer) for lane_group_id in self.predecessor_ids + ] + return predecessors + + @property + def successor_ids(self) -> List[MapObjectIDType]: + """List of IDs of the successor lane groups.""" + return self._successor_ids + + @property + def successors(self) -> List[LaneGroup]: + """List of successor :class:`LaneGroup` instances.""" + successors: Optional[List[LaneGroup]] = None + if self._map_api is not None: + successors = [ + self._map_api.get_map_object(lane_group_id, self.layer) for lane_group_id in self.successor_ids + ] + return successors + + @property + def trimesh_mesh(self) -> Trimesh: + """The trimesh mesh representation of the lane group.""" + return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) + + +class Intersection(BaseMapSurfaceObject): + """Class representing an intersection in a map, which consists of multiple lane groups.""" + + __slots__ = ("_lane_group_ids", "_map_api") + + def __init__( + self, + object_id: MapObjectIDType, + lane_group_ids: List[MapObjectIDType], + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + map_api: Optional["MapAPI"] = None, + ): + """Initialize an :class:`Intersection` instance. + + Notes + ----- + If the map_api is provided, lane groups can be accessed through the properties. + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the intersection. + :param lane_group_ids: The IDs of the lane groups that belong to the intersection. + :param outline: The outline of the intersection, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the intersection, defaults to None. + :param map_api: The MapAPI instance, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + self._lane_group_ids = lane_group_ids + self._map_api = map_api + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.INTERSECTION + + @property + def lane_group_ids(self) -> List[MapObjectIDType]: + """List of IDs of the lane groups that belong to the intersection.""" + return self._lane_group_ids + + @property + def lane_groups(self) -> List[LaneGroup]: + """List of :class:`LaneGroup` instances that belong to the intersection.""" + lane_groups: Optional[List[LaneGroup]] = None + if self._map_api is not None: + lane_groups = [ + self._map_api.get_map_object(lane_group_id, MapLayer.LANE_GROUP) + for lane_group_id in self.lane_group_ids + ] + return lane_groups + + +class Crosswalk(BaseMapSurfaceObject): + """Class representing a crosswalk in a map.""" + + __slots__ = () + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ): + """Initialize a Crosswalk instance. + + Notes + ----- + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the crosswalk. + :param outline: The outline of the crosswalk, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the crosswalk, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.CROSSWALK + + +class Carpark(BaseMapSurfaceObject): + """Class representing a carpark or driveway in a map.""" + + __slots__ = () + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ): + """Initialize a Carpark instance. + + Notes + ----- + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the carpark. + :param outline: The outline of the carpark, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the carpark, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.CARPARK + + +class Walkway(BaseMapSurfaceObject): + __slots__ = () + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ): + """Initialize a Walkway instance. + + Notes + ----- + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the walkway. + :param outline: The outline of the walkway, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the walkway, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + + @property + def layer(self) -> MapLayer: + return MapLayer.WALKWAY + + +class GenericDrivable(BaseMapSurfaceObject): + """Class representing a generic drivable area in a map. + Can overlap with other drivable areas, depending on the dataset. + """ + + __slots__ = () + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ): + """Initialize a GenericDrivable instance. + + Notes + ----- + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the walkway. + :param outline: The outline of the walkway, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the walkway, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.GENERIC_DRIVABLE + + +class StopZone(BaseMapSurfaceObject): + """Placeholder class representing a stop zone in a map. Requires further implementation based on dataset specifics.""" + + __slots__ = () + + def __init__( + self, + object_id: MapObjectIDType, + outline: Optional[Union[Polyline2D, Polyline3D]] = None, + shapely_polygon: Optional[geom.Polygon] = None, + ): + """Initialize a StopZone instance. + + Notes + ----- + Either outline or shapely_polygon must be provided. + + :param object_id: The ID of the stop zone. + :param outline: The outline of the stop zone, defaults to None. + :param shapely_polygon: The Shapely polygon representation of the stop zone, defaults to None. + """ + super().__init__(object_id, outline, shapely_polygon) + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.STOP_ZONE + + +class RoadEdge(BaseMapLineObject): + """Class representing a road edge in a map.""" + + __slots__ = ("_road_edge_type",) + + def __init__( + self, + object_id: MapObjectIDType, + road_edge_type: RoadEdgeType, + polyline: Union[Polyline2D, Polyline3D], + ): + """Initialize a RoadEdge instance. + + :param object_id: The ID of the road edge. + :param road_edge_type: The type of the road edge. + :param polyline: The polyline representation of the road edge. + """ + super().__init__(object_id, polyline) + self._road_edge_type = road_edge_type + + @property + def layer(self) -> MapLayer: + return MapLayer.ROAD_EDGE + + @property + def road_edge_type(self) -> RoadEdgeType: + """The type of road edge, according to :class:`~py123d.datatypes.map_objects.map_layer_types.RoadEdgeType`.""" + return self._road_edge_type + + +class RoadLine(BaseMapLineObject): + """Class representing a road line in a map.""" + + __slots__ = ("_road_line_type",) + + def __init__( + self, + object_id: MapObjectIDType, + road_line_type: RoadLineType, + polyline: Union[Polyline2D, Polyline3D], + ): + """Initialize a RoadLine instance. + + :param object_id: The ID of the road line. + :param road_line_type: The type of the road line. + :param polyline: The polyline representation of the road line. + """ + super().__init__(object_id, polyline) + self._road_line_type = road_line_type + + @property + def layer(self) -> MapLayer: + """The :class:`~py123d.datatypes.map_objects.map_layer_types.MapLayer` of the map object.""" + return MapLayer.ROAD_LINE + + @property + def road_line_type(self) -> RoadLineType: + """The type of road edge, according to :class:`~py123d.datatypes.map_objects.map_layer_types.RoadLineType`.""" + return self._road_line_type diff --git a/src/py123d/datatypes/map_objects/utils.py b/src/py123d/datatypes/map_objects/utils.py new file mode 100644 index 00000000..7d5557d1 --- /dev/null +++ b/src/py123d/datatypes/map_objects/utils.py @@ -0,0 +1,57 @@ +import numpy as np +import numpy.typing as npt +import trimesh + +from py123d.geometry import Polyline3D + + +def get_trimesh_from_boundaries( + left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 0.25 +) -> trimesh.Trimesh: + """Helper function to create a trimesh from two lane boundaries. + + :param left_boundary: The left boundary polyline. + :param right_boundary: The right boundary polyline. + :param resolution: The resolution for the mesh, defaults to 0.25. + :return: A trimesh representation of the lane. + """ + + def _interpolate_polyline(polyline_3d: Polyline3D, num_samples: int) -> npt.NDArray[np.float64]: + """Helper function to interpolate a polyline to a fixed number of samples.""" + num_samples = max(num_samples, 2) + distances = np.linspace(0, polyline_3d.length, num=num_samples, endpoint=True, dtype=np.float64) + return polyline_3d.interpolate(distances) + + average_length = (left_boundary.length + right_boundary.length) / 2 + num_samples = int(average_length // resolution) + 1 + left_boundary_array = _interpolate_polyline(left_boundary, num_samples) + right_boundary_array = _interpolate_polyline(right_boundary, num_samples) + return _create_lane_mesh_from_boundary_arrays(left_boundary_array, right_boundary_array) + + +def _create_lane_mesh_from_boundary_arrays( + left_boundary_array: npt.NDArray[np.float64], right_boundary_array: npt.NDArray[np.float64] +) -> trimesh.Trimesh: + """Helper function to create a trimesh from two boundary arrays. + + :param left_boundary_array: The left boundary array. + :param right_boundary_array: The right boundary array. + :raises ValueError: If the boundary arrays do not have the same number of points. + :return: A trimesh representation of the lane. + """ + + # Ensure both polylines have the same number of points + if left_boundary_array.shape[0] != right_boundary_array.shape[0]: + raise ValueError("Both polylines must have the same number of points") + + n_points = left_boundary_array.shape[0] + vertices = np.vstack([left_boundary_array, right_boundary_array]) + + faces = [] + for i in range(n_points - 1): + faces.append([i, i + n_points, i + 1]) + faces.append([i + 1, i + n_points, i + n_points + 1]) + + faces = np.array(faces) + mesh = trimesh.Trimesh(vertices=vertices, faces=faces) + return mesh diff --git a/src/py123d/datatypes/maps/abstract_map.py b/src/py123d/datatypes/maps/abstract_map.py deleted file mode 100644 index 3aee0e99..00000000 --- a/src/py123d/datatypes/maps/abstract_map.py +++ /dev/null @@ -1,88 +0,0 @@ -from __future__ import annotations - -import abc -from typing import Dict, Iterable, List, Optional, Union - -import shapely - -from py123d.datatypes.maps.abstract_map_objects import AbstractMapObject -from py123d.datatypes.maps.map_datatypes import MapLayer -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.geometry import Point2D - -# TODO: -# - add docstrings -# - rename methods? -# - Combine query and query_object_ids into one method with an additional parameter to specify whether to return objects or IDs? -# - Add stop pads or stop lines. - - -class AbstractMap(abc.ABC): - - @abc.abstractmethod - def get_map_metadata(self) -> MapMetadata: - pass - - @abc.abstractmethod - def initialize(self) -> None: - pass - - @abc.abstractmethod - def get_available_map_objects(self) -> List[MapLayer]: - pass - - @abc.abstractmethod - def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: - pass - - @abc.abstractmethod - def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: - pass - - @abc.abstractmethod - def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: - pass - - @abc.abstractmethod - def get_proximal_map_objects( - self, point: Point2D, radius: float, layers: List[MapLayer] - ) -> Dict[MapLayer, List[AbstractMapObject]]: - pass - - @abc.abstractmethod - def query( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: - pass - - @abc.abstractmethod - def query_object_ids( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: - pass - - @abc.abstractmethod - def query_nearest( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - return_all: bool = True, - max_distance: Optional[float] = None, - return_distance: bool = False, - exclusive: bool = False, - ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: - pass - - @property - def location(self) -> str: - return self.get_map_metadata().location diff --git a/src/py123d/datatypes/maps/abstract_map_objects.py b/src/py123d/datatypes/maps/abstract_map_objects.py deleted file mode 100644 index baea8d87..00000000 --- a/src/py123d/datatypes/maps/abstract_map_objects.py +++ /dev/null @@ -1,449 +0,0 @@ -from __future__ import annotations - -import abc -from typing import List, Optional, Tuple, Union - -import shapely.geometry as geom -import trimesh -from typing_extensions import TypeAlias - -from py123d.datatypes.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType -from py123d.geometry import Polyline2D, Polyline3D, PolylineSE2 - -# TODO: Refactor and just use int -# type MapObjectIDType = Union[str, int] for Python >= 3.12 -MapObjectIDType: TypeAlias = Union[str, int] - - -class AbstractMapObject(abc.ABC): - """ - Base interface representation of all map objects. - """ - - def __init__(self, object_id: MapObjectIDType): - """ - Constructor of the base map object type. - :param object_id: unique identifier of the map object. - """ - self.object_id: MapObjectIDType = object_id - - @property - @abc.abstractmethod - def layer(self) -> MapLayer: - """ - Returns map layer type, e.g. LANE, ROAD_EDGE. - :return: map layer type - """ - - -class AbstractSurfaceMapObject(AbstractMapObject): - """ - Base interface representation of all map objects. - """ - - @property - @abc.abstractmethod - def shapely_polygon(self) -> geom.Polygon: - """ - Returns the 2D shapely polygon of the map object. - :return: shapely polygon - """ - - @property - @abc.abstractmethod - def outline(self) -> Union[Polyline2D, Polyline3D]: - """ - Returns the 2D or 3D outline of the map surface, if available. - :return: 2D or 3D polyline - """ - - @property - @abc.abstractmethod - def trimesh_mesh(self) -> trimesh.Trimesh: - """ - Returns a triangle mesh of the map surface. - :return: Trimesh - """ - - @property - def outline_3d(self) -> Polyline3D: - """Returns the 3D outline of the map surface, or converts 2D to 3D if necessary. - - :return: 3D polyline - """ - if isinstance(self.outline, Polyline3D): - return self.outline - # Converts 2D polyline to 3D by adding a default (zero) z-coordinate - return Polyline3D.from_linestring(self.outline.linestring) - - @property - def outline_2d(self) -> Polyline2D: - """Returns the 2D outline of the map surface, or converts 3D to 2D if necessary. - - :return: 2D polyline - """ - if isinstance(self.outline, Polyline2D): - return self.outline - # Converts 3D polyline to 2D by dropping the z-coordinate - return self.outline.polyline_2d - - -class AbstractLineMapObject(AbstractMapObject): - - @property - @abc.abstractmethod - def polyline(self) -> Union[Polyline2D, Polyline3D]: - """ - Returns the polyline of the road edge, either 2D or 3D. - :return: polyline - """ - - @property - def polyline_3d(self) -> Polyline3D: - """ - Returns the 3D polyline of the road edge. - :return: 3D polyline - """ - if isinstance(self.polyline, Polyline3D): - return self.polyline - # Converts 2D polyline to 3D by adding a default (zero) z-coordinate - return Polyline3D.from_linestring(self.polyline.linestring) - - @property - def polyline_2d(self) -> Polyline2D: - """ - Returns the 2D polyline of the road line. - :return: 2D polyline - """ - if isinstance(self.polyline, Polyline2D): - return self.polyline - # Converts 3D polyline to 2D by dropping the z-coordinate - return self.polyline.polyline_2d - - @property - def polyline_se2(self) -> PolylineSE2: - """ - Returns the 2D polyline of the road line in SE(2) coordinates. - :return: 2D polyline in SE(2) - """ - return self.polyline_2d.polyline_se2 - - @property - def shapely_linestring(self) -> geom.LineString: - """ - Returns the shapely linestring of the line, either 2D or 3D. - :return: shapely linestring - """ - return self.polyline.linestring - - -class AbstractLane(AbstractSurfaceMapObject): - """Abstract interface for lane objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.LANE - - @property - @abc.abstractmethod - def speed_limit_mps(self) -> Optional[float]: - """ - Property of lanes speed limit in m/s, if available. - :return: float or none - """ - - @property - @abc.abstractmethod - def successor_ids(self) -> List[MapObjectIDType]: - """ - Property of succeeding lane object ids (front). - :return: list of lane ids - """ - - @property - @abc.abstractmethod - def successors(self) -> List[AbstractLane]: - """ - Property of succeeding lane objects (front). - :return: list of lane class - """ - - @property - @abc.abstractmethod - def predecessor_ids(self) -> List[MapObjectIDType]: - """ - Property of preceding lane object ids (behind). - :return: list of lane ids - """ - - @property - @abc.abstractmethod - def predecessors(self) -> List[AbstractLane]: - """ - Property of preceding lane objects (behind). - :return: list of lane class - """ - - @property - @abc.abstractmethod - def left_boundary(self) -> Polyline3D: - """ - Property of left boundary of lane. - :return: returns 3D polyline - """ - - @property - @abc.abstractmethod - def right_boundary(self) -> Polyline3D: - """ - Property of right boundary of lane. - :return: returns 3D polyline - """ - - @property - @abc.abstractmethod - def left_lane_id(self) -> Optional[MapObjectIDType]: - """ - Property of left lane id of lane. - :return: returns left lane id or none, if no left lane - """ - - @property - @abc.abstractmethod - def left_lane(self) -> Optional[AbstractLane]: - """ - Property of left lane of lane. - :return: returns left lane or none, if no left lane - """ - - @property - @abc.abstractmethod - def right_lane_id(self) -> Optional[MapObjectIDType]: - """ - Property of right lane id of lane. - :return: returns right lane id or none, if no right lane - """ - - @property - @abc.abstractmethod - def right_lane(self) -> Optional[AbstractLane]: - """ - Property of right lane of lane. - :return: returns right lane or none, if no right lane - """ - - @property - @abc.abstractmethod - def centerline(self) -> Polyline3D: - """ - Property of centerline of lane. - :return: returns 3D polyline - """ - - @property - @abc.abstractmethod - def lane_group_id(self) -> AbstractLaneGroup: - """ - Property of lane group id of lane. - :return: returns lane group id - """ - - @property - @abc.abstractmethod - def lane_group(self) -> AbstractLaneGroup: - """ - Property of lane group of lane. - :return: returns lane group - """ - - @property - def boundaries(self) -> Tuple[Polyline3D, Polyline3D]: - """ - Property of left and right boundary. - :return: returns tuple of left and right boundary polylines - """ - return self.left_boundary, self.right_boundary - - -class AbstractLaneGroup(AbstractSurfaceMapObject): - """Abstract interface lane groups (nearby lanes going in the same direction).""" - - @property - def layer(self) -> MapLayer: - return MapLayer.LANE_GROUP - - @property - @abc.abstractmethod - def successor_ids(self) -> List[MapObjectIDType]: - """ - Property of succeeding lane object ids (front). - :return: list of lane group ids - """ - - @property - @abc.abstractmethod - def successors(self) -> List[AbstractLaneGroup]: - """ - Property of succeeding lane group objects (front). - :return: list of lane group class - """ - - @property - @abc.abstractmethod - def predecessor_ids(self) -> List[MapObjectIDType]: - """ - Property of preceding lane object ids (behind). - :return: list of lane group ids - """ - - @property - @abc.abstractmethod - def predecessors(self) -> List[AbstractLaneGroup]: - """ - Property of preceding lane group objects (behind). - :return: list of lane group class - """ - - @property - @abc.abstractmethod - def left_boundary(self) -> Polyline3D: - """ - Property of left boundary of lane group. - :return: returns 3D polyline - """ - - @property - @abc.abstractmethod - def right_boundary(self) -> Polyline3D: - """ - Property of right boundary of lane group. - :return: returns 3D polyline - """ - - @property - @abc.abstractmethod - def lane_ids(self) -> List[MapObjectIDType]: - """ - Property of interior lane ids of a lane group. - :return: returns list of lane ids - """ - - @property - @abc.abstractmethod - def lanes(self) -> List[AbstractLane]: - """ - Property of interior lanes of a lane group. - :return: returns list of lanes - """ - - @property - @abc.abstractmethod - def intersection_id(self) -> Optional[MapObjectIDType]: - """ - Property of intersection id of a lane group. - :return: returns intersection id or none, if lane group not in intersection - """ - - @property - @abc.abstractmethod - def intersection(self) -> Optional[AbstractIntersection]: - """ - Property of intersection of a lane group. - :return: returns intersection or none, if lane group not in intersection - """ - - -class AbstractIntersection(AbstractSurfaceMapObject): - """Abstract interface for intersection objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.INTERSECTION - - @property - @abc.abstractmethod - def lane_group_ids(self) -> List[MapObjectIDType]: - """ - Property of lane group ids of intersection. - :return: returns list of lane group ids - """ - - @property - @abc.abstractmethod - def lane_groups(self) -> List[AbstractLaneGroup]: - """ - Property of lane groups of intersection. - :return: returns list of lane groups - """ - - -class AbstractCrosswalk(AbstractSurfaceMapObject): - """Abstract interface for crosswalk objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.CROSSWALK - - -class AbstractWalkway(AbstractSurfaceMapObject): - """Abstract interface for walkway objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.WALKWAY - - -class AbstractCarpark(AbstractSurfaceMapObject): - """Abstract interface for carpark objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.CARPARK - - -class AbstractGenericDrivable(AbstractSurfaceMapObject): - """Abstract interface for generic drivable objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.GENERIC_DRIVABLE - - -class AbstractStopLine(AbstractSurfaceMapObject): - """Abstract interface for stop line objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.STOP_LINE - - -class AbstractRoadEdge(AbstractLineMapObject): - """Abstract interface for road edge objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.ROAD_EDGE - - @property - @abc.abstractmethod - def road_edge_type(self) -> RoadEdgeType: - """ - Returns the road edge type. - :return: RoadEdgeType - """ - - -class AbstractRoadLine(AbstractLineMapObject): - """Abstract interface for road line objects.""" - - @property - def layer(self) -> MapLayer: - return MapLayer.ROAD_LINE - - @property - @abc.abstractmethod - def road_line_type(self) -> RoadLineType: - """ - Returns the road line type. - :return: RoadLineType - """ diff --git a/src/py123d/datatypes/maps/cache/cache_map_objects.py b/src/py123d/datatypes/maps/cache/cache_map_objects.py deleted file mode 100644 index 498c01c1..00000000 --- a/src/py123d/datatypes/maps/cache/cache_map_objects.py +++ /dev/null @@ -1,311 +0,0 @@ -from __future__ import annotations - -from typing import List, Optional, Union - -import numpy as np -import shapely.geometry as geom -import trimesh - -from py123d.datatypes.maps.abstract_map_objects import ( - AbstractCarpark, - AbstractCrosswalk, - AbstractGenericDrivable, - AbstractIntersection, - AbstractLane, - AbstractLaneGroup, - AbstractLineMapObject, - AbstractRoadEdge, - AbstractRoadLine, - AbstractSurfaceMapObject, - AbstractWalkway, - MapObjectIDType, -) -from py123d.datatypes.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType -from py123d.geometry import Polyline3D -from py123d.geometry.polyline import Polyline2D - - -class CacheSurfaceObject(AbstractSurfaceMapObject): - """ - Base interface representation of all map objects. - """ - - def __init__( - self, - object_id: MapObjectIDType, - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ) -> None: - super().__init__(object_id) - - assert outline is not None or geometry is not None, "Either outline or geometry must be provided." - - if outline is None: - outline = Polyline3D.from_linestring(geometry.exterior) - - if geometry is None: - geometry = geom.Polygon(outline.array[:, :2]) - - self._outline = outline - self._geometry = geometry - - outline = property(lambda self: self._outline) - - @property - def shapely_polygon(self) -> geom.Polygon: - """Inherited, see superclass.""" - return self._geometry - - @property - def outline_3d(self) -> Polyline3D: - """Inherited, see superclass.""" - if isinstance(self.outline, Polyline3D): - return self.outline - # Converts 2D polyline to 3D by adding a default (zero) z-coordinate - return Polyline3D.from_linestring(self.outline.linestring) - - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - raise NotImplementedError - - -class CacheLineObject(AbstractLineMapObject): - - def __init__(self, object_id: MapObjectIDType, polyline: Union[Polyline2D, Polyline3D]) -> None: - """ - Constructor of the base line map object type. - :param object_id: unique identifier of a line map object. - """ - super().__init__(object_id) - self._polyline = polyline - - polyline = property(lambda self: self._polyline) - - -class CacheLane(CacheSurfaceObject, AbstractLane): - - def __init__( - self, - object_id: MapObjectIDType, - lane_group_id: MapObjectIDType, - left_boundary: Polyline3D, - right_boundary: Polyline3D, - centerline: Polyline3D, - left_lane_id: Optional[MapObjectIDType] = None, - right_lane_id: Optional[MapObjectIDType] = None, - predecessor_ids: List[MapObjectIDType] = [], - successor_ids: List[MapObjectIDType] = [], - speed_limit_mps: Optional[float] = None, - outline: Optional[Polyline3D] = None, - geometry: Optional[geom.Polygon] = None, - ) -> None: - - if outline is None: - outline_array = np.vstack( - ( - left_boundary.array, - right_boundary.array[::-1], - left_boundary.array[0], - ) - ) - outline = Polyline3D.from_linestring(geom.LineString(outline_array)) - - super().__init__(object_id, outline, geometry) - - self._lane_group_id = lane_group_id - self._left_boundary = left_boundary - self._right_boundary = right_boundary - self._centerline = centerline - self._left_lane_id = left_lane_id - self._right_lane_id = right_lane_id - self._predecessor_ids = predecessor_ids - self._successor_ids = successor_ids - self._speed_limit_mps = speed_limit_mps - - lane_group_id = property(lambda self: self._lane_group_id) - left_boundary = property(lambda self: self._left_boundary) - right_boundary = property(lambda self: self._right_boundary) - centerline = property(lambda self: self._centerline) - left_lane_id = property(lambda self: self._left_lane_id) - right_lane_id = property(lambda self: self._right_lane_id) - predecessor_ids = property(lambda self: self._predecessor_ids) - successor_ids = property(lambda self: self._successor_ids) - speed_limit_mps = property(lambda self: self._speed_limit_mps) - - @property - def layer(self) -> MapLayer: - """Inherited, see superclass.""" - return MapLayer.LANE - - @property - def successors(self) -> List[AbstractLane]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def predecessors(self) -> List[AbstractLane]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def left_lane(self) -> Optional[AbstractLane]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def right_lane(self) -> Optional[AbstractLane]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def lane_group(self) -> AbstractLaneGroup: - """Inherited, see superclass.""" - raise NotImplementedError - - -class CacheLaneGroup(CacheSurfaceObject, AbstractLaneGroup): - def __init__( - self, - object_id: MapObjectIDType, - lane_ids: List[MapObjectIDType], - left_boundary: Polyline3D, - right_boundary: Polyline3D, - intersection_id: Optional[MapObjectIDType] = None, - predecessor_ids: List[MapObjectIDType] = [], - successor_ids: List[MapObjectIDType] = [], - outline: Optional[Polyline3D] = None, - geometry: Optional[geom.Polygon] = None, - ): - if outline is None: - outline_array = np.vstack( - ( - left_boundary.array, - right_boundary.array[::-1], - left_boundary.array[0], - ) - ) - outline = Polyline3D.from_linestring(geom.LineString(outline_array)) - super().__init__(object_id, outline, geometry) - - self._lane_ids = lane_ids - self._left_boundary = left_boundary - self._right_boundary = right_boundary - self._intersection_id = intersection_id - self._predecessor_ids = predecessor_ids - self._successor_ids = successor_ids - - layer = property(lambda self: MapLayer.LANE_GROUP) - lane_ids = property(lambda self: self._lane_ids) - intersection_id = property(lambda self: self._intersection_id) - predecessor_ids = property(lambda self: self._predecessor_ids) - successor_ids = property(lambda self: self._successor_ids) - left_boundary = property(lambda self: self._left_boundary) - right_boundary = property(lambda self: self._right_boundary) - - @property - def successors(self) -> List[AbstractLaneGroup]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def predecessors(self) -> List[AbstractLaneGroup]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def lanes(self) -> List[AbstractLane]: - """Inherited, see superclass.""" - raise NotImplementedError - - @property - def intersection(self) -> Optional[AbstractIntersection]: - """Inherited, see superclass.""" - raise NotImplementedError - - -class CacheIntersection(CacheSurfaceObject, AbstractIntersection): - def __init__( - self, - object_id: MapObjectIDType, - lane_group_ids: List[MapObjectIDType], - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ): - - super().__init__(object_id, outline, geometry) - self._lane_group_ids = lane_group_ids - - layer = property(lambda self: MapLayer.INTERSECTION) - lane_group_ids = property(lambda self: self._lane_group_ids) - - @property - def lane_groups(self) -> List[CacheLaneGroup]: - """Inherited, see superclass.""" - raise NotImplementedError - - -class CacheCrosswalk(CacheSurfaceObject, AbstractCrosswalk): - def __init__( - self, - object_id: MapObjectIDType, - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ): - super().__init__(object_id, outline, geometry) - - -class CacheCarpark(CacheSurfaceObject, AbstractCarpark): - def __init__( - self, - object_id: MapObjectIDType, - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ): - super().__init__(object_id, outline, geometry) - - -class CacheWalkway(CacheSurfaceObject, AbstractWalkway): - def __init__( - self, - object_id: MapObjectIDType, - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ): - super().__init__(object_id, outline, geometry) - - -class CacheGenericDrivable(CacheSurfaceObject, AbstractGenericDrivable): - def __init__( - self, - object_id: MapObjectIDType, - outline: Optional[Union[Polyline2D, Polyline3D]] = None, - geometry: Optional[geom.Polygon] = None, - ): - super().__init__(object_id, outline, geometry) - - -class CacheRoadEdge(CacheLineObject, AbstractRoadEdge): - def __init__( - self, - object_id: MapObjectIDType, - road_edge_type: RoadEdgeType, - polyline: Union[Polyline2D, Polyline3D], - ): - super().__init__(object_id, polyline) - self._road_edge_type = road_edge_type - - road_edge_type = property(lambda self: self._road_edge_type) - - -class CacheRoadLine(CacheLineObject, AbstractRoadLine): - def __init__( - self, - object_id: MapObjectIDType, - road_line_type: RoadLineType, - polyline: Union[Polyline2D, Polyline3D], - ): - super().__init__(object_id, polyline) - self._road_line_type = road_line_type - - road_line_type = property(lambda self: self._road_line_type) diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_map.py b/src/py123d/datatypes/maps/gpkg/gpkg_map.py deleted file mode 100644 index acc1cfbb..00000000 --- a/src/py123d/datatypes/maps/gpkg/gpkg_map.py +++ /dev/null @@ -1,394 +0,0 @@ -from __future__ import annotations - -import warnings -from collections import defaultdict -from functools import lru_cache -from pathlib import Path -from typing import Callable, Dict, Final, Iterable, List, Optional, Tuple, Union - -import geopandas as gpd -import shapely -import shapely.geometry as geom - -from py123d.datatypes.maps.abstract_map import AbstractMap -from py123d.datatypes.maps.abstract_map_objects import AbstractMapObject -from py123d.datatypes.maps.gpkg.gpkg_map_objects import ( - GPKGCarpark, - GPKGCrosswalk, - GPKGGenericDrivable, - GPKGIntersection, - GPKGLane, - GPKGLaneGroup, - GPKGRoadEdge, - GPKGRoadLine, - GPKGWalkway, -) -from py123d.datatypes.maps.gpkg.gpkg_utils import load_gdf_with_geometry_columns -from py123d.datatypes.maps.map_datatypes import MapLayer -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.geometry import Point2D -from py123d.script.utils.dataset_path_utils import get_dataset_paths - -USE_ARROW: bool = True -MAX_LRU_CACHED_TABLES: Final[int] = 128 # TODO: add to some configs - - -class GPKGMap(AbstractMap): - def __init__(self, file_path: Path) -> None: - - self._file_path = file_path - self._map_object_getter: Dict[MapLayer, Callable[[str], Optional[AbstractMapObject]]] = { - MapLayer.LANE: self._get_lane, - MapLayer.LANE_GROUP: self._get_lane_group, - MapLayer.INTERSECTION: self._get_intersection, - MapLayer.CROSSWALK: self._get_crosswalk, - MapLayer.WALKWAY: self._get_walkway, - MapLayer.CARPARK: self._get_carpark, - MapLayer.GENERIC_DRIVABLE: self._get_generic_drivable, - MapLayer.ROAD_EDGE: self._get_road_edge, - MapLayer.ROAD_LINE: self._get_road_line, - } - - # loaded during `.initialize()` - self._gpd_dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} - self._map_metadata: Optional[MapMetadata] = None - - def initialize(self) -> None: - """Inherited, see superclass.""" - if len(self._gpd_dataframes) == 0: - available_layers = list(gpd.list_layers(self._file_path).name) - for map_layer in list(MapLayer): - map_layer_name = map_layer.serialize() - if map_layer_name in available_layers: - self._gpd_dataframes[map_layer] = gpd.read_file( - self._file_path, layer=map_layer_name, use_arrow=USE_ARROW - ) - load_gdf_with_geometry_columns( - self._gpd_dataframes[map_layer], - geometry_column_names=["centerline", "right_boundary", "left_boundary", "outline"], - ) - # TODO: remove the temporary fix and enforce consistent id types in the GPKG files - if "id" in self._gpd_dataframes[map_layer].columns: - self._gpd_dataframes[map_layer]["id"] = self._gpd_dataframes[map_layer]["id"].astype(str) - else: - warnings.warn(f"GPKGMap: {map_layer_name} not available in {str(self._file_path)}") - self._gpd_dataframes[map_layer] = None - - assert "map_metadata" in list(gpd.list_layers(self._file_path).name) - metadata_gdf = gpd.read_file(self._file_path, layer="map_metadata", use_arrow=USE_ARROW) - self._map_metadata = MapMetadata.from_dict(metadata_gdf.iloc[0].to_dict()) - - def _assert_initialize(self) -> None: - "Checks if `.initialize()` was called, before retrieving data." - assert len(self._gpd_dataframes) > 0, "GPKGMap: Call `.initialize()` before retrieving data!" - - def _assert_layer_available(self, layer: MapLayer) -> None: - "Checks if layer is available." - assert layer in self.get_available_map_objects(), f"GPKGMap: MapLayer {layer.name} is unavailable." - - def get_map_metadata(self): - return self._map_metadata - - def get_available_map_objects(self) -> List[MapLayer]: - """Inherited, see superclass.""" - self._assert_initialize() - return list(self._gpd_dataframes.keys()) - - def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: - """Inherited, see superclass.""" - - self._assert_initialize() - self._assert_layer_available(layer) - try: - return self._map_object_getter[layer](object_id) - except KeyError: - raise ValueError(f"Object representation for layer: {layer.name} object: {object_id} is unavailable") - - def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: - """Inherited, see superclass.""" - raise NotImplementedError - - def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: - """Inherited, see superclass.""" - raise NotImplementedError - - def get_proximal_map_objects( - self, point_2d: Point2D, radius: float, layers: List[MapLayer] - ) -> Dict[MapLayer, List[AbstractMapObject]]: - """Inherited, see superclass.""" - center_point = geom.Point(point_2d.x, point_2d.y) - patch = center_point.buffer(radius) - return self.query(geometry=patch, layers=layers, predicate="intersects") - - def query( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: - supported_layers = self.get_available_map_objects() - unsupported_layers = [layer for layer in layers if layer not in supported_layers] - assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]] = defaultdict( - list - ) - for layer in layers: - object_map[layer] = self._query_layer(geometry, layer, predicate, sort, distance) - return object_map - - def query_object_ids( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: - supported_layers = self.get_available_map_objects() - unsupported_layers = [layer for layer in layers if layer not in supported_layers] - assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) - for layer in layers: - object_map[layer] = self._query_layer_objects_ids(geometry, layer, predicate, sort, distance) - return object_map - - def query_nearest( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapLayer], - return_all: bool = True, - max_distance: Optional[float] = None, - return_distance: bool = False, - exclusive: bool = False, - ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: - supported_layers = self.get_available_map_objects() - unsupported_layers = [layer for layer in layers if layer not in supported_layers] - assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) - for layer in layers: - object_map[layer] = self._query_layer_nearest( - geometry, layer, return_all, max_distance, return_distance, exclusive - ) - return object_map - - def _query_layer( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapLayer, - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]: - queried_indices = self._gpd_dataframes[layer].sindex.query( - geometry, predicate=predicate, sort=sort, distance=distance - ) - - if queried_indices.ndim == 2: - query_dict: Dict[int, List[AbstractMapObject]] = defaultdict(list) - for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): - map_object_id = self._gpd_dataframes[layer]["id"].iloc[map_object_idx] - query_dict[int(geometry_idx)].append(self.get_map_object(map_object_id, layer)) - return query_dict - else: - map_object_ids = self._gpd_dataframes[layer]["id"].iloc[queried_indices] - return [self.get_map_object(map_object_id, layer) for map_object_id in map_object_ids] - - def _query_layer_objects_ids( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapLayer, - predicate: Optional[str] = None, - sort: bool = False, - distance: Optional[float] = None, - ) -> Union[List[str], Dict[int, List[str]]]: - # Use numpy for fast indexing and avoid .iloc in a loop - - queried_indices = self._gpd_dataframes[layer].sindex.query( - geometry, predicate=predicate, sort=sort, distance=distance - ) - ids = self._gpd_dataframes[layer]["id"].values # numpy array for fast access - - if queried_indices.ndim == 2: - query_dict: Dict[int, List[str]] = defaultdict(list) - for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): - query_dict[int(geometry_idx)].append(ids[map_object_idx]) - return query_dict - else: - return list(ids[queried_indices]) - - def _query_layer_nearest( - self, - geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapLayer, - return_all: bool = True, - max_distance: Optional[float] = None, - return_distance: bool = False, - exclusive: bool = False, - ) -> Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]: - # Use numpy for fast indexing and avoid .iloc in a loop - - queried_indices = self._gpd_dataframes[layer].sindex.nearest( - geometry, - return_all=return_all, - max_distance=max_distance, - return_distance=return_distance, - exclusive=exclusive, - ) - ids = self._gpd_dataframes[layer]["id"].values # numpy array for fast access - - if return_distance: - queried_indices, distances = queried_indices - query_dict: Dict[int, List[Tuple[str, float]]] = defaultdict(list) - for geometry_idx, map_object_idx, distance in zip(queried_indices[0], queried_indices[1], distances): - query_dict[int(geometry_idx)].append((ids[map_object_idx], float(distance))) - return query_dict - - elif queried_indices.ndim == 2: - query_dict: Dict[int, List[str]] = defaultdict(list) - for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): - query_dict[int(geometry_idx)].append(ids[map_object_idx]) - return query_dict - else: - return list(ids[queried_indices]) - - def _get_lane(self, id: str) -> Optional[GPKGLane]: - return ( - GPKGLane( - id, - self._gpd_dataframes[MapLayer.LANE], - self._gpd_dataframes[MapLayer.LANE_GROUP], - self._gpd_dataframes[MapLayer.INTERSECTION], - ) - if id in self._gpd_dataframes[MapLayer.LANE]["id"].tolist() - else None - ) - - def _get_lane_group(self, id: str) -> Optional[GPKGLaneGroup]: - return ( - GPKGLaneGroup( - id, - self._gpd_dataframes[MapLayer.LANE_GROUP], - self._gpd_dataframes[MapLayer.LANE], - self._gpd_dataframes[MapLayer.INTERSECTION], - ) - if id in self._gpd_dataframes[MapLayer.LANE_GROUP]["id"].tolist() - else None - ) - - def _get_intersection(self, id: str) -> Optional[GPKGIntersection]: - return ( - GPKGIntersection( - id, - self._gpd_dataframes[MapLayer.INTERSECTION], - self._gpd_dataframes[MapLayer.LANE], - self._gpd_dataframes[MapLayer.LANE_GROUP], - ) - if id in self._gpd_dataframes[MapLayer.INTERSECTION]["id"].tolist() - else None - ) - - def _get_crosswalk(self, id: str) -> Optional[GPKGCrosswalk]: - return ( - GPKGCrosswalk(id, self._gpd_dataframes[MapLayer.CROSSWALK]) - if id in self._gpd_dataframes[MapLayer.CROSSWALK]["id"].tolist() - else None - ) - - def _get_walkway(self, id: str) -> Optional[GPKGWalkway]: - return ( - GPKGWalkway(id, self._gpd_dataframes[MapLayer.WALKWAY]) - if id in self._gpd_dataframes[MapLayer.WALKWAY]["id"].tolist() - else None - ) - - def _get_carpark(self, id: str) -> Optional[GPKGCarpark]: - return ( - GPKGCarpark(id, self._gpd_dataframes[MapLayer.CARPARK]) - if id in self._gpd_dataframes[MapLayer.CARPARK]["id"].tolist() - else None - ) - - def _get_generic_drivable(self, id: str) -> Optional[GPKGGenericDrivable]: - return ( - GPKGGenericDrivable(id, self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]) - if id in self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]["id"].tolist() - else None - ) - - def _get_road_edge(self, id: str) -> Optional[GPKGRoadEdge]: - return ( - GPKGRoadEdge(id, self._gpd_dataframes[MapLayer.ROAD_EDGE]) - if id in self._gpd_dataframes[MapLayer.ROAD_EDGE]["id"].tolist() - else None - ) - - def _get_road_line(self, id: str) -> Optional[GPKGRoadLine]: - return ( - GPKGRoadLine(id, self._gpd_dataframes[MapLayer.ROAD_LINE]) - if id in self._gpd_dataframes[MapLayer.ROAD_LINE]["id"].tolist() - else None - ) - - # def _query_layer( - # self, - # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapLayer, - # predicate: Optional[str] = None, - # sort: bool = False, - # distance: Optional[float] = None, - # ) -> Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]: - # queried_indices = self._gpd_dataframes[layer].sindex.query( - # geometry, predicate=predicate, sort=sort, distance=distance - # ) - # ids = self._gpd_dataframes[layer]["id"].values # numpy array for fast access - # if queried_indices.ndim == 2: - # query_dict: Dict[int, List[AbstractMapObject]] = defaultdict(list) - # for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): - # map_object_id = ids[map_object_idx] - # query_dict[int(geometry_idx)].append(self.get_map_object(map_object_id, layer)) - # return query_dict - # else: - # map_object_ids = ids[queried_indices] - # return [self.get_map_object(map_object_id, layer) for map_object_id in map_object_ids] - - # def _query_layer_objects_ids( - # self, - # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapLayer, - # predicate: Optional[str] = None, - # sort: bool = False, - # distance: Optional[float] = None, - # ) -> Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]: - # queried_indices = self._gpd_dataframes[layer].sindex.query( - # geometry, predicate=predicate, sort=sort, distance=distance - # ) - # if queried_indices.ndim == 2: - # query_dict: Dict[int, List[AbstractMapObject]] = defaultdict(list) - # for geometry_idx, map_object_idx in zip(queried_indices[0], queried_indices[1]): - # map_object_id = self._gpd_dataframes[layer]["id"].iloc[map_object_idx] - # query_dict[int(geometry_idx)].append(map_object_id) - # return query_dict - # else: - # map_object_ids = self._gpd_dataframes[layer]["id"].iloc[queried_indices] - # return list(map_object_ids) - - -@lru_cache(maxsize=MAX_LRU_CACHED_TABLES) -def get_global_map_api(dataset: str, location: str) -> GPKGMap: - PY123D_MAPS_ROOT: Path = Path(get_dataset_paths().py123d_maps_root) - gpkg_path = PY123D_MAPS_ROOT / dataset / f"{dataset}_{location}.gpkg" - assert gpkg_path.is_file(), f"{dataset}_{location}.gpkg not found in {str(PY123D_MAPS_ROOT)}." - map_api = GPKGMap(gpkg_path) - map_api.initialize() - return map_api - - -def get_local_map_api(split_name: str, log_name: str) -> GPKGMap: - PY123D_MAPS_ROOT: Path = Path(get_dataset_paths().py123d_maps_root) - gpkg_path = PY123D_MAPS_ROOT / split_name / f"{log_name}.gpkg" - assert gpkg_path.is_file(), f"{log_name}.gpkg not found in {str(PY123D_MAPS_ROOT)}." - map_api = GPKGMap(gpkg_path) - map_api.initialize() - return map_api diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py b/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py deleted file mode 100644 index 97b11d73..00000000 --- a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py +++ /dev/null @@ -1,385 +0,0 @@ -from __future__ import annotations - -import ast -from functools import cached_property -from typing import List, Optional, Union - -import geopandas as gpd -import numpy as np -import pandas as pd -import shapely.geometry as geom -import trimesh - -from py123d.datatypes.maps.abstract_map_objects import ( - AbstractCarpark, - AbstractCrosswalk, - AbstractGenericDrivable, - AbstractIntersection, - AbstractLane, - AbstractLaneGroup, - AbstractLineMapObject, - AbstractRoadEdge, - AbstractRoadLine, - AbstractSurfaceMapObject, - AbstractWalkway, - MapObjectIDType, -) -from py123d.datatypes.maps.gpkg.gpkg_utils import get_row_with_value, get_trimesh_from_boundaries -from py123d.datatypes.maps.map_datatypes import RoadEdgeType, RoadLineType -from py123d.geometry import Point3DIndex, Polyline3D -from py123d.geometry.polyline import Polyline2D - - -class GPKGSurfaceObject(AbstractSurfaceMapObject): - """ - Base interface representation of all map objects. - """ - - def __init__(self, object_id: MapObjectIDType, surface_df: gpd.GeoDataFrame) -> None: - """ - Constructor of the base surface map object type. - :param object_id: unique identifier of a surface map object. - """ - super().__init__(object_id) - # TODO: add assertion if columns are available - self._object_df = surface_df - - @property - def shapely_polygon(self) -> geom.Polygon: - """Inherited, see superclass.""" - return self._object_row.geometry - - @cached_property - def _object_row(self) -> gpd.GeoSeries: - return get_row_with_value(self._object_df, "id", self.object_id) - - @property - def outline(self) -> Polyline3D: - """Inherited, see superclass.""" - outline_3d: Optional[Polyline3D] = None - if "outline" in self._object_df.columns: - outline_3d = Polyline3D.from_linestring(self._object_row.outline) - else: - outline_3d = Polyline3D.from_linestring(geom.LineString(self.shapely_polygon.exterior.coords)) - return outline_3d - - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - - trimesh_mesh: Optional[trimesh.Trimesh] = None - if "right_boundary" in self._object_df.columns and "left_boundary" in self._object_df.columns: - left_boundary = Polyline3D.from_linestring(self._object_row.left_boundary) - right_boundary = Polyline3D.from_linestring(self._object_row.right_boundary) - trimesh_mesh = get_trimesh_from_boundaries(left_boundary, right_boundary) - else: - # Fallback to geometry if no boundaries are available - outline_3d_array = self.outline_3d.array - vertices_2d, faces = trimesh.creation.triangulate_polygon( - geom.Polygon(outline_3d_array[:, Point3DIndex.XY]) - ) - if len(vertices_2d) == len(outline_3d_array): - # Regular case, where vertices match outline_3d_array - vertices_3d = outline_3d_array - elif len(vertices_2d) == len(outline_3d_array) + 1: - # outline array was not closed, so we need to add the first vertex again - vertices_3d = np.vstack((outline_3d_array, outline_3d_array[0])) - else: - raise ValueError("No vertices found for triangulation.") - trimesh_mesh = trimesh.Trimesh(vertices=vertices_3d, faces=faces) - return trimesh_mesh - - -class GPKGLineObject(AbstractLineMapObject): - - def __init__(self, object_id: MapObjectIDType, line_df: gpd.GeoDataFrame) -> None: - """ - Constructor of the base line map object type. - :param object_id: unique identifier of a line map object. - """ - super().__init__(object_id) - # TODO: add assertion if columns are available - self._object_df = line_df - - @cached_property - def _object_row(self) -> gpd.GeoSeries: - return get_row_with_value(self._object_df, "id", self.object_id) - - @property - def polyline(self) -> Union[Polyline2D, Polyline3D]: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.geometry) - - -class GPKGLane(GPKGSurfaceObject, AbstractLane): - def __init__( - self, - object_id: MapObjectIDType, - object_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - intersection_df: gpd.GeoDataFrame, - ) -> None: - super().__init__(object_id, object_df) - self._lane_group_df = lane_group_df - self._intersection_df = intersection_df - - @property - def speed_limit_mps(self) -> Optional[float]: - """Inherited, see superclass.""" - return self._object_row.speed_limit_mps - - @property - def successor_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.successor_ids) - - @property - def successors(self) -> List[GPKGLane]: - """Inherited, see superclass.""" - return [GPKGLane(lane_id, self._object_df) for lane_id in self.successor_ids] - - @property - def predecessor_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.predecessor_ids) - - @property - def predecessors(self) -> List[GPKGLane]: - """Inherited, see superclass.""" - return [GPKGLane(lane_id, self._object_df) for lane_id in self.predecessor_ids] - - @property - def left_boundary(self) -> Polyline3D: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.left_boundary) - - @property - def right_boundary(self) -> Polyline3D: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.right_boundary) - - @property - def left_lane_id(self) -> Optional[MapObjectIDType]: - """ "Inherited, see superclass.""" - return self._object_row.left_lane_id - - @property - def left_lane(self) -> Optional[GPKGLane]: - """Inherited, see superclass.""" - return ( - GPKGLane(self.left_lane_id, self._object_df, self._lane_group_df, self._intersection_df) - if self.left_lane_id is not None and not pd.isna(self.left_lane_id) - else None - ) - - @property - def right_lane_id(self) -> Optional[MapObjectIDType]: - """Inherited, see superclass.""" - return self._object_row.right_lane_id - - @property - def right_lane(self) -> Optional[GPKGLane]: - """Inherited, see superclass.""" - return ( - GPKGLane(self.right_lane_id, self._object_df, self._lane_group_df, self._intersection_df) - if self.right_lane_id is not None and not pd.isna(self.right_lane_id) - else None - ) - - @property - def centerline(self) -> Polyline3D: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.centerline) - - @property - def outline_3d(self) -> Polyline3D: - """Inherited, see superclass.""" - outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) - outline_array = np.vstack((outline_array, outline_array[0])) - return Polyline3D.from_linestring(geom.LineString(outline_array)) - - @property - def lane_group_id(self) -> MapObjectIDType: - """Inherited, see superclass.""" - return self._object_row.lane_group_id - - @property - def lane_group(self) -> GPKGLaneGroup: - """Inherited, see superclass.""" - return GPKGLaneGroup( - self.lane_group_id, - self._lane_group_df, - self._object_df, - self._intersection_df, - ) - - -class GPKGLaneGroup(GPKGSurfaceObject, AbstractLaneGroup): - def __init__( - self, - object_id: MapObjectIDType, - object_df: gpd.GeoDataFrame, - lane_df: gpd.GeoDataFrame, - intersection_df: gpd.GeoDataFrame, - ): - super().__init__(object_id, object_df) - self._lane_df = lane_df - self._intersection_df = intersection_df - - @property - def successor_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.successor_ids) - - @property - def successors(self) -> List[GPKGLaneGroup]: - """Inherited, see superclass.""" - return [ - GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) - for lane_group_id in self.successor_ids - ] - - @property - def predecessor_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.predecessor_ids) - - @property - def predecessors(self) -> List[GPKGLaneGroup]: - """Inherited, see superclass.""" - return [ - GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) - for lane_group_id in self.predecessor_ids - ] - - @property - def left_boundary(self) -> Polyline3D: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.left_boundary) - - @property - def right_boundary(self) -> Polyline3D: - """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.right_boundary) - - @property - def outline_3d(self) -> Polyline3D: - """Inherited, see superclass.""" - outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) - return Polyline3D.from_linestring(geom.LineString(outline_array)) - - @property - def lane_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.lane_ids) - - @property - def lanes(self) -> List[GPKGLane]: - """Inherited, see superclass.""" - return [ - GPKGLane( - lane_id, - self._lane_df, - self._object_df, - self._intersection_df, - ) - for lane_id in self.lane_ids - ] - - @property - def intersection_id(self) -> Optional[MapObjectIDType]: - """Inherited, see superclass.""" - return self._object_row.intersection_id - - @property - def intersection(self) -> Optional[GPKGIntersection]: - """Inherited, see superclass.""" - return ( - GPKGIntersection( - self.intersection_id, - self._intersection_df, - self._lane_df, - self._object_df, - ) - if self.intersection_id is not None and not pd.isna(self.intersection_id) - else None - ) - - -class GPKGIntersection(GPKGSurfaceObject, AbstractIntersection): - def __init__( - self, - object_id: MapObjectIDType, - object_df: gpd.GeoDataFrame, - lane_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - ): - super().__init__(object_id, object_df) - self._lane_df = lane_df - self._lane_group_df = lane_group_df - - @property - def lane_group_ids(self) -> List[MapObjectIDType]: - """Inherited, see superclass.""" - return ast.literal_eval(self._object_row.lane_group_ids) - - @property - def lane_groups(self) -> List[GPKGLaneGroup]: - """Inherited, see superclass.""" - return [ - GPKGLaneGroup( - lane_group_id, - self._lane_group_df, - self._lane_df, - self._object_df, - ) - for lane_group_id in self.lane_group_ids - ] - - -class GPKGCrosswalk(GPKGSurfaceObject, AbstractCrosswalk): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - -class GPKGCarpark(GPKGSurfaceObject, AbstractCarpark): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - -class GPKGWalkway(GPKGSurfaceObject, AbstractWalkway): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - -class GPKGGenericDrivable(GPKGSurfaceObject, AbstractGenericDrivable): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - -class GPKGRoadEdge(GPKGLineObject, AbstractRoadEdge): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - @cached_property - def _object_row(self) -> gpd.GeoSeries: - return get_row_with_value(self._object_df, "id", self.object_id) - - @property - def road_edge_type(self) -> RoadEdgeType: - """Inherited, see superclass.""" - return RoadEdgeType(int(self._object_row.road_edge_type)) - - -class GPKGRoadLine(GPKGLineObject, AbstractRoadLine): - def __init__(self, object_id: MapObjectIDType, object_df: gpd.GeoDataFrame): - super().__init__(object_id, object_df) - - @cached_property - def _object_row(self) -> gpd.GeoSeries: - return get_row_with_value(self._object_df, "id", self.object_id) - - @property - def road_line_type(self) -> RoadLineType: - """Inherited, see superclass.""" - return RoadLineType(int(self._object_row.road_line_type)) diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py b/src/py123d/datatypes/maps/gpkg/gpkg_utils.py deleted file mode 100644 index 54dd93e6..00000000 --- a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py +++ /dev/null @@ -1,90 +0,0 @@ -from typing import List - -import geopandas as gpd -import numpy as np -import numpy.typing as npt -import trimesh -from shapely import wkt - -from py123d.geometry.polyline import Polyline3D - - -def load_gdf_with_geometry_columns(gdf: gpd.GeoDataFrame, geometry_column_names: List[str] = []): - # TODO: refactor - # Convert string geometry columns back to shapely objects - for col in geometry_column_names: - if col in gdf.columns and len(gdf) > 0 and isinstance(gdf[col].iloc[0], str): - try: - gdf[col] = gdf[col].apply(lambda x: wkt.loads(x) if isinstance(x, str) else x) - except Exception as e: - print(f"Warning: Could not convert column {col} to geometry: {str(e)}") - - -def get_all_rows_with_value( - elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str -) -> gpd.geodataframe.GeoDataFrame: - """ - Extract all matching elements. Note, if no matching desired_key is found and empty list is returned. - :param elements: data frame from MapsDb. - :param column_label: key to extract from a column. - :param desired_value: key which is compared with the values of column_label entry. - :return: a subset of the original GeoDataFrame containing the matching key. - """ - return elements.iloc[np.where(elements[column_label].to_numpy().astype(int) == int(desired_value))] - - -def get_row_with_value(elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str) -> gpd.GeoSeries: - """ - Extract a matching element. - :param elements: data frame from MapsDb. - :param column_label: key to extract from a column. - :param desired_value: key which is compared with the values of column_label entry. - :return row from GeoDataFrame. - """ - if column_label == "fid": - return elements.loc[desired_value] - - matching_rows = get_all_rows_with_value(elements, column_label, desired_value) - assert len(matching_rows) > 0, f"Could not find the desired key = {desired_value}" - assert len(matching_rows) == 1, ( - f"{len(matching_rows)} matching keys found. Expected to only find one." "Try using get_all_rows_with_value" - ) - return matching_rows.iloc[0] - - -def get_trimesh_from_boundaries( - left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 0.25 -) -> trimesh.Trimesh: - - def _interpolate_polyline(polyline_3d: Polyline3D, num_samples: int) -> npt.NDArray[np.float64]: - if num_samples < 2: - num_samples = 2 - distances = np.linspace(0, polyline_3d.length, num=num_samples, endpoint=True, dtype=np.float64) - return polyline_3d.interpolate(distances) - - average_length = (left_boundary.length + right_boundary.length) / 2 - num_samples = int(average_length // resolution) + 1 - left_boundary_array = _interpolate_polyline(left_boundary, num_samples) - right_boundary_array = _interpolate_polyline(right_boundary, num_samples) - return _create_lane_mesh_from_boundary_arrays(left_boundary_array, right_boundary_array) - - -def _create_lane_mesh_from_boundary_arrays( - left_boundary_array: npt.NDArray[np.float64], right_boundary_array: npt.NDArray[np.float64] -) -> trimesh.Trimesh: - - # Ensure both polylines have the same number of points - if left_boundary_array.shape[0] != right_boundary_array.shape[0]: - raise ValueError("Both polylines must have the same number of points") - - n_points = left_boundary_array.shape[0] - vertices = np.vstack([left_boundary_array, right_boundary_array]) - - faces = [] - for i in range(n_points - 1): - faces.append([i, i + n_points, i + 1]) - faces.append([i + 1, i + n_points, i + n_points + 1]) - - faces = np.array(faces) - mesh = trimesh.Trimesh(vertices=vertices, faces=faces) - return mesh diff --git a/src/py123d/datatypes/maps/map_datatypes.py b/src/py123d/datatypes/maps/map_datatypes.py deleted file mode 100644 index dc9e4820..00000000 --- a/src/py123d/datatypes/maps/map_datatypes.py +++ /dev/null @@ -1,75 +0,0 @@ -from __future__ import annotations - -from py123d.common.utils.enums import SerialIntEnum - -# TODO: Add stop pads or stop lines. -# - Add type for stop zones. -# - Add type for carparks, e.g. outline, driveway (Waymo), or other types. -# - Check if intersections should have types. -# - Use consistent naming conventions unknown, undefined, none, etc. - - -class MapLayer(SerialIntEnum): - """ - Enum for AbstractMapSurface. - """ - - LANE = 0 - LANE_GROUP = 1 - INTERSECTION = 2 - CROSSWALK = 3 - WALKWAY = 4 - CARPARK = 5 - GENERIC_DRIVABLE = 6 - STOP_LINE = 7 - ROAD_EDGE = 8 - ROAD_LINE = 9 - - -class LaneType(SerialIntEnum): - """ - Enum for LaneType. - NOTE: We use the lane types from Waymo. - https://github.com/waymo-research/waymo-open-dataset/blob/99a4cb3ff07e2fe06c2ce73da001f850f628e45a/src/waymo_open_dataset/protos/map.proto#L147 - """ - - UNDEFINED = 0 - FREEWAY = 1 - SURFACE_STREET = 2 - BIKE_LANE = 3 - - -class RoadEdgeType(SerialIntEnum): - """ - Enum for RoadEdgeType. - NOTE: We use the road line types from Waymo. - https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L188 - """ - - UNKNOWN = 0 - ROAD_EDGE_BOUNDARY = 1 - ROAD_EDGE_MEDIAN = 2 - - -class RoadLineType(SerialIntEnum): - """ - Enum for RoadLineType. - TODO: Use the Argoverse 2 road line types. - https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L208 - """ - - NONE = 0 - UNKNOWN = 1 - DASH_SOLID_YELLOW = 2 - DASH_SOLID_WHITE = 3 - DASHED_WHITE = 4 - DASHED_YELLOW = 5 - DOUBLE_SOLID_YELLOW = 6 - DOUBLE_SOLID_WHITE = 7 - DOUBLE_DASH_YELLOW = 8 - DOUBLE_DASH_WHITE = 9 - SOLID_YELLOW = 10 - SOLID_WHITE = 11 - SOLID_DASH_WHITE = 12 - SOLID_DASH_YELLOW = 13 - SOLID_BLUE = 14 diff --git a/src/py123d/datatypes/maps/map_metadata.py b/src/py123d/datatypes/maps/map_metadata.py deleted file mode 100644 index 14fd13c8..00000000 --- a/src/py123d/datatypes/maps/map_metadata.py +++ /dev/null @@ -1,27 +0,0 @@ -from __future__ import annotations - -from dataclasses import asdict, dataclass -from typing import Any, Dict, Optional - -import py123d - -# TODO: Refactor the usage of the map map metadata in this repo. - - -@dataclass -class MapMetadata: - """Class to hold metadata information about a map.""" - - dataset: str - split: Optional[str] # None, if map is not per log - log_name: Optional[str] # None, if map is per log - location: str - map_has_z: bool - map_is_local: bool # True, if map is per log - version: str = str(py123d.__version__) - - def to_dict(self) -> dict: - return asdict(self) - - def from_dict(data_dict: Dict[str, Any]) -> MapMetadata: - return MapMetadata(**data_dict) diff --git a/src/py123d/datatypes/metadata/__init__.py b/src/py123d/datatypes/metadata/__init__.py new file mode 100644 index 00000000..89685f72 --- /dev/null +++ b/src/py123d/datatypes/metadata/__init__.py @@ -0,0 +1,2 @@ +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.datatypes.metadata.log_metadata import LogMetadata diff --git a/src/py123d/datatypes/metadata/log_metadata.py b/src/py123d/datatypes/metadata/log_metadata.py new file mode 100644 index 00000000..d78b5317 --- /dev/null +++ b/src/py123d/datatypes/metadata/log_metadata.py @@ -0,0 +1,225 @@ +from __future__ import annotations + +from typing import Dict, Optional, Type + +import py123d +from py123d.conversion.registry.box_detection_label_registry import BOX_DETECTION_LABEL_REGISTRY, BoxDetectionLabel +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICameraMetadata, FisheyeMEICameraType +from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType +from py123d.datatypes.sensors.pinhole_camera import PinholeCameraMetadata, PinholeCameraType +from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters + + +class LogMetadata: + """Class to hold metadata information about a log.""" + + __slots__ = ( + "_dataset", + "_split", + "_log_name", + "_location", + "_timestep_seconds", + "_vehicle_parameters", + "_box_detection_label_class", + "_pinhole_camera_metadata", + "_fisheye_mei_camera_metadata", + "_lidar_metadata", + "_map_metadata", + "_version", + ) + + def __init__( + self, + dataset: str, + split: str, + log_name: str, + location: str, + timestep_seconds: float, + vehicle_parameters: Optional[VehicleParameters] = None, + box_detection_label_class: Optional[Type[BoxDetectionLabel]] = None, + pinhole_camera_metadata: Optional[Dict[PinholeCameraType, PinholeCameraMetadata]] = {}, + fisheye_mei_camera_metadata: Optional[Dict[FisheyeMEICameraType, FisheyeMEICameraMetadata]] = {}, + lidar_metadata: Optional[Dict[LiDARType, LiDARMetadata]] = {}, + map_metadata: Optional[MapMetadata] = None, + version: str = str(py123d.__version__), + ): + """Create a :class:`LogMetadata` instance from a dictionary. + + :param dataset: The dataset name in lowercase. + :param split: Data split name, typically ``{dataset_name}_{train/val/test}``. + :param log_name: Name of the log file. + :param location: Location of the log data. + :param timestep_seconds: The time interval between consecutive frames in seconds. + :param vehicle_parameters: The :class:`~py123d.datatypes.vehicle_state.VehicleParameters` + of the ego vehicle, if available. + :param box_detection_label_class: The box detection label class specific to the dataset, if available. + :param pinhole_camera_metadata: Dictionary of :class:`~py123d.datatypes.sensors.PinholeCameraType` + to :class:`~py123d.datatypes.sensors.PinholeCameraMetadata`, defaults to {} + :param fisheye_mei_camera_metadata: Dictionary of :class:`~py123d.datatypes.sensors.FisheyeMEICameraType` + to :class:`~py123d.datatypes.sensors.FisheyeMEICameraMetadata`, defaults to {} + :param lidar_metadata: Dictionary of :class:`~py123d.datatypes.sensors.LiDARType` + to :class:`~py123d.datatypes.sensors.LiDARMetadata`, defaults to {} + :param map_metadata: The :class:`~py123d.datatypes.metadata.MapMetadata` for the log, if available, defaults to None + :param version: The version of the log metadata, defaults to str(py123d.__version__) + """ + self._dataset = dataset + self._split = split + self._log_name = log_name + self._location = location + self._timestep_seconds = timestep_seconds + self._vehicle_parameters = vehicle_parameters + self._box_detection_label_class = box_detection_label_class + self._pinhole_camera_metadata = pinhole_camera_metadata + self._fisheye_mei_camera_metadata = fisheye_mei_camera_metadata + self._lidar_metadata = lidar_metadata + self._map_metadata = map_metadata + self._version = version + + @property + def dataset(self) -> str: + """The dataset name in lowercase.""" + return self._dataset + + @property + def split(self) -> str: + """Data split name, typically ``{dataset_name}_{train/val/test}``.""" + return self._split + + @property + def log_name(self) -> str: + """Name of the log file.""" + return self._log_name + + @property + def location(self) -> str: + """Location of the log data.""" + return self._location + + @property + def timestep_seconds(self) -> float: + """The time interval between consecutive frames in seconds.""" + return self._timestep_seconds + + @property + def vehicle_parameters(self) -> Optional[VehicleParameters]: + """The :class:`~py123d.datatypes.vehicle_state.VehicleParameters` of the ego vehicle, if available.""" + return self._vehicle_parameters + + @property + def box_detection_label_class(self) -> Optional[Type[BoxDetectionLabel]]: + """The box detection label class specific to the dataset, if available.""" + return self._box_detection_label_class + + @property + def pinhole_camera_metadata(self) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + """Dictionary of :class:`~py123d.datatypes.sensors.PinholeCameraType` + to :class:`~py123d.datatypes.sensors.PinholeCameraMetadata`. + """ + return self._pinhole_camera_metadata + + @property + def fisheye_mei_camera_metadata(self) -> Dict[FisheyeMEICameraType, FisheyeMEICameraMetadata]: + """Dictionary of :class:`~py123d.datatypes.sensors.FisheyeMEICameraType` + to :class:`~py123d.datatypes.sensors.FisheyeMEICameraMetadata`. + """ + return self._fisheye_mei_camera_metadata + + @property + def lidar_metadata(self) -> Dict[LiDARType, LiDARMetadata]: + """Dictionary of :class:`~py123d.datatypes.sensors.LiDARType` + to :class:`~py123d.datatypes.sensors.LiDARMetadata`. + """ + return self._lidar_metadata + + @property + def map_metadata(self) -> Optional[MapMetadata]: + """The :class:`~py123d.datatypes.metadata.MapMetadata` associated with the log, if available.""" + return self._map_metadata + + @property + def version(self) -> str: + """Version of the py123d library used to create this log metadata (not used currently).""" + return self._version + + @classmethod + def from_dict(cls, data_dict: Dict) -> LogMetadata: + """Create a :class:`LogMetadata` instance from a Python dictionary. + + :param data_dict: Dictionary containing log metadata. + :raises ValueError: If the dictionary is missing required fields. + :return: A :class:`LogMetadata` instance. + """ + + # Ego Vehicle Parameters + if data_dict["vehicle_parameters"] is not None: + data_dict["vehicle_parameters"] = VehicleParameters.from_dict(data_dict["vehicle_parameters"]) + + # Box detection label class specific to the dataset + if data_dict["box_detection_label_class"] in BOX_DETECTION_LABEL_REGISTRY: + data_dict["box_detection_label_class"] = BOX_DETECTION_LABEL_REGISTRY[ + data_dict["box_detection_label_class"] + ] + elif data_dict["box_detection_label_class"] is None: + data_dict["box_detection_label_class"] = None + else: + raise ValueError(f"Unknown box detection label class: {data_dict['box_detection_label_class']}") + + # Pinhole Camera Metadata + pinhole_camera_metadata = {} + for key, value in data_dict.get("pinhole_camera_metadata", {}).items(): + pinhole_camera_metadata[PinholeCameraType.deserialize(key)] = PinholeCameraMetadata.from_dict(value) + data_dict["pinhole_camera_metadata"] = pinhole_camera_metadata + + # Fisheye MEI Camera Metadata + fisheye_mei_camera_metadata = {} + for key, value in data_dict.get("fisheye_mei_camera_metadata", {}).items(): + fisheye_mei_camera_metadata[FisheyeMEICameraType.deserialize(key)] = FisheyeMEICameraMetadata.from_dict( + value + ) + data_dict["fisheye_mei_camera_metadata"] = fisheye_mei_camera_metadata + + # LiDAR Metadata + data_dict["lidar_metadata"] = { + LiDARType.deserialize(key): LiDARMetadata.from_dict(value) + for key, value in data_dict.get("lidar_metadata", {}).items() + } + + # Map Metadata + if data_dict["map_metadata"] is not None: + data_dict["map_metadata"] = MapMetadata.from_dict(data_dict["map_metadata"]) + + return LogMetadata(**data_dict) + + def to_dict(self) -> Dict: + """Convert the :class:`LogMetadata` instance to a Python dictionary. + + :return: A dictionary representation of the log metadata. + """ + data_dict = {slot.lstrip("_"): getattr(self, slot) for slot in self.__slots__} + + # Override complex types with their dictionary representations + data_dict["vehicle_parameters"] = self.vehicle_parameters.to_dict() if self.vehicle_parameters else None + if self.box_detection_label_class is not None: + data_dict["box_detection_label_class"] = self.box_detection_label_class.__name__ + data_dict["pinhole_camera_metadata"] = { + key.serialize(): value.to_dict() for key, value in self.pinhole_camera_metadata.items() + } + data_dict["fisheye_mei_camera_metadata"] = { + key.serialize(): value.to_dict() for key, value in self.fisheye_mei_camera_metadata.items() + } + data_dict["lidar_metadata"] = {key.serialize(): value.to_dict() for key, value in self.lidar_metadata.items()} + data_dict["map_metadata"] = self.map_metadata.to_dict() if self.map_metadata else None + return data_dict + + def __repr__(self) -> str: + return ( + f"LogMetadata(dataset={self.dataset}, split={self.split}, log_name={self.log_name}, " + f"location={self.location}, timestep_seconds={self.timestep_seconds}, " + f"vehicle_parameters={self.vehicle_parameters}, " + f"box_detection_label_class={self.box_detection_label_class}, " + f"pinhole_camera_metadata={self.pinhole_camera_metadata}, " + f"fisheye_mei_camera_metadata={self.fisheye_mei_camera_metadata}, " + f"lidar_metadata={self.lidar_metadata}, map_metadata={self.map_metadata}, " + f"version={self.version})" + ) diff --git a/src/py123d/datatypes/metadata/map_metadata.py b/src/py123d/datatypes/metadata/map_metadata.py new file mode 100644 index 00000000..9cd6511a --- /dev/null +++ b/src/py123d/datatypes/metadata/map_metadata.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from typing import Any, Dict, Optional + +import py123d + + +class MapMetadata: + """Class to hold metadata information about a map.""" + + __slots__ = ("_dataset", "_split", "_log_name", "_location", "_map_has_z", "_map_is_local", "_version") + + def __init__( + self, + dataset: str, + location: str, + map_has_z: bool, + map_is_local: bool, + split: Optional[str] = None, + log_name: Optional[str] = None, + version: str = str(py123d.__version__), + ): + """Initialize a MapMetadata instance. + + :param dataset: The dataset name in lowercase. + :param location: The location of the map data. + :param map_has_z: Indicates if the map includes Z (elevation) data. + :param map_is_local: Indicates if the map is local (map for each log) or + global (map for multiple logs in dataset). + :param split: Data split name, typically ``{dataset_name}_{train/val/test}``, defaults to None + :param log_name: Name of the log file, defaults to None + :param version: Version of the py123d library used to create this map metadata, + defaults to str(py123d.__version__) + """ + self._dataset = dataset + self._split = split + self._log_name = log_name + self._location = location + self._map_has_z = map_has_z + self._map_is_local = map_is_local + self._version = version + + @property + def dataset(self) -> str: + """The dataset name in lowercase.""" + return self._dataset + + @property + def split(self) -> Optional[str]: + """Data split name, typically ``{dataset_name}_{train/val/test}``.""" + return self._split + + @property + def log_name(self) -> Optional[str]: + """Name of the log file.""" + return self._log_name + + @property + def location(self) -> str: + """Location of the map data.""" + return self._location + + @property + def map_has_z(self) -> bool: + """Indicates if the map includes Z (elevation) data.""" + return self._map_has_z + + @property + def map_is_local(self) -> bool: + """Indicates if the map is local (map for each log) or global (map for multiple logs in dataset).""" + return self._map_is_local + + @property + def version(self) -> str: + """Version of the py123d library used to create this map metadata.""" + return self._version + + @classmethod + def from_dict(cls, data_dict: Dict[str, Any]) -> MapMetadata: + """Create a MapMetadata instance from a dictionary. + + :param data_dict: A dictionary representation of a MapMetadata instance. + :return: A MapMetadata instance. + """ + return MapMetadata(**data_dict) + + def to_dict(self) -> Dict[str, Any]: + """Convert the MapMetadata instance to a dictionary. + + :return: A dictionary representation of the MapMetadata instance. + """ + return {slot.lstrip("_"): getattr(self, slot) for slot in self.__slots__} + + def __repr__(self) -> str: + return ( + f"MapMetadata(" + f"dataset={self.dataset!r}, " + f"split={self.split!r}, " + f"log_name={self.log_name!r}, " + f"location={self.location!r}, " + f"map_has_z={self.map_has_z}, " + f"map_is_local={self.map_is_local}, " + f"version={self.version!r}" + f")" + ) diff --git a/src/py123d/datatypes/scene/abstract_scene.py b/src/py123d/datatypes/scene/abstract_scene.py deleted file mode 100644 index 33611539..00000000 --- a/src/py123d/datatypes/scene/abstract_scene.py +++ /dev/null @@ -1,116 +0,0 @@ -from __future__ import annotations - -import abc -from typing import List, Optional - -from py123d.datatypes.detections.box_detections import BoxDetectionWrapper -from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper -from py123d.datatypes.maps.abstract_map import AbstractMap -from py123d.datatypes.scene.scene_metadata import LogMetadata, SceneExtractionMetadata -from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICamera, FisheyeMEICameraType -from py123d.datatypes.sensors.lidar import LiDAR, LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeCameraType -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters - - -class AbstractScene(abc.ABC): - - #################################################################################################################### - # Abstract Methods, to be implemented by subclasses - #################################################################################################################### - - @abc.abstractmethod - def get_log_metadata(self) -> LogMetadata: - raise NotImplementedError - - @abc.abstractmethod - def get_scene_extraction_metadata(self) -> SceneExtractionMetadata: - raise NotImplementedError - - @abc.abstractmethod - def get_map_api(self) -> Optional[AbstractMap]: - raise NotImplementedError - - @abc.abstractmethod - def get_timepoint_at_iteration(self, iteration: int) -> TimePoint: - raise NotImplementedError - - @abc.abstractmethod - def get_ego_state_at_iteration(self, iteration: int) -> Optional[EgoStateSE3]: - raise NotImplementedError - - @abc.abstractmethod - def get_box_detections_at_iteration(self, iteration: int) -> Optional[BoxDetectionWrapper]: - raise NotImplementedError - - @abc.abstractmethod - def get_traffic_light_detections_at_iteration(self, iteration: int) -> Optional[TrafficLightDetectionWrapper]: - raise NotImplementedError - - @abc.abstractmethod - def get_route_lane_group_ids(self, iteration: int) -> Optional[List[int]]: - raise NotImplementedError - - @abc.abstractmethod - def get_pinhole_camera_at_iteration( - self, iteration: int, camera_type: PinholeCameraType - ) -> Optional[PinholeCamera]: - raise NotImplementedError - - @abc.abstractmethod - def get_fisheye_mei_camera_at_iteration( - self, iteration: int, camera_type: FisheyeMEICameraType - ) -> Optional[FisheyeMEICamera]: - raise NotImplementedError - - @abc.abstractmethod - def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> Optional[LiDAR]: - raise NotImplementedError - - #################################################################################################################### - # Syntactic Sugar / Properties, for easier access to common attributes - #################################################################################################################### - - # 1. Log Metadata properties - @property - def log_metadata(self) -> LogMetadata: - return self.get_log_metadata() - - @property - def log_name(self) -> str: - return self.log_metadata.log_name - - @property - def vehicle_parameters(self) -> VehicleParameters: - return self.log_metadata.vehicle_parameters - - @property - def available_pinhole_camera_types(self) -> List[PinholeCameraType]: - return list(self.log_metadata.pinhole_camera_metadata.keys()) - - @property - def available_fisheye_mei_camera_types(self) -> List[FisheyeMEICameraType]: - return list(self.log_metadata.fisheye_mei_camera_metadata.keys()) - - @property - def available_lidar_types(self) -> List[LiDARType]: - return list(self.log_metadata.lidar_metadata.keys()) - - # 2. Scene Extraction Metadata properties - @property - def scene_extraction_metadata(self) -> SceneExtractionMetadata: - return self.get_scene_extraction_metadata() - - @property - def uuid(self) -> str: - return self.scene_extraction_metadata.initial_uuid - - @property - def number_of_iterations(self) -> int: - return self.scene_extraction_metadata.number_of_iterations - - @property - def number_of_history_iterations(self) -> int: - return self.scene_extraction_metadata.number_of_history_iterations diff --git a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py b/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py deleted file mode 100644 index 71c08154..00000000 --- a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py +++ /dev/null @@ -1,230 +0,0 @@ -from pathlib import Path -from typing import Dict, List, Optional, Union - -import cv2 -import numpy as np -import numpy.typing as npt -import pyarrow as pa -from omegaconf import DictConfig - -from py123d.conversion.registry.lidar_index_registry import DefaultLiDARIndex -from py123d.conversion.sensor_io.camera.jpeg_camera_io import decode_image_from_jpeg_binary -from py123d.conversion.sensor_io.camera.mp4_camera_io import get_mp4_reader_from_path -from py123d.conversion.sensor_io.lidar.draco_lidar_io import load_lidar_from_draco_binary -from py123d.conversion.sensor_io.lidar.file_lidar_io import load_lidar_pcs_from_file -from py123d.conversion.sensor_io.lidar.laz_lidar_io import load_lidar_from_laz_binary -from py123d.datatypes.detections.box_detections import ( - BoxDetection, - BoxDetectionMetadata, - BoxDetectionSE3, - BoxDetectionWrapper, -) -from py123d.datatypes.detections.traffic_light_detections import ( - TrafficLightDetection, - TrafficLightDetectionWrapper, - TrafficLightStatus, -) -from py123d.datatypes.scene.scene_metadata import LogMetadata -from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICamera, FisheyeMEICameraType -from py123d.datatypes.sensors.lidar import LiDAR, LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeCameraType -from py123d.datatypes.time.time_point import TimePoint -from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters -from py123d.geometry import BoundingBoxSE3, StateSE3, Vector3D -from py123d.script.utils.dataset_path_utils import get_dataset_paths - -DATASET_PATHS: DictConfig = get_dataset_paths() -DATASET_SENSOR_ROOT: Dict[str, Path] = { - "av2-sensor": DATASET_PATHS.av2_sensor_data_root, - "nuplan": DATASET_PATHS.nuplan_sensor_root, - "nuscenes": DATASET_PATHS.nuscenes_data_root, - "wopd": DATASET_PATHS.wopd_data_root, - "pandaset": DATASET_PATHS.pandaset_data_root, - "kitti360": DATASET_PATHS.kitti360_data_root, -} - - -def get_timepoint_from_arrow_table(arrow_table: pa.Table, index: int) -> TimePoint: - return TimePoint.from_us(arrow_table["timestamp"][index].as_py()) - - -def get_ego_vehicle_state_from_arrow_table( - arrow_table: pa.Table, index: int, vehicle_parameters: VehicleParameters -) -> EgoStateSE3: - timepoint = get_timepoint_from_arrow_table(arrow_table, index) - return EgoStateSE3.from_array( - array=pa.array(arrow_table["ego_state"][index]).to_numpy(), - vehicle_parameters=vehicle_parameters, - timepoint=timepoint, - ) - - -def get_box_detections_from_arrow_table( - arrow_table: pa.Table, - index: int, - log_metadata: LogMetadata, -) -> BoxDetectionWrapper: - timepoint = get_timepoint_from_arrow_table(arrow_table, index) - box_detections: List[BoxDetection] = [] - box_detection_label_class = log_metadata.box_detection_label_class - - for detection_state, detection_velocity, detection_token, detection_label in zip( - arrow_table["box_detection_state"][index].as_py(), - arrow_table["box_detection_velocity"][index].as_py(), - arrow_table["box_detection_token"][index].as_py(), - arrow_table["box_detection_label"][index].as_py(), - ): - box_detection = BoxDetectionSE3( - metadata=BoxDetectionMetadata( - label=box_detection_label_class(detection_label), - timepoint=timepoint, - track_token=detection_token, - confidence=None, - ), - bounding_box_se3=BoundingBoxSE3.from_array(np.array(detection_state)), - velocity=Vector3D.from_array(np.array(detection_velocity)) if detection_velocity else None, - ) - box_detections.append(box_detection) - return BoxDetectionWrapper(box_detections=box_detections) - - -def get_traffic_light_detections_from_arrow_table(arrow_table: pa.Table, index: int) -> TrafficLightDetectionWrapper: - timepoint = get_timepoint_from_arrow_table(arrow_table, index) - traffic_light_detections: Optional[List[TrafficLightDetection]] = None - - if "traffic_light_ids" in arrow_table.schema.names and "traffic_light_types" in arrow_table.schema.names: - traffic_light_detections: List[TrafficLightDetection] = [] - for lane_id, status in zip( - arrow_table["traffic_light_ids"][index].as_py(), - arrow_table["traffic_light_types"][index].as_py(), - ): - traffic_light_detection = TrafficLightDetection( - timepoint=timepoint, - lane_id=lane_id, - status=TrafficLightStatus(status), - ) - traffic_light_detections.append(traffic_light_detection) - - traffic_light_detections = TrafficLightDetectionWrapper(traffic_light_detections=traffic_light_detections) - - return traffic_light_detections - - -def get_camera_from_arrow_table( - arrow_table: pa.Table, - index: int, - camera_type: Union[PinholeCameraType, FisheyeMEICameraType], - log_metadata: LogMetadata, -) -> Union[PinholeCamera, FisheyeMEICamera]: - - camera_name = camera_type.serialize() - table_data = arrow_table[f"{camera_name}_data"][index].as_py() - extrinsic_values = arrow_table[f"{camera_name}_extrinsic"][index].as_py() - extrinsic = StateSE3.from_list(extrinsic_values) if extrinsic_values is not None else None - - if table_data is None or extrinsic is None: - return None - - image: Optional[npt.NDArray[np.uint8]] = None - - if isinstance(table_data, str): - sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] - assert sensor_root is not None, f"Dataset path for sensor loading not found for dataset: {log_metadata.dataset}" - full_image_path = Path(sensor_root) / table_data - assert full_image_path.exists(), f"Camera file not found: {full_image_path}" - - image = cv2.imread(str(full_image_path), cv2.IMREAD_COLOR) - image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) - - elif isinstance(table_data, bytes): - image = decode_image_from_jpeg_binary(table_data) - elif isinstance(table_data, int): - image = _unoptimized_demo_mp4_read(log_metadata, camera_name, table_data) - else: - raise NotImplementedError( - f"Only string file paths, bytes, or int frame indices are supported for camera data, got {type(table_data)}" - ) - - if camera_name.startswith("fcam"): - camera_metadata = log_metadata.fisheye_mei_camera_metadata[camera_type] - return FisheyeMEICamera( - metadata=camera_metadata, - image=image, - extrinsic=extrinsic, - ) - else: - camera_metadata = log_metadata.pinhole_camera_metadata[camera_type] - return PinholeCamera( - metadata=camera_metadata, - image=image, - extrinsic=extrinsic, - ) - - -def get_lidar_from_arrow_table( - arrow_table: pa.Table, - index: int, - lidar_type: LiDARType, - log_metadata: LogMetadata, -) -> LiDAR: - - lidar: Optional[LiDAR] = None - lidar_column_name = f"{lidar_type.serialize()}_data" - lidar_column_name = ( - f"{LiDARType.LIDAR_MERGED.serialize()}_data" - if lidar_column_name not in arrow_table.schema.names - else lidar_column_name - ) - if lidar_column_name in arrow_table.schema.names: - - lidar_data = arrow_table[lidar_column_name][index].as_py() - if isinstance(lidar_data, str): - lidar_pc_dict = load_lidar_pcs_from_file(relative_path=lidar_data, log_metadata=log_metadata, index=index) - if lidar_type == LiDARType.LIDAR_MERGED: - # Merge all available LiDAR point clouds into one - merged_pc = np.vstack(list(lidar_pc_dict.values())) - lidar = LiDAR( - metadata=LiDARMetadata( - lidar_type=LiDARType.LIDAR_MERGED, - lidar_index=DefaultLiDARIndex, - extrinsic=None, - ), - point_cloud=merged_pc, - ) - elif lidar_type in lidar_pc_dict: - lidar = LiDAR( - metadata=log_metadata.lidar_metadata[lidar_type], - point_cloud=lidar_pc_dict[lidar_type], - ) - elif isinstance(lidar_data, bytes): - lidar_metadata = log_metadata.lidar_metadata[lidar_type] - if lidar_data.startswith(b"DRACO"): - # NOTE: DRACO only allows XYZ compression, so we need to override the lidar index here. - lidar_metadata.lidar_index = DefaultLiDARIndex - - lidar = load_lidar_from_draco_binary(lidar_data, lidar_metadata) - elif lidar_data.startswith(b"LASF"): - - lidar = load_lidar_from_laz_binary(lidar_data, lidar_metadata) - elif lidar_data is None: - lidar = None - else: - raise NotImplementedError( - f"Only string file paths or bytes for LiDAR data are supported, got {type(lidar_data)}" - ) - - return lidar - - -def _unoptimized_demo_mp4_read(log_metadata: LogMetadata, camera_name: str, frame_index: int) -> Optional[np.ndarray]: - """A quick and dirty MP4 reader for testing purposes only. Not optimized for performance.""" - image: Optional[npt.NDArray[np.uint8]] = None - - py123d_sensor_root = Path(DATASET_PATHS.py123d_sensors_root) - mp4_path = py123d_sensor_root / log_metadata.split / log_metadata.log_name / f"{camera_name}.mp4" - if mp4_path.exists(): - reader = get_mp4_reader_from_path(str(mp4_path)) - image = reader.get_frame(frame_index) - - return image diff --git a/src/py123d/datatypes/scene/arrow/utils/arrow_metadata_utils.py b/src/py123d/datatypes/scene/arrow/utils/arrow_metadata_utils.py deleted file mode 100644 index 3f264b62..00000000 --- a/src/py123d/datatypes/scene/arrow/utils/arrow_metadata_utils.py +++ /dev/null @@ -1,21 +0,0 @@ -import json -from functools import lru_cache -from pathlib import Path -from typing import Union - -import pyarrow as pa - -from py123d.common.utils.arrow_helper import get_lru_cached_arrow_table -from py123d.datatypes.scene.scene_metadata import LogMetadata - - -@lru_cache(maxsize=10000) -def get_log_metadata_from_arrow(arrow_file_path: Union[Path, str]) -> LogMetadata: - table = get_lru_cached_arrow_table(arrow_file_path) - log_metadata = LogMetadata.from_dict(json.loads(table.schema.metadata[b"log_metadata"].decode())) - return log_metadata - - -def add_log_metadata_to_arrow_schema(schema: pa.schema, log_metadata: LogMetadata) -> pa.schema: - schema = schema.with_metadata({"log_metadata": json.dumps(log_metadata.to_dict())}) - return schema diff --git a/src/py123d/datatypes/scene/scene_metadata.py b/src/py123d/datatypes/scene/scene_metadata.py deleted file mode 100644 index 2bc271f1..00000000 --- a/src/py123d/datatypes/scene/scene_metadata.py +++ /dev/null @@ -1,111 +0,0 @@ -from __future__ import annotations - -from dataclasses import asdict, dataclass, field -from typing import Dict, Optional, Type - -import py123d -from py123d.conversion.registry.box_detection_label_registry import BOX_DETECTION_LABEL_REGISTRY, BoxDetectionLabel -from py123d.datatypes.maps.map_metadata import MapMetadata -from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICameraMetadata, FisheyeMEICameraType -from py123d.datatypes.sensors.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCameraMetadata, PinholeCameraType -from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters - - -@dataclass -class LogMetadata: - - dataset: str - split: str - log_name: str - location: str - timestep_seconds: float - - vehicle_parameters: Optional[VehicleParameters] = None - box_detection_label_class: Optional[Type[BoxDetectionLabel]] = None - pinhole_camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = field(default_factory=dict) - fisheye_mei_camera_metadata: Dict[FisheyeMEICameraType, FisheyeMEICameraMetadata] = field(default_factory=dict) - lidar_metadata: Dict[LiDARType, LiDARMetadata] = field(default_factory=dict) - - map_metadata: Optional[MapMetadata] = None - version: str = str(py123d.__version__) - - @classmethod - def from_dict(cls, data_dict: Dict) -> LogMetadata: - - # Ego Vehicle Parameters - if data_dict["vehicle_parameters"] is not None: - data_dict["vehicle_parameters"] = VehicleParameters.from_dict(data_dict["vehicle_parameters"]) - - # Box detection label class specific to the dataset - if data_dict["box_detection_label_class"] in BOX_DETECTION_LABEL_REGISTRY: - data_dict["box_detection_label_class"] = BOX_DETECTION_LABEL_REGISTRY[ - data_dict["box_detection_label_class"] - ] - elif data_dict["box_detection_label_class"] is None: - data_dict["box_detection_label_class"] = None - else: - raise ValueError(f"Unknown box detection label class: {data_dict['box_detection_label_class']}") - - # Pinhole Camera Metadata - pinhole_camera_metadata = {} - for key, value in data_dict.get("pinhole_camera_metadata", {}).items(): - pinhole_camera_metadata[PinholeCameraType.deserialize(key)] = PinholeCameraMetadata.from_dict(value) - data_dict["pinhole_camera_metadata"] = pinhole_camera_metadata - - # Fisheye MEI Camera Metadata - fisheye_mei_camera_metadata = {} - for key, value in data_dict.get("fisheye_mei_camera_metadata", {}).items(): - fisheye_mei_camera_metadata[FisheyeMEICameraType.deserialize(key)] = FisheyeMEICameraMetadata.from_dict( - value - ) - data_dict["fisheye_mei_camera_metadata"] = fisheye_mei_camera_metadata - - # LiDAR Metadata - data_dict["lidar_metadata"] = { - LiDARType.deserialize(key): LiDARMetadata.from_dict(value) - for key, value in data_dict.get("lidar_metadata", {}).items() - } - - # Map Metadata - if data_dict["map_metadata"] is not None: - data_dict["map_metadata"] = MapMetadata.from_dict(data_dict["map_metadata"]) - - return LogMetadata(**data_dict) - - def to_dict(self) -> Dict: - data_dict = asdict(self) - data_dict["vehicle_parameters"] = self.vehicle_parameters.to_dict() if self.vehicle_parameters else None - if self.box_detection_label_class is not None: - data_dict["box_detection_label_class"] = self.box_detection_label_class.__name__ - data_dict["pinhole_camera_metadata"] = { - key.serialize(): value.to_dict() for key, value in self.pinhole_camera_metadata.items() - } - data_dict["fisheye_mei_camera_metadata"] = { - key.serialize(): value.to_dict() for key, value in self.fisheye_mei_camera_metadata.items() - } - data_dict["lidar_metadata"] = {key.serialize(): value.to_dict() for key, value in self.lidar_metadata.items()} - data_dict["map_metadata"] = self.map_metadata.to_dict() if self.map_metadata else None - return data_dict - - -@dataclass(frozen=True) -class SceneExtractionMetadata: - - initial_uuid: str - initial_idx: int - duration_s: float - history_s: float - iteration_duration_s: float - - @property - def number_of_iterations(self) -> int: - return round(self.duration_s / self.iteration_duration_s) - - @property - def number_of_history_iterations(self) -> int: - return round(self.history_s / self.iteration_duration_s) - - @property - def end_idx(self) -> int: - return self.initial_idx + self.number_of_iterations diff --git a/src/py123d/datatypes/sensors/__init__.py b/src/py123d/datatypes/sensors/__init__.py index 54cd70a1..6ce7fd0c 100644 --- a/src/py123d/datatypes/sensors/__init__.py +++ b/src/py123d/datatypes/sensors/__init__.py @@ -11,8 +11,9 @@ FisheyeMEICameraType, FisheyeMEICamera, FisheyeMEIDistortionIndex, - FisheyeMEIProjectionIndex, + FisheyeMEIDistortion, FisheyeMEIProjection, + FisheyeMEIProjectionIndex, FisheyeMEICameraMetadata, ) from py123d.datatypes.sensors.lidar import ( diff --git a/src/py123d/datatypes/sensors/fisheye_mei_camera.py b/src/py123d/datatypes/sensors/fisheye_mei_camera.py index 4a98afad..32e70ced 100644 --- a/src/py123d/datatypes/sensors/fisheye_mei_camera.py +++ b/src/py123d/datatypes/sensors/fisheye_mei_camera.py @@ -1,46 +1,92 @@ from __future__ import annotations -from dataclasses import asdict, dataclass +from dataclasses import dataclass +from enum import IntEnum from typing import Any, Dict, Optional import numpy as np import numpy.typing as npt -from zmq import IntEnum from py123d.common.utils.enums import SerialIntEnum -from py123d.common.utils.mixin import ArrayMixin -from py123d.geometry.se import StateSE3 +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr +from py123d.geometry.pose import PoseSE3 class FisheyeMEICameraType(SerialIntEnum): - """ - Enum for fisheye cameras in d123. - """ + """Enumeration of fisheye MEI camera types in multi-sensor setups.""" FCAM_L = 0 + """Left-facing fisheye MEI camera.""" + FCAM_R = 1 + """Right-facing fisheye MEI camera.""" -@dataclass class FisheyeMEICamera: + """Fisheye MEI camera data structure.""" + + __slots__ = ("_metadata", "_image", "_extrinsic") + + def __init__( + self, + metadata: FisheyeMEICameraMetadata, + image: npt.NDArray[np.uint8], + extrinsic: PoseSE3, + ) -> None: + """Initialize a Fisheye MEI camera. + + :param metadata: Metadata for the camera. + :param image: Image captured by the camera. + :param extrinsic: Extrinsic pose of the camera. + """ + self._metadata = metadata + self._image = image + self._extrinsic = extrinsic + + @property + def metadata(self) -> FisheyeMEICameraMetadata: + """The :class:`FisheyeMEICameraMetadata` object for the camera.""" + return self._metadata - metadata: FisheyeMEICameraMetadata - image: npt.NDArray[np.uint8] - extrinsic: StateSE3 + @property + def image(self) -> npt.NDArray[np.uint8]: + """Image captured by the camera, as a NumPy array.""" + return self._image + + @property + def extrinsic(self) -> PoseSE3: + """Extrinsic :class:`~py123d.geometry.PoseSE3` of the camera.""" + return self._extrinsic class FisheyeMEIDistortionIndex(IntEnum): + """Indexing for fisheye MEI distortion parameters.""" K1 = 0 + """Radial distortion coefficient k1.""" + K2 = 1 + """Radial distortion coefficient k2.""" + P1 = 2 + """Tangential distortion coefficient p1.""" + P2 = 3 + """Tangential distortion coefficient p2.""" class FisheyeMEIDistortion(ArrayMixin): + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, k1: float, k2: float, p1: float, p2: float) -> None: + """Initialize the fisheye MEI distortion parameters. + + :param k1: Radial distortion coefficient k1. + :param k2: Radial distortion coefficient k2. + :param p1: Tangential distortion coefficient p1. + :param p2: Tangential distortion coefficient p2. + """ array = np.zeros(len(FisheyeMEIDistortionIndex), dtype=np.float64) array[FisheyeMEIDistortionIndex.K1] = k1 array[FisheyeMEIDistortionIndex.K2] = k2 @@ -50,6 +96,13 @@ def __init__(self, k1: float, k2: float, p1: float, p2: float) -> None: @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> FisheyeMEIDistortion: + """Creates a :class:`FisheyeMEIDistortion` instance from a NumPy array, + indexing according to :class:`FisheyeMEIDistortionIndex`. + + :param array: Input array containing distortion parameters. + :param copy: Whether to copy the array data, defaults to True. + :return: A new instance of :class:`FisheyeMEIDistortion`. + """ assert array.ndim == 1 assert array.shape[-1] == len(FisheyeMEIDistortionIndex) instance = object.__new__(cls) @@ -58,37 +111,64 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Fishey @property def array(self) -> npt.NDArray[np.float64]: + """Underlying NumPy array of distortion parameters, indexed by :class:`FisheyeMEIDistortionIndex`.""" return self._array @property def k1(self) -> float: + """Radial distortion coefficient.""" return self._array[FisheyeMEIDistortionIndex.K1] @property def k2(self) -> float: + """Radial distortion coefficient.""" return self._array[FisheyeMEIDistortionIndex.K2] @property def p1(self) -> float: + """Tangential distortion coefficient.""" return self._array[FisheyeMEIDistortionIndex.P1] @property def p2(self) -> float: + """Tangential distortion coefficient.""" return self._array[FisheyeMEIDistortionIndex.P2] + def __repr__(self) -> str: + """String representation of :class:`FisheyeMEIDistortion`.""" + return indexed_array_repr(self, FisheyeMEIDistortionIndex) + class FisheyeMEIProjectionIndex(IntEnum): + """Indexing for fisheye MEI projection parameters.""" GAMMA1 = 0 + """Generalized focal length gamma1.""" + GAMMA2 = 1 + """Generalized focal length gamma2.""" + U0 = 2 + """Principal point x-coordinate.""" + V0 = 3 + """Principal point y-coordinate.""" class FisheyeMEIProjection(ArrayMixin): + """Fisheye MEI projection parameters.""" + + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, gamma1: float, gamma2: float, u0: float, v0: float) -> None: + """Initialize the fisheye MEI projection parameters. + + :param gamma1: Generalized focal length gamma1. + :param gamma2: Generalized focal length gamma2. + :param u0: Principal point x-coordinate. + :param v0: Principal point y-coordinate. + """ array = np.zeros(len(FisheyeMEIProjectionIndex), dtype=np.float64) array[FisheyeMEIProjectionIndex.GAMMA1] = gamma1 array[FisheyeMEIProjectionIndex.GAMMA2] = gamma2 @@ -98,6 +178,13 @@ def __init__(self, gamma1: float, gamma2: float, u0: float, v0: float) -> None: @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> FisheyeMEIProjection: + """Intializes a :class:`FisheyeMEIProjection` from a NumPy array, + indexing according to :class:`FisheyeMEIProjectionIndex`. + + :param array: Input array containing projection parameters. + :param copy: Whether to copy the array data, defaults to True. + :return: A new instance of :class:`FisheyeMEIProjection`. + """ assert array.ndim == 1 assert array.shape[-1] == len(FisheyeMEIProjectionIndex) instance = object.__new__(cls) @@ -106,37 +193,72 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Fishey @property def array(self) -> npt.NDArray[np.float64]: + """Underlying NumPy array of projection parameters, indexed by :class:`FisheyeMEIProjectionIndex`.""" return self._array @property def gamma1(self) -> float: + """Generalized focal length gamma1.""" return self._array[FisheyeMEIProjectionIndex.GAMMA1] @property def gamma2(self) -> float: + """Generalized focal length gamma2.""" return self._array[FisheyeMEIProjectionIndex.GAMMA2] @property def u0(self) -> float: + """Principal point x-coordinate.""" return self._array[FisheyeMEIProjectionIndex.U0] @property def v0(self) -> float: + """Principal point y-coordinate.""" return self._array[FisheyeMEIProjectionIndex.V0] + def __repr__(self) -> str: + """String representation of :class:`FisheyeMEIProjection`.""" + return indexed_array_repr(self, FisheyeMEIProjectionIndex) + @dataclass class FisheyeMEICameraMetadata: - - camera_type: FisheyeMEICameraType - mirror_parameter: Optional[float] - distortion: Optional[FisheyeMEIDistortion] - projection: Optional[FisheyeMEIProjection] - width: int - height: int + """Metadata for a fisheye MEI camera.""" + + __slots__ = ("_camera_type", "_mirror_parameter", "_distortion", "_projection", "_width", "_height") + + def __init__( + self, + camera_type: FisheyeMEICameraType, + mirror_parameter: Optional[float], + distortion: Optional[FisheyeMEIDistortion], + projection: Optional[FisheyeMEIProjection], + width: int, + height: int, + ) -> None: + """Initialize the fisheye MEI camera metadata. + + :param camera_type: Type of the fisheye MEI camera. + :param mirror_parameter: Mirror parameter of the camera model. + :param distortion: Distortion parameters of the camera. + :param projection: Projection parameters of the camera. + :param width: Width of the camera image in pixels. + :param height: Height of the camera image in pixels. + """ + self._camera_type = camera_type + self._mirror_parameter = mirror_parameter + self._distortion = distortion + self._projection = projection + self._width = width + self._height = height @classmethod def from_dict(cls, data_dict: Dict[str, Any]) -> FisheyeMEICameraMetadata: + """Create a :class:`FisheyeMEICameraMetadata` instance from a dictionary. + + :param data_dict: Dictionary containing camera metadata. + :return: A new instance of :class:`FisheyeMEICameraMetadata`. + """ data_dict["camera_type"] = FisheyeMEICameraType(data_dict["camera_type"]) data_dict["distortion"] = ( FisheyeMEIDistortion.from_array(np.array(data_dict["distortion"])) @@ -150,15 +272,43 @@ def from_dict(cls, data_dict: Dict[str, Any]) -> FisheyeMEICameraMetadata: ) return FisheyeMEICameraMetadata(**data_dict) + @property + def camera_type(self) -> FisheyeMEICameraType: + """The type of the fisheye MEI camera.""" + return self._camera_type + + @property + def mirror_parameter(self) -> Optional[float]: + """The mirror parameter of the fisheye MEI camera.""" + return self._mirror_parameter + + @property + def distortion(self) -> Optional[FisheyeMEIDistortion]: + """The distortion parameters of the fisheye MEI camera, if available.""" + return self._distortion + + @property + def projection(self) -> Optional[FisheyeMEIProjection]: + """The projection parameters of the fisheye MEI camera, if available.""" + return self._projection + @property def aspect_ratio(self) -> float: - return self.width / self.height + """The aspect ratio of the fisheye MEI camera.""" + return self._width / self._height def to_dict(self) -> Dict[str, Any]: - data_dict = asdict(self) - data_dict["camera_type"] = int(self.camera_type) - data_dict["distortion"] = self.distortion.array.tolist() if self.distortion is not None else None - data_dict["projection"] = self.projection.array.tolist() if self.projection is not None else None + """Converts the :class:`FisheyeMEICameraMetadata` instance to a Python dictionary. + + :return: A dictionary representation of the camera metadata. + """ + data_dict: Dict[str, Any] = {} + data_dict["mirror_parameter"] = self._mirror_parameter + data_dict["camera_type"] = int(self._camera_type) + data_dict["distortion"] = self._distortion.array.tolist() if self._distortion is not None else None + data_dict["projection"] = self._projection.array.tolist() if self._projection is not None else None + data_dict["width"] = self._width + data_dict["height"] = self._height return data_dict def cam2image(self, points_3d: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: diff --git a/src/py123d/datatypes/sensors/lidar.py b/src/py123d/datatypes/sensors/lidar.py index 0950c77c..4167cb2d 100644 --- a/src/py123d/datatypes/sensors/lidar.py +++ b/src/py123d/datatypes/sensors/lidar.py @@ -1,99 +1,172 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Optional, Type +from typing import Any, Dict, Optional, Type import numpy as np import numpy.typing as npt from py123d.common.utils.enums import SerialIntEnum -from py123d.conversion.registry.lidar_index_registry import LIDAR_INDEX_REGISTRY, LiDARIndex -from py123d.geometry import StateSE3 +from py123d.conversion.registry import LIDAR_INDEX_REGISTRY, LiDARIndex +from py123d.geometry import PoseSE3 class LiDARType(SerialIntEnum): + """Enumeration of LiDAR sensors, in multi-sensor setups.""" LIDAR_UNKNOWN = 0 + """Unknown LiDAR type.""" + LIDAR_MERGED = 1 + """Merged LiDAR type.""" + LIDAR_TOP = 2 + """Top-facing LiDAR type.""" + LIDAR_FRONT = 3 + """Front-facing LiDAR type.""" + LIDAR_SIDE_LEFT = 4 + """Left-side LiDAR type.""" + LIDAR_SIDE_RIGHT = 5 + """Right-side LiDAR type.""" + LIDAR_BACK = 6 + """Back-facing LiDAR type.""" + LIDAR_DOWN = 7 + """Down-facing LiDAR type.""" -@dataclass class LiDARMetadata: + """Metadata for LiDAR sensor, static for a given sensor.""" - lidar_type: LiDARType - lidar_index: Type[LiDARIndex] - extrinsic: Optional[StateSE3] = None - # TODO: add identifier if point cloud is returned in lidar or ego frame. + __slots__ = ("_lidar_type", "_lidar_index", "_extrinsic") - def to_dict(self) -> dict: - return { - "lidar_type": self.lidar_type.name, - "lidar_index": self.lidar_index.__name__, - "extrinsic": self.extrinsic.tolist() if self.extrinsic is not None else None, - } + def __init__( + self, + lidar_type: LiDARType, + lidar_index: Type[LiDARIndex], + extrinsic: Optional[PoseSE3] = None, + ): + """Initialize LiDAR metadata. + + :param lidar_type: The type of the LiDAR sensor. + :param lidar_index: The indexing schema of the LiDAR point cloud. + :param extrinsic: The extrinsic pose of the LiDAR sensor, defaults to None + """ + self._lidar_type = lidar_type + self._lidar_index = lidar_index + self._extrinsic = extrinsic + + @property + def lidar_type(self) -> LiDARType: + """The type of the LiDAR sensor.""" + return self._lidar_type + + @property + def lidar_index(self) -> LiDARIndex: + """The indexing schema of the LiDAR point cloud.""" + return self._lidar_index + + @property + def extrinsic(self) -> Optional[PoseSE3]: + """The extrinsic :class:`~py123d.geometry.PoseSE3` of the LiDAR sensor, relative to the vehicle frame.""" + return self._extrinsic @classmethod def from_dict(cls, data_dict: dict) -> LiDARMetadata: + """Construct the LiDAR metadata from a dictionary. + + :param data_dict: A dictionary containing LiDAR metadata. + :raises ValueError: If the dictionary is missing required fields or contains invalid data. + :return: An instance of LiDARMetadata. + """ lidar_type = LiDARType[data_dict["lidar_type"]] if data_dict["lidar_index"] not in LIDAR_INDEX_REGISTRY: raise ValueError(f"Unknown lidar index: {data_dict['lidar_index']}") lidar_index_class = LIDAR_INDEX_REGISTRY[data_dict["lidar_index"]] - extrinsic = StateSE3.from_list(data_dict["extrinsic"]) if data_dict["extrinsic"] is not None else None + extrinsic = PoseSE3.from_list(data_dict["extrinsic"]) if data_dict["extrinsic"] is not None else None return cls(lidar_type=lidar_type, lidar_index=lidar_index_class, extrinsic=extrinsic) + def to_dict(self) -> Dict[str, Any]: + """Convert the LiDAR metadata to a dictionary. + + :return: A dictionary representation of the LiDAR metadata. + """ + return { + "lidar_type": self.lidar_type.name, + "lidar_index": self.lidar_index.__name__, + "extrinsic": self.extrinsic.tolist() if self.extrinsic is not None else None, + } + -@dataclass class LiDAR: + """Data structure for LiDAR point cloud data and associated metadata.""" - metadata: LiDARMetadata - point_cloud: npt.NDArray[np.float32] + __slots__ = ("_metadata", "_point_cloud") - @property - def xyz(self) -> npt.NDArray[np.float32]: + def __init__(self, metadata: LiDARMetadata, point_cloud: npt.NDArray[np.float32]) -> None: + """Initialize LiDAR data structure. + + :param metadata: LiDAR metadata. + :param point_cloud: LiDAR point cloud as an NxM numpy array, where N is the number of points + and M is the number of attributes per point as defined by the :class:`~py123d.conversion.registry.LiDARIndex`. """ - Returns the point cloud as an Nx3 array of x, y, z coordinates. + self._metadata = metadata + self._point_cloud = point_cloud + + @property + def metadata(self) -> LiDARMetadata: + """The :class:`LiDARMetadata` associated with this LiDAR recording.""" + return self._metadata + + @property + def point_cloud(self) -> npt.NDArray[np.float32]: + """The raw point cloud as an NxM numpy array, + where N is the number of points and M is the number of attributes per point, + as defined by the :class:`~py123d.conversion.registry.LiDARIndex`. Point cloud in vehicle frame. """ - return self.point_cloud[:, self.metadata.lidar_index.XYZ] + return self._point_cloud + + @property + def xyz(self) -> npt.NDArray[np.float32]: + """The point cloud as an Nx3 array of x, y, z coordinates.""" + return self._point_cloud[:, self.metadata.lidar_index.XYZ] @property def xy(self) -> npt.NDArray[np.float32]: - """ - Returns the point cloud as an Nx2 array of x, y coordinates. - """ - return self.point_cloud[:, self.metadata.lidar_index.XY] + """The point cloud as an Nx2 array of x, y coordinates.""" + return self._point_cloud[:, self.metadata.lidar_index.XY] @property def intensity(self) -> Optional[npt.NDArray[np.float32]]: - """ - Returns the intensity values of the LiDAR point cloud if available. - Returns None if intensity is not part of the point cloud. - """ - if hasattr(self.metadata.lidar_index, "INTENSITY"): - return self.point_cloud[:, self.metadata.lidar_index.INTENSITY] - return None + """The point cloud as an Nx1 array of intensity values, if available.""" + intensity: Optional[npt.NDArray[np.float32]] = None + if hasattr(self._metadata.lidar_index, "INTENSITY"): + intensity = self._point_cloud[:, self._metadata.lidar_index.INTENSITY] + return intensity @property def range(self) -> Optional[npt.NDArray[np.float32]]: - """ - Returns the range values of the LiDAR point cloud if available. - Returns None if range is not part of the point cloud. - """ - if hasattr(self.metadata.lidar_index, "RANGE"): - return self.point_cloud[:, self.metadata.lidar_index.RANGE] - return None + """The point cloud as an Nx1 array of range values, if available.""" + range: Optional[npt.NDArray[np.float32]] = None + if hasattr(self._metadata.lidar_index, "RANGE"): + range = self._point_cloud[:, self._metadata.lidar_index.RANGE] + return range @property def elongation(self) -> Optional[npt.NDArray[np.float32]]: - """ - Returns the elongation values of the LiDAR point cloud if available. - Returns None if elongation is not part of the point cloud. - """ - if hasattr(self.metadata.lidar_index, "ELONGATION"): - return self.point_cloud[:, self.metadata.lidar_index.ELONGATION] - return None + """The point cloud as an Nx1 array of elongation values, if available.""" + elongation: Optional[npt.NDArray[np.float32]] = None + if hasattr(self._metadata.lidar_index, "ELONGATION"): + elongation = self._point_cloud[:, self._metadata.lidar_index.ELONGATION] + return elongation + + @property + def ring(self) -> Optional[npt.NDArray[np.int32]]: + """The point cloud as an Nx1 array of ring values, if available.""" + ring: Optional[npt.NDArray[np.int32]] = None + if hasattr(self._metadata.lidar_index, "RING"): + ring = self._point_cloud[:, self._metadata.lidar_index.RING] + return ring diff --git a/src/py123d/datatypes/sensors/pinhole_camera.py b/src/py123d/datatypes/sensors/pinhole_camera.py index beefa883..5755c944 100644 --- a/src/py123d/datatypes/sensors/pinhole_camera.py +++ b/src/py123d/datatypes/sensors/pinhole_camera.py @@ -1,53 +1,121 @@ from __future__ import annotations -from dataclasses import asdict, dataclass +from enum import IntEnum from typing import Any, Dict, Optional import numpy as np import numpy.typing as npt -from zmq import IntEnum from py123d.common.utils.enums import SerialIntEnum -from py123d.common.utils.mixin import ArrayMixin -from py123d.geometry.se import StateSE3 +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr +from py123d.geometry import PoseSE3 class PinholeCameraType(SerialIntEnum): + """Enumeration of pinhole camera types.""" PCAM_F0 = 0 + """Front camera.""" + PCAM_B0 = 1 + """Back camera.""" + PCAM_L0 = 2 + """Left camera, first from front to back.""" + PCAM_L1 = 3 + """Left camera, second from front to back.""" + PCAM_L2 = 4 + """Left camera, third from front to back.""" + PCAM_R0 = 5 + """Right camera, first from front to back.""" + PCAM_R1 = 6 + """Right camera, second from front to back.""" + PCAM_R2 = 7 + """Right camera, third from front to back.""" + PCAM_STEREO_L = 8 + """Left stereo camera.""" + PCAM_STEREO_R = 9 + """Right stereo camera.""" -@dataclass class PinholeCamera: + """Represents the recording of a pinhole camera including its metadata, image, and extrinsic pose.""" + + __slots__ = ("_metadata", "_image", "_extrinsic") + + def __init__( + self, + metadata: PinholeCameraMetadata, + image: npt.NDArray[np.uint8], + extrinsic: PoseSE3, + ) -> None: + """Initialize a PinholeCamera instance. + + :param metadata: The metadata associated with the camera. + :param image: The image captured by the camera. + :param extrinsic: The extrinsic pose of the camera. + """ + self._metadata = metadata + self._image = image + self._extrinsic = extrinsic - metadata: PinholeCameraMetadata - image: npt.NDArray[np.uint8] - extrinsic: StateSE3 + @property + def metadata(self) -> PinholeCameraMetadata: + """The static :class:`PinholeCameraMetadata` associated with the pinhole camera.""" + return self._metadata + + @property + def image(self) -> npt.NDArray[np.uint8]: + """The image captured by the pinhole camera, as a numpy array.""" + return self._image + + @property + def extrinsic(self) -> PoseSE3: + """The extrinsic :class:`~py123d.geometry.PoseSE3` of the pinhole camera, relative to the ego vehicle frame.""" + return self._extrinsic class PinholeIntrinsicsIndex(IntEnum): + """Enumeration of pinhole camera intrinsic parameters.""" FX = 0 + """Focal length in x direction.""" + FY = 1 + """Focal length in y direction.""" + CX = 2 + """Optical center x coordinate.""" + CY = 3 - SKEW = 4 # NOTE: not used, but added for completeness + """Optical center y coordinate.""" + + SKEW = 4 + """Skew coefficient. Not used in most cases.""" class PinholeIntrinsics(ArrayMixin): + """Pinhole camera intrinsics representation.""" + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, fx: float, fy: float, cx: float, cy: float, skew: float = 0.0) -> None: + """Initialize PinholeIntrinsics. + + :param fx: Focal length in x direction. + :param fy: Focal length in y direction. + :param cx: Optical center x coordinate. + :param cy: Optical center y coordinate. + :param skew: Skew coefficient. Not used in most cases, defaults to 0.0 + """ array = np.zeros(len(PinholeIntrinsicsIndex), dtype=np.float64) array[PinholeIntrinsicsIndex.FX] = fx array[PinholeIntrinsicsIndex.FY] = fy @@ -58,6 +126,12 @@ def __init__(self, fx: float, fy: float, cx: float, cy: float, skew: float = 0.0 @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> PinholeIntrinsics: + """Creates a PinholeIntrinsics from a numpy array, indexed by :class:`PinholeIntrinsicsIndex`. + + :param array: A 1D numpy array containing the intrinsic parameters. + :param copy: Whether to copy the array, defaults to True + :return: A :class:`PinholeIntrinsics` instance. + """ assert array.ndim == 1 assert array.shape[-1] == len(PinholeIntrinsicsIndex) instance = object.__new__(cls) @@ -66,10 +140,10 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Pinhol @classmethod def from_camera_matrix(cls, intrinsic: npt.NDArray[np.float64]) -> PinholeIntrinsics: - """ - Create a PinholeIntrinsics from a 3x3 intrinsic matrix. + """Create a PinholeIntrinsics from a 3x3 intrinsic matrix. + :param intrinsic: A 3x3 numpy array representing the intrinsic matrix. - :return: A PinholeIntrinsics instance. + :return: A :class:`PinholeIntrinsics` instance. """ assert intrinsic.shape == (3, 3) fx = intrinsic[0, 0] @@ -82,34 +156,37 @@ def from_camera_matrix(cls, intrinsic: npt.NDArray[np.float64]) -> PinholeIntrin @property def array(self) -> npt.NDArray[np.float64]: + """A numpy array representation of the pinhole intrinsics, indexed by :class:`PinholeIntrinsicsIndex`.""" return self._array @property def fx(self) -> float: + """Focal length in x direction.""" return self._array[PinholeIntrinsicsIndex.FX] @property def fy(self) -> float: + """Focal length in y direction.""" return self._array[PinholeIntrinsicsIndex.FY] @property def cx(self) -> float: + """Optical center x coordinate.""" return self._array[PinholeIntrinsicsIndex.CX] @property def cy(self) -> float: + """Optical center y coordinate.""" return self._array[PinholeIntrinsicsIndex.CY] @property def skew(self) -> float: + """Skew coefficient. Not used in most cases.""" return self._array[PinholeIntrinsicsIndex.SKEW] @property def camera_matrix(self) -> npt.NDArray[np.float64]: - """ - Returns the intrinsic matrix. - :return: A 3x3 numpy array representing the intrinsic matrix. - """ + """The 3x3 camera intrinsic matrix K.""" K = np.array( [ [self.fx, self.skew, self.cx], @@ -120,19 +197,45 @@ def camera_matrix(self) -> npt.NDArray[np.float64]: ) return K + def __repr__(self) -> str: + """String representation of :class:`PinholeIntrinsics`.""" + return indexed_array_repr(self, PinholeIntrinsicsIndex) + class PinholeDistortionIndex(IntEnum): + """Enumeration of pinhole camera distortion parameters.""" + K1 = 0 + """Radial distortion coefficient k1.""" + K2 = 1 + """Radial distortion coefficient k2.""" + P1 = 2 + """Tangential distortion coefficient p1.""" + P2 = 3 + """Tangential distortion coefficient p2.""" + K3 = 4 + """Radial distortion coefficient k3.""" class PinholeDistortion(ArrayMixin): + """Pinhole camera distortion representation.""" + + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, k1: float, k2: float, p1: float, p2: float, k3: float) -> None: + """Initialize :class:`:PinholeDistortion`. + + :param k1: Radial distortion coefficient k1. + :param k2: Radial distortion coefficient k2. + :param p1: Tangential distortion coefficient p1. + :param p2: Tangential distortion coefficient p2. + :param k3: Radial distortion coefficient k3. + """ array = np.zeros(len(PinholeDistortionIndex), dtype=np.float64) array[PinholeDistortionIndex.K1] = k1 array[PinholeDistortionIndex.K2] = k2 @@ -143,6 +246,12 @@ def __init__(self, k1: float, k2: float, p1: float, p2: float, k3: float) -> Non @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> PinholeDistortion: + """Creates a PinholeDistortion from a numpy array, indexed by :class:`PinholeDistortionIndex`. + + :param array: A 1D numpy array containing the distortion parameters. + :param copy: Whether to copy the array, defaults to True + :return: A :class:`PinholeDistortion` instance. + """ assert array.ndim == 1 assert array.shape[-1] == len(PinholeDistortionIndex) instance = object.__new__(cls) @@ -151,40 +260,73 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Pinhol @property def array(self) -> npt.NDArray[np.float64]: + """A numpy array representation of the pinhole distortion, indexed by :class:`PinholeDistortionIndex`.""" return self._array @property def k1(self) -> float: + """Radial distortion coefficient k1.""" return self._array[PinholeDistortionIndex.K1] @property def k2(self) -> float: + """Radial distortion coefficient k2.""" return self._array[PinholeDistortionIndex.K2] @property def p1(self) -> float: + """Tangential distortion coefficient p1.""" return self._array[PinholeDistortionIndex.P1] @property def p2(self) -> float: + """Tangential distortion coefficient p2.""" return self._array[PinholeDistortionIndex.P2] @property def k3(self) -> float: + """Radial distortion coefficient k3.""" return self._array[PinholeDistortionIndex.K3] + def __repr__(self) -> str: + """String representation of :class:`PinholeDistortion`.""" + return indexed_array_repr(self, PinholeDistortionIndex) -@dataclass -class PinholeCameraMetadata: - camera_type: PinholeCameraType - intrinsics: Optional[PinholeIntrinsics] - distortion: Optional[PinholeDistortion] - width: int - height: int +class PinholeCameraMetadata: + """Static metadata for a pinhole camera, stored in a log.""" + + __slots__ = ("_camera_type", "_intrinsics", "_distortion", "_width", "_height") + + def __init__( + self, + camera_type: PinholeCameraType, + intrinsics: Optional[PinholeIntrinsics], + distortion: Optional[PinholeDistortion], + width: int, + height: int, + ) -> None: + """Initialize a :class:`PinholeCameraMetadata` instance. + + :param camera_type: The type of the pinhole camera. + :param intrinsics: The :class:`PinholeIntrinsics` of the pinhole camera. + :param distortion: The :class:`PinholeDistortion` of the pinhole camera. + :param width: The image width in pixels. + :param height: The image height in pixels. + """ + self._camera_type = camera_type + self._intrinsics = intrinsics + self._distortion = distortion + self._width = width + self._height = height @classmethod def from_dict(cls, data_dict: Dict[str, Any]) -> PinholeCameraMetadata: + """Create a :class:`PinholeCameraMetadata` from a dictionary. + + :param data_dict: A dictionary containing the metadata. + :return: A PinholeCameraMetadata instance. + """ data_dict["camera_type"] = PinholeCameraType(data_dict["camera_type"]) data_dict["intrinsics"] = ( PinholeIntrinsics.from_list(data_dict["intrinsics"]) if data_dict["intrinsics"] is not None else None @@ -195,28 +337,56 @@ def from_dict(cls, data_dict: Dict[str, Any]) -> PinholeCameraMetadata: return PinholeCameraMetadata(**data_dict) def to_dict(self) -> Dict[str, Any]: - data_dict = asdict(self) + """Converts the :class:`PinholeCameraMetadata` to a dictionary. + + :return: A dictionary representation of the PinholeCameraMetadata instance, with default Python types. + """ + data_dict = {} data_dict["camera_type"] = int(self.camera_type) data_dict["intrinsics"] = self.intrinsics.tolist() if self.intrinsics is not None else None data_dict["distortion"] = self.distortion.tolist() if self.distortion is not None else None + data_dict["width"] = self.width + data_dict["height"] = self.height return data_dict + @property + def camera_type(self) -> PinholeCameraType: + """The :class:`PinholeCameraType` of the pinhole camera.""" + return self._camera_type + + @property + def intrinsics(self) -> Optional[PinholeIntrinsics]: + """The :class:`PinholeIntrinsics` of the pinhole camera.""" + return self._intrinsics + + @property + def distortion(self) -> Optional[PinholeDistortion]: + """The :class:`PinholeDistortion` of the pinhole camera.""" + return self._distortion + + @property + def width(self) -> int: + """The image width in pixels.""" + return self._width + + @property + def height(self) -> int: + """The image height in pixels.""" + return self._height + @property def aspect_ratio(self) -> float: + """The aspect ratio (width / height) of the pinhole camera.""" return self.width / self.height @property def fov_x(self) -> float: - """ - Calculates the horizontal field of view (FOV) in radian. - """ + """The horizontal field of view (FOV) of the pinhole camera in radians.""" fov_x_rad = 2 * np.arctan(self.width / (2 * self.intrinsics.fx)) return fov_x_rad @property def fov_y(self) -> float: - """ - Calculates the vertical field of view (FOV) in radian. - """ + """The vertical field of view (FOV) of the pinhole camera in radians.""" fov_y_rad = 2 * np.arctan(self.height / (2 * self.intrinsics.fy)) return fov_y_rad diff --git a/src/py123d/datatypes/time/__init__.py b/src/py123d/datatypes/time/__init__.py index e69de29b..5ca914d0 100644 --- a/src/py123d/datatypes/time/__init__.py +++ b/src/py123d/datatypes/time/__init__.py @@ -0,0 +1 @@ +from py123d.datatypes.time.time_point import TimePoint diff --git a/src/py123d/datatypes/time/time_point.py b/src/py123d/datatypes/time/time_point.py index 437bf967..2162fb02 100644 --- a/src/py123d/datatypes/time/time_point.py +++ b/src/py123d/datatypes/time/time_point.py @@ -1,353 +1,78 @@ from __future__ import annotations -from dataclasses import dataclass -# TODO: Refactor -# NOTE: Taken from nuplan and adapted. Generally, these types are handy, when handling time discrete data. - - -class TimeDuration: - """Class representing a time delta, with a microsecond resolution.""" - - __slots__ = "_time_us" - - def __init__(self, *, time_us: int, _direct: bool = True) -> None: - """Constructor, should not be called directly. Raises if the keyword parameter _direct is not set to false.""" - if _direct: - raise RuntimeError("Don't initialize this class directly, use one of the constructors instead!") - - self._time_us = time_us - - @classmethod - def from_us(cls, t_us: int) -> TimeDuration: - """ - Constructs a TimeDuration from a value in microseconds. - :param t_us: Time in microseconds. - :return: TimeDuration. - """ - assert isinstance(t_us, int), "Microseconds must be an integer!" - return cls(time_us=t_us, _direct=False) - - @classmethod - def from_ms(cls, t_ms: float) -> TimeDuration: - """ - Constructs a TimeDuration from a value in milliseconds. - :param t_ms: Time in milliseconds. - :return: TimeDuration. - """ - return cls(time_us=int(t_ms * int(1e3)), _direct=False) - - @classmethod - def from_s(cls, t_s: float) -> TimeDuration: - """ - Constructs a TimeDuration from a value in seconds. - :param t_s: Time in seconds. - :return: TimeDuration. - """ - return cls(time_us=int(t_s * int(1e6)), _direct=False) - - @property - def time_us(self) -> int: - """ - :return: TimeDuration in microseconds. - """ - return self._time_us - - @property - def time_ms(self) -> float: - """ - :return: TimeDuration in milliseconds. - """ - return self._time_us / 1e3 - - @property - def time_s(self) -> float: - """ - :return: TimeDuration in seconds. - """ - return self._time_us / 1e6 - - def __add__(self, other: object) -> TimeDuration: - """ - Adds a time duration to a time duration. - :param other: time duration. - :return: self + other if other is a TimeDuration. - """ - if isinstance(other, TimeDuration): - return TimeDuration.from_us(self.time_us + other.time_us) - return NotImplemented - - def __sub__(self, other: object) -> TimeDuration: - """ - Subtract a time duration from a time duration. - :param other: time duration. - :return: self - other if other is a TimeDuration. - """ - if isinstance(other, TimeDuration): - return TimeDuration.from_us(self.time_us - other.time_us) - return NotImplemented - - def __mul__(self, other: object) -> TimeDuration: - """ - Multiply a time duration by a scalar value. - :param other: value to multiply. - :return: self * other if other is a scalar. - """ - if isinstance(other, (int, float)): - return TimeDuration.from_s(self.time_s * other) - return NotImplemented - - def __rmul__(self, other: object) -> TimeDuration: - """ - Multiply a time duration by a scalar value. - :param other: value to multiply. - :return: self * other if other is a scalar. - """ - if isinstance(other, (int, float)): - return self * other - return NotImplemented - - def __truediv__(self, other: object) -> TimeDuration: - """ - Divides a time duration by a scalar value. - :param other: value to divide for. - :return: self / other if other is a scalar. - """ - if isinstance(other, (int, float)): - return TimeDuration.from_s(self.time_s / other) - return NotImplemented - - def __floordiv__(self, other: object) -> TimeDuration: - """ - Floor divides a time duration by a scalar value. - :param other: value to divide for. - :return: self // other if other is a scalar. - """ - if isinstance(other, (int, float)): - return TimeDuration.from_s(self.time_s // other) - return NotImplemented - - def __gt__(self, other: TimeDuration) -> bool: - """ - Self is greater than other. - :param other: TimeDuration. - :return: True if self > other, False otherwise. - """ - if isinstance(other, TimeDuration): - return self.time_us > other.time_us - return NotImplemented - - def __ge__(self, other: object) -> bool: - """ - Self is greater or equal than other. - :param other: TimeDuration. - :return: True if self >= other, False otherwise. - """ - if isinstance(other, TimeDuration): - return self.time_us >= other.time_us - return NotImplemented - - def __lt__(self, other: TimeDuration) -> bool: - """ - Self is less than other. - :param other: TimeDuration. - :return: True if self < other, False otherwise. - """ - if isinstance(other, TimeDuration): - return self.time_us < other.time_us - return NotImplemented - - def __le__(self, other: TimeDuration) -> bool: - """ - Self is less or equal than other. - :param other: TimeDuration. - :return: True if self <= other, False otherwise. - """ - if isinstance(other, TimeDuration): - return self.time_us <= other.time_us - return NotImplemented - - def __eq__(self, other: object) -> bool: - """ - Self is equal to other. - :param other: TimeDuration. - :return: True if self == other, False otherwise. - """ - if not isinstance(other, TimeDuration): - return NotImplemented - - return self.time_us == other.time_us - - def __hash__(self) -> int: - """ - :return: hash for this object. - """ - return hash(self.time_us) - - def __repr__(self) -> str: - """ - :return: String representation. - """ - return "TimeDuration({}s)".format(self.time_s) - - -@dataclass class TimePoint: - """ - Time instance in a time series. - """ - - time_us: int # [micro seconds] time since epoch in micro seconds - __slots__ = "time_us" + """Time instance in a time series.""" - def __post_init__(self) -> None: - """ - Validate class after creation. - """ - assert self.time_us >= 0, "Time point has to be positive!" + __slots__ = ("_time_us",) + _time_us: int # [micro seconds] time since epoch in micro seconds @classmethod def from_ns(cls, t_ns: int) -> TimePoint: - """ - Constructs a TimePoint from a value in nanoseconds. + """Constructs a TimePoint from a value in nanoseconds. + :param t_ns: Time in nanoseconds. :return: TimePoint. """ assert isinstance(t_ns, int), "Nanoseconds must be an integer!" - return TimePoint(time_us=t_ns // 1000) + instance = object.__new__(cls) + object.__setattr__(instance, "_time_us", t_ns // 1000) + return instance @classmethod def from_us(cls, t_us: int) -> TimePoint: - """ - Constructs a TimePoint from a value in microseconds. + """Constructs a TimePoint from a value in microseconds. + :param t_us: Time in microseconds. :return: TimePoint. """ assert isinstance(t_us, int), "Microseconds must be an integer!" - return TimePoint(time_us=t_us) + instance = object.__new__(cls) + object.__setattr__(instance, "_time_us", t_us) + return instance @classmethod def from_ms(cls, t_ms: float) -> TimePoint: - """ - Constructs a TimePoint from a value in milliseconds. + """Constructs a TimePoint from a value in milliseconds. + :param t_ms: Time in milliseconds. :return: TimePoint. """ - return TimePoint(time_us=int(t_ms * int(1e3))) + instance = object.__new__(cls) + object.__setattr__(instance, "_time_us", int(t_ms * int(1e3))) + return instance @classmethod def from_s(cls, t_s: float) -> TimePoint: - """ - Constructs a TimePoint from a value in seconds. + """Constructs a TimePoint from a value in seconds. + :param t_s: Time in seconds. :return: TimePoint. """ - return TimePoint(time_us=int(t_s * int(1e6))) + instance = object.__new__(cls) + object.__setattr__(instance, "_time_us", int(t_s * int(1e6))) + return instance @property - def time_ms(self) -> float: - """ - :return: TimePoint in milliseconds. - """ - return self.time_us / 1e3 + def time_ns(self) -> int: + """The timepoint in nanoseconds [ns].""" + return self._time_us * 1000 @property - def time_s(self) -> float: - """ - :return: TimePoint in seconds. - """ - return self.time_us / 1e6 - - def __add__(self, other: object) -> TimePoint: - """ - Adds a TimeDuration to generate a new TimePoint. - :param other: time point. - :return: self + other. - """ - if isinstance(other, (TimeDuration, TimePoint)): - return TimePoint(self.time_us + other.time_us) - return NotImplemented - - def __radd__(self, other: object) -> TimePoint: - """ - :param other: Right addition target. - :return: Addition with other if other is a TimeDuration. - """ - if isinstance(other, TimeDuration): - return self.__add__(other) - return NotImplemented - - def __sub__(self, other: object) -> TimePoint: - """ - Subtract a time duration from a time point. - :param other: time duration. - :return: self - other if other is a TimeDuration. - """ - if isinstance(other, (TimeDuration, TimePoint)): - return TimePoint(self.time_us - other.time_us) - return NotImplemented - - def __gt__(self, other: TimePoint) -> bool: - """ - Self is greater than other. - :param other: time point. - :return: True if self > other, False otherwise. - """ - if isinstance(other, TimePoint): - return self.time_us > other.time_us - return NotImplemented - - def __ge__(self, other: TimePoint) -> bool: - """ - Self is greater or equal than other. - :param other: time point. - :return: True if self >= other, False otherwise. - """ - if isinstance(other, TimePoint): - return self.time_us >= other.time_us - return NotImplemented - - def __lt__(self, other: TimePoint) -> bool: - """ - Self is less than other. - :param other: time point. - :return: True if self < other, False otherwise. - """ - if isinstance(other, TimePoint): - return self.time_us < other.time_us - return NotImplemented - - def __le__(self, other: TimePoint) -> bool: - """ - Self is less or equal than other. - :param other: time point. - :return: True if self <= other, False otherwise. - """ - if isinstance(other, TimePoint): - return self.time_us <= other.time_us - return NotImplemented - - def __eq__(self, other: object) -> bool: - """ - Self is equal to other - :param other: time point - :return: True if self == other, False otherwise - """ - if not isinstance(other, TimePoint): - return NotImplemented + def time_us(self) -> int: + """The timepoint in microseconds [μs].""" + return self._time_us - return self.time_us == other.time_us + @property + def time_ms(self) -> float: + """The timepoint in milliseconds [ms].""" + return self._time_us / 1e3 - def __hash__(self) -> int: - """ - :return: hash for this object - """ - return hash(self.time_us) + @property + def time_s(self) -> float: + """The timepoint in seconds [s].""" + return self._time_us / 1e6 - def diff(self, time_point: TimePoint) -> TimeDuration: - """ - Computes the TimeDuration between self and another TimePoint. - :param time_point: The other time point. - :return: The TimeDuration between the two TimePoints. - """ - return TimeDuration.from_us(int(self.time_us - time_point.time_us)) + def __repr__(self): + """String representation of :class:`TimePoint`.""" + return f"TimePoint(time_us={self._time_us})" diff --git a/src/py123d/datatypes/vehicle_state/__init__.py b/src/py123d/datatypes/vehicle_state/__init__.py index e69de29b..64b7a273 100644 --- a/src/py123d/datatypes/vehicle_state/__init__.py +++ b/src/py123d/datatypes/vehicle_state/__init__.py @@ -0,0 +1,8 @@ +from py123d.datatypes.vehicle_state.dynamic_state import ( + DynamicStateSE2, + DynamicStateSE2Index, + DynamicStateSE3, + DynamicStateSE3Index, +) +from py123d.datatypes.vehicle_state.ego_state import EgoStateSE2, EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters diff --git a/src/py123d/datatypes/vehicle_state/dynamic_state.py b/src/py123d/datatypes/vehicle_state/dynamic_state.py new file mode 100644 index 00000000..942fd7ef --- /dev/null +++ b/src/py123d/datatypes/vehicle_state/dynamic_state.py @@ -0,0 +1,239 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import IntEnum + +import numpy as np +import numpy.typing as npt + +from py123d.common.utils.enums import classproperty +from py123d.common.utils.mixin import ArrayMixin +from py123d.geometry import Vector2D, Vector3D + + +class DynamicStateSE3Index(IntEnum): + """The indices for the dynamic state in SE3.""" + + VELOCITY_X = 0 + """Velocity in the X direction (forward).""" + + VELOCITY_Y = 1 + """Velocity in the Y direction (left).""" + + VELOCITY_Z = 2 + """Velocity in the Z direction (up).""" + + ACCELERATION_X = 3 + """Acceleration in the X direction (forward).""" + + ACCELERATION_Y = 4 + """Acceleration in the Y direction (left).""" + + ACCELERATION_Z = 5 + """Acceleration in the Z direction (up).""" + + ANGULAR_VELOCITY_X = 6 + """Angular velocity around the X axis (roll).""" + + ANGULAR_VELOCITY_Y = 7 + """Angular velocity around the Y axis (pitch).""" + + ANGULAR_VELOCITY_Z = 8 + """Angular velocity around the Z axis (yaw).""" + + @classproperty + def VELOCITY_3D(cls) -> slice: + """Slice for the 3D velocity components (x,y,z).""" + return slice(cls.VELOCITY_X, cls.VELOCITY_Z + 1) + + @classproperty + def VELOCITY_2D(cls) -> slice: + """Slice for the 2D velocity components (x,y).""" + return slice(cls.VELOCITY_X, cls.VELOCITY_Y + 1) + + @classproperty + def ACCELERATION_3D(cls) -> slice: + """Slice for the 3D acceleration components (x,y,z).""" + return slice(cls.ACCELERATION_X, cls.ACCELERATION_Z + 1) + + @classproperty + def ACCELERATION_2D(cls) -> slice: + """Slice for the 2D acceleration components (x,y).""" + return slice(cls.ACCELERATION_X, cls.ACCELERATION_Y + 1) + + @classproperty + def ANGULAR_VELOCITY_3D(cls) -> slice: + """Slice for the 3D angular velocity components (x,y,z).""" + return slice(cls.ANGULAR_VELOCITY_X, cls.ANGULAR_VELOCITY_Z + 1) + + +class DynamicStateSE3(ArrayMixin): + """The dynamic state of a vehicle in SE3 (3D space).""" + + __slots__ = ("_array",) + _array: npt.NDArray[np.float64] + + def __init__( + self, + velocity: Vector3D, + acceleration: Vector3D, + angular_velocity: Vector3D, + ): + """Initialize a :class:`DynamicStateSE3` instance. + + :param velocity: 3D velocity vector. + :param acceleration: 3D acceleration vector. + :param angular_velocity: 3D angular velocity vector. + """ + array = np.zeros(len(DynamicStateSE3Index), dtype=np.float64) + array[DynamicStateSE3Index.VELOCITY_3D] = velocity.array + array[DynamicStateSE3Index.ACCELERATION_3D] = acceleration.array + array[DynamicStateSE3Index.ANGULAR_VELOCITY_3D] = angular_velocity.array + self._array = array + + @classmethod + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> DynamicStateSE3: + """Create a :class:`DynamicStateSE3` from NumPy array of shape (9,), indexed by :class:`DynamicStateSE3Index`. + + :param array: The array containing the dynamic state information. + :param copy: Whether to copy the array data. + :return: A :class:`DynamicStateSE3` instance. + """ + assert array.ndim == 1 + assert array.shape[0] == len(DynamicStateSE3Index) + instance = object.__new__(cls) + object.__setattr__(instance, "_array", array.copy() if copy else array) + return instance + + @property + def velocity_3d(self) -> Vector3D: + """3D velocity vector.""" + return Vector3D.from_array(self._array[DynamicStateSE3Index.VELOCITY_3D], copy=False) + + @property + def velocity_2d(self) -> Vector2D: + """2D velocity vector.""" + return Vector2D.from_array(self._array[DynamicStateSE3Index.VELOCITY_2D], copy=False) + + @property + def acceleration_3d(self) -> Vector3D: + """3D acceleration vector.""" + return Vector3D.from_array(self._array[DynamicStateSE3Index.ACCELERATION_3D], copy=False) + + @property + def acceleration_2d(self) -> Vector2D: + """2D acceleration vector.""" + return Vector2D.from_array(self._array[DynamicStateSE3Index.ACCELERATION_2D], copy=False) + + @property + def angular_velocity(self) -> Vector3D: + """3D angular velocity vector.""" + return Vector3D.from_array(self._array[DynamicStateSE3Index.ANGULAR_VELOCITY_3D], copy=False) + + @property + def array(self) -> npt.NDArray[np.float64]: + """NumPy array representation of shape (9,), indexed by :class:`DynamicStateSE3Index`.""" + return self._array + + @property + def dynamic_state_se2(self) -> DynamicStateSE2: + """The :class:`DynamicStateSE2` projection of this SE3 dynamic state.""" + _array = np.zeros(len(DynamicStateSE2Index), dtype=np.float64) + _array[DynamicStateSE2Index.VELOCITY_2D] = self._array[DynamicStateSE3Index.VELOCITY_2D] + _array[DynamicStateSE2Index.ACCELERATION_2D] = self._array[DynamicStateSE3Index.ACCELERATION_2D] + _array[DynamicStateSE2Index.ANGULAR_VELOCITY_Z] = self._array[DynamicStateSE3Index.ANGULAR_VELOCITY_Z] + return DynamicStateSE2.from_array(_array, copy=False) + + +class DynamicStateSE2Index(IntEnum): + """The indices for the dynamic state in SE2.""" + + VELOCITY_X = 0 + """Velocity in the X direction (forward).""" + + VELOCITY_Y = 1 + """Velocity in the Y direction (left).""" + + ACCELERATION_X = 2 + """Acceleration in the X direction (forward).""" + + ACCELERATION_Y = 3 + """Acceleration in the Y direction (left).""" + + ANGULAR_VELOCITY_Z = 4 + """Angular velocity around the Z axis (yaw).""" + + @classproperty + def VELOCITY_2D(cls) -> slice: + """Slice for the 2D velocity components (x,y).""" + return slice(cls.VELOCITY_X, cls.VELOCITY_Y + 1) + + @classproperty + def ACCELERATION_2D(cls) -> slice: + """Slice for the 2D acceleration components (x,y).""" + return slice(cls.ACCELERATION_X, cls.ACCELERATION_Y + 1) + + @classproperty + def ANGULAR_VELOCITY(cls) -> int: + """Index for the angular velocity component (yaw).""" + return cls.ANGULAR_VELOCITY_Z + + +@dataclass +class DynamicStateSE2(ArrayMixin): + """The dynamic state of a vehicle in SE2 (2D plane).""" + + __slots__ = ("_array",) + _array: npt.NDArray[np.float64] + + def __init__( + self, + velocity: Vector2D, + acceleration: Vector2D, + angular_velocity: float, + ): + """Initialize a :class:`DynamicStateSE2` instance. + + :param velocity: 2D velocity vector. + :param acceleration: 2D acceleration vector. + :param angular_velocity: Angular velocity around the Z axis (yaw). + """ + array = np.zeros(len(DynamicStateSE2Index), dtype=np.float64) + array[DynamicStateSE2Index.VELOCITY_2D] = velocity.array + array[DynamicStateSE2Index.ACCELERATION_2D] = acceleration.array + array[DynamicStateSE2Index.ANGULAR_VELOCITY_Z] = angular_velocity + self._array = array + + @classmethod + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> DynamicStateSE2: + """Create a :class:`DynamicStateSE2` from NumPy array of shape (5,), indexed by :class:`DynamicStateSE2Index`. + + :param array: The array containing the dynamic state information. + :param copy: Whether to copy the array data. + :return: A :class:`DynamicStateSE2` instance. + """ + assert array.ndim == 1 + assert array.shape[0] == len(DynamicStateSE2Index) + instance = object.__new__(cls) + object.__setattr__(instance, "_array", array.copy() if copy else array) + return instance + + @property + def velocity_2d(self) -> Vector2D: + """2D velocity vector.""" + return Vector2D.from_array(self._array[DynamicStateSE2Index.VELOCITY_2D], copy=False) + + @property + def acceleration_2d(self) -> Vector2D: + """2D acceleration vector.""" + return Vector2D.from_array(self._array[DynamicStateSE2Index.ACCELERATION_2D], copy=False) + + @property + def angular_velocity(self) -> float: + """Angular velocity around the Z axis (yaw).""" + return self._array[DynamicStateSE2Index.ANGULAR_VELOCITY_Z] + + @property + def array(self) -> npt.NDArray[np.float64]: + """NumPy array representation of shape (5,), indexed by :class:`DynamicStateSE2Index`.""" + return self._array diff --git a/src/py123d/datatypes/vehicle_state/ego_state.py b/src/py123d/datatypes/vehicle_state/ego_state.py index 3ddc09a5..f6b40497 100644 --- a/src/py123d/datatypes/vehicle_state/ego_state.py +++ b/src/py123d/datatypes/vehicle_state/ego_state.py @@ -1,16 +1,11 @@ from __future__ import annotations -from dataclasses import dataclass -from enum import IntEnum from typing import Final, Optional -import numpy as np -import numpy.typing as npt - -from py123d.common.utils.enums import classproperty from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE2, BoxDetectionSE3 from py123d.datatypes.time.time_point import TimePoint +from py123d.datatypes.vehicle_state.dynamic_state import DynamicStateSE2, DynamicStateSE3 from py123d.datatypes.vehicle_state.vehicle_parameters import ( VehicleParameters, center_se2_to_rear_axle_se2, @@ -18,331 +13,324 @@ rear_axle_se2_to_center_se2, rear_axle_se3_to_center_se3, ) -from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, StateSE2, StateSE3, Vector2D, Vector3D +from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, PoseSE2, PoseSE3 EGO_TRACK_TOKEN: Final[str] = "ego_vehicle" -class EgoStateSE3Index(IntEnum): - - X = 0 - Y = 1 - Z = 2 - QW = 3 - QX = 4 - QY = 5 - QZ = 6 - VELOCITY_X = 7 - VELOCITY_Y = 8 - VELOCITY_Z = 9 - ACCELERATION_X = 10 - ACCELERATION_Y = 11 - ACCELERATION_Z = 12 - ANGULAR_VELOCITY_X = 13 - ANGULAR_VELOCITY_Y = 14 - ANGULAR_VELOCITY_Z = 15 - - @classproperty - def STATE_SE3(cls) -> slice: - return slice(cls.X, cls.QZ + 1) - - @classproperty - def DYNAMIC_VEHICLE_STATE(cls) -> slice: - return slice(cls.VELOCITY_X, cls.ANGULAR_VELOCITY_Z + 1) - - @classproperty - def SCALAR(cls) -> slice: - return slice(cls.QW, cls.QW + 1) - - @classproperty - def VECTOR(cls) -> slice: - return slice(cls.QX, cls.QZ + 1) - - -@dataclass class EgoStateSE3: - - center_se3: StateSE3 - dynamic_state_se3: DynamicStateSE3 - vehicle_parameters: VehicleParameters - timepoint: Optional[TimePoint] = None - tire_steering_angle: float = 0.0 + """The EgoStateSE3 represents the state of the ego vehicle in SE3 (3D space). + It includes the rear axle pose, vehicle parameters, optional dynamic state, + optional timepoint, and optional tire steering angle. + """ + + def __init__( + self, + rear_axle_se3: PoseSE3, + vehicle_parameters: VehicleParameters, + dynamic_state_se3: Optional[DynamicStateSE3] = None, + timepoint: Optional[TimePoint] = None, + tire_steering_angle: Optional[float] = 0.0, + ): + """Initialize an :class:`EgoStateSE3` instance. + + :param rear_axle_se3: The pose of the rear axle in SE3. + :param vehicle_parameters: The parameters of the vehicle. + :param dynamic_state_se3: The dynamic state of the vehicle, defaults to None. + :param timepoint: The timepoint of the state, defaults to None. + :param tire_steering_angle: The tire steering angle, defaults to 0.0. + """ + self._rear_axle_se3 = rear_axle_se3 + self._vehicle_parameters = vehicle_parameters + self._dynamic_state_se3 = dynamic_state_se3 + self._timepoint: Optional[TimePoint] = timepoint + self._tire_steering_angle: Optional[float] = tire_steering_angle @classmethod - def from_array( + def from_center( cls, - array: npt.NDArray[np.float64], + center_se3: PoseSE3, vehicle_parameters: VehicleParameters, + dynamic_state_se3: Optional[DynamicStateSE3] = None, timepoint: Optional[TimePoint] = None, + tire_steering_angle: float = 0.0, ) -> EgoStateSE3: - state_se3 = StateSE3.from_array(array[EgoStateSE3Index.STATE_SE3]) - dynamic_state = DynamicStateSE3.from_array(array[EgoStateSE3Index.DYNAMIC_VEHICLE_STATE]) - return EgoStateSE3(state_se3, dynamic_state, vehicle_parameters, timepoint) + """Create an :class:`EgoStateSE3` from the center pose. + + :param center_se3: The center pose in SE3. + :param vehicle_parameters: The parameters of the vehicle. + :param dynamic_state_se3: The dynamic state of the vehicle, defaults to None. + :param timepoint: The timepoint of the state, defaults to None. + :param tire_steering_angle: The tire steering angle, defaults to 0.0. + :return: An :class:`EgoStateSE3` instance. + """ + + rear_axle_se3 = center_se3_to_rear_axle_se3( + center_se3=center_se3, + vehicle_parameters=vehicle_parameters, + ) + + # TODO @DanielDauner: Adapt dynamic state from center to rear-axle + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_se3, + vehicle_parameters=vehicle_parameters, + dynamic_state_se3=dynamic_state_se3, + timepoint=timepoint, + tire_steering_angle=tire_steering_angle, + ) @classmethod def from_rear_axle( cls, - rear_axle_se3: StateSE3, - dynamic_state_se3: DynamicStateSE3, + rear_axle_se3: PoseSE3, vehicle_parameters: VehicleParameters, - time_point: TimePoint, + dynamic_state_se3: Optional[DynamicStateSE3] = None, + timepoint: Optional[TimePoint] = None, tire_steering_angle: float = 0.0, ) -> EgoStateSE3: + """Create an :class:`EgoStateSE3` from the rear axle pose. + + :param rear_axle_se3: The pose of the rear axle in SE3. + :param vehicle_parameters: The parameters of the vehicle. + :param dynamic_state_se3: The dynamic state of the vehicle, defaults to None. + :param timepoint: The timepoint of the state, defaults to None. + :param tire_steering_angle: The tire steering angle, defaults to 0.0. + :return: An :class:`EgoStateSE3` instance. + """ return EgoStateSE3( - center_se3=rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_se3, vehicle_parameters=vehicle_parameters), - dynamic_state_se3=dynamic_state_se3, # TODO: Adapt dynamic state rear-axle to center + rear_axle_se3=rear_axle_se3, vehicle_parameters=vehicle_parameters, - timepoint=time_point, + dynamic_state_se3=dynamic_state_se3, + timepoint=timepoint, tire_steering_angle=tire_steering_angle, ) @property - def array(self) -> npt.NDArray[np.float64]: - """ - Convert the EgoVehicleState to an array. - :return: An array containing the bounding box and dynamic state information. - """ - assert isinstance(self.center_se3, StateSE3) - assert isinstance(self.dynamic_state_se3, DynamicStateSE3) + def rear_axle_se3(self) -> PoseSE3: + """The :class:`~py123d.geometry.PoseSE3` of the rear axle in SE3.""" + return self._rear_axle_se3 - center_array = self.center_se3.array - dynamic_array = self.dynamic_state_se3.array + @property + def rear_axle_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` of the rear axle in SE2.""" + return self._rear_axle_se3.pose_se2 - return np.concatenate((center_array, dynamic_array), axis=0) + @property + def vehicle_parameters(self) -> VehicleParameters: + """The :class:`~py123d.datatypes.vehicle_state.VehicleParameters` of the vehicle.""" + return self._vehicle_parameters @property - def center(self) -> StateSE3: - return self.center_se3 + def dynamic_state_se3(self) -> Optional[DynamicStateSE3]: + """The :class:`~py123d.datatypes.vehicle_state.DynamicStateSE3` of the vehicle.""" + return self._dynamic_state_se3 @property - def rear_axle_se3(self) -> StateSE3: - return center_se3_to_rear_axle_se3(center_se3=self.center_se3, vehicle_parameters=self.vehicle_parameters) + def timepoint(self) -> Optional[TimePoint]: + """The :class:`~py123d.datatypes.time.TimePoint` of the ego state, if available.""" + return self._timepoint @property - def rear_axle_se2(self) -> StateSE2: - return self.rear_axle_se3.state_se2 + def tire_steering_angle(self) -> Optional[float]: + """The tire steering angle of the ego state, if available.""" + return self._tire_steering_angle @property - def rear_axle(self) -> StateSE3: - return self.rear_axle_se3 + def center_se3(self) -> PoseSE3: + """The :class:`~py123d.geometry.PoseSE3` of the vehicle center in SE3.""" + return rear_axle_se3_to_center_se3( + rear_axle_se3=self._rear_axle_se3, + vehicle_parameters=self._vehicle_parameters, + ) + + @property + def center_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` of the vehicle center in SE2.""" + return self.center_se3.pose_se2 @property - def bounding_box(self) -> BoundingBoxSE3: + def bounding_box_se3(self) -> BoundingBoxSE3: + """The :class:`~py123d.geometry.BoundingBoxSE3` of the ego vehicle.""" return BoundingBoxSE3( - center=self.center_se3, + center_se3=self.center_se3, length=self.vehicle_parameters.length, width=self.vehicle_parameters.width, height=self.vehicle_parameters.height, ) - @property - def bounding_box_se3(self) -> BoundingBoxSE3: - return self.bounding_box - @property def bounding_box_se2(self) -> BoundingBoxSE2: - return self.bounding_box.bounding_box_se2 + """The :class:`~py123d.geometry.BoundingBoxSE2` of the ego vehicle.""" + return self.bounding_box_se3.bounding_box_se2 @property - def box_detection(self) -> BoxDetectionSE3: + def box_detection_se3(self) -> BoxDetectionSE3: + """The :class:`~py123d.datatypes.detections.BoxDetectionSE3` projection of the ego vehicle.""" return BoxDetectionSE3( metadata=BoxDetectionMetadata( label=DefaultBoxDetectionLabel.EGO, timepoint=self.timepoint, track_token=EGO_TRACK_TOKEN, - confidence=1.0, + num_lidar_points=None, ), - bounding_box_se3=self.bounding_box, - velocity=self.dynamic_state_se3.velocity, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=self.dynamic_state_se3.velocity_3d if self.dynamic_state_se3 else None, ) - @property - def box_detection_se3(self) -> BoxDetectionSE3: - return self.box_detection - @property def box_detection_se2(self) -> BoxDetectionSE2: - return self.box_detection.box_detection_se2 + """The :class:`~py123d.datatypes.detections.BoxDetectionSE2` projection of the ego vehicle.""" + return self.box_detection_se3.box_detection_se2 @property def ego_state_se2(self) -> EgoStateSE2: - return EgoStateSE2( - center_se2=self.center_se3.state_se2, - dynamic_state_se2=self.dynamic_state_se3.dynamic_state_se2, + """The :class:`EgoStateSE2` projection of this SE3 ego state.""" + return EgoStateSE2.from_rear_axle( + rear_axle_se2=self.rear_axle_se2, vehicle_parameters=self.vehicle_parameters, + dynamic_state_se2=self.dynamic_state_se3.dynamic_state_se2 if self.dynamic_state_se3 else None, timepoint=self.timepoint, tire_steering_angle=self.tire_steering_angle, ) -@dataclass class EgoStateSE2: + """The EgoStateSE2 represents the state of the ego vehicle in SE2 (2D space). + It includes the rear axle pose, vehicle parameters, optional dynamic state, and optional timepoint. + """ - center_se2: StateSE2 - dynamic_state_se2: DynamicStateSE2 - vehicle_parameters: VehicleParameters - timepoint: Optional[TimePoint] = None - tire_steering_angle: float = 0.0 + def __init__( + self, + rear_axle_se2: PoseSE2, + vehicle_parameters: VehicleParameters, + dynamic_state_se2: Optional[DynamicStateSE2] = None, + timepoint: Optional[TimePoint] = None, + tire_steering_angle: Optional[float] = 0.0, + ): + """Initialize an :class:`EgoStateSE2` instance. + + :param rear_axle_se2: The pose of the rear axle in SE2. + :param vehicle_parameters: The parameters of the vehicle. + :param dynamic_state_se2: The dynamic state of the vehicle in SE2, defaults to None. + :param timepoint: The timepoint of the state, defaults to None. + :param tire_steering_angle: The tire steering angle, defaults to 0.0 + """ + self._rear_axle_se2: PoseSE2 = rear_axle_se2 + self._vehicle_parameters: VehicleParameters = vehicle_parameters + self._dynamic_state_se2: Optional[DynamicStateSE2] = dynamic_state_se2 + self._timepoint: Optional[TimePoint] = timepoint + self._tire_steering_angle: Optional[float] = tire_steering_angle @classmethod def from_rear_axle( cls, - rear_axle_se2: StateSE2, + rear_axle_se2: PoseSE2, dynamic_state_se2: DynamicStateSE2, vehicle_parameters: VehicleParameters, - time_point: TimePoint, + timepoint: TimePoint, tire_steering_angle: float = 0.0, ) -> EgoStateSE2: + """Create an :class:`EgoStateSE2` from the rear axle pose. + + :param rear_axle_se2: The pose of the rear axle in SE2. + :param dynamic_state_se2: The dynamic state of the vehicle in SE2. + :param vehicle_parameters: The parameters of the vehicle. + :param timepoint: The timepoint of the state. + :param tire_steering_angle: The tire steering angle, defaults to 0.0. + :return: An instance of :class:`EgoStateSE2`. + """ return EgoStateSE2( - center_se2=rear_axle_se2_to_center_se2(rear_axle_se2=rear_axle_se2, vehicle_parameters=vehicle_parameters), - dynamic_state_se2=dynamic_state_se2, # TODO: Adapt dynamic state rear-axle to center + rear_axle_se2=rear_axle_se2, + dynamic_state_se2=dynamic_state_se2, vehicle_parameters=vehicle_parameters, - timepoint=time_point, + timepoint=timepoint, tire_steering_angle=tire_steering_angle, ) + @classmethod + def from_center( + cls, + center_se2: PoseSE2, + dynamic_state_se2: DynamicStateSE2, + vehicle_parameters: VehicleParameters, + timepoint: TimePoint, + tire_steering_angle: float = 0.0, + ) -> EgoStateSE2: + """Create an :class:`EgoStateSE2` from the center pose. + + :param center_se2: The pose of the center in SE2. + :param dynamic_state_se2: The dynamic state of the vehicle in SE2. + :param vehicle_parameters: The parameters of the vehicle. + :param timepoint: The timepoint of the state. + :param tire_steering_angle: The tire steering angle, defaults to 0.0. + :return: An instance of :class:`EgoStateSE2`. + """ + + rear_axle_se2 = center_se2_to_rear_axle_se2( + center_se2=center_se2, + vehicle_parameters=vehicle_parameters, + ) + + # TODO @DanielDauner: Adapt dynamic state from center to rear-axle + return EgoStateSE2.from_rear_axle( + rear_axle_se2=rear_axle_se2, + dynamic_state_se2=dynamic_state_se2, + vehicle_parameters=vehicle_parameters, + timepoint=timepoint, + tire_steering_angle=tire_steering_angle, + ) + + @property + def rear_axle_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` of the rear axle in SE2.""" + return self._rear_axle_se2 + @property - def center(self) -> StateSE2: - return self.center_se2 + def vehicle_parameters(self) -> VehicleParameters: + """The :class:`~py123d.datatypes.vehicle_state.VehicleParameters` of the vehicle.""" + return self._vehicle_parameters @property - def rear_axle_se2(self) -> StateSE2: - return center_se2_to_rear_axle_se2(center_se2=self.center_se2, vehicle_parameters=self.vehicle_parameters) + def dynamic_state_se2(self) -> Optional[DynamicStateSE2]: + """The :class:`~py123d.datatypes.vehicle_state.DynamicStateSE2` of the vehicle.""" + return self._dynamic_state_se2 @property - def rear_axle(self) -> StateSE2: - return self.rear_axle_se2 + def timepoint(self) -> Optional[TimePoint]: + """The :class:`~py123d.datatypes.time.TimePoint` of the ego state, if available.""" + return self._timepoint @property - def bounding_box(self) -> BoundingBoxSE2: + def tire_steering_angle(self) -> Optional[float]: + """The tire steering angle of the ego state, if available.""" + return self._tire_steering_angle + + @property + def center_se2(self) -> PoseSE2: + """The :class:`~py123d.geometry.PoseSE2` of the center in SE2.""" + return rear_axle_se2_to_center_se2(rear_axle_se2=self.rear_axle_se2, vehicle_parameters=self.vehicle_parameters) + + @property + def bounding_box_se2(self) -> BoundingBoxSE2: + """The :class:`~py123d.geometry.BoundingBoxSE2` of the ego vehicle.""" return BoundingBoxSE2( - center=self.center_se2, + center_se2=self.center_se2, length=self.vehicle_parameters.length, width=self.vehicle_parameters.width, ) @property - def bounding_box_se2(self) -> BoundingBoxSE2: - return self.bounding_box - - @property - def box_detection(self) -> BoxDetectionSE2: + def box_detection_se2(self) -> BoxDetectionSE2: + """The :class:`~py123d.datatypes.detections.BoxDetectionSE2` projection of the ego vehicle.""" return BoxDetectionSE2( metadata=BoxDetectionMetadata( label=DefaultBoxDetectionLabel.EGO, timepoint=self.timepoint, track_token=EGO_TRACK_TOKEN, - confidence=1.0, + num_lidar_points=None, ), bounding_box_se2=self.bounding_box_se2, - velocity=self.dynamic_state_se2.velocity, - ) - - @property - def box_detection_se2(self) -> BoxDetectionSE2: - return self.box_detection - - -class DynamicStateSE3Index(IntEnum): - - VELOCITY_X = 0 - VELOCITY_Y = 1 - VELOCITY_Z = 2 - ACCELERATION_X = 3 - ACCELERATION_Y = 4 - ACCELERATION_Z = 5 - ANGULAR_VELOCITY_X = 6 - ANGULAR_VELOCITY_Y = 7 - ANGULAR_VELOCITY_Z = 8 - - @classproperty - def VELOCITY(cls) -> slice: - return slice(cls.VELOCITY_X, cls.VELOCITY_Z + 1) - - @classproperty - def ACCELERATION(cls) -> slice: - return slice(cls.ACCELERATION_X, cls.ACCELERATION_Z + 1) - - @classproperty - def ANGULAR_VELOCITY(cls) -> slice: - return slice(cls.ANGULAR_VELOCITY_X, cls.ANGULAR_VELOCITY_Z + 1) - - -@dataclass -class DynamicStateSE3: - # TODO: Make class array like - - velocity: Vector3D - acceleration: Vector3D - angular_velocity: Vector3D - - tire_steering_rate: float = 0.0 - angular_acceleration: float = 0.0 - - @classmethod - def from_array(cls, array: npt.NDArray[np.float64]) -> DynamicStateSE3: - """ - Create a DynamicVehicleState from an array. - :param array: The array containing the dynamic state information. - :return: A DynamicVehicleState instance. - """ - assert array.ndim == 1 - assert array.shape[0] == len(DynamicStateSE3Index) - velocity = Vector3D.from_array(array[DynamicStateSE3Index.VELOCITY]) - acceleration = Vector3D.from_array(array[DynamicStateSE3Index.ACCELERATION]) - angular_velocity = Vector3D.from_array(array[DynamicStateSE3Index.ANGULAR_VELOCITY]) - return DynamicStateSE3(velocity, acceleration, angular_velocity) - - @property - def array(self) -> npt.NDArray[np.float64]: - """ - Convert the DynamicVehicleState to an array. - :return: An array containing the velocity, acceleration, and angular velocity. - """ - assert isinstance(self.velocity, Vector3D) - assert isinstance(self.acceleration, Vector3D) - assert isinstance(self.angular_velocity, Vector3D) - - return np.concatenate( - ( - self.velocity.array, - self.acceleration.array, - self.angular_velocity.array, - ), - axis=0, + velocity_2d=self.dynamic_state_se2.velocity_2d if self.dynamic_state_se2 else None, ) - - @property - def dynamic_state_se2(self) -> DynamicStateSE2: - """ - Convert the DynamicVehicleState to a 2D dynamic state. - :return: A DynamicStateSE2 instance. - """ - return DynamicStateSE2( - velocity=self.velocity.vector_2d, - acceleration=self.acceleration.vector_2d, - angular_velocity=self.angular_velocity.z, - tire_steering_rate=self.tire_steering_rate, - angular_acceleration=self.angular_acceleration, - ) - - -@dataclass -class DynamicStateSE2: - - velocity: Vector2D - acceleration: Vector2D - angular_velocity: float - - tire_steering_rate: float = 0.0 - angular_acceleration: float = 0.0 - - @property - def array(self) -> npt.NDArray[np.float64]: - """ - Convert the DynamicVehicleState to an array. - :return: An array containing the velocity, acceleration, and angular velocity. - """ - return np.concatenate((self.velocity.array, self.acceleration.array, np.array([self.angular_velocity])), axis=0) diff --git a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py index ca2a1944..92ae035c 100644 --- a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py +++ b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py @@ -2,72 +2,102 @@ from dataclasses import asdict, dataclass -from py123d.geometry import StateSE2, StateSE3, Vector2D, Vector3D -from py123d.geometry.transform.transform_se2 import translate_se2_along_body_frame -from py123d.geometry.transform.transform_se3 import translate_se3_along_body_frame +from py123d.geometry import PoseSE2, PoseSE3, Vector2D, Vector3D +from py123d.geometry.transform import translate_se2_along_body_frame, translate_se3_along_body_frame @dataclass class VehicleParameters: + """Parameters that describe the physical dimensions of a vehicle.""" vehicle_name: str + """Name of the vehicle model.""" width: float + """Width of the vehicle.""" + length: float + """Length of the vehicle.""" + height: float + """Height of the vehicle.""" wheel_base: float + """Wheel base of the vehicle (longitudinal distance between front and rear axles).""" + rear_axle_to_center_vertical: float + """Distance from the rear axle to the center of the vehicle (vertical).""" + rear_axle_to_center_longitudinal: float + """Distance from the rear axle to the center of the vehicle (longitudinal).""" @classmethod def from_dict(cls, data_dict: dict) -> VehicleParameters: - return VehicleParameters(**data_dict) + """Creates a VehicleParameters instance from a dictionary. - def to_dict(self) -> dict: - return asdict(self) + :param data_dict: Dictionary containing vehicle parameters. + :return: VehicleParameters instance. + """ + return VehicleParameters(**data_dict) @property def half_width(self) -> float: + """Half the width of the vehicle.""" return self.width / 2.0 @property def half_length(self) -> float: + """Half the length of the vehicle.""" return self.length / 2.0 @property def half_height(self) -> float: + """Half the height of the vehicle.""" return self.height / 2.0 + def to_dict(self) -> dict: + """Converts the :class:`VehicleParameters` instance to a dictionary. + + :return: Dictionary representation of the vehicle parameters. + """ + return asdict(self) + def get_nuplan_chrysler_pacifica_parameters() -> VehicleParameters: - # NOTE: use parameters from nuPlan dataset + """Helper function to get nuPlan Chrysler Pacifica vehicle parameters.""" + # NOTE: These parameters are mostly available in nuPlan, except for the rear_axle_to_center_vertical. + # The value is estimated based the LiDAR point cloud. + # [1] https://en.wikipedia.org/wiki/Chrysler_Pacifica_(minivan) return VehicleParameters( vehicle_name="nuplan_chrysler_pacifica", width=2.297, length=5.176, height=1.777, wheel_base=3.089, - rear_axle_to_center_vertical=0.45, # NOTE: missing in nuPlan, TODO: find more accurate value + rear_axle_to_center_vertical=0.45, rear_axle_to_center_longitudinal=1.461, ) def get_nuscenes_renault_zoe_parameters() -> VehicleParameters: - # https://en.wikipedia.org/wiki/Renault_Zoe + """Helper function to get nuScenes Renault Zoe vehicle parameters.""" + # NOTE: The parameters in nuScenes are estimates, and partially taken from the Renault Zoe model [1]. + # [1] https://en.wikipedia.org/wiki/Renault_Zoe return VehicleParameters( vehicle_name="nuscenes_renault_zoe", width=1.730, length=4.084, height=1.562, wheel_base=2.588, - rear_axle_to_center_vertical=1.562 / 2, # NOTE: missing in nuscenes, TODO: find more accurate value + rear_axle_to_center_vertical=1.562 / 2, rear_axle_to_center_longitudinal=1.385, ) def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: - # NOTE: values are extracted from CARLA + """Helper function to get CARLA Lincoln MKZ 2020 vehicle parameters.""" + # NOTE: These parameters are taken from the CARLA simulator vehicle model. The rear axles to center transform + # parameters are calculated based on parameters from the CARLA simulator. return VehicleParameters( vehicle_name="carla_lincoln_mkz_2020", width=1.83671, @@ -80,8 +110,10 @@ def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: def get_wopd_chrysler_pacifica_parameters() -> VehicleParameters: - # NOTE: use parameters from nuPlan dataset - # Find better parameters for WOPD ego vehicle + """Helper function to get WOPD Chrysler Pacifica vehicle parameters.""" + # NOTE: These parameters are estimates based on the vehicle model used in the WOPD dataset. + # The vehicle should be the same (or a similar) vehicle model to nuPlan and PandaSet [1]. + # [1] https://en.wikipedia.org/wiki/Chrysler_Pacifica_(minivan) return VehicleParameters( vehicle_name="wopd_chrysler_pacifica", width=2.297, @@ -94,11 +126,13 @@ def get_wopd_chrysler_pacifica_parameters() -> VehicleParameters: def get_kitti360_vw_passat_parameters() -> VehicleParameters: - # The KITTI-360 dataset uses a 2006 VW Passat Variant B6. - # https://en.wikipedia.org/wiki/Volkswagen_Passat_(B6) - # [1] https://scispace.com/pdf/team-annieway-s-autonomous-system-18ql8b7kki.pdf - # NOTE: Parameters are estimated from the vehicle model. - # https://www.cvlibs.net/datasets/kitti-360/documentation.php + """Helper function to get KITTI-360 VW Passat vehicle parameters.""" + # NOTE: The parameters in KITTI-360 are estimates based on the vehicle model used in the dataset + # Uses a 2006 VW Passat Variant B6 [1]. Vertical distance is estimated based on the LiDAR. + # KITTI-360 is currently the only dataset where the IMU has a lateral offset to the rear axle [2] + # We do account for such offsets, but the overall estimations are not perfect. + # [1] https://en.wikipedia.org/wiki/Volkswagen_Passat_(B6) + # [2] https://www.cvlibs.net/datasets/kitti-360/documentation.php return VehicleParameters( vehicle_name="kitti360_vw_passat", width=1.820, @@ -111,8 +145,9 @@ def get_kitti360_vw_passat_parameters() -> VehicleParameters: def get_av2_ford_fusion_hybrid_parameters() -> VehicleParameters: - # NOTE: Parameters are estimated from the vehicle model. - # https://en.wikipedia.org/wiki/Ford_Fusion_Hybrid#Second_generation + """Helper function to get Argoverse 2 Ford Fusion Hybrid vehicle parameters.""" + # NOTE: Parameters are estimated from the vehicle model [1] and LiDAR point cloud. + # [1] https://en.wikipedia.org/wiki/Ford_Fusion_Hybrid#Second_generation # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/tests/unit/map/test_map_api.py#L375 return VehicleParameters( vehicle_name="av2_ford_fusion_hybrid", @@ -126,6 +161,10 @@ def get_av2_ford_fusion_hybrid_parameters() -> VehicleParameters: def get_pandaset_chrysler_pacifica_parameters() -> VehicleParameters: + """Helper function to get PandaSet Chrysler Pacifica vehicle parameters.""" + # NOTE: Some parameters are available in PandaSet [1], others are estimated based on the vehicle model [2]. + # [1] https://arxiv.org/pdf/2112.12610 (Figure 3 (a)) + # [2] https://en.wikipedia.org/wiki/Chrysler_Pacifica_(minivan) return VehicleParameters( vehicle_name="pandaset_chrysler_pacifica", width=2.297, @@ -137,9 +176,9 @@ def get_pandaset_chrysler_pacifica_parameters() -> VehicleParameters: ) -def center_se3_to_rear_axle_se3(center_se3: StateSE3, vehicle_parameters: VehicleParameters) -> StateSE3: - """ - Converts a center state to a rear axle state. +def center_se3_to_rear_axle_se3(center_se3: PoseSE3, vehicle_parameters: VehicleParameters) -> PoseSE3: + """Converts a center state to a rear axle state. + :param center_se3: The center state. :param vehicle_parameters: The vehicle parameters. :return: The rear axle state. @@ -154,9 +193,9 @@ def center_se3_to_rear_axle_se3(center_se3: StateSE3, vehicle_parameters: Vehicl ) -def rear_axle_se3_to_center_se3(rear_axle_se3: StateSE3, vehicle_parameters: VehicleParameters) -> StateSE3: - """ - Converts a rear axle state to a center state. +def rear_axle_se3_to_center_se3(rear_axle_se3: PoseSE3, vehicle_parameters: VehicleParameters) -> PoseSE3: + """Converts a rear axle state to a center state. + :param rear_axle_se3: The rear axle state. :param vehicle_parameters: The vehicle parameters. :return: The center state. @@ -171,9 +210,9 @@ def rear_axle_se3_to_center_se3(rear_axle_se3: StateSE3, vehicle_parameters: Veh ) -def center_se2_to_rear_axle_se2(center_se2: StateSE2, vehicle_parameters: VehicleParameters) -> StateSE2: - """ - Converts a center state to a rear axle state in 2D. +def center_se2_to_rear_axle_se2(center_se2: PoseSE2, vehicle_parameters: VehicleParameters) -> PoseSE2: + """Converts a center state to a rear axle state in 2D. + :param center_se2: The center state in 2D. :param vehicle_parameters: The vehicle parameters. :return: The rear axle state in 2D. @@ -181,9 +220,9 @@ def center_se2_to_rear_axle_se2(center_se2: StateSE2, vehicle_parameters: Vehicl return translate_se2_along_body_frame(center_se2, Vector2D(-vehicle_parameters.rear_axle_to_center_longitudinal, 0)) -def rear_axle_se2_to_center_se2(rear_axle_se2: StateSE2, vehicle_parameters: VehicleParameters) -> StateSE2: - """ - Converts a rear axle state to a center state in 2D. +def rear_axle_se2_to_center_se2(rear_axle_se2: PoseSE2, vehicle_parameters: VehicleParameters) -> PoseSE2: + """Converts a rear axle state to a center state in 2D. + :param rear_axle_se2: The rear axle state in 2D. :param vehicle_parameters: The vehicle parameters. :return: The center state in 2D. diff --git a/src/py123d/geometry/__init__.py b/src/py123d/geometry/__init__.py index c391b29a..0de0f86f 100644 --- a/src/py123d/geometry/__init__.py +++ b/src/py123d/geometry/__init__.py @@ -6,17 +6,17 @@ Corners2DIndex, Corners3DIndex, EulerAnglesIndex, - EulerStateSE3Index, + EulerPoseSE3Index, QuaternionIndex, - StateSE2Index, - StateSE3Index, + PoseSE2Index, + PoseSE3Index, Vector2DIndex, Vector3DIndex, ) from py123d.geometry.point import Point2D, Point3D from py123d.geometry.vector import Vector2D, Vector3D from py123d.geometry.rotation import EulerAngles, Quaternion -from py123d.geometry.se import EulerStateSE3, StateSE2, StateSE3 +from py123d.geometry.pose import EulerPoseSE3, PoseSE2, PoseSE3 from py123d.geometry.bounding_box import BoundingBoxSE2, BoundingBoxSE3 from py123d.geometry.polyline import Polyline2D, Polyline3D, PolylineSE2 from py123d.geometry.occupancy_map import OccupancyMap2D diff --git a/src/py123d/geometry/bounding_box.py b/src/py123d/geometry/bounding_box.py index bc3e1b73..3ac375fc 100644 --- a/src/py123d/geometry/bounding_box.py +++ b/src/py123d/geometry/bounding_box.py @@ -1,26 +1,25 @@ from __future__ import annotations -from ast import Dict -from typing import Union +from typing import Dict, Union import numpy as np import numpy.typing as npt import shapely.geometry as geom -from py123d.common.utils.mixin import ArrayMixin +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr from py123d.geometry.geometry_index import BoundingBoxSE2Index, BoundingBoxSE3Index, Corners2DIndex, Corners3DIndex from py123d.geometry.point import Point2D, Point3D -from py123d.geometry.se import StateSE2, StateSE3 +from py123d.geometry.pose import PoseSE2, PoseSE3 from py123d.geometry.utils.bounding_box_utils import bbse2_array_to_corners_array, bbse3_array_to_corners_array class BoundingBoxSE2(ArrayMixin): """ - Rotated bounding box in 2D defined by center (StateSE2), length and width. + Rotated bounding box in 2D defined by a center :class:`~py123d.geometry.PoseSE2`, length and width. Example: - >>> from py123d.geometry import StateSE2 - >>> bbox = BoundingBoxSE2(center=StateSE2(1.0, 2.0, 0.5), length=4.0, width=2.0) + >>> from py123d.geometry import PoseSE2, BoundingBoxSE2 + >>> bbox = BoundingBoxSE2(center_se2=PoseSE2(1.0, 2.0, 0.5), length=4.0, width=2.0) >>> bbox.array array([1. , 2. , 0.5, 4. , 2. ]) >>> bbox.corners_array.shape @@ -29,29 +28,30 @@ class BoundingBoxSE2(ArrayMixin): 8.0 """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] - def __init__(self, center: StateSE2, length: float, width: float): - """Initialize BoundingBoxSE2 with center (StateSE2), length and width. + def __init__(self, center_se2: PoseSE2, length: float, width: float): + """Initialize :class:`BoundingBoxSE2` with :class:`~py123d.geometry.PoseSE2` center, length and width. - :param center: Center of the bounding box as a StateSE2 instance. + :param center_se2: Center of the bounding box as a :class:`~py123d.geometry.PoseSE2` instance. :param length: Length of the bounding box along the x-axis in the local frame. :param width: Width of the bounding box along the y-axis in the local frame. """ array = np.zeros(len(BoundingBoxSE2Index), dtype=np.float64) - array[BoundingBoxSE2Index.SE2] = center.array + array[BoundingBoxSE2Index.SE2] = center_se2.array array[BoundingBoxSE2Index.LENGTH] = length array[BoundingBoxSE2Index.WIDTH] = width object.__setattr__(self, "_array", array) @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> BoundingBoxSE2: - """Create a BoundingBoxSE2 from a numpy array. + """Create a :class:`BoundingBoxSE2` from a (5,) numpy array, \ + indexed by :class:`~py123d.geometry.BoundingBoxSE2Index`. - :param array: A 1D numpy array containing the bounding box parameters, indexed by \ - :class:`~py123d.geometry.BoundingBoxSE2Index`. + :param array: A 1D numpy array containing the bounding box parameters. :param copy: Whether to copy the input array. Defaults to True. - :return: A BoundingBoxSE2 instance. + :return: A :class:`BoundingBoxSE2` instance. """ assert array.ndim == 1 assert array.shape[-1] == len(BoundingBoxSE2Index) @@ -60,87 +60,62 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Boundi return instance @property - def center(self) -> StateSE2: - """The center of the bounding box as a StateSE2 instance. - - :return: The center of the bounding box as a StateSE2 instance. - """ - return StateSE2.from_array(self._array[BoundingBoxSE2Index.SE2]) - - @property - def center_se2(self) -> StateSE2: - """The center of the bounding box as a StateSE2 instance. - - :return: The center of the bounding box as a StateSE2 instance. - """ - return self.center + def center_se2(self) -> PoseSE2: + """The center of the bounding box as a :class:`~py123d.geometry.PoseSE2` instance.""" + return PoseSE2.from_array(self._array[BoundingBoxSE2Index.SE2]) @property def length(self) -> float: - """The length of the bounding box along the x-axis in the local frame. - - :return: The length of the bounding box. - """ + """Length of the bounding box along the x-axis in the local frame.""" return self._array[BoundingBoxSE2Index.LENGTH] @property def width(self) -> float: - """The width of the bounding box along the y-axis in the local frame. - - :return: The width of the bounding box. - """ + """Width of the bounding box along the y-axis in the local frame.""" return self._array[BoundingBoxSE2Index.WIDTH] @property def array(self) -> npt.NDArray[np.float64]: - """Converts the BoundingBoxSE2 instance to a numpy array, indexed by :class:`~py123d.geometry.BoundingBoxSE2Index`. - - :return: A numpy array of shape (5,) containing the bounding box parameters [x, y, yaw, length, width]. - """ + """The numpy array representation of shape (5,), indexed by :class:`~py123d.geometry.BoundingBoxSE2Index`.""" return self._array - @property - def shapely_polygon(self) -> geom.Polygon: - """Return a Shapely polygon representation of the bounding box. - - :return: A Shapely polygon representing the bounding box. - """ - return geom.Polygon(self.corners_array) - - @property - def bounding_box_se2(self) -> BoundingBoxSE2: - """Returns bounding box itself for polymorphism. - - :return: A BoundingBoxSE2 instance representing the 2D bounding box. - """ - return self - @property def corners_array(self) -> npt.NDArray[np.float64]: - """Returns the corner points of the bounding box as a numpy array. - - :return: A numpy array of shape (4, 2) containing the corner points of the bounding box, \ - indexed by :class:`~py123d.geometry.Corners2DIndex` and :class:`~py123d.geometry.Point2DIndex`. + """The corner points of the bounding box as a numpy array of shape (4, 2), indexed by \ + :class:`~py123d.geometry.Corners2DIndex` and :class:`~py123d.geometry.Point2DIndex`, respectively. """ return bbse2_array_to_corners_array(self.array) @property def corners_dict(self) -> Dict[Corners2DIndex, Point2D]: - """Returns the corner points of the bounding box as a dictionary. - - :return: A dictionary mapping :class:`~py123d.geometry.Corners2DIndex` to :class:`~py123d.geometry.Point2D` instances. + """Dictionary of corner points of the bounding box, mapping :class:`~py123d.geometry.Corners2DIndex` to \ + :class:`~py123d.geometry.Point2D` instances. """ corners_array = self.corners_array return {index: Point2D.from_array(corners_array[index]) for index in Corners2DIndex} + @property + def shapely_polygon(self) -> geom.Polygon: + """The shapely polygon representation of the bounding box.""" + return geom.Polygon(self.corners_array) + + @property + def bounding_box_se2(self) -> BoundingBoxSE2: + """The :class:`BoundingBoxSE2` instance itself.""" + return self + + def __repr__(self) -> str: + """String representation of :class:`BoundingBoxSE2`.""" + return indexed_array_repr(self, BoundingBoxSE2Index) + class BoundingBoxSE3(ArrayMixin): """ - Rotated bounding box in 3D defined by center with quaternion rotation (StateSE3), length, width and height. + Rotated bounding box in 3D defined by center with quaternion rotation (PoseSE3), length, width and height. Example: - >>> from py123d.geometry import StateSE3 - >>> bbox = BoundingBoxSE3(center=StateSE3(1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0), length=4.0, width=2.0, height=1.5) + >>> from py123d.geometry import PoseSE3, BoundingBoxSE3 + >>> bbox = BoundingBoxSE3(center_se3=PoseSE3(1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0), length=4.0, width=2.0, height=1.5) >>> bbox.array array([1. , 2. , 3. , 1. , 0. , 0. , 0. , 4. , 2. , 1.5]) >>> bbox.bounding_box_se2.array @@ -149,18 +124,19 @@ class BoundingBoxSE3(ArrayMixin): 8.0 """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] - def __init__(self, center: StateSE3, length: float, width: float, height: float): - """Initialize BoundingBoxSE3 with center (StateSE3), length, width and height. + def __init__(self, center_se3: PoseSE3, length: float, width: float, height: float): + """Initialize :class:`BoundingBoxSE3` with :class:`~py123d.geometry.PoseSE3` center, length, width and height. - :param center: Center of the bounding box as a StateSE3 instance. + :param center_se3: Center of the bounding box as a :class:`~py123d.geometry.PoseSE3` instance. :param length: Length of the bounding box along the x-axis in the local frame. :param width: Width of the bounding box along the y-axis in the local frame. :param height: Height of the bounding box along the z-axis in the local frame. """ array = np.zeros(len(BoundingBoxSE3Index), dtype=np.float64) - array[BoundingBoxSE3Index.STATE_SE3] = center.array + array[BoundingBoxSE3Index.SE3] = center_se3.array array[BoundingBoxSE3Index.LENGTH] = length array[BoundingBoxSE3Index.WIDTH] = width array[BoundingBoxSE3Index.HEIGHT] = height @@ -168,10 +144,10 @@ def __init__(self, center: StateSE3, length: float, width: float, height: float) @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> BoundingBoxSE3: - """Create a BoundingBoxSE3 from a numpy array. + """Create a :class:`BoundingBoxSE3` from a (10,) numpy array, \ + indexed by :class:`~py123d.geometry.BoundingBoxSE3Index`. - :param array: A 1D numpy array containing the bounding box parameters, indexed by \ - :class:`~py123d.geometry.BoundingBoxSE3Index`. + :param array: A (10,) numpy array containing the bounding box parameters. :param copy: Whether to copy the input array. Defaults to True. :return: A BoundingBoxSE3 instance. """ @@ -182,100 +158,67 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Boundi return instance @property - def center(self) -> StateSE3: - """The center of the bounding box as a StateSE3 instance. - - :return: The center of the bounding box as a StateSE3 instance. - """ - return StateSE3.from_array(self._array[BoundingBoxSE3Index.STATE_SE3]) + def center_se3(self) -> PoseSE3: + """The center of the bounding box as a :class:`~py123d.geometry.PoseSE3` instance.""" + return PoseSE3.from_array(self._array[BoundingBoxSE3Index.SE3]) @property - def center_se3(self) -> StateSE3: - """The center of the bounding box as a StateSE3 instance. - - :return: The center of the bounding box as a StateSE3 instance. - """ - return self.center - - @property - def center_se2(self) -> StateSE2: - """The center of the bounding box as a StateSE2 instance. - - :return: The center of the bounding box as a StateSE2 instance. - """ - return self.center_se3.state_se2 + def center_se2(self) -> PoseSE2: + """The center of the bounding box as a :class:`~py123d.geometry.PoseSE2` instance.""" + return self.center_se3.pose_se2 @property def length(self) -> float: - """The length of the bounding box along the x-axis in the local frame. - - :return: The length of the bounding box. - """ + """The length of the bounding box along the x-axis in the local frame.""" return self._array[BoundingBoxSE3Index.LENGTH] @property def width(self) -> float: - """The width of the bounding box along the y-axis in the local frame. - - :return: The width of the bounding box. - """ + """The width of the bounding box along the y-axis in the local frame.""" return self._array[BoundingBoxSE3Index.WIDTH] @property def height(self) -> float: - """The height of the bounding box along the z-axis in the local frame. - - :return: The height of the bounding box. - """ + """The height of the bounding box along the z-axis in the local frame.""" return self._array[BoundingBoxSE3Index.HEIGHT] @property def array(self) -> npt.NDArray[np.float64]: - """Convert the BoundingBoxSE3 instance to a numpy array. - - :return: A 1D numpy array containing the bounding box parameters, indexed by \ - :class:`~py123d.geometry.BoundingBoxSE3Index`. - """ + """The numpy array representation of shape (10,), indexed by :class:`~py123d.geometry.BoundingBoxSE3Index`.""" return self._array @property def bounding_box_se2(self) -> BoundingBoxSE2: - """Converts the 3D bounding box to a 2D bounding box by dropping the z, roll and pitch components. - - :return: A BoundingBoxSE2 instance. - """ + """The SE2 projection :class:`~py123d.geometry.BoundingBoxSE2` of the bounding box.""" return BoundingBoxSE2( - center=self.center_se2, + center_se2=self.center_se2, length=self.length, width=self.width, ) - @property - def shapely_polygon(self) -> geom.Polygon: - """Return a Shapely polygon representation of the 2D projection of the bounding box. - - :return: A shapely polygon representing the 2D bounding box. - """ - return self.bounding_box_se2.shapely_polygon - @property def corners_array(self) -> npt.NDArray[np.float64]: - """Returns the corner points of the bounding box as a numpy array, shape (8, 3). - - :return: A numpy array of shape (8, 3) containing the corner points of the bounding box, \ - indexed by :class:`~py123d.geometry.Corners3DIndex` and :class:`~py123d.geometry.Point3DIndex`. + """The corner points of the bounding box as a numpy array of shape (8, 3), indexed by \ + :class:`~py123d.geometry.Corners3DIndex` and :class:`~py123d.geometry.Point3DIndex`, respectively. """ return bbse3_array_to_corners_array(self.array) @property def corners_dict(self) -> Dict[Corners3DIndex, Point3D]: - """Returns the corner points of the bounding box as a dictionary. - - :return: A dictionary mapping :class:`~py123d.geometry.Corners3DIndex` to \ + """Dictionary of corner points of the bounding box, mapping :class:`~py123d.geometry.Corners3DIndex` to \ :class:`~py123d.geometry.Point3D` instances. """ corners_array = self.corners_array return {index: Point3D.from_array(corners_array[index]) for index in Corners3DIndex} + @property + def shapely_polygon(self) -> geom.Polygon: + """The shapely polygon representation of the SE2 projection of the bounding box.""" + return self.bounding_box_se2.shapely_polygon + + def __repr__(self) -> str: + """String representation of :class:`BoundingBoxSE3`.""" + return indexed_array_repr(self, BoundingBoxSE3Index) + BoundingBox = Union[BoundingBoxSE2, BoundingBoxSE3] diff --git a/src/py123d/geometry/geometry_index.py b/src/py123d/geometry/geometry_index.py index 5d596f77..2acb5f07 100644 --- a/src/py123d/geometry/geometry_index.py +++ b/src/py123d/geometry/geometry_index.py @@ -4,35 +4,31 @@ class Point2DIndex(IntEnum): - """ - Indexes array-like representations of 2D points (x,y). - """ + """Indexing enum for array-like representations of 2D points (x,y).""" X = 0 Y = 1 @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) class Vector2DIndex(IntEnum): - """ - Indexes array-like representations of 2D vectors (x,y). - """ + """Indexing enum for array-like representations of 2D vectors (x,y).""" X = 0 Y = 1 @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) vector components.""" return slice(cls.X, cls.Y + 1) -class StateSE2Index(IntEnum): - """ - Indexes array-like representations of SE2 states (x,y,yaw). - """ +class PoseSE2Index(IntEnum): + """Indexing enum for array-like representations of SE2 poses (x,y,yaw).""" X = 0 Y = 1 @@ -40,13 +36,17 @@ class StateSE2Index(IntEnum): @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) + @classproperty + def SE2(cls) -> slice: + """Slice for accessing (x,y,yaw) pose components.""" + return slice(cls.X, cls.YAW + 1) + class Point3DIndex(IntEnum): - """ - Indexes array-like representations of 3D points (x,y,z). - """ + """Indexing enum for array-like representations of 3D points (x,y,z).""" X = 0 Y = 1 @@ -54,17 +54,17 @@ class Point3DIndex(IntEnum): @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) @classproperty def XYZ(cls) -> slice: + """Slice for accessing (x,y,z) coordinates.""" return slice(cls.X, cls.Z + 1) class Vector3DIndex(IntEnum): - """ - Indexes array-like representations of 3D vectors (x,y,z). - """ + """Indexing enum for array-like representations of 3D vectors (x,y,z).""" X = 0 Y = 1 @@ -72,13 +72,12 @@ class Vector3DIndex(IntEnum): @classproperty def XYZ(cls) -> slice: + """Slice for accessing (x,y,z) vector components.""" return slice(cls.X, cls.Z + 1) class EulerAnglesIndex(IntEnum): - """ - Indexes array-like representations of Euler angles (roll,pitch,yaw). - """ + """Indexing enum for array-like representations of Euler angles (roll,pitch,yaw).""" ROLL = 0 PITCH = 1 @@ -86,9 +85,7 @@ class EulerAnglesIndex(IntEnum): class QuaternionIndex(IntEnum): - """ - Indexes array-like representations of quaternions (qw,qx,qy,qz). - """ + """Indexing enum for array-like representations of quaternions (qw,qx,qy,qz), scalar-first.""" QW = 0 QX = 1 @@ -97,17 +94,22 @@ class QuaternionIndex(IntEnum): @classproperty def SCALAR(cls) -> int: + """Index for the scalar part of the quaternion.""" return cls.QW @classproperty def VECTOR(cls) -> slice: + """Slice for accessing the imaginary vector part of the quaternion.""" return slice(cls.QX, cls.QZ + 1) -class EulerStateSE3Index(IntEnum): - """ - Indexes array-like representations of SE3 states (x,y,z,roll,pitch,yaw). - TODO: Use quaternions for rotation. +class EulerPoseSE3Index(IntEnum): + """Indexing enum for array-like representations of SE3 states with Euler angles (x,y,z,roll,pitch,yaw). + + Notes + ----- + Representing a pose with Euler angles is deprecated but left in for testing purposes. + """ X = 0 @@ -119,21 +121,22 @@ class EulerStateSE3Index(IntEnum): @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) @classproperty def XYZ(cls) -> slice: + """Slice for accessing (x,y,z) coordinates.""" return slice(cls.X, cls.Z + 1) @classproperty def EULER_ANGLES(cls) -> slice: + """Slice for accessing (roll,pitch,yaw) Euler angles.""" return slice(cls.ROLL, cls.YAW + 1) -class StateSE3Index(IntEnum): - """ - Indexes array-like representations of SE3 states with quaternions (x,y,z,qw,qx,qy,qz). - """ +class PoseSE3Index(IntEnum): + """Indexing enum for array-like representations of SE3 poses (x,y,z,qw,qx,qy,qz).""" X = 0 Y = 1 @@ -145,28 +148,35 @@ class StateSE3Index(IntEnum): @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) @classproperty def XYZ(cls) -> slice: + """Slice for accessing (x,y,z) coordinates.""" return slice(cls.X, cls.Z + 1) @classproperty def QUATERNION(cls) -> slice: + """Slice for accessing (qw,qx,qy,qz) quaternion components.""" return slice(cls.QW, cls.QZ + 1) @classproperty def SCALAR(cls) -> slice: + """Slice for accessing the scalar part of the quaternion.""" return slice(cls.QW, cls.QW + 1) @classproperty def VECTOR(cls) -> slice: + """Slice for accessing the vector part of the quaternion.""" return slice(cls.QX, cls.QZ + 1) class BoundingBoxSE2Index(IntEnum): - """ - Indexes array-like representations of rotated 2D bounding boxes (x,y,yaw,length,width). + """Indexing enum for array-like representations of bounding boxes in SE2 + - center point (x,y). + - yaw rotation. + - extent (length,width). """ X = 0 @@ -177,21 +187,22 @@ class BoundingBoxSE2Index(IntEnum): @classproperty def XY(cls) -> slice: + """Slice for accessing (x,y) coordinates.""" return slice(cls.X, cls.Y + 1) @classproperty def SE2(cls) -> slice: + """Slice for accessing (x,y,yaw) SE2 representation.""" return slice(cls.X, cls.YAW + 1) @classproperty def EXTENT(cls) -> slice: + """Slice for accessing (length,width) extent.""" return slice(cls.LENGTH, cls.WIDTH + 1) class Corners2DIndex(IntEnum): - """ - Indexes the corners of a BoundingBoxSE2 in the order: front-left, front-right, back-right, back-left. - """ + """Indexes the corners of a bounding boxes in SE2 in the order: front-left, front-right, back-right, back-left.""" FRONT_LEFT = 0 FRONT_RIGHT = 1 @@ -202,8 +213,8 @@ class Corners2DIndex(IntEnum): class BoundingBoxSE3Index(IntEnum): """ Indexes array-like representations of rotated 3D bounding boxes - - center (x,y,z). - - rotation (qw,qx,qy,qz). + - center point (x,y,z). + - quaternion rotation (qw,qx,qy,qz). - extent (length,width,height). """ @@ -220,34 +231,40 @@ class BoundingBoxSE3Index(IntEnum): @classproperty def XYZ(cls) -> slice: + """Slice for accessing (x,y,z) coordinates.""" return slice(cls.X, cls.Z + 1) @classproperty - def STATE_SE3(cls) -> slice: + def SE3(cls) -> slice: + """Slice for accessing the full SE3 pose representation.""" return slice(cls.X, cls.QZ + 1) @classproperty def QUATERNION(cls) -> slice: + """Slice for accessing (qw,qx,qy,qz) quaternion components.""" return slice(cls.QW, cls.QZ + 1) @classproperty def EXTENT(cls) -> slice: + """Slice for accessing (length,width,height) extent.""" return slice(cls.LENGTH, cls.HEIGHT + 1) @classproperty def SCALAR(cls) -> slice: + """Slice for accessing the scalar part of the quaternion.""" return slice(cls.QW, cls.QW + 1) @classproperty def VECTOR(cls) -> slice: + """Slice for accessing the vector part of the quaternion.""" return slice(cls.QX, cls.QZ + 1) class Corners3DIndex(IntEnum): """ - Indexes the corners of a BoundingBoxSE3 in the order: - front-left-bottom, front-right-bottom, back-right-bottom, back-left-bottom, - front-left-top, front-right-top, back-right-top, back-left-top. + Indexes the corners of a BoundingBoxSE3 in the order: \ + front-left-bottom, front-right-bottom, back-right-bottom, back-left-bottom,\ + front-left-top, front-right-top, back-right-top, back-left-top. """ FRONT_LEFT_BOTTOM = 0 @@ -261,8 +278,10 @@ class Corners3DIndex(IntEnum): @classproperty def BOTTOM(cls) -> slice: + """Slice for accessing the four bottom corners.""" return slice(cls.FRONT_LEFT_BOTTOM, cls.BACK_LEFT_BOTTOM + 1) @classproperty def TOP(cls) -> slice: + """Slice for accessing the four top corners.""" return slice(cls.FRONT_LEFT_TOP, cls.BACK_LEFT_TOP + 1) diff --git a/src/py123d/geometry/occupancy_map.py b/src/py123d/geometry/occupancy_map.py index a0e4021a..8ee6ed32 100644 --- a/src/py123d/geometry/occupancy_map.py +++ b/src/py123d/geometry/occupancy_map.py @@ -12,10 +12,12 @@ class OccupancyMap2D: + """Class to represent a 2D occupancy map of shapely geometries using an str-tree for efficient spatial queries.""" + def __init__( self, geometries: Sequence[BaseGeometry], - ids: Optional[Union[List[str], List[int]]] = None, + ids: Optional[Union[Sequence[str], Sequence[int]]] = None, node_capacity: int = 10, ): """Constructs a 2D occupancy map of shapely geometries using an str-tree for efficient spatial queries. @@ -26,10 +28,10 @@ def __init__( """ assert ids is None or len(ids) == len(geometries), "Length of ids must match length of geometries" + if ids is not None: + assert all(isinstance(id, (str, int)) for id in ids), "IDs must be either strings or integers" - self._ids: Union[List[str], List[int]] = ( - ids if ids is not None else [str(idx) for idx in range(len(geometries))] - ) + self._ids: Sequence[Union[str, int]] = ids if ids is not None else [str(idx) for idx in range(len(geometries))] self._id_to_idx: Dict[Union[str, int], int] = {id: idx for idx, id in enumerate(self._ids)} self._geometries = geometries @@ -37,7 +39,11 @@ def __init__( self._str_tree = STRtree(self._geometries, node_capacity) @classmethod - def from_dict(cls, geometry_dict: Dict[Union[str, int], BaseGeometry], node_capacity: int = 10) -> OccupancyMap2D: + def from_dict( + cls, + geometry_dict: Union[Dict[str, BaseGeometry], Dict[int, BaseGeometry]], + node_capacity: int = 10, + ) -> OccupancyMap2D: """Constructs a 2D occupancy map from a dictionary of geometries. :param geometry_dict: Dictionary mapping geometry identifiers to shapely geometries @@ -102,6 +108,7 @@ def query( Literal[ "intersects", "within", + "dwithin", "contains", "overlaps", "crosses", @@ -153,7 +160,10 @@ def query_nearest( def contains_vectorized(self, points: npt.NDArray[np.float64]) -> npt.NDArray[np.bool_]: """Determines wether input-points are in geometries (i.e. polygons) of the occupancy map. - NOTE: This function can be significantly faster than using the str-tree, if the number of geometries is + + Notes + ----- + This function can be significantly faster than using the str-tree, if the number of geometries is relatively small compared to the number of input-points. :param points: array of shape (num_points, 2), indexed by :class:`~py123d.geometry.Point2DIndex`. diff --git a/src/py123d/geometry/point.py b/src/py123d/geometry/point.py index 571567be..14cf74c1 100644 --- a/src/py123d/geometry/point.py +++ b/src/py123d/geometry/point.py @@ -1,22 +1,34 @@ from __future__ import annotations -from typing import Iterable - import numpy as np import numpy.typing as npt import shapely.geometry as geom -from py123d.common.utils.mixin import ArrayMixin +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr from py123d.geometry.geometry_index import Point2DIndex, Point3DIndex class Point2D(ArrayMixin): - """Class to represents 2D points.""" - + """Class presenting a 2D point. + + Example: + >>> from py123d.geometry import Point2D + >>> point_2d = Point2D(1.0, 2.0) + >>> point_2d.x, point_2d.y + (1.0, 2.0) + >>> point_2d.array + array([1., 2.]) + """ + + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, x: float, y: float): - """Initialize StateSE2 with x, y, yaw coordinates.""" + """Initializes :class:`Point2D` with x, y coordinates. + + :param x: The x coordinate. + :param y: The y coordinate. + """ array = np.zeros(len(Point2DIndex), dtype=np.float64) array[Point2DIndex.X] = x array[Point2DIndex.Y] = y @@ -24,10 +36,9 @@ def __init__(self, x: float, y: float): @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Point2D: - """Constructs a Point2D from a numpy array. + """Creates a :class:`Point2D` from a (2,) shaped numpy array, indexed by :class:`~py123d.geometry.Point2DIndex`. - :param array: Array of shape (2,) representing the point coordinates [x, y], indexed by \ - :class:`~py123d.geometry.Point2DIndex`. + :param array: A (2,) shaped numpy array representing the point coordinates (x,y). :param copy: Whether to copy the input array. Defaults to True. :return: A Point2D instance. """ @@ -39,53 +50,57 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Point2 @property def x(self) -> float: - """The x coordinate of the point. - - :return: The x coordinate of the point. - """ + """The x coordinate of the point.""" return self._array[Point2DIndex.X] @property def y(self) -> float: - """The y coordinate of the point. - - :return: The y coordinate of the point. - """ + """The y coordinate of the point.""" return self._array[Point2DIndex.Y] @property def array(self) -> npt.NDArray[np.float64]: - """The array representation of the point. - - :return: A numpy array of shape (2,) containing the point coordinates [x, y], indexed by \ - :class:`~py123d.geometry.Point2DIndex`. - """ + """The array representation of shape (2,), indexed by :class:`~py123d.geometry.Point2DIndex`.""" return self._array @property def shapely_point(self) -> geom.Point: - """The Shapely Point representation of the 2D point. - - :return: A Shapely Point representation of the 2D point. - """ + """The shapely point representation of the 2D point.""" return geom.Point(self.x, self.y) - def __iter__(self) -> Iterable[float]: - """Iterator over point coordinates.""" - return iter((self.x, self.y)) + @property + def point_2d(self) -> Point2D: + """Returns the :class:`Point2D` instance itself.""" + return self - def __hash__(self) -> int: - """Hash method""" - return hash((self.x, self.y)) + def __repr__(self) -> str: + """String representation of :class:`Point2D`.""" + return indexed_array_repr(self, Point2DIndex) class Point3D(ArrayMixin): - """Class to represents 3D points.""" + """Class presenting a 3D point. + + Example: + >>> from py123d.geometry import Point3D + >>> point_3d = Point3D(1.0, 2.0, 3.0) + >>> point_3d.x, point_3d.y, point_3d.z + (1.0, 2.0, 3.0) + >>> point_3d.array + array([1., 2., 3.]) + + """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, x: float, y: float, z: float): - """Initialize Point3D with x, y, z coordinates.""" + """Initializes :class:`Point3D` with x, y, z coordinates. + + :param x: The x coordinate. + :param y: The y coordinate. + :param z: The z coordinate. + """ array = np.zeros(len(Point3DIndex), dtype=np.float64) array[Point3DIndex.X] = x array[Point3DIndex.Y] = y @@ -94,12 +109,10 @@ def __init__(self, x: float, y: float, z: float): @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Point3D: - """Constructs a Point3D from a numpy array. + """Creates a :class:`Point3D` from a (3,) shaped numpy array, indexed by :class:`~py123d.geometry.Point3DIndex`. - :param array: Array of shape (3,) representing the point coordinates [x, y, z], indexed by \ - :class:`~py123d.geometry.Point3DIndex`. + :param array: A (3,) shaped numpy array representing the point coordinates (x,y,z). :param copy: Whether to copy the input array. Defaults to True. - :return: A Point3D instance. """ assert array.ndim == 1 assert array.shape[0] == len(Point3DIndex) @@ -107,60 +120,41 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Point3 object.__setattr__(instance, "_array", array.copy() if copy else array) return instance - @property - def array(self) -> npt.NDArray[np.float64]: - """The array representation of the point. - - :return: A numpy array of shape (3,) containing the point coordinates [x, y, z], indexed by \ - :class:`~py123d.geometry.Point3DIndex`. - """ - return self._array - @property def x(self) -> float: - """The x coordinate of the point. - - :return: The x coordinate of the point. - """ + """The x coordinate of the point.""" return self._array[Point3DIndex.X] @property def y(self) -> float: - """The y coordinate of the point. - - :return: The y coordinate of the point. - """ + """The y coordinate of the point.""" return self._array[Point3DIndex.Y] @property def z(self) -> float: - """The z coordinate of the point. - - :return: The z coordinate of the point. - """ + """The z coordinate of the point.""" return self._array[Point3DIndex.Z] @property - def point_2d(self) -> Point2D: - """The 2D projection of the 3D point. + def array(self) -> npt.NDArray[np.float64]: + """The array representation of shape (3,), indexed by :class:`~py123d.geometry.Point3DIndex`.""" + return self._array - :return: A Point2D instance representing the 2D projection of the 3D point. - """ + @property + def point_3d(self) -> Point3D: + """Returns the :class:`Point3D` instance itself.""" + return self + + @property + def point_2d(self) -> Point2D: + """The 2D projection of the 3D point as a :class:`Point2D` instance.""" return Point2D.from_array(self.array[Point3DIndex.XY], copy=False) @property def shapely_point(self) -> geom.Point: - """The Shapely Point representation of the 3D point. \ - This geometry contains the z-coordinate, but many Shapely operations ignore it. - - :return: A Shapely Point representation of the 3D point. - """ + """The shapely point representation of the 3D point.""" return geom.Point(self.x, self.y, self.z) - def __iter__(self) -> Iterable[float]: - """Iterator over the point coordinates (x, y, z).""" - return iter((self.x, self.y, self.z)) - - def __hash__(self) -> int: - """Hash method""" - return hash((self.x, self.y, self.z)) + def __repr__(self) -> str: + """String representation of :class:`Point3D`.""" + return indexed_array_repr(self, Point3DIndex) diff --git a/src/py123d/geometry/polyline.py b/src/py123d/geometry/polyline.py index 3f3b7f65..5bd61eb7 100644 --- a/src/py123d/geometry/polyline.py +++ b/src/py123d/geometry/polyline.py @@ -1,7 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import List, Optional, Union +from typing import Optional, Union import numpy as np import numpy.typing as npt @@ -10,69 +9,78 @@ from scipy.interpolate import interp1d from py123d.common.utils.mixin import ArrayMixin -from py123d.geometry.geometry_index import Point2DIndex, Point3DIndex, StateSE2Index +from py123d.geometry.geometry_index import Point2DIndex, Point3DIndex, PoseSE2Index from py123d.geometry.point import Point2D, Point3D -from py123d.geometry.se import StateSE2 +from py123d.geometry.pose import PoseSE2 from py123d.geometry.utils.constants import DEFAULT_Z -from py123d.geometry.utils.polyline_utils import get_linestring_yaws, get_path_progress +from py123d.geometry.utils.polyline_utils import get_linestring_yaws, get_path_progress_2d, get_path_progress_3d from py123d.geometry.utils.rotation_utils import normalize_angle -# TODO: Implement PolylineSE3 -# TODO: Benchmark interpolation performance and reconsider reliance on LineString - -@dataclass class Polyline2D(ArrayMixin): - """Represents a interpolatable 2D polyline.""" + """Represents a interpolatable 2D polyline. + + Example: + >>> import numpy as np + >>> from py123d.geometry import Polyline2D + >>> polyline = Polyline2D.from_array(np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]])) + >>> polyline.length + 2.8284271247461903 + >>> polyline.interpolate(np.sqrt(2)) + Point2D(array=[1. 1.]) - linestring: geom.LineString + """ + + __slots__ = ("_linestring",) + _linestring: geom.LineString @classmethod def from_linestring(cls, linestring: geom.LineString) -> Polyline2D: - """Creates a Polyline2D from a Shapely LineString. If the LineString has Z-coordinates, they are ignored. + """Creates a :class:`Polyline2D` from a Shapely LineString. If the LineString has Z-coordinates, they are ignored. - :param linestring: A Shapely LineString object. + :param linestring: A shapely LineString object. :return: A Polyline2D instance. """ if linestring.has_z: - linestring_ = geom_creation.linestrings(*linestring.xy) + linestring_ = geom_creation.linestrings(*linestring.xy) # pyright: ignore[reportUnknownMemberType] else: linestring_ = linestring - return Polyline2D(linestring_) + + instance = object.__new__(cls) + object.__setattr__(instance, "_linestring", linestring_) + return instance @classmethod - def from_array(cls, polyline_array: npt.NDArray[np.float32]) -> Polyline2D: - """Creates a Polyline2D from a numpy array. + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Polyline2D: + """Creates a :class:`Polyline2D` from a (N, 2) or (N, 3) shaped numpy array. \ + Assumes [...,:2] slices are XY coordinates. :param polyline_array: A numpy array of shape (N, 2) or (N, 3), e.g. indexed by \ :class:`~py123d.geometry.Point2DIndex` or :class:`~py123d.geometry.Point3DIndex`. :raises ValueError: If the input array is not of the expected shape. - :return: A Polyline2D instance. + :return: A :class:`Polyline2D` instance. """ - assert polyline_array.ndim == 2 - linestring: Optional[geom.LineString] = None - if polyline_array.shape[-1] == len(Point2DIndex): - linestring = geom_creation.linestrings(polyline_array) - elif polyline_array.shape[-1] == len(Point3DIndex): - linestring = geom_creation.linestrings(polyline_array[:, Point3DIndex.XY]) + assert array.ndim == 2 + linestring_: Optional[geom.LineString] = None + if array.shape[-1] == len(Point2DIndex): + linestring_ = geom.LineString(array) + elif array.shape[-1] == len(Point3DIndex): + linestring_ = geom.LineString(array[:, Point3DIndex.XY]) # pyright: ignore[reportUnknownMemberType] else: raise ValueError("Array must have shape (N, 2) or (N, 3) for Point2D or Point3D respectively.") - return Polyline2D(linestring) - def from_discrete_points(cls, discrete_points: List[Point2D]) -> Polyline2D: - """Creates a Polyline2D from a list of discrete 2D points. + instance = object.__new__(cls) + object.__setattr__(instance, "_linestring", linestring_) + return instance - :param discrete_points: A list of Point2D instances. - :return: A Polyline2D instance. - """ - return Polyline2D.from_array(np.array(discrete_points, dtype=np.float64)) + @property + def linestring(self) -> geom.LineString: + """The shapely LineString representation of the polyline.""" + return self._linestring @property def array(self) -> npt.NDArray[np.float64]: - """Converts the polyline to a numpy array, indexed by :class:`~py123d.geometry.Point2DIndex`. - - :return: A numpy array of shape (N, 2) representing the polyline. - """ + """The numpy array representation of shape (N, 2), indexed by :class:`~py123d.geometry.Point2DIndex`.""" x, y = self.linestring.xy array = np.zeros((len(x), len(Point2DIndex)), dtype=np.float64) array[:, Point2DIndex.X] = x @@ -81,18 +89,12 @@ def array(self) -> npt.NDArray[np.float64]: @property def polyline_se2(self) -> PolylineSE2: - """Converts the 2D polyline to a 2D SE(2) polyline and retrieves the yaw angles. - - :return: A PolylineSE2 instance representing the 2D polyline. - """ - return PolylineSE2.from_linestring(self.linestring) + """The :class:`~py123d.geometry.PolylineSE2` representation of the polyline, with inferred yaw angles.""" + return PolylineSE2.from_linestring(self._linestring) @property def length(self) -> float: - """Returns the length of the polyline. - - :return: The length of the polyline. - """ + """Returns the length of the polyline.""" return self.linestring.length def interpolate( @@ -100,9 +102,9 @@ def interpolate( distances: Union[float, npt.NDArray[np.float64]], normalized: bool = False, ) -> Union[Point2D, npt.NDArray[np.float64]]: - """Interpolates the polyline at the given distances. + """Interpolates the :class:`Polyline2D` at the given distances. - :param distances: The distances at which to interpolate the polyline. + :param distances: Array-like or float distances along the polyline to interpolate. :return: The interpolated point(s) on the polyline. """ @@ -110,13 +112,17 @@ def interpolate( point = self.linestring.interpolate(distances, normalized=normalized) return Point2D(point.x, point.y) else: - distances_ = np.asarray(distances, dtype=np.float64) points = self.linestring.interpolate(distances, normalized=normalized) return np.array([[p.x, p.y] for p in points], dtype=np.float64) def project( self, - point: Union[geom.Point, Point2D, StateSE2, npt.NDArray[np.float64]], + point: Union[ + geom.Point, + Point2D, + PoseSE2, + npt.NDArray[np.float64], + ], normalized: bool = False, ) -> npt.NDArray[np.float64]: """Projects a point onto the polyline and returns the distance along the polyline to the closest point. @@ -125,119 +131,130 @@ def project( :param normalized: Whether to return the normalized distance, defaults to False. :return: The distance along the polyline to the closest point. """ - if isinstance(point, Point2D) or isinstance(point, StateSE2): + if isinstance(point, Point2D) or isinstance(point, PoseSE2): point_ = point.shapely_point elif isinstance(point, geom.Point): point_ = point else: point_ = np.array(point, dtype=np.float64) - return self.linestring.project(point_, normalized=normalized) + return self._linestring.project(point_, normalized=normalized) # type: ignore -@dataclass class PolylineSE2(ArrayMixin): - """Represents a interpolatable SE2 polyline.""" + """Represents a interpolatable SE2 polyline. - _array: npt.NDArray[np.float64] - linestring: Optional[geom.LineString] = None + Example: + >>> import numpy as np + >>> from py123d.geometry import PolylineSE2 + >>> polyline_se2 = PolylineSE2.from_array(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, np.pi/4], [2.0, 0.0, 0.0]])) + >>> polyline_se2.length + 2.8284271247461903 + >>> polyline_se2.interpolate(np.sqrt(2)) + PoseSE2(array=[1. 1. 0.78539816]) - _progress: Optional[npt.NDArray[np.float64]] = None - _interpolator: Optional[interp1d] = None + """ - def __post_init__(self): - assert self._array is not None + __slots__ = ("_array", "_progress", "_linestring") - if self.linestring is None: - self.linestring = geom_creation.linestrings(self._array[..., StateSE2Index.XY]) + def __init__( + self, + array: npt.NDArray[np.float64], + linestring: Optional[geom.LineString] = None, + ): + """Initializes :class:`PolylineSE2` with a numpy array of SE2 states. + + :param array: A numpy array of shape (N, 3) representing SE2 states, indexed by \ + :class:`~py123d.geometry.PoseSE2Index`. + :param linestring: Optional shapely LineString representing the XY path. If not provided,\ + it will be created from the array. + """ - self._array[:, StateSE2Index.YAW] = np.unwrap(self._array[:, StateSE2Index.YAW], axis=0) - self._progress = get_path_progress(self._array) - self._interpolator = interp1d(self._progress, self._array, axis=0, bounds_error=False, fill_value=0.0) + self._array = array + self._array[:, PoseSE2Index.YAW] = np.unwrap(self._array[:, PoseSE2Index.YAW], axis=0) + self._progress = get_path_progress_2d(self._array) + self._linestring = geom.LineString(self._array[..., PoseSE2Index.XY]) if linestring is None else linestring @classmethod def from_linestring(cls, linestring: geom.LineString) -> PolylineSE2: - """Creates a PolylineSE2 from a LineString. This requires computing the yaw angles along the path. + """Creates a :class:`PolylineSE2` from a shapely LineString. \ + The yaw angles are inferred from the LineString coordinates. :param linestring: The LineString to convert. - :return: A PolylineSE2 representing the same path as the LineString. + :return: A :class:`PolylineSE2` representing the same path as the LineString. """ - points_2d = np.array(linestring.coords, dtype=np.float64)[..., StateSE2Index.XY] - se2_array = np.zeros((len(points_2d), len(StateSE2Index)), dtype=np.float64) - se2_array[:, StateSE2Index.XY] = points_2d - se2_array[:, StateSE2Index.YAW] = get_linestring_yaws(linestring) + points_2d_array = np.array(linestring.coords, dtype=np.float64)[..., PoseSE2Index.XY] + se2_array = np.zeros((len(points_2d_array), len(PoseSE2Index)), dtype=np.float64) + se2_array[:, PoseSE2Index.XY] = points_2d_array + se2_array[:, PoseSE2Index.YAW] = get_linestring_yaws(linestring) return PolylineSE2(se2_array, linestring) @classmethod def from_array(cls, polyline_array: npt.NDArray[np.float32]) -> PolylineSE2: - """Creates a PolylineSE2 from a numpy array. + """Creates a :class:`PolylineSE2` from a numpy array. :param polyline_array: The input numpy array representing, either indexed by \ - :class:`~py123d.geometry.Point2DIndex` or :class:`~py123d.geometry.StateSE2Index`. + :class:`~py123d.geometry.Point2DIndex` or :class:`~py123d.geometry.PoseSE2Index`. :raises ValueError: If the input array is not of the expected shape. - :return: A PolylineSE2 representing the same path as the input array. + :return: A :class:`PolylineSE2` representing the same path as the input array. """ assert polyline_array.ndim == 2 if polyline_array.shape[-1] == len(Point2DIndex): - se2_array = np.zeros((len(polyline_array), len(StateSE2Index)), dtype=np.float64) - se2_array[:, StateSE2Index.XY] = polyline_array - se2_array[:, StateSE2Index.YAW] = get_linestring_yaws(geom_creation.linestrings(*polyline_array.T)) - elif polyline_array.shape[-1] == len(StateSE2Index): + se2_array = np.zeros((len(polyline_array), len(PoseSE2Index)), dtype=np.float64) + se2_array[:, PoseSE2Index.XY] = polyline_array + se2_array[:, PoseSE2Index.YAW] = get_linestring_yaws(geom_creation.linestrings(*polyline_array.T)) + elif polyline_array.shape[-1] == len(PoseSE2Index): se2_array = np.array(polyline_array, dtype=np.float64) else: - raise ValueError("Invalid polyline array shape.") + raise ValueError(f"Invalid polyline array shape, expected (N, 2) or (N, 3), got {polyline_array.shape}.") return PolylineSE2(se2_array) - @classmethod - def from_discrete_se2(cls, discrete_se2: List[StateSE2]) -> PolylineSE2: - """Creates a PolylineSE2 from a list of discrete SE2 states. - - :param discrete_se2: The list of discrete SE2 states. - :return: A PolylineSE2 representing the same path as the discrete SE2 states. - """ - return PolylineSE2.from_array(np.array(discrete_se2, dtype=np.float64)) + @property + def linestring(self) -> geom.LineString: + """The shapely LineString representation of the polyline.""" + return self._linestring @property def array(self) -> npt.NDArray[np.float64]: - """Converts the polyline to a numpy array, indexed by :class:`~py123d.geometry.StateSE2Index`. - - :return: A numpy array of shape (N, 3) representing the polyline. - """ + """The numpy array representation of shape (N, 3), indexed by :class:`~py123d.geometry.PoseSE2Index`.""" return self._array @property def length(self) -> float: - """Returns the length of the polyline. - - :return: The length of the polyline. - """ + """Returns the length of the polyline.""" + assert self._progress is not None return float(self._progress[-1]) def interpolate( self, distances: Union[float, npt.NDArray[np.float64]], normalized: bool = False, - ) -> Union[StateSE2, npt.NDArray[np.float64]]: + ) -> Union[PoseSE2, npt.NDArray[np.float64]]: """Interpolates the polyline at the given distances. :param distances: The distances along the polyline to interpolate. :param normalized: Whether the distances are normalized (0 to 1), defaults to False :return: The interpolated StateSE2 or an array of interpolated states, according to """ - + _interpolator = interp1d(self._progress, self._array, axis=0, bounds_error=False, fill_value=0.0) distances_ = distances * self.length if normalized else distances clipped_distances = np.clip(distances_, 1e-8, self.length) - interpolated_se2_array = self._interpolator(clipped_distances) - interpolated_se2_array[..., StateSE2Index.YAW] = normalize_angle(interpolated_se2_array[..., StateSE2Index.YAW]) + interpolated_se2_array = _interpolator(clipped_distances) + interpolated_se2_array[..., PoseSE2Index.YAW] = normalize_angle(interpolated_se2_array[..., PoseSE2Index.YAW]) if clipped_distances.ndim == 0: - return StateSE2(*interpolated_se2_array) + return PoseSE2(*interpolated_se2_array) else: return interpolated_se2_array def project( self, - point: Union[geom.Point, Point2D, StateSE2, npt.NDArray[np.float64]], + point: Union[ + geom.Point, + Point2D, + PoseSE2, + npt.NDArray[np.float64], + ], normalized: bool = False, ) -> npt.NDArray[np.float64]: """Projects a point onto the polyline and returns the distance along the polyline to the closest point. @@ -246,78 +263,102 @@ def project( :param normalized: Whether to return the normalized distance, defaults to False. :return: The distance along the polyline to the closest point. """ - if isinstance(point, Point2D) or isinstance(point, StateSE2): + if isinstance(point, Point2D) or isinstance(point, PoseSE2): point_ = point.shapely_point elif isinstance(point, geom.Point): point_ = point else: point_ = np.array(point, dtype=np.float64) - return self.linestring.project(point_, normalized=normalized) + return self.linestring.project(point_, normalized=normalized) # type: ignore -@dataclass class Polyline3D(ArrayMixin): - """Represents a interpolatable 3D polyline.""" + """Represents a interpolatable 3D polyline. - linestring: geom.LineString + Example: + >>> import numpy as np + >>> from py123d.geometry import Polyline3D + >>> polyline_3d = Polyline3D.from_array(np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [2.0, 0.0, 0.0]])) + >>> polyline_3d.length + 3.4641016151377544 + >>> polyline_3d.interpolate(np.sqrt(3)) + Point3D(array=[1. 1. 1.]) + + """ + + __slots__ = ("_array", "_progress", "_linestring") + + def __init__(self, array: npt.NDArray[np.float64], linestring: Optional[geom.LineString] = None): + """Initializes :class:`Polyline3D` with a numpy array of 3D points. + + :param array: A numpy array of shape (N, 3) representing 3D points, e.g. indexed by \ + :class:`~py123d.geometry.Point3DIndex`. + :param linestring: Optional shapely LineString representing the 3D path. If not provided,\ + it will be created from the array. + """ + assert len(array.shape) == 2 and array.shape[1] == len(Point3DIndex) + self._array = array + self._progress = get_path_progress_3d(self._array[:, Point3DIndex.XYZ]) + self._linestring = geom.LineString(self._array) if linestring is None else linestring @classmethod def from_linestring(cls, linestring: geom.LineString) -> Polyline3D: - """Creates a Polyline3D from a Shapely LineString. If the LineString does not have Z-coordinates, \ - a default Z-value is added. + """Creates a :class:`Polyline3D` from a shapely LineString. If the LineString does not have Z-coordinates, \ + the coordinate is zero-padded. :param linestring: The input LineString. - :return: A Polyline3D instance. + :return: A :class:`Polyline3D` instance. """ - return ( - Polyline3D(linestring) - if linestring.has_z - else Polyline3D(geom_creation.linestrings(*linestring.xy, z=DEFAULT_Z)) - ) + if linestring.has_z: + linestring_ = linestring + else: + linestring_ = geom_creation.linestrings(*linestring.xy, z=DEFAULT_Z) # type: ignore + array = np.array(linestring_.coords, dtype=np.float64) + return Polyline3D(array, linestring_) @classmethod def from_array(cls, array: npt.NDArray[np.float64]) -> Polyline3D: - """Creates a Polyline3D from a numpy array. + """Creates a :class:`Polyline3D` from a numpy array. :param array: A numpy array of shape (N, 3) representing 3D points, e.g. indexed by \ :class:`~py123d.geometry.Point3DIndex`. - :return: A Polyline3D instance. + :return: A :class:`Polyline3D` instance. """ - assert array.ndim == 2 and array.shape[1] == len(Point3DIndex), "Array must be 3D with shape (N, 3)" - linestring = geom_creation.linestrings(*array.T) - return Polyline3D(linestring) + assert array.ndim == 2, "Array must be 2D with shape (N, 3) or (N, 2)." + if array.shape[1] == len(Point2DIndex): + array = np.hstack((array, np.full((array.shape[0], 1), DEFAULT_Z))) + elif array.shape[1] != len(Point3DIndex): + raise ValueError("Array must have shape (N, 3) for Point3D.") + return Polyline3D(array) @property - def polyline_2d(self) -> Polyline2D: - """Converts the 3D polyline to a 2D polyline by dropping the Z-coordinates. - - :return: A Polyline2D instance. - """ - return Polyline2D(geom_creation.linestrings(*self.linestring.xy)) + def linestring(self) -> geom.LineString: + """The shapely LineString representation of the 3D polyline.""" + if not self._linestring.has_z: + linestring_ = geom_creation.linestrings(*self._array.T) # type: ignore + object.__setattr__(self, "_linestring", linestring_) + return self._linestring @property - def polyline_se2(self) -> PolylineSE2: - """Converts the 3D polyline to a 2D SE(2) polyline. - - :return: A PolylineSE2 instance. - """ - return PolylineSE2.from_linestring(self.linestring) + def array(self) -> npt.NDArray[np.float64]: + """The numpy array representation of shape (N, 3), indexed by :class:`~py123d.geometry.Point3DIndex`.""" + return np.array(self.linestring.coords, dtype=np.float64) @property - def array(self) -> npt.NDArray[np.float64]: - """Converts the 3D polyline to the discrete 3D points. + def polyline_2d(self) -> Polyline2D: + """The :class:`~py123d.geometry.Polyline2D` representation of the 3D polyline.""" + return Polyline2D.from_linestring(geom_creation.linestrings(*self.linestring.xy)) # type: ignore - :return: A numpy array of shape (N, 3), indexed by :class:`~py123d.geometry.Point3DIndex`. - """ - return np.array(self.linestring.coords, dtype=np.float64) + @property + def polyline_se2(self) -> PolylineSE2: + """The :class:`~py123d.geometry.PolylineSE2` representation of the 3D polyline.""" + return PolylineSE2.from_linestring(self.linestring) # type: ignore\ @property def length(self) -> float: - """Returns the length of the 3D polyline. - - :return: The length of the polyline. - """ - return self.linestring.length + """Returns the length of the 3D polyline.""" + array = self.array + return np.linalg.norm(array[:-1, :] - array[1:, :], axis=1).sum() def interpolate( self, @@ -331,13 +372,15 @@ def interpolate( :return: A Point3D instance or a numpy array of shape (N, 3) representing the interpolated points. """ - if isinstance(distances, float) or isinstance(distances, int): - point = self.linestring.interpolate(distances, normalized=normalized) - return Point3D(point.x, point.y, point.z) + _interpolator = interp1d(self._progress, self._array, axis=0, bounds_error=False, fill_value="extrapolate") + distances_ = distances * self.length if normalized else distances + clipped_distances = np.clip(distances_, 1e-8, self.length) + + interpolated_3d_array = _interpolator(clipped_distances) + if clipped_distances.ndim == 0: + return Point3D(*interpolated_3d_array) else: - distances = np.asarray(distances, dtype=np.float64) - points = self.linestring.interpolate(distances, normalized=normalized) - return np.array([[p.x, p.y, p.z] for p in points], dtype=np.float64) + return interpolated_3d_array def project( self, @@ -350,21 +393,10 @@ def project( :param normalized: Whether to return normalized distances, defaults to False. :return: The distance along the polyline to the closest point. """ - if isinstance(point, Point2D) or isinstance(point, StateSE2) or isinstance(point, Point3D): + if isinstance(point, Point2D) or isinstance(point, PoseSE2) or isinstance(point, Point3D): point_ = point.shapely_point elif isinstance(point, geom.Point): point_ = point else: point_ = np.array(point, dtype=np.float64) - return self.linestring.project(point_, normalized=normalized) - - -@dataclass -class PolylineSE3: - # TODO: Implement PolylineSE3 once quaternions are used in StateSE3 - # Interpolating along SE3 states (i.e., 3D position + orientation) is meaningful, - # but more complex than SE2 due to 3D rotations (quaternions or rotation matrices). - # Linear interpolation of positions is straightforward, but orientation interpolation - # should use SLERP (spherical linear interpolation) for quaternions. - # This is commonly needed in robotics, animation, and path planning. - pass + return self.linestring.project(point_, normalized=normalized) # type: ignore diff --git a/src/py123d/geometry/pose.py b/src/py123d/geometry/pose.py new file mode 100644 index 00000000..802ee116 --- /dev/null +++ b/src/py123d/geometry/pose.py @@ -0,0 +1,492 @@ +from __future__ import annotations + +import numpy as np +import numpy.typing as npt +import shapely.geometry as geom + +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr +from py123d.geometry.geometry_index import EulerPoseSE3Index, Point3DIndex, PoseSE2Index, PoseSE3Index +from py123d.geometry.point import Point2D, Point3D +from py123d.geometry.rotation import EulerAngles, Quaternion + + +class PoseSE2(ArrayMixin): + """Class to represents a 2D pose as SE2 (x, y, yaw). + + Examples: + >>> from py123d.geometry import PoseSE2 + >>> pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + >>> print(pose.x, pose.y, pose.yaw) + 1.0 2.0 0.5 + >>> print(pose.rotation_matrix) + [[ 0.87758256 -0.47942554] + [ 0.47942554 0.87758256]] + + """ + + __slots__ = ("_array",) + _array: npt.NDArray[np.float64] + + def __init__(self, x: float, y: float, yaw: float): + """Init :class:`PoseSE2` with x, y, yaw coordinates. + + :param x: The x-coordinate. + :param y: The y-coordinate. + :param yaw: The yaw angle in radians. + """ + array = np.zeros(len(PoseSE2Index), dtype=np.float64) + array[PoseSE2Index.X] = x + array[PoseSE2Index.Y] = y + array[PoseSE2Index.YAW] = yaw + object.__setattr__(self, "_array", array) + + @classmethod + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> PoseSE2: + """Constructs a PoseSE2 from a numpy array. + + :param array: Array of shape (3,) representing the state [x, y, yaw], indexed by \ + :class:`~py123d.geometry.geometry_index.PoseSE2Index`. + :param copy: Whether to copy the input array. Defaults to True. + :return: A PoseSE2 instance. + """ + assert array.ndim == 1 + assert array.shape[0] == len(PoseSE2Index) + instance = object.__new__(cls) + object.__setattr__(instance, "_array", array.copy() if copy else array) + return instance + + @property + def x(self) -> float: + """The x-coordinate of the pose.""" + return self._array[PoseSE2Index.X] + + @property + def y(self) -> float: + """The y-coordinate of the pose.""" + return self._array[PoseSE2Index.Y] + + @property + def yaw(self) -> float: + """The yaw angle of the pose.""" + return self._array[PoseSE2Index.YAW] + + @property + def array(self) -> npt.NDArray[np.float64]: + """Pose as numpy array of shape (3,), indexed by :class:`~py123d.geometry.geometry_index.PoseSE2Index`.""" + return self._array + + @property + def pose_se2(self) -> PoseSE2: + """Returns self to match interface of other pose classes.""" + return self + + @property + def point_2d(self) -> Point2D: + """The :class:`~py123d.geometry.Point2D` of the pose, i.e. the translation part.""" + return Point2D.from_array(self.array[PoseSE2Index.XY]) + + @property + def rotation_matrix(self) -> npt.NDArray[np.float64]: + """The 2x2 rotation matrix representation of the pose.""" + cos_yaw = np.cos(self.yaw) + sin_yaw = np.sin(self.yaw) + return np.array([[cos_yaw, -sin_yaw], [sin_yaw, cos_yaw]], dtype=np.float64) + + @property + def transformation_matrix(self) -> npt.NDArray[np.float64]: + """The 3x3 transformation matrix representation of the pose.""" + matrix = np.zeros((3, 3), dtype=np.float64) + matrix[:2, :2] = self.rotation_matrix + matrix[0, 2] = self.x + matrix[1, 2] = self.y + return matrix + + @property + def shapely_point(self) -> geom.Point: + """The Shapely point representation of the pose.""" + return geom.Point(self.x, self.y) + + def __repr__(self) -> str: + """String representation of :class:`PoseSE2`.""" + return indexed_array_repr(self, PoseSE2Index) + + +class PoseSE3(ArrayMixin): + """Class representing a quaternion in SE3 space + + Examples: + >>> from py123d.geometry import PoseSE3 + >>> pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + >>> pose.point_3d + Point3D(array=[1. 2. 3.]) + >>> pose.transformation_matrix + array([[1., 0., 0., 1.], + [0., 1., 0., 2.], + [0., 0., 1., 3.], + [0., 0., 0., 1.]]) + >>> PoseSE3.from_transformation_matrix(pose.transformation_matrix) == pose + True + >>> print(pose.yaw, pose.pitch, pose.roll) + 0.0 0.0 0.0 + """ + + __slots__ = ("_array",) + _array: npt.NDArray[np.float64] + + def __init__(self, x: float, y: float, z: float, qw: float, qx: float, qy: float, qz: float): + """Initialize :class:`PoseSE3` with x, y, z, qw, qx, qy, qz coordinates. + + :param x: The x-coordinate. + :param y: The y-coordinate. + :param z: The z-coordinate. + :param qw: The w-coordinate of the quaternion, representing the scalar part. + :param qx: The x-coordinate of the quaternion, representing the first component of the vector part. + :param qy: The y-coordinate of the quaternion, representing the second component of the vector part. + :param qz: The z-coordinate of the quaternion, representing the third component of the vector part. + """ + array = np.zeros(len(PoseSE3Index), dtype=np.float64) + array[PoseSE3Index.X] = x + array[PoseSE3Index.Y] = y + array[PoseSE3Index.Z] = z + array[PoseSE3Index.QW] = qw + array[PoseSE3Index.QX] = qx + array[PoseSE3Index.QY] = qy + array[PoseSE3Index.QZ] = qz + object.__setattr__(self, "_array", array) + + @classmethod + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> PoseSE3: + """Constructs a :class:`PoseSE3` from a numpy array of shape (7,), \ + indexed by :class:`~py123d.geometry.geometry_index.PoseSE3Index`. + + :param array: Array of shape (7,) representing the state [x, y, z, qw, qx, qy, qz]. + :param copy: Whether to copy the input array. Defaults to True. + :return: A :class:`PoseSE3` instance. + """ + assert array.ndim == 1 + assert array.shape[0] == len(PoseSE3Index) + instance = object.__new__(cls) + object.__setattr__(instance, "_array", array.copy() if copy else array) + return instance + + @classmethod + def from_transformation_matrix(cls, transformation_matrix: npt.NDArray[np.float64]) -> PoseSE3: + """Constructs a :class:`PoseSE3` from a 4x4 transformation matrix. + + :param transformation_matrix: A 4x4 numpy array representing the transformation matrix. + :return: A :class:`PoseSE3` instance. + """ + assert transformation_matrix.ndim == 2 + assert transformation_matrix.shape == (4, 4) + array = np.zeros(len(PoseSE3Index), dtype=np.float64) + array[PoseSE3Index.XYZ] = transformation_matrix[:3, 3] + array[PoseSE3Index.QUATERNION] = Quaternion.from_rotation_matrix(transformation_matrix[:3, :3]) + return PoseSE3.from_array(array, copy=False) + + @property + def x(self) -> float: + """The x-coordinate of the pose.""" + return self._array[PoseSE3Index.X] + + @property + def y(self) -> float: + """The y-coordinate of the pose.""" + return self._array[PoseSE3Index.Y] + + @property + def z(self) -> float: + """The z-coordinate of the pose.""" + return self._array[PoseSE3Index.Z] + + @property + def qw(self) -> float: + """The w-coordinate of the quaternion, representing the scalar part.""" + return self._array[PoseSE3Index.QW] + + @property + def qx(self) -> float: + """The x-coordinate of the quaternion, representing the first component of the vector part.""" + return self._array[PoseSE3Index.QX] + + @property + def qy(self) -> float: + """The y-coordinate of the quaternion, representing the second component of the vector part.""" + return self._array[PoseSE3Index.QY] + + @property + def qz(self) -> float: + """The z-coordinate of the quaternion, representing the third component of the vector part.""" + return self._array[PoseSE3Index.QZ] + + @property + def array(self) -> npt.NDArray[np.float64]: + """The numpy array representation of the pose with shape (7,), \ + indexed by :class:`~py123d.geometry.geometry_index.PoseSE3Index`""" + return self._array + + @property + def pose_se3(self) -> PoseSE3: + """The :class:`PoseSE3` itself.""" + return self + + @property + def pose_se2(self) -> PoseSE2: + """The :class:`PoseSE2` representation of the SE3 pose.""" + return PoseSE2(self.x, self.y, self.yaw) + + @property + def point_3d(self) -> Point3D: + """The :class:`Point3D` representation of the SE3 pose, i.e. the translation part.""" + return Point3D(self.x, self.y, self.z) + + @property + def point_2d(self) -> Point2D: + """The :class:`Point2D` representation of the SE3 pose, i.e. the translation part.""" + return Point2D(self.x, self.y) + + @property + def shapely_point(self) -> geom.Point: + """The Shapely point representation, of the translation part of the SE3 pose.""" + return self.point_3d.shapely_point + + @property + def quaternion(self) -> Quaternion: + """The :class:`~py123d.geometry.Quaternion` representation of the state's orientation.""" + return Quaternion.from_array(self.array[PoseSE3Index.QUATERNION]) + + @property + def euler_angles(self) -> EulerAngles: + """The :class:`~py123d.geometry.EulerAngles` representation of the state's orientation.""" + return self.quaternion.euler_angles + + @property + def roll(self) -> float: + """The roll (x-axis rotation) angle in radians.""" + return self.euler_angles.roll + + @property + def pitch(self) -> float: + """The pitch (y-axis rotation) angle in radians.""" + return self.euler_angles.pitch + + @property + def yaw(self) -> float: + """The yaw (z-axis rotation) angle in radians.""" + return self.euler_angles.yaw + + @property + def rotation_matrix(self) -> npt.NDArray[np.float64]: + """Returns the 3x3 rotation matrix representation of the state's orientation.""" + return self.quaternion.rotation_matrix + + @property + def transformation_matrix(self) -> npt.NDArray[np.float64]: + """Returns the 4x4 transformation matrix representation of the state.""" + transformation_matrix = np.eye(4, dtype=np.float64) + transformation_matrix[:3, :3] = self.rotation_matrix + transformation_matrix[:3, 3] = self.array[PoseSE3Index.XYZ] + return transformation_matrix + + def __repr__(self) -> str: + """String representation of :class:`PoseSE3`.""" + return indexed_array_repr(self, PoseSE3Index) + + +class EulerPoseSE3(ArrayMixin): + """ + Class to represents a 3D pose as SE3 (x, y, z, roll, pitch, yaw). + + Notes + ----- + This class is deprecated, use :class:`~py123d.geometry.PoseSE3` instead (quaternion based). + """ + + __slots__ = ("_array",) + _array: npt.NDArray[np.float64] + + def __init__(self, x: float, y: float, z: float, roll: float, pitch: float, yaw: float): + """Initialize PoseSE3 with x, y, z, roll, pitch, yaw coordinates.""" + array = np.zeros(len(EulerPoseSE3Index), dtype=np.float64) + array[EulerPoseSE3Index.X] = x + array[EulerPoseSE3Index.Y] = y + array[EulerPoseSE3Index.Z] = z + array[EulerPoseSE3Index.ROLL] = roll + array[EulerPoseSE3Index.PITCH] = pitch + array[EulerPoseSE3Index.YAW] = yaw + object.__setattr__(self, "_array", array) + + @classmethod + def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> EulerPoseSE3: + """Constructs a PoseSE3 from a numpy array. + + :param array: Array of shape (6,) representing the state [x, y, z, roll, pitch, yaw], indexed by \ + :class:`~py123d.geometry.geometry_index.PoseSE3Index`. + :param copy: Whether to copy the input array. Defaults to True. + :return: A PoseSE3 instance. + """ + assert array.ndim == 1 + assert array.shape[0] == len(EulerPoseSE3Index) + instance = object.__new__(cls) + object.__setattr__(instance, "_array", array.copy() if copy else array) + return instance + + @classmethod + def from_transformation_matrix(cls, transformation_matrix: npt.NDArray[np.float64]) -> EulerPoseSE3: + """Constructs a EulerPoseSE3 from a 4x4 transformation matrix. + + :param array: A 4x4 numpy array representing the transformation matrix. + :return: A EulerPoseSE3 instance. + """ + assert transformation_matrix.ndim == 2 + assert transformation_matrix.shape == (4, 4) + translation = transformation_matrix[:3, 3] + rotation = transformation_matrix[:3, :3] + roll, pitch, yaw = EulerAngles.from_rotation_matrix(rotation) + return EulerPoseSE3( + x=translation[Point3DIndex.X], + y=translation[Point3DIndex.Y], + z=translation[Point3DIndex.Z], + roll=roll, + pitch=pitch, + yaw=yaw, + ) + + @property + def x(self) -> float: + """Returns the x-coordinate of the 3D state. + + :return: The x-coordinate. + """ + return self._array[EulerPoseSE3Index.X] + + @property + def y(self) -> float: + """Returns the y-coordinate of the 3D state. + + :return: The y-coordinate. + """ + return self._array[EulerPoseSE3Index.Y] + + @property + def z(self) -> float: + """Returns the z-coordinate of the 3D state. + + :return: The z-coordinate. + """ + return self._array[EulerPoseSE3Index.Z] + + @property + def roll(self) -> float: + """Returns the roll (x-axis rotation) of the 3D state. + + :return: The roll angle. + """ + return self._array[EulerPoseSE3Index.ROLL] + + @property + def pitch(self) -> float: + """Returns the pitch (y-axis rotation) of the 3D state. + + :return: The pitch angle. + """ + return self._array[EulerPoseSE3Index.PITCH] + + @property + def yaw(self) -> float: + """Returns the yaw (z-axis rotation) of the 3D state. + + :return: The yaw angle. + """ + return self._array[EulerPoseSE3Index.YAW] + + @property + def array(self) -> npt.NDArray[np.float64]: + """Returns the PoseSE3 instance as a numpy array. + + :return: A numpy array of shape (6,), indexed by \ + :class:`~py123d.geometry.geometry_index.PoseSE3Index`. + """ + return self._array + + @property + def pose_se2(self) -> PoseSE2: + """Returns the 3D state as a 2D state by ignoring the z-axis. + + :return: A StateSE2 instance representing the 2D projection of the 3D state. + """ + return PoseSE2(self.x, self.y, self.yaw) + + @property + def point_3d(self) -> Point3D: + """Returns the 3D point representation of the state. + + :return: A Point3D instance representing the 3D point. + """ + return Point3D(self.x, self.y, self.z) + + @property + def point_2d(self) -> Point2D: + """Returns the 2D point representation of the state. + + :return: A Point2D instance representing the 2D point. + """ + return Point2D(self.x, self.y) + + @property + def shapely_point(self) -> geom.Point: + """Returns the Shapely point representation of the state. + + :return: A Shapely Point instance representing the 3D point. + """ + return self.point_3d.shapely_point + + @property + def rotation_matrix(self) -> npt.NDArray[np.float64]: + """Returns the 3x3 rotation matrix representation of the state's orientation. + + :return: A 3x3 numpy array representing the rotation matrix. + """ + return self.euler_angles.rotation_matrix + + @property + def transformation_matrix(self) -> npt.NDArray[np.float64]: + """Returns the 4x4 transformation matrix representation of the state. + + :return: A 4x4 numpy array representing the transformation matrix. + """ + rotation_matrix = self.rotation_matrix + transformation_matrix = np.eye(4, dtype=np.float64) + transformation_matrix[:3, :3] = rotation_matrix + transformation_matrix[:3, 3] = self.array[EulerPoseSE3Index.XYZ] + return transformation_matrix + + @property + def euler_angles(self) -> EulerAngles: + """Returns the :class:`~py123d.geometry.EulerAngles` representation of the state's orientation. + + :return: An EulerAngles instance representing the state's orientation. + """ + return EulerAngles.from_array(self.array[EulerPoseSE3Index.EULER_ANGLES]) + + @property + def pose_se3(self) -> PoseSE3: + """Returns the :class:`~py123d.geometry.PoseSE3` representation of the state. + + :return: A PoseSE3 instance representing the state. + """ + quaternion_se3_array = np.zeros(len(PoseSE3Index), dtype=np.float64) + quaternion_se3_array[PoseSE3Index.XYZ] = self.array[EulerPoseSE3Index.XYZ] + quaternion_se3_array[PoseSE3Index.QUATERNION] = Quaternion.from_euler_angles(self.euler_angles) + return PoseSE3.from_array(quaternion_se3_array, copy=False) + + @property + def quaternion(self) -> Quaternion: + """Returns the :class:`~py123d.geometry.Quaternion` representation of the state's orientation. + + :return: A Quaternion instance representing the state's orientation. + """ + return Quaternion.from_euler_angles(self.euler_angles) + + def __repr__(self) -> str: + """String representation of :class:`EulerPoseSE3`.""" + return indexed_array_repr(self, EulerPoseSE3Index) diff --git a/src/py123d/geometry/rotation.py b/src/py123d/geometry/rotation.py index 1f54431a..29607098 100644 --- a/src/py123d/geometry/rotation.py +++ b/src/py123d/geometry/rotation.py @@ -4,7 +4,7 @@ import numpy.typing as npt import pyquaternion -from py123d.common.utils.mixin import ArrayMixin +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr from py123d.geometry.geometry_index import EulerAnglesIndex, QuaternionIndex from py123d.geometry.utils.rotation_utils import ( get_euler_array_from_quaternion_array, @@ -18,14 +18,41 @@ class EulerAngles(ArrayMixin): """Class to represent 3D rotation using Euler angles (roll, pitch, yaw) in radians. - NOTE: The rotation order is intrinsic Z-Y'-X'' (yaw-pitch-roll). - See https://en.wikipedia.org/wiki/Euler_angles for more details. + + Examples + -------- + >>> import numpy as np + >>> from py123d.geometry import EulerAngles + >>> euler_angles = EulerAngles(roll=0.0, pitch=0.0, yaw=np.pi) + >>> euler_angles.roll + 0.0 + >>> euler_angles.yaw + 3.141592653589793 + >>> euler_angles.array + array([0.0, 0.0, 3.14159265]) + >>> EulerAngles.from_rotation_matrix(euler_angles.rotation_matrix).yaw + 3.141592653589793 + + Notes + ----- + The rotation order is intrinsic Z-Y'-X'' (yaw-pitch-roll) [1]_. + + References + ---------- + .. [1] https://en.wikipedia.org/wiki/Euler_angles + """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, roll: float, pitch: float, yaw: float): - """Initialize EulerAngles with roll, pitch, yaw coordinates.""" + """Initialize EulerAngles with roll, pitch, yaw angles in radians. + + :param roll: The roll (x-axis rotation) angle in radians. + :param pitch: The pitch (y-axis rotation) angle in radians. + :param yaw: The yaw (z-axis rotation) angle in radians. + """ array = np.zeros(len(EulerAnglesIndex), dtype=np.float64) array[EulerAnglesIndex.ROLL] = roll array[EulerAnglesIndex.PITCH] = pitch @@ -34,12 +61,12 @@ def __init__(self, roll: float, pitch: float, yaw: float): @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> EulerAngles: - """Constructs a EulerAngles from a numpy array. - - :param array: Array of shape (3,) representing the euler angles [roll, pitch, yaw], indexed by \ + """Constructs a :class:`EulerAngles` from a numpy array of shape (3,) representing, indexed by \ :class:`~py123d.geometry.EulerAnglesIndex`. + + :param array: Array of shape (3,) representing the euler angles [roll, pitch, yaw]. :param copy: Whether to copy the input array. Defaults to True. - :return: A EulerAngles instance. + :return: A :class:`EulerAngles` instance. """ assert array.ndim == 1 assert array.shape[0] == len(EulerAnglesIndex) @@ -49,11 +76,10 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> EulerA @classmethod def from_rotation_matrix(cls, rotation_matrix: npt.NDArray[np.float64]) -> EulerAngles: - """Constructs a EulerAngles from a 3x3 rotation matrix. - NOTE: The rotation order is intrinsic Z-Y'-X'' (yaw-pitch-roll). + """Constructs a :class:`EulerAngles` from a 3x3 rotation matrix. :param rotation_matrix: A 3x3 numpy array representing the rotation matrix. - :return: A EulerAngles instance. + :return: A :class:`EulerAngles` instance. """ assert rotation_matrix.ndim == 2 assert rotation_matrix.shape == (3, 3) @@ -61,68 +87,74 @@ def from_rotation_matrix(cls, rotation_matrix: npt.NDArray[np.float64]) -> Euler @property def roll(self) -> float: - """The roll (x-axis rotation) angle in radians. - - :return: The roll angle in radians. - """ + """The roll (x-axis rotation) angle in radians.""" return self._array[EulerAnglesIndex.ROLL] @property def pitch(self) -> float: - """The pitch (y-axis rotation) angle in radians. - - :return: The pitch angle in radians. - """ + """The pitch (y-axis rotation) angle in radians.""" return self._array[EulerAnglesIndex.PITCH] @property def yaw(self) -> float: - """The yaw (z-axis rotation) angle in radians. - - :return: The yaw angle in radians. - """ + """The yaw (z-axis rotation) angle in radians.""" return self._array[EulerAnglesIndex.YAW] @property def array(self) -> npt.NDArray[np.float64]: - """Converts the EulerAngles instance to a numpy array. - - :return: A numpy array of shape (3,) containing the Euler angles [roll, pitch, yaw], indexed by \ - :class:`~py123d.geometry.EulerAnglesIndex`. + """Converts the EulerAngles instance to a numpy array of shape (3,),\ + indexed by :class:`~py123d.geometry.EulerAnglesIndex`. """ return self._array @property def quaternion(self) -> Quaternion: + """The :class:`Quaternion` representation of the Euler angles.""" return Quaternion.from_euler_angles(self) @property def rotation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 3x3 rotation matrix representation of the Euler angles. - NOTE: The rotation order is intrinsic Z-Y'-X'' (yaw-pitch-roll). - - :return: A 3x3 numpy array representing the rotation matrix. - """ + """Returns the 3x3 rotation matrix representation of the Euler angles.""" return get_rotation_matrix_from_euler_array(self.array) - def __iter__(self): - """Iterator over euler angles.""" - return iter((self.roll, self.pitch, self.yaw)) - - def __hash__(self): - """Hash function for euler angles.""" - return hash((self.roll, self.pitch, self.yaw)) + def __repr__(self) -> str: + """String representation of :class:`EulerAngles`.""" + return indexed_array_repr(self, EulerAnglesIndex) class Quaternion(ArrayMixin): """ Represents a quaternion for 3D rotations. + + Examples + -------- + >>> import numpy as np + >>> from py123d.geometry import Quaternion + >>> quat = Quaternion(1.0, 0.0, 0.0, 0.0) + >>> quat.qw + 1.0 + >>> quat.qx + 0.0 + >>> quat.array + array([1.0, 0.0, 0.0, 0.0]) + >>> quat.rotation_matrix + array([[1., 0., 0.], + [0., 1., 0.], + [0., 0., 1.]]) + """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, qw: float, qx: float, qy: float, qz: float): - """Initialize Quaternion with qw, qx, qy, qz components.""" + """Initialize Quaternion with components. + + :param qw: The scalar component of the quaternion. + :param qx: The x component of the quaternion. + :param qy: The y component of the quaternion. + :param qz: The z component of the quaternion. + """ array = np.zeros(len(QuaternionIndex), dtype=np.float64) array[QuaternionIndex.QW] = qw array[QuaternionIndex.QX] = qx @@ -167,74 +199,46 @@ def from_euler_angles(cls, euler_angles: EulerAngles) -> Quaternion: @property def qw(self) -> float: - """The scalar part of the quaternion. - - :return: The qw component. - """ + """The scalar component of the quaternion.""" return self._array[QuaternionIndex.QW] @property def qx(self) -> float: - """The x component of the quaternion. - - :return: The qx component. - """ + """The x component of the quaternion.""" return self._array[QuaternionIndex.QX] @property def qy(self) -> float: - """The y component of the quaternion. - - :return: The qy component. - """ + """The y component of the quaternion.""" return self._array[QuaternionIndex.QY] @property def qz(self) -> float: - """The z component of the quaternion. - - :return: The qz component. - """ + """The z component of the quaternion.""" return self._array[QuaternionIndex.QZ] @property def array(self) -> npt.NDArray[np.float64]: - """Converts the Quaternion instance to a numpy array. - - :return: A numpy array of shape (4,) containing the quaternion [qw, qx, qy, qz], indexed by \ + """The numpy array of shape (4,) containing the quaternion [qw, qx, qy, qz], indexed by \ :class:`~py123d.geometry.QuaternionIndex`. """ return self._array @property def pyquaternion(self) -> pyquaternion.Quaternion: - """Returns the pyquaternion.Quaternion representation of the quaternion. - - :return: A pyquaternion.Quaternion representation of the quaternion. - """ + """The pyquaternion.Quaternion representation of the quaternion.""" return pyquaternion.Quaternion(array=self.array) @property def euler_angles(self) -> EulerAngles: - """Returns the Euler angles (roll, pitch, yaw) representation of the quaternion. - NOTE: The rotation order is intrinsic Z-Y'-X'' (yaw-pitch-roll). - - :return: An EulerAngles instance representing the Euler angles. - """ + """The :class:`EulerAngles` representation of the quaternion.""" return EulerAngles.from_array(get_euler_array_from_quaternion_array(self.array), copy=False) @property def rotation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 3x3 rotation matrix representation of the quaternion. - - :return: A 3x3 numpy array representing the rotation matrix. - """ + """Returns the 3x3 rotation matrix representation of the quaternion.""" return get_rotation_matrix_from_quaternion_array(self.array) - def __iter__(self): - """Iterator over quaternion components.""" - return iter((self.qw, self.qx, self.qy, self.qz)) - - def __hash__(self): - """Hash function for quaternion.""" - return hash((self.qw, self.qx, self.qy, self.qz)) + def __repr__(self) -> str: + """String representation of :class:`Quaternion`.""" + return indexed_array_repr(self, QuaternionIndex) diff --git a/src/py123d/geometry/se.py b/src/py123d/geometry/se.py deleted file mode 100644 index b8b30cc8..00000000 --- a/src/py123d/geometry/se.py +++ /dev/null @@ -1,495 +0,0 @@ -from __future__ import annotations - -from typing import Iterable - -import numpy as np -import numpy.typing as npt -import shapely.geometry as geom - -from py123d.common.utils.mixin import ArrayMixin -from py123d.geometry.geometry_index import EulerStateSE3Index, Point3DIndex, StateSE2Index, StateSE3Index -from py123d.geometry.point import Point2D, Point3D -from py123d.geometry.rotation import EulerAngles, Quaternion - - -class StateSE2(ArrayMixin): - """Class to represents a 2D pose as SE2 (x, y, yaw).""" - - _array: npt.NDArray[np.float64] - - def __init__(self, x: float, y: float, yaw: float): - """Initialize StateSE2 with x, y, yaw coordinates.""" - array = np.zeros(len(StateSE2Index), dtype=np.float64) - array[StateSE2Index.X] = x - array[StateSE2Index.Y] = y - array[StateSE2Index.YAW] = yaw - object.__setattr__(self, "_array", array) - - @classmethod - def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> StateSE2: - """Constructs a StateSE2 from a numpy array. - - :param array: Array of shape (3,) representing the state [x, y, yaw], indexed by \ - :class:`~py123d.geometry.geometry_index.StateSE2Index`. - :param copy: Whether to copy the input array. Defaults to True. - :return: A StateSE2 instance. - """ - assert array.ndim == 1 - assert array.shape[0] == len(StateSE2Index) - instance = object.__new__(cls) - object.__setattr__(instance, "_array", array.copy() if copy else array) - return instance - - @property - def x(self) -> float: - return self._array[StateSE2Index.X] - - @property - def y(self) -> float: - return self._array[StateSE2Index.Y] - - @property - def yaw(self) -> float: - return self._array[StateSE2Index.YAW] - - @property - def array(self) -> npt.NDArray[np.float64]: - """Converts the StateSE2 instance to a numpy array - - :return: A numpy array of shape (3,) containing the state, indexed by \ - :class:`~py123d.geometry.geometry_index.StateSE2Index`. - """ - return self._array - - @property - def state_se2(self) -> StateSE2: - """The 2D pose itself. Helpful for polymorphism. - - :return: A StateSE2 instance representing the 2D pose. - """ - return self - - @property - def point_2d(self) -> Point2D: - """The 2D projection of the 2D pose. - - :return: A Point2D instance representing the 2D projection of the 2D pose. - """ - return Point2D.from_array(self.array[StateSE2Index.XY]) - - @property - def rotation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 2x2 rotation matrix representation of the state's orientation. - - :return: A 2x2 numpy array representing the rotation matrix. - """ - cos_yaw = np.cos(self.yaw) - sin_yaw = np.sin(self.yaw) - return np.array([[cos_yaw, -sin_yaw], [sin_yaw, cos_yaw]], dtype=np.float64) - - @property - def transformation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 3x3 transformation matrix representation of the state. - - :return: A 3x3 numpy array representing the transformation matrix. - """ - matrix = np.zeros((3, 3), dtype=np.float64) - matrix[:2, :2] = self.rotation_matrix - matrix[0, 2] = self.x - matrix[1, 2] = self.y - return matrix - - @property - def shapely_point(self) -> geom.Point: - return geom.Point(self.x, self.y) - - -class StateSE3(ArrayMixin): - """Class representing a quaternion in SE3 space.""" - - _array: npt.NDArray[np.float64] - - def __init__(self, x: float, y: float, z: float, qw: float, qx: float, qy: float, qz: float): - """Initialize QuaternionSE3 with x, y, z, qw, qx, qy, qz coordinates.""" - array = np.zeros(len(StateSE3Index), dtype=np.float64) - array[StateSE3Index.X] = x - array[StateSE3Index.Y] = y - array[StateSE3Index.Z] = z - array[StateSE3Index.QW] = qw - array[StateSE3Index.QX] = qx - array[StateSE3Index.QY] = qy - array[StateSE3Index.QZ] = qz - object.__setattr__(self, "_array", array) - - @classmethod - def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> StateSE3: - """Constructs a QuaternionSE3 from a numpy array. - - :param array: Array of shape (7,), indexed by :class:`~py123d.geometry.geometry_index.QuaternionSE3Index`. - :param copy: Whether to copy the input array. Defaults to True. - :return: A QuaternionSE3 instance. - """ - assert array.ndim == 1 - assert array.shape[0] == len(StateSE3Index) - instance = object.__new__(cls) - object.__setattr__(instance, "_array", array.copy() if copy else array) - return instance - - @classmethod - def from_transformation_matrix(cls, transformation_matrix: npt.NDArray[np.float64]) -> StateSE3: - """Constructs a StateSE3 from a 4x4 transformation matrix. - - :param transformation_matrix: A 4x4 numpy array representing the transformation matrix. - :return: A StateSE3 instance. - """ - assert transformation_matrix.ndim == 2 - assert transformation_matrix.shape == (4, 4) - array = np.zeros(len(StateSE3Index), dtype=np.float64) - array[StateSE3Index.XYZ] = transformation_matrix[:3, 3] - array[StateSE3Index.QUATERNION] = Quaternion.from_rotation_matrix(transformation_matrix[:3, :3]) - return StateSE3.from_array(array, copy=False) - - @property - def x(self) -> float: - """Returns the x-coordinate of the quaternion. - - :return: The x-coordinate. - """ - return self._array[StateSE3Index.X] - - @property - def y(self) -> float: - """Returns the y-coordinate of the quaternion. - - :return: The y-coordinate. - """ - return self._array[StateSE3Index.Y] - - @property - def z(self) -> float: - """Returns the z-coordinate of the quaternion. - - :return: The z-coordinate. - """ - return self._array[StateSE3Index.Z] - - @property - def qw(self) -> float: - """Returns the w-coordinate of the quaternion. - - :return: The w-coordinate. - """ - return self._array[StateSE3Index.QW] - - @property - def qx(self) -> float: - """Returns the x-coordinate of the quaternion. - - :return: The x-coordinate. - """ - return self._array[StateSE3Index.QX] - - @property - def qy(self) -> float: - """Returns the y-coordinate of the quaternion. - - :return: The y-coordinate. - """ - return self._array[StateSE3Index.QY] - - @property - def qz(self) -> float: - """Returns the z-coordinate of the quaternion. - - :return: The z-coordinate. - """ - return self._array[StateSE3Index.QZ] - - @property - def array(self) -> npt.NDArray[np.float64]: - """Converts the QuaternionSE3 instance to a numpy array. - - :return: A numpy array of shape (7,), indexed by :class:`~py123d.geometry.geometry_index.QuaternionSE3Index`. - """ - return self._array - - @property - def state_se2(self) -> StateSE2: - """Returns the quaternion state as a 2D state by ignoring the z-axis. - - :return: A StateSE2 instance representing the 2D projection of the 3D state. - """ - # Convert quaternion to yaw angle - yaw = self.quaternion.euler_angles.yaw - return StateSE2(self.x, self.y, yaw) - - @property - def point_3d(self) -> Point3D: - """Returns the 3D point representation of the state. - - :return: A Point3D instance representing the 3D point. - """ - return Point3D(self.x, self.y, self.z) - - @property - def point_2d(self) -> Point2D: - """Returns the 2D point representation of the state. - - :return: A Point2D instance representing the 2D point. - """ - return Point2D(self.x, self.y) - - @property - def shapely_point(self) -> geom.Point: - """Returns the Shapely point representation of the state. - - :return: A Shapely Point instance representing the 3D point. - """ - return self.point_3d.shapely_point - - @property - def quaternion(self) -> Quaternion: - """Returns the quaternion (w, x, y, z) representation of the state's orientation. - - :return: A Quaternion instance representing the quaternion. - """ - return Quaternion.from_array(self.array[StateSE3Index.QUATERNION]) - - @property - def euler_angles(self) -> EulerAngles: - """Returns the Euler angles (roll, pitch, yaw) representation of the state's orientation. - - :return: An EulerAngles instance representing the Euler angles. - """ - return self.quaternion.euler_angles - - @property - def roll(self) -> float: - """The roll (x-axis rotation) angle in radians. - - :return: The roll angle in radians. - """ - return self.euler_angles.roll - - @property - def pitch(self) -> float: - """The pitch (y-axis rotation) angle in radians. - - :return: The pitch angle in radians. - """ - return self.euler_angles.pitch - - @property - def yaw(self) -> float: - """The yaw (z-axis rotation) angle in radians. - - :return: The yaw angle in radians. - """ - return self.euler_angles.yaw - - @property - def rotation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 3x3 rotation matrix representation of the state's orientation. - - :return: A 3x3 numpy array representing the rotation matrix. - """ - return self.quaternion.rotation_matrix - - @property - def transformation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 4x4 transformation matrix representation of the state. - - :return: A 4x4 numpy array representing the transformation matrix. - """ - transformation_matrix = np.eye(4, dtype=np.float64) - transformation_matrix[:3, :3] = self.rotation_matrix - transformation_matrix[:3, 3] = self.array[StateSE3Index.XYZ] - return transformation_matrix - - -class EulerStateSE3(ArrayMixin): - """ - Class to represents a 3D pose as SE3 (x, y, z, roll, pitch, yaw). - NOTE: This class is deprecated, use :class:`~py123d.geometry.StateSE3` instead (quaternion based). - """ - - _array: npt.NDArray[np.float64] - - def __init__(self, x: float, y: float, z: float, roll: float, pitch: float, yaw: float): - """Initialize StateSE3 with x, y, z, roll, pitch, yaw coordinates.""" - array = np.zeros(len(EulerStateSE3Index), dtype=np.float64) - array[EulerStateSE3Index.X] = x - array[EulerStateSE3Index.Y] = y - array[EulerStateSE3Index.Z] = z - array[EulerStateSE3Index.ROLL] = roll - array[EulerStateSE3Index.PITCH] = pitch - array[EulerStateSE3Index.YAW] = yaw - object.__setattr__(self, "_array", array) - - @classmethod - def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> EulerStateSE3: - """Constructs a StateSE3 from a numpy array. - - :param array: Array of shape (6,) representing the state [x, y, z, roll, pitch, yaw], indexed by \ - :class:`~py123d.geometry.geometry_index.StateSE3Index`. - :param copy: Whether to copy the input array. Defaults to True. - :return: A StateSE3 instance. - """ - assert array.ndim == 1 - assert array.shape[0] == len(EulerStateSE3Index) - instance = object.__new__(cls) - object.__setattr__(instance, "_array", array.copy() if copy else array) - return instance - - @classmethod - def from_transformation_matrix(cls, transformation_matrix: npt.NDArray[np.float64]) -> EulerStateSE3: - """Constructs a EulerStateSE3 from a 4x4 transformation matrix. - - :param array: A 4x4 numpy array representing the transformation matrix. - :return: A EulerStateSE3 instance. - """ - assert transformation_matrix.ndim == 2 - assert transformation_matrix.shape == (4, 4) - translation = transformation_matrix[:3, 3] - rotation = transformation_matrix[:3, :3] - roll, pitch, yaw = EulerAngles.from_rotation_matrix(rotation) - return EulerStateSE3( - x=translation[Point3DIndex.X], - y=translation[Point3DIndex.Y], - z=translation[Point3DIndex.Z], - roll=roll, - pitch=pitch, - yaw=yaw, - ) - - @property - def x(self) -> float: - """Returns the x-coordinate of the 3D state. - - :return: The x-coordinate. - """ - return self._array[EulerStateSE3Index.X] - - @property - def y(self) -> float: - """Returns the y-coordinate of the 3D state. - - :return: The y-coordinate. - """ - return self._array[EulerStateSE3Index.Y] - - @property - def z(self) -> float: - """Returns the z-coordinate of the 3D state. - - :return: The z-coordinate. - """ - return self._array[EulerStateSE3Index.Z] - - @property - def roll(self) -> float: - """Returns the roll (x-axis rotation) of the 3D state. - - :return: The roll angle. - """ - return self._array[EulerStateSE3Index.ROLL] - - @property - def pitch(self) -> float: - """Returns the pitch (y-axis rotation) of the 3D state. - - :return: The pitch angle. - """ - return self._array[EulerStateSE3Index.PITCH] - - @property - def yaw(self) -> float: - """Returns the yaw (z-axis rotation) of the 3D state. - - :return: The yaw angle. - """ - return self._array[EulerStateSE3Index.YAW] - - @property - def array(self) -> npt.NDArray[np.float64]: - """Returns the StateSE3 instance as a numpy array. - - :return: A numpy array of shape (6,), indexed by \ - :class:`~py123d.geometry.geometry_index.StateSE3Index`. - """ - return self._array - - @property - def state_se2(self) -> StateSE2: - """Returns the 3D state as a 2D state by ignoring the z-axis. - - :return: A StateSE2 instance representing the 2D projection of the 3D state. - """ - return StateSE2(self.x, self.y, self.yaw) - - @property - def point_3d(self) -> Point3D: - """Returns the 3D point representation of the state. - - :return: A Point3D instance representing the 3D point. - """ - return Point3D(self.x, self.y, self.z) - - @property - def point_2d(self) -> Point2D: - """Returns the 2D point representation of the state. - - :return: A Point2D instance representing the 2D point. - """ - return Point2D(self.x, self.y) - - @property - def shapely_point(self) -> geom.Point: - """Returns the Shapely point representation of the state. - - :return: A Shapely Point instance representing the 3D point. - """ - return self.point_3d.shapely_point - - @property - def rotation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 3x3 rotation matrix representation of the state's orientation. - - :return: A 3x3 numpy array representing the rotation matrix. - """ - return self.euler_angles.rotation_matrix - - @property - def transformation_matrix(self) -> npt.NDArray[np.float64]: - """Returns the 4x4 transformation matrix representation of the state. - - :return: A 4x4 numpy array representing the transformation matrix. - """ - rotation_matrix = self.rotation_matrix - transformation_matrix = np.eye(4, dtype=np.float64) - transformation_matrix[:3, :3] = rotation_matrix - transformation_matrix[:3, 3] = self.array[EulerStateSE3Index.XYZ] - return transformation_matrix - - @property - def euler_angles(self) -> EulerAngles: - return EulerAngles.from_array(self.array[EulerStateSE3Index.EULER_ANGLES]) - - @property - def state_se3(self) -> StateSE3: - quaternion_se3_array = np.zeros(len(StateSE3Index), dtype=np.float64) - quaternion_se3_array[StateSE3Index.XYZ] = self.array[EulerStateSE3Index.XYZ] - quaternion_se3_array[StateSE3Index.QUATERNION] = Quaternion.from_euler_angles(self.euler_angles) - return StateSE3.from_array(quaternion_se3_array, copy=False) - - @property - def quaternion(self) -> Quaternion: - return Quaternion.from_euler_angles(self.euler_angles) - - def __iter__(self) -> Iterable[float]: - """Iterator over the state coordinates (x, y, z, roll, pitch, yaw).""" - return iter((self.x, self.y, self.z, self.roll, self.pitch, self.yaw)) - - def __hash__(self) -> int: - """Hash method""" - return hash((self.x, self.y, self.z, self.roll, self.pitch, self.yaw)) diff --git a/src/py123d/geometry/transform/__init__.py b/src/py123d/geometry/transform/__init__.py index cf22b848..d464c7cb 100644 --- a/src/py123d/geometry/transform/__init__.py +++ b/src/py123d/geometry/transform/__init__.py @@ -1,8 +1,10 @@ from py123d.geometry.transform.transform_se2 import ( - convert_absolute_to_relative_point_2d_array, + convert_absolute_to_relative_points_2d_array, convert_absolute_to_relative_se2_array, - convert_relative_to_absolute_point_2d_array, + convert_relative_to_absolute_points_2d_array, convert_relative_to_absolute_se2_array, + convert_points_2d_array_between_origins, + convert_se2_array_between_origins, translate_se2_along_body_frame, translate_se2_along_x, translate_se2_along_y, diff --git a/src/py123d/geometry/transform/transform_euler_se3.py b/src/py123d/geometry/transform/transform_euler_se3.py index 398b2af4..5d6e02a5 100644 --- a/src/py123d/geometry/transform/transform_euler_se3.py +++ b/src/py123d/geometry/transform/transform_euler_se3.py @@ -3,7 +3,7 @@ import numpy as np import numpy.typing as npt -from py123d.geometry import EulerAngles, EulerStateSE3, EulerStateSE3Index, Point3DIndex, Vector3D, Vector3DIndex +from py123d.geometry import EulerAngles, EulerPoseSE3, EulerPoseSE3Index, Point3DIndex, Vector3D, Vector3DIndex from py123d.geometry.utils.rotation_utils import ( get_euler_array_from_rotation_matrices, get_rotation_matrices_from_euler_array, @@ -11,81 +11,74 @@ ) -def translate_euler_se3_along_z(state_se3: EulerStateSE3, distance: float) -> EulerStateSE3: - - R = state_se3.rotation_matrix +def translate_euler_se3_along_z(pose_se3: EulerPoseSE3, distance: float) -> EulerPoseSE3: + R = pose_se3.rotation_matrix z_axis = R[:, 2] - state_se3_array = state_se3.array.copy() - state_se3_array[EulerStateSE3Index.XYZ] += distance * z_axis[Vector3DIndex.XYZ] - return EulerStateSE3.from_array(state_se3_array, copy=False) - + pose_se3_array = pose_se3.array.copy() + pose_se3_array[EulerPoseSE3Index.XYZ] += distance * z_axis[Vector3DIndex.XYZ] + return EulerPoseSE3.from_array(pose_se3_array, copy=False) -def translate_euler_se3_along_y(state_se3: EulerStateSE3, distance: float) -> EulerStateSE3: - R = state_se3.rotation_matrix +def translate_euler_se3_along_y(pose_se3: EulerPoseSE3, distance: float) -> EulerPoseSE3: + R = pose_se3.rotation_matrix y_axis = R[:, 1] - state_se3_array = state_se3.array.copy() - state_se3_array[EulerStateSE3Index.XYZ] += distance * y_axis[Vector3DIndex.XYZ] - return EulerStateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[EulerPoseSE3Index.XYZ] += distance * y_axis[Vector3DIndex.XYZ] + return EulerPoseSE3.from_array(pose_se3_array, copy=False) -def translate_euler_se3_along_x(state_se3: EulerStateSE3, distance: float) -> EulerStateSE3: - - R = state_se3.rotation_matrix +def translate_euler_se3_along_x(pose_se3: EulerPoseSE3, distance: float) -> EulerPoseSE3: + R = pose_se3.rotation_matrix x_axis = R[:, 0] - state_se3_array = state_se3.array.copy() - state_se3_array[EulerStateSE3Index.XYZ] += distance * x_axis[Vector3DIndex.XYZ] - return EulerStateSE3.from_array(state_se3_array, copy=False) - + pose_se3_array = pose_se3.array.copy() + pose_se3_array[EulerPoseSE3Index.XYZ] += distance * x_axis[Vector3DIndex.XYZ] + return EulerPoseSE3.from_array(pose_se3_array, copy=False) -def translate_euler_se3_along_body_frame(state_se3: EulerStateSE3, vector_3d: Vector3D) -> EulerStateSE3: - R = state_se3.rotation_matrix +def translate_euler_se3_along_body_frame(pose_se3: EulerPoseSE3, vector_3d: Vector3D) -> EulerPoseSE3: + R = pose_se3.rotation_matrix world_translation = R @ vector_3d.array - state_se3_array = state_se3.array.copy() - state_se3_array[EulerStateSE3Index.XYZ] += world_translation[Vector3DIndex.XYZ] - return EulerStateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[EulerPoseSE3Index.XYZ] += world_translation[Vector3DIndex.XYZ] + return EulerPoseSE3.from_array(pose_se3_array, copy=False) def convert_absolute_to_relative_euler_se3_array( - origin: Union[EulerStateSE3, npt.NDArray[np.float64]], se3_array: npt.NDArray[np.float64] + origin: Union[EulerPoseSE3, npt.NDArray[np.float64]], se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: - - if isinstance(origin, EulerStateSE3): + if isinstance(origin, EulerPoseSE3): origin_array = origin.array t_origin = origin.point_3d.array R_origin = origin.rotation_matrix elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(EulerStateSE3Index) + assert origin.ndim == 1 and origin.shape[-1] == len(EulerPoseSE3Index) origin_array = origin - t_origin = origin_array[EulerStateSE3Index.XYZ] - R_origin = get_rotation_matrix_from_euler_array(origin_array[EulerStateSE3Index.EULER_ANGLES]) + t_origin = origin_array[EulerPoseSE3Index.XYZ] + R_origin = get_rotation_matrix_from_euler_array(origin_array[EulerPoseSE3Index.EULER_ANGLES]) else: - raise TypeError(f"Expected StateSE3 or np.ndarray, got {type(origin)}") + raise TypeError(f"Expected PoseSE3 or np.ndarray, got {type(origin)}") assert se3_array.ndim >= 1 - assert se3_array.shape[-1] == len(EulerStateSE3Index) + assert se3_array.shape[-1] == len(EulerPoseSE3Index) # Prepare output array rel_se3_array = se3_array.copy() # Vectorized relative position calculation - abs_positions = se3_array[..., EulerStateSE3Index.XYZ] + abs_positions = se3_array[..., EulerPoseSE3Index.XYZ] rel_positions = (abs_positions - t_origin) @ R_origin - rel_se3_array[..., EulerStateSE3Index.XYZ] = rel_positions + rel_se3_array[..., EulerPoseSE3Index.XYZ] = rel_positions # Convert absolute rotation matrices to relative rotation matrices - abs_rotation_matrices = get_rotation_matrices_from_euler_array(se3_array[..., EulerStateSE3Index.EULER_ANGLES]) + abs_rotation_matrices = get_rotation_matrices_from_euler_array(se3_array[..., EulerPoseSE3Index.EULER_ANGLES]) rel_rotation_matrices = np.einsum("ij,...jk->...ik", R_origin.T, abs_rotation_matrices) if se3_array.shape[0] != 0: - # rel_euler_angles = np.array([EulerAngles.from_rotation_matrix(R).array for R in rel_rotation_matrices]) - # rel_se3_array[..., EulerStateSE3Index.EULER_ANGLES] = normalize_angle(rel_euler_angles) - rel_se3_array[..., EulerStateSE3Index.EULER_ANGLES] = get_euler_array_from_rotation_matrices( + rel_se3_array[..., EulerPoseSE3Index.EULER_ANGLES] = get_euler_array_from_rotation_matrices( rel_rotation_matrices ) @@ -93,55 +86,53 @@ def convert_absolute_to_relative_euler_se3_array( def convert_relative_to_absolute_euler_se3_array( - origin: EulerStateSE3, se3_array: npt.NDArray[np.float64] + origin: EulerPoseSE3, se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: - - if isinstance(origin, EulerStateSE3): + if isinstance(origin, EulerPoseSE3): origin_array = origin.array t_origin = origin.point_3d.array R_origin = origin.rotation_matrix elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(EulerStateSE3Index) + assert origin.ndim == 1 and origin.shape[-1] == len(EulerPoseSE3Index) origin_array = origin - t_origin = origin_array[EulerStateSE3Index.XYZ] - R_origin = get_rotation_matrix_from_euler_array(origin_array[EulerStateSE3Index.EULER_ANGLES]) + t_origin = origin_array[EulerPoseSE3Index.XYZ] + R_origin = get_rotation_matrix_from_euler_array(origin_array[EulerPoseSE3Index.EULER_ANGLES]) else: - raise TypeError(f"Expected StateSE3 or np.ndarray, got {type(origin)}") + raise TypeError(f"Expected PoseSE3 or np.ndarray, got {type(origin)}") assert se3_array.ndim >= 1 - assert se3_array.shape[-1] == len(EulerStateSE3Index) + assert se3_array.shape[-1] == len(EulerPoseSE3Index) # Prepare output array abs_se3_array = se3_array.copy() # Vectorized absolute position calculation: rotate and translate - rel_positions = se3_array[..., EulerStateSE3Index.XYZ] + rel_positions = se3_array[..., EulerPoseSE3Index.XYZ] abs_positions = (rel_positions @ R_origin.T) + t_origin - abs_se3_array[..., EulerStateSE3Index.XYZ] = abs_positions + abs_se3_array[..., EulerPoseSE3Index.XYZ] = abs_positions # Convert relative rotation matrices to absolute rotation matrices - rel_rotation_matrices = get_rotation_matrices_from_euler_array(se3_array[..., EulerStateSE3Index.EULER_ANGLES]) + rel_rotation_matrices = get_rotation_matrices_from_euler_array(se3_array[..., EulerPoseSE3Index.EULER_ANGLES]) abs_rotation_matrices = np.einsum("ij,...jk->...ik", R_origin, rel_rotation_matrices) if se3_array.shape[0] != 0: - abs_se3_array[..., EulerStateSE3Index.EULER_ANGLES] = get_euler_array_from_rotation_matrices( + abs_se3_array[..., EulerPoseSE3Index.EULER_ANGLES] = get_euler_array_from_rotation_matrices( abs_rotation_matrices ) return abs_se3_array def convert_absolute_to_relative_points_3d_array( - origin: Union[EulerStateSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] + origin: Union[EulerPoseSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: - - if isinstance(origin, EulerStateSE3): + if isinstance(origin, EulerPoseSE3): t_origin = origin.point_3d.array R_origin = origin.rotation_matrix elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(EulerStateSE3Index) - t_origin = origin[EulerStateSE3Index.XYZ] - R_origin = get_rotation_matrix_from_euler_array(origin[EulerStateSE3Index.EULER_ANGLES]) + assert origin.ndim == 1 and origin.shape[-1] == len(EulerPoseSE3Index) + t_origin = origin[EulerPoseSE3Index.XYZ] + R_origin = get_rotation_matrix_from_euler_array(origin[EulerPoseSE3Index.EULER_ANGLES]) else: - raise TypeError(f"Expected StateSE3 or np.ndarray, got {type(origin)}") + raise TypeError(f"Expected EulerPoseSE3 or np.ndarray, got {type(origin)}") assert points_3d_array.ndim >= 1 assert points_3d_array.shape[-1] == len(Point3DIndex) @@ -152,19 +143,18 @@ def convert_absolute_to_relative_points_3d_array( def convert_relative_to_absolute_points_3d_array( - origin: Union[EulerStateSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] + origin: Union[EulerPoseSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: - - if isinstance(origin, EulerStateSE3): + if isinstance(origin, EulerPoseSE3): origin_array = origin.array elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(EulerStateSE3Index) + assert origin.ndim == 1 and origin.shape[-1] == len(EulerPoseSE3Index) origin_array = origin else: - raise TypeError(f"Expected EulerStateSE3 or np.ndarray, got {type(origin)}") + raise TypeError(f"Expected EulerPoseSE3 or np.ndarray, got {type(origin)}") assert points_3d_array.shape[-1] == len(Point3DIndex) - R = EulerAngles.from_array(origin_array[EulerStateSE3Index.EULER_ANGLES]).rotation_matrix - absolute_points = points_3d_array @ R.T + origin.point_3d.array + R = EulerAngles.from_array(origin_array[EulerPoseSE3Index.EULER_ANGLES]).rotation_matrix + absolute_points = points_3d_array @ R.T + origin_array[EulerPoseSE3Index.XYZ] return absolute_points diff --git a/src/py123d/geometry/transform/transform_se2.py b/src/py123d/geometry/transform/transform_se2.py index 0ee2d5b5..5d021df4 100644 --- a/src/py123d/geometry/transform/transform_se2.py +++ b/src/py123d/geometry/transform/transform_se2.py @@ -3,136 +3,188 @@ import numpy as np import numpy.typing as npt -from py123d.geometry import Point2DIndex, StateSE2, StateSE2Index, Vector2D, Vector2DIndex +from py123d.geometry import Point2DIndex, PoseSE2, PoseSE2Index, Vector2D, Vector2DIndex from py123d.geometry.utils.rotation_utils import normalize_angle +def _extract_pose_se2_array(pose: Union[PoseSE2, npt.NDArray[np.float64]]) -> npt.NDArray[np.float64]: + """Helper function to extract SE2 pose array from a PoseSE2 or np.ndarray. + + :param pose: Input pose, either a PoseSE2 instance or a 1D numpy array. + :raises TypeError: If the input is neither a PoseSE2 nor a 1D numpy array. + :return: A 1D numpy array representing the SE2 pose. + """ + if isinstance(pose, PoseSE2): + pose_array = pose.array + elif isinstance(pose, np.ndarray): + assert pose.ndim == 1 and pose.shape[-1] == len(PoseSE2Index) + pose_array = pose + else: + raise TypeError(f"Expected PoseSE2 or np.ndarray, got {type(pose)}") + return pose_array + + def convert_absolute_to_relative_se2_array( - origin: Union[StateSE2, npt.NDArray[np.float64]], state_se2_array: npt.NDArray[np.float64] + origin: Union[PoseSE2, npt.NDArray[np.float64]], pose_se2_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts an StateSE2 array from global to relative coordinates. :param origin: origin pose of relative coords system - :param state_se2_array: array of SE2 states with (x,y,yaw), indexed by \ - :class:`~py123d.geometry.geometry_index.StateSE2Index`, in last dim + :param pose_se2_array: array of SE2 poses with (x,y,yaw), indexed by \ + :class:`~py123d.geometry.geometry_index.PoseSE2Index`, in last dim :return: SE2 array, index by \ - :class:`~py123d.geometry.geometry_index.StateSE2Index`, in last dim + :class:`~py123d.geometry.geometry_index.PoseSE2Index`, in last dim """ - if isinstance(origin, StateSE2): - origin_array = origin.array - elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(StateSE2Index) - origin_array = origin - else: - raise TypeError(f"Expected StateSE2 or np.ndarray, got {type(origin)}") - - assert len(StateSE2Index) == state_se2_array.shape[-1] + assert len(PoseSE2Index) == pose_se2_array.shape[-1] + origin_array = _extract_pose_se2_array(origin) - rotate_rad = -origin_array[StateSE2Index.YAW] + rotate_rad = -origin_array[PoseSE2Index.YAW] cos, sin = np.cos(rotate_rad), np.sin(rotate_rad) R_inv = np.array([[cos, -sin], [sin, cos]]) - state_se2_rel = state_se2_array - origin_array - state_se2_rel[..., StateSE2Index.XY] = state_se2_rel[..., StateSE2Index.XY] @ R_inv.T - state_se2_rel[..., StateSE2Index.YAW] = normalize_angle(state_se2_rel[..., StateSE2Index.YAW]) + pose_se2_rel = pose_se2_array - origin_array + pose_se2_rel[..., PoseSE2Index.XY] = pose_se2_rel[..., PoseSE2Index.XY] @ R_inv.T + pose_se2_rel[..., PoseSE2Index.YAW] = normalize_angle(pose_se2_rel[..., PoseSE2Index.YAW]) - return state_se2_rel + return pose_se2_rel def convert_relative_to_absolute_se2_array( - origin: Union[StateSE2, npt.NDArray[np.float64]], state_se2_array: npt.NDArray[np.float64] + origin: Union[PoseSE2, npt.NDArray[np.float64]], pose_se2_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: - """ - Converts an StateSE2 array from global to relative coordinates. + """Converts an StateSE2 array from global to relative coordinates. + :param origin: origin pose of relative coords system - :param state_se2_array: array of SE2 states with (x,y,θ) in last dim + :param pose_se2_array: array of SE2 poses with (x,y,θ) in last dim :return: SE2 coords array in relative coordinates """ - if isinstance(origin, StateSE2): - origin_array = origin.array - elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(StateSE2Index) - origin_array = origin - else: - raise TypeError(f"Expected StateSE2 or np.ndarray, got {type(origin)}") + assert len(PoseSE2Index) == pose_se2_array.shape[-1] + origin_array = _extract_pose_se2_array(origin) - assert len(StateSE2Index) == state_se2_array.shape[-1] - - rotate_rad = origin_array[StateSE2Index.YAW] + rotate_rad = origin_array[PoseSE2Index.YAW] cos, sin = np.cos(rotate_rad), np.sin(rotate_rad) R = np.array([[cos, -sin], [sin, cos]]) - state_se2_abs = np.zeros_like(state_se2_array, dtype=np.float64) - state_se2_abs[..., StateSE2Index.XY] = state_se2_array[..., StateSE2Index.XY] @ R.T - state_se2_abs[..., StateSE2Index.XY] += origin_array[..., StateSE2Index.XY] - state_se2_abs[..., StateSE2Index.YAW] = normalize_angle( - state_se2_array[..., StateSE2Index.YAW] + origin_array[..., StateSE2Index.YAW] + pose_se2_abs = np.zeros_like(pose_se2_array, dtype=np.float64) + pose_se2_abs[..., PoseSE2Index.XY] = pose_se2_array[..., PoseSE2Index.XY] @ R.T + pose_se2_abs[..., PoseSE2Index.XY] += origin_array[..., PoseSE2Index.XY] + pose_se2_abs[..., PoseSE2Index.YAW] = normalize_angle( + pose_se2_array[..., PoseSE2Index.YAW] + origin_array[..., PoseSE2Index.YAW] ) - return state_se2_abs + return pose_se2_abs + + +def convert_se2_array_between_origins( + from_origin: Union[PoseSE2, npt.NDArray[np.float64]], + to_origin: Union[PoseSE2, npt.NDArray[np.float64]], + se2_array: npt.NDArray[np.float64], +) -> npt.NDArray[np.float64]: + """Converts an SE2 array from one origin frame to another origin frame. + + :param from_origin: The source origin state in the absolute frame, as a PoseSE2 or np.ndarray. + :param to_origin: The target origin state in the absolute frame, as a PoseSE2 or np.ndarray. + :param se2_array: The SE2 array in the source origin frame. + :raises TypeError: If the origins are not PoseSE2 or np.ndarray. + :return: The SE2 array in the target origin frame, indexed by :class:`~py123d.geometry.PoseSE2Index`. + """ + # Parse from_origin & to_origin + from_origin_array = _extract_pose_se2_array(from_origin) + to_origin_array = _extract_pose_se2_array(to_origin) + + assert se2_array.ndim >= 1 + assert se2_array.shape[-1] == len(PoseSE2Index) + + # TODO: Re-write withouts transforming to absolute frame intermediate step + abs_array = convert_relative_to_absolute_se2_array(from_origin_array, se2_array) + result_se2_array = convert_absolute_to_relative_se2_array(to_origin_array, abs_array) + + return result_se2_array -def convert_absolute_to_relative_point_2d_array( - origin: Union[StateSE2, npt.NDArray[np.float64]], point_2d_array: npt.NDArray[np.float64] +def convert_absolute_to_relative_points_2d_array( + origin: Union[PoseSE2, npt.NDArray[np.float64]], points_2d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts an absolute 2D point array from global to relative coordinates. :param origin: origin pose of relative coords system - :param point_2d_array: array of 2D points with (x,y) in last dim + :param points_2d_array: array of 2D points with (x,y) in last dim :return: 2D points array in relative coordinates """ - if isinstance(origin, StateSE2): - origin_array = origin.array - elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(StateSE2Index) - origin_array = origin - else: - raise TypeError(f"Expected StateSE2 or np.ndarray, got {type(origin)}") + assert points_2d_array.ndim >= 1 + assert points_2d_array.shape[-1] == len(Point2DIndex) + origin_array = _extract_pose_se2_array(origin) - rotate_rad = -origin_array[StateSE2Index.YAW] + rotate_rad = -origin_array[PoseSE2Index.YAW] cos, sin = np.cos(rotate_rad), np.sin(rotate_rad) R = np.array([[cos, -sin], [sin, cos]], dtype=np.float64) - point_2d_rel = point_2d_array - origin_array[..., StateSE2Index.XY] + point_2d_rel = points_2d_array - origin_array[..., PoseSE2Index.XY] point_2d_rel = point_2d_rel @ R.T return point_2d_rel -def convert_relative_to_absolute_point_2d_array( - origin: Union[StateSE2, npt.NDArray[np.float64]], point_2d_array: npt.NDArray[np.float64] +def convert_relative_to_absolute_points_2d_array( + origin: Union[PoseSE2, npt.NDArray[np.float64]], points_2d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: + """Converts relative 2D point array to absolute coordinates. - if isinstance(origin, StateSE2): - origin_array = origin.array - elif isinstance(origin, np.ndarray): - assert origin.ndim == 1 and origin.shape[-1] == len(StateSE2Index) - origin_array = origin - else: - raise TypeError(f"Expected StateSE2 or np.ndarray, got {type(origin)}") + :param origin: origin pose of relative coords system + :param points_2d_array: array of 2D points with (x,y) in last dim + :return: 2D points array in absolute coordinates + """ - rotate_rad = origin_array[StateSE2Index.YAW] + origin_array = _extract_pose_se2_array(origin) + + rotate_rad = origin_array[PoseSE2Index.YAW] cos, sin = np.cos(rotate_rad), np.sin(rotate_rad) R = np.array([[cos, -sin], [sin, cos]], dtype=np.float64) - point_2d_abs = point_2d_array @ R.T - point_2d_abs = point_2d_abs + origin_array[..., StateSE2Index.XY] + points_2d_abs = points_2d_array @ R.T + points_2d_abs = points_2d_abs + origin_array[..., PoseSE2Index.XY] + + return points_2d_abs + + +def convert_points_2d_array_between_origins( + from_origin: Union[PoseSE2, npt.NDArray[np.float64]], + to_origin: Union[PoseSE2, npt.NDArray[np.float64]], + points_2d_array: npt.NDArray[np.float64], +) -> npt.NDArray[np.float64]: + """Converts 2D points from one origin frame to another origin frame. + + :param from_origin: The source origin state in the absolute frame, as a PoseSE2 or np.ndarray. + :param to_origin: The target origin state in the absolute frame, as a PoseSE2 or np.ndarray. + :param points_2d_array: The 2D points in the source origin frame. + :raises TypeError: If the origins are not PoseSE2 or np.ndarray. + :return: The 2D points in the target origin frame, indexed by :class:`~py123d.geometry.Point2DIndex`. + """ + + assert points_2d_array.ndim >= 1 + assert points_2d_array.shape[-1] == len(Point2DIndex) + + from_origin_array = _extract_pose_se2_array(from_origin) + to_origin_array = _extract_pose_se2_array(to_origin) + + abs_points_array = convert_relative_to_absolute_points_2d_array(from_origin_array, points_2d_array) + result_points_array = convert_absolute_to_relative_points_2d_array(to_origin_array, abs_points_array) - return point_2d_abs + return result_points_array def translate_se2_array_along_body_frame( - state_se2_array: npt.NDArray[np.float64], translation: Vector2D + pose_se2_array: npt.NDArray[np.float64], translation: Vector2D ) -> npt.NDArray[np.float64]: """Translate an array of SE2 states along their respective local coordinate frames. - :param state_se2_array: array of SE2 states with (x,y,yaw) in last dim + :param pose_se2_array: array of SE2 states with (x,y,yaw) in last dim :param translation: 2D translation in local frame (x: forward, y: left) :return: translated SE2 array """ - assert len(StateSE2Index) == state_se2_array.shape[-1] - result = state_se2_array.copy() - yaws = state_se2_array[..., StateSE2Index.YAW] + assert len(PoseSE2Index) == pose_se2_array.shape[-1] + result = pose_se2_array.copy() + yaws = pose_se2_array[..., PoseSE2Index.YAW] cos_yaws, sin_yaws = np.cos(yaws), np.sin(yaws) # Transform translation from local to global frame for each state @@ -143,41 +195,41 @@ def translate_se2_array_along_body_frame( translation_vector = translation.array[Vector2DIndex.XY] # [x, y] global_translation = np.einsum("...ij,...j->...i", R, translation_vector) - result[..., StateSE2Index.XY] += global_translation + result[..., PoseSE2Index.XY] += global_translation return result -def translate_se2_along_body_frame(state_se2: StateSE2, translation: Vector2D) -> StateSE2: +def translate_se2_along_body_frame(pose_se2: PoseSE2, translation: Vector2D) -> PoseSE2: """Translate a single SE2 state along its local coordinate frame. - :param state_se2: SE2 state to translate + :param pose_se2: SE2 state to translate :param translation: 2D translation in local frame (x: forward, y: left) :return: translated SE2 state """ - return StateSE2.from_array(translate_se2_array_along_body_frame(state_se2.array, translation), copy=False) + return PoseSE2.from_array(translate_se2_array_along_body_frame(pose_se2.array, translation), copy=False) -def translate_se2_along_x(state_se2: StateSE2, distance: float) -> StateSE2: +def translate_se2_along_x(pose_se2: PoseSE2, distance: float) -> PoseSE2: """Translate a single SE2 state along its local X-axis. - :param state_se2: SE2 state to translate + :param pose_se2: SE2 state to translate :param distance: distance to translate along the local X-axis :return: translated SE2 state """ translation = Vector2D.from_array(np.array([distance, 0.0], dtype=np.float64)) - return StateSE2.from_array(translate_se2_array_along_body_frame(state_se2.array, translation), copy=False) + return PoseSE2.from_array(translate_se2_array_along_body_frame(pose_se2.array, translation), copy=False) -def translate_se2_along_y(state_se2: StateSE2, distance: float) -> StateSE2: +def translate_se2_along_y(pose_se2: PoseSE2, distance: float) -> PoseSE2: """Translate a single SE2 state along its local Y-axis. - :param state_se2: SE2 state to translate + :param pose_se2: SE2 state to translate :param distance: distance to translate along the local Y-axis :return: translated SE2 state """ translation = Vector2D.from_array(np.array([0.0, distance], dtype=np.float64)) - return StateSE2.from_array(translate_se2_array_along_body_frame(state_se2.array, translation), copy=False) + return PoseSE2.from_array(translate_se2_array_along_body_frame(pose_se2.array, translation), copy=False) def translate_2d_along_body_frame( diff --git a/src/py123d/geometry/transform/transform_se3.py b/src/py123d/geometry/transform/transform_se3.py index 8f394772..1274030a 100644 --- a/src/py123d/geometry/transform/transform_se3.py +++ b/src/py123d/geometry/transform/transform_se3.py @@ -3,7 +3,7 @@ import numpy as np import numpy.typing as npt -from py123d.geometry import Point3DIndex, QuaternionIndex, StateSE3, StateSE3Index, Vector3D, Vector3DIndex +from py123d.geometry import Point3DIndex, PoseSE3, PoseSE3Index, QuaternionIndex, Vector3D, Vector3DIndex from py123d.geometry.utils.rotation_utils import ( conjugate_quaternion_array, get_rotation_matrices_from_quaternion_array, @@ -13,37 +13,37 @@ def _extract_rotation_translation_pose_arrays( - pose: Union[StateSE3, npt.NDArray[np.float64]], + pose: Union[PoseSE3, npt.NDArray[np.float64]], ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.float64]]: - """Helper function to extract rotation matrix and translation vector from a StateSE3 or np.ndarray. + """Helper function to extract rotation matrix and translation vector from a PoseSE3 or np.ndarray. - :param pose: A StateSE3 pose or np.ndarray, indexed by :class:`~py123d.geometry.StateSE3Index`. - :raises TypeError: If the pose is not a StateSE3 or np.ndarray. + :param pose: A PoseSE3 pose or np.ndarray, indexed by :class:`~py123d.geometry.PoseSE3Index`. + :raises TypeError: If the pose is not a PoseSE3 or np.ndarray. :return: A tuple containing the rotation matrix, translation vector, and pose array. """ - if isinstance(pose, StateSE3): + if isinstance(pose, PoseSE3): translation = pose.point_3d.array rotation = pose.rotation_matrix pose_array = pose.array elif isinstance(pose, np.ndarray): - assert pose.ndim == 1 and pose.shape[-1] == len(StateSE3Index) - translation = pose[StateSE3Index.XYZ] - rotation = get_rotation_matrix_from_quaternion_array(pose[StateSE3Index.QUATERNION]) + assert pose.ndim == 1 and pose.shape[-1] == len(PoseSE3Index) + translation = pose[PoseSE3Index.XYZ] + rotation = get_rotation_matrix_from_quaternion_array(pose[PoseSE3Index.QUATERNION]) pose_array = pose else: - raise TypeError(f"Expected StateSE3 or np.ndarray, got {type(pose)}") + raise TypeError(f"Expected PoseSE3 or np.ndarray, got {type(pose)}") return rotation, translation, pose_array def convert_absolute_to_relative_points_3d_array( - origin: Union[StateSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] + origin: Union[PoseSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts 3D points from the absolute frame to the relative frame. - :param origin: The origin state in the absolute frame, as a StateSE3 or np.ndarray. + :param origin: The origin state in the absolute frame, as a PoseSE3 or np.ndarray. :param points_3d_array: The 3D points in the absolute frame. - :raises TypeError: If the origin is not a StateSE3 or np.ndarray. + :raises TypeError: If the origin is not a PoseSE3 or np.ndarray. :return: The 3D points in the relative frame, indexed by :class:`~py123d.geometry.Point3DIndex`. """ @@ -58,46 +58,46 @@ def convert_absolute_to_relative_points_3d_array( def convert_absolute_to_relative_se3_array( - origin: Union[StateSE3, npt.NDArray[np.float64]], se3_array: npt.NDArray[np.float64] + origin: Union[PoseSE3, npt.NDArray[np.float64]], se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts an SE3 array from the absolute frame to the relative frame. - :param origin: The origin state in the absolute frame, as a StateSE3 or np.ndarray. + :param origin: The origin state in the absolute frame, as a PoseSE3 or np.ndarray. :param se3_array: The SE3 array in the absolute frame. - :raises TypeError: If the origin is not a StateSE3 or np.ndarray. - :return: The SE3 array in the relative frame, indexed by :class:`~py123d.geometry.StateSE3Index`. + :raises TypeError: If the origin is not a PoseSE3 or np.ndarray. + :return: The SE3 array in the relative frame, indexed by :class:`~py123d.geometry.PoseSE3Index`. """ R_origin, t_origin, origin_array = _extract_rotation_translation_pose_arrays(origin) assert se3_array.ndim >= 1 - assert se3_array.shape[-1] == len(StateSE3Index) + assert se3_array.shape[-1] == len(PoseSE3Index) - abs_positions = se3_array[..., StateSE3Index.XYZ] - abs_quaternions = se3_array[..., StateSE3Index.QUATERNION] + abs_positions = se3_array[..., PoseSE3Index.XYZ] + abs_quaternions = se3_array[..., PoseSE3Index.QUATERNION] rel_se3_array = np.zeros_like(se3_array) # 1. Vectorized relative position calculation: translate and rotate rel_positions = (abs_positions - t_origin) @ R_origin - rel_se3_array[..., StateSE3Index.XYZ] = rel_positions + rel_se3_array[..., PoseSE3Index.XYZ] = rel_positions # 2. Vectorized relative orientation calculation: quaternion multiplication with conjugate - q_origin_conj = conjugate_quaternion_array(origin_array[StateSE3Index.QUATERNION]) + q_origin_conj = conjugate_quaternion_array(origin_array[PoseSE3Index.QUATERNION]) rel_quaternions = multiply_quaternion_arrays(q_origin_conj, abs_quaternions) - rel_se3_array[..., StateSE3Index.QUATERNION] = rel_quaternions + rel_se3_array[..., PoseSE3Index.QUATERNION] = rel_quaternions return rel_se3_array def convert_relative_to_absolute_points_3d_array( - origin: Union[StateSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] + origin: Union[PoseSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts 3D points from the relative frame to the absolute frame. - :param origin: The origin state in the absolute frame, as a StateSE3 or np.ndarray. + :param origin: The origin state in the absolute frame, as a PoseSE3 or np.ndarray. :param points_3d_array: The 3D points in the relative frame, indexed by :class:`~py123d.geometry.Point3DIndex`. - :raises TypeError: If the origin is not a StateSE3 or np.ndarray. + :raises TypeError: If the origin is not a PoseSE3 or np.ndarray. :return: The 3D points in the absolute frame, indexed by :class:`~py123d.geometry.Point3DIndex`. """ R_origin, t_origin, _ = _extract_rotation_translation_pose_arrays(origin) @@ -109,67 +109,67 @@ def convert_relative_to_absolute_points_3d_array( def convert_relative_to_absolute_se3_array( - origin: StateSE3, se3_array: npt.NDArray[np.float64] + origin: PoseSE3, se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Converts an SE3 array from the relative frame to the absolute frame. - :param origin: The origin state in the relative frame, as a StateSE3 or np.ndarray. + :param origin: The origin state in the relative frame, as a PoseSE3 or np.ndarray. :param se3_array: The SE3 array in the relative frame. - :raises TypeError: If the origin is not a StateSE3 or np.ndarray. - :return: The SE3 array in the absolute frame, indexed by :class:`~py123d.geometry.StateSE3Index`. + :raises TypeError: If the origin is not a PoseSE3 or np.ndarray. + :return: The SE3 array in the absolute frame, indexed by :class:`~py123d.geometry.PoseSE3Index`. """ R_origin, t_origin, origin_array = _extract_rotation_translation_pose_arrays(origin) assert se3_array.ndim >= 1 - assert se3_array.shape[-1] == len(StateSE3Index) + assert se3_array.shape[-1] == len(PoseSE3Index) # Extract relative positions and orientations - rel_positions = se3_array[..., StateSE3Index.XYZ] - rel_quaternions = se3_array[..., StateSE3Index.QUATERNION] + rel_positions = se3_array[..., PoseSE3Index.XYZ] + rel_quaternions = se3_array[..., PoseSE3Index.QUATERNION] # Vectorized absolute position calculation: rotate and translate abs_positions = (R_origin @ rel_positions.T).T + t_origin - abs_quaternions = multiply_quaternion_arrays(origin_array[StateSE3Index.QUATERNION], rel_quaternions) + abs_quaternions = multiply_quaternion_arrays(origin_array[PoseSE3Index.QUATERNION], rel_quaternions) # Prepare output array abs_se3_array = se3_array.copy() - abs_se3_array[..., StateSE3Index.XYZ] = abs_positions - abs_se3_array[..., StateSE3Index.QUATERNION] = abs_quaternions + abs_se3_array[..., PoseSE3Index.XYZ] = abs_positions + abs_se3_array[..., PoseSE3Index.QUATERNION] = abs_quaternions return abs_se3_array def convert_se3_array_between_origins( - from_origin: Union[StateSE3, npt.NDArray[np.float64]], - to_origin: Union[StateSE3, npt.NDArray[np.float64]], + from_origin: Union[PoseSE3, npt.NDArray[np.float64]], + to_origin: Union[PoseSE3, npt.NDArray[np.float64]], se3_array: npt.NDArray[np.float64], ) -> npt.NDArray[np.float64]: """Converts an SE3 array from one origin frame to another origin frame. - :param from_origin: The source origin state in the absolute frame, as a StateSE3 or np.ndarray. - :param to_origin: The target origin state in the absolute frame, as a StateSE3 or np.ndarray. + :param from_origin: The source origin state in the absolute frame, as a PoseSE3 or np.ndarray. + :param to_origin: The target origin state in the absolute frame, as a PoseSE3 or np.ndarray. :param se3_array: The SE3 array in the source origin frame. - :raises TypeError: If the origins are not StateSE3 or np.ndarray. - :return: The SE3 array in the target origin frame, indexed by :class:`~py123d.geometry.StateSE3Index`. + :raises TypeError: If the origins are not PoseSE3 or np.ndarray. + :return: The SE3 array in the target origin frame, indexed by :class:`~py123d.geometry.PoseSE3Index`. """ # Parse from_origin & to_origin R_from, t_from, from_origin_array = _extract_rotation_translation_pose_arrays(from_origin) R_to, t_to, to_origin_array = _extract_rotation_translation_pose_arrays(to_origin) assert se3_array.ndim >= 1 - assert se3_array.shape[-1] == len(StateSE3Index) + assert se3_array.shape[-1] == len(PoseSE3Index) - rel_positions = se3_array[..., StateSE3Index.XYZ] - rel_quaternions = se3_array[..., StateSE3Index.QUATERNION] + rel_positions = se3_array[..., PoseSE3Index.XYZ] + rel_quaternions = se3_array[..., PoseSE3Index.QUATERNION] # Compute relative transformation: T_to^-1 * T_from R_rel = R_to.T @ R_from # Relative rotation matrix t_rel = R_to.T @ (t_from - t_to) # Relative translation q_rel = multiply_quaternion_arrays( - conjugate_quaternion_array(to_origin_array[StateSE3Index.QUATERNION]), - from_origin_array[StateSE3Index.QUATERNION], + conjugate_quaternion_array(to_origin_array[PoseSE3Index.QUATERNION]), + from_origin_array[PoseSE3Index.QUATERNION], ) # Transform positions: rotate and translate @@ -180,23 +180,23 @@ def convert_se3_array_between_origins( # Prepare output array result_se3_array = np.zeros_like(se3_array) - result_se3_array[..., StateSE3Index.XYZ] = new_rel_positions - result_se3_array[..., StateSE3Index.QUATERNION] = new_rel_quaternions + result_se3_array[..., PoseSE3Index.XYZ] = new_rel_positions + result_se3_array[..., PoseSE3Index.QUATERNION] = new_rel_quaternions return result_se3_array def convert_points_3d_array_between_origins( - from_origin: Union[StateSE3, npt.NDArray[np.float64]], - to_origin: Union[StateSE3, npt.NDArray[np.float64]], + from_origin: Union[PoseSE3, npt.NDArray[np.float64]], + to_origin: Union[PoseSE3, npt.NDArray[np.float64]], points_3d_array: npt.NDArray[np.float64], ) -> npt.NDArray[np.float64]: """Converts 3D points from one origin frame to another origin frame. - :param from_origin: The source origin state in the absolute frame, as a StateSE3 or np.ndarray. - :param to_origin: The target origin state in the absolute frame, as a StateSE3 or np.ndarray. + :param from_origin: The source origin state in the absolute frame, as a PoseSE3 or np.ndarray. + :param to_origin: The target origin state in the absolute frame, as a PoseSE3 or np.ndarray. :param points_3d_array: The 3D points in the source origin frame. - :raises TypeError: If the origins are not StateSE3 or np.ndarray. + :raises TypeError: If the origins are not PoseSE3 or np.ndarray. :return: The 3D points in the target origin frame, indexed by :class:`~py123d.geometry.Point3DIndex`. """ # Parse from_origin & to_origin @@ -213,64 +213,64 @@ def convert_points_3d_array_between_origins( return conv_points_3d_array -def translate_se3_along_z(state_se3: StateSE3, distance: float) -> StateSE3: +def translate_se3_along_z(pose_se3: PoseSE3, distance: float) -> PoseSE3: """Translates an SE3 state along the Z-axis. - :param state_se3: The SE3 state to translate. + :param pose_se3: The SE3 state to translate. :param distance: The distance to translate along the Z-axis. :return: The translated SE3 state. """ - R = state_se3.rotation_matrix + R = pose_se3.rotation_matrix z_axis = R[:, 2] - state_se3_array = state_se3.array.copy() - state_se3_array[StateSE3Index.XYZ] += distance * z_axis[Vector3DIndex.XYZ] - return StateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[PoseSE3Index.XYZ] += distance * z_axis[Vector3DIndex.XYZ] + return PoseSE3.from_array(pose_se3_array, copy=False) -def translate_se3_along_y(state_se3: StateSE3, distance: float) -> StateSE3: +def translate_se3_along_y(pose_se3: PoseSE3, distance: float) -> PoseSE3: """Translates a SE3 state along the Y-axis. - :param state_se3: The SE3 state to translate. + :param pose_se3: The SE3 state to translate. :param distance: The distance to translate along the Y-axis. :return: The translated SE3 state. """ - R = state_se3.rotation_matrix + R = pose_se3.rotation_matrix y_axis = R[:, 1] - state_se3_array = state_se3.array.copy() - state_se3_array[StateSE3Index.XYZ] += distance * y_axis[Vector3DIndex.XYZ] - return StateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[PoseSE3Index.XYZ] += distance * y_axis[Vector3DIndex.XYZ] + return PoseSE3.from_array(pose_se3_array, copy=False) -def translate_se3_along_x(state_se3: StateSE3, distance: float) -> StateSE3: +def translate_se3_along_x(pose_se3: PoseSE3, distance: float) -> PoseSE3: """Translates a SE3 state along the X-axis. - :param state_se3: The SE3 state to translate. + :param pose_se3: The SE3 state to translate. :param distance: The distance to translate along the X-axis. :return: The translated SE3 state. """ - R = state_se3.rotation_matrix + R = pose_se3.rotation_matrix x_axis = R[:, 0] - state_se3_array = state_se3.array.copy() - state_se3_array[StateSE3Index.XYZ] += distance * x_axis[Vector3DIndex.XYZ] - return StateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[PoseSE3Index.XYZ] += distance * x_axis[Vector3DIndex.XYZ] + return PoseSE3.from_array(pose_se3_array, copy=False) -def translate_se3_along_body_frame(state_se3: StateSE3, vector_3d: Vector3D) -> StateSE3: +def translate_se3_along_body_frame(pose_se3: PoseSE3, vector_3d: Vector3D) -> PoseSE3: """Translates a SE3 state along a vector in the body frame. - :param state_se3: The SE3 state to translate. + :param pose_se3: The SE3 state to translate. :param vector_3d: The vector to translate along in the body frame. :return: The translated SE3 state. """ - R = state_se3.rotation_matrix + R = pose_se3.rotation_matrix world_translation = R @ vector_3d.array - state_se3_array = state_se3.array.copy() - state_se3_array[StateSE3Index.XYZ] += world_translation - return StateSE3.from_array(state_se3_array, copy=False) + pose_se3_array = pose_se3.array.copy() + pose_se3_array[PoseSE3Index.XYZ] += world_translation + return PoseSE3.from_array(pose_se3_array, copy=False) def translate_3d_along_body_frame( diff --git a/src/py123d/geometry/utils/bounding_box_utils.py b/src/py123d/geometry/utils/bounding_box_utils.py index d31058bd..bedaba37 100644 --- a/src/py123d/geometry/utils/bounding_box_utils.py +++ b/src/py123d/geometry/utils/bounding_box_utils.py @@ -97,8 +97,7 @@ def corners_2d_array_to_polygon_array(corners_array: npt.NDArray[np.float64]) -> :param corners_array: Array of shape (..., 4, 2) where 4 is the number of corners. :return: Array of shapely Polygons. """ - polygons = shapely.creation.polygons(corners_array) - return polygons + return shapely.creation.polygons(corners_array) # type: ignore def bbse2_array_to_polygon_array(bbse2: npt.NDArray[np.float64]) -> npt.NDArray[np.object_]: @@ -144,6 +143,11 @@ def bbse3_array_to_corners_array(bbse3_array: npt.NDArray[np.float64]) -> npt.ND def corners_array_to_3d_mesh( corners_array: npt.NDArray[np.float64], ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.int32]]: + """Creates a triangular mesh representation of boxes defined by their corner points. + + :param corners_array: An array of shape (..., 8, 3) representing the corners of the boxes. + :return: A tuple containing the vertices and faces of the mesh. + """ num_boxes = corners_array.shape[0] vertices = corners_array.reshape(-1, 3) @@ -180,6 +184,12 @@ def corners_array_to_3d_mesh( def corners_array_to_edge_lines(corners_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Creates line segments representing the edges of boxes defined by their corner points. + + :param corners_array: An array of shape (..., 8, 3) representing the corners of the boxes. + :return: An array of shape (..., 12, 2, 3) representing the edge lines of the boxes. + """ + assert corners_array.shape[-1] == len(Point3DIndex) assert corners_array.shape[-2] == len(Corners3DIndex) assert corners_array.ndim >= 2 diff --git a/src/py123d/geometry/utils/polyline_utils.py b/src/py123d/geometry/utils/polyline_utils.py index d66628e8..9a65b3af 100644 --- a/src/py123d/geometry/utils/polyline_utils.py +++ b/src/py123d/geometry/utils/polyline_utils.py @@ -2,14 +2,14 @@ import numpy.typing as npt from shapely.geometry import LineString -from py123d.geometry.geometry_index import Point2DIndex, StateSE2Index +from py123d.geometry.geometry_index import Point2DIndex, Point3DIndex, PoseSE2Index from py123d.geometry.transform.transform_se2 import translate_2d_along_body_frame def get_linestring_yaws(linestring: LineString) -> npt.NDArray[np.float64]: - """ - Compute the heading of each coordinate to its successor coordinate. The last coordinate will have the same heading - as the second last coordinate. + """Compute the heading of each coordinate to its successor coordinate. The last coordinate \ + will have the same heading as the second last coordinate. + :param linestring: linestring as a shapely LineString. :return: a list of headings associated to each starting coordinate. """ @@ -18,6 +18,12 @@ def get_linestring_yaws(linestring: LineString) -> npt.NDArray[np.float64]: def get_points_2d_yaws(points_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Compute the heading of each 2D point to its successor point. The last point \ + will have the same heading as the second last point. + + :param points_array: Array of shape (..., 2) representing 2D points. + :return: Array of shape (...,) representing the yaw angles of the points. + """ assert points_array.ndim == 2 assert points_array.shape[-1] == len(Point2DIndex) vectors = np.diff(points_array, axis=0) @@ -27,39 +33,73 @@ def get_points_2d_yaws(points_array: npt.NDArray[np.float64]) -> npt.NDArray[np. return yaws -def get_path_progress(points_array: npt.NDArray[np.float64]) -> list[float]: +def get_path_progress_2d(points_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Compute the cumulative path progress along a series of 2D points. + + :param points_array: Array of shape (..., 2) representing 2D points. + :raises ValueError: If the input points_array is not valid. + :return: Array of shape (...) representing the cumulative path progress. + """ if points_array.shape[-1] == len(Point2DIndex): x_diff = np.diff(points_array[..., Point2DIndex.X]) y_diff = np.diff(points_array[..., Point2DIndex.X]) - elif points_array.shape[-1] == len(StateSE2Index): - x_diff = np.diff(points_array[..., StateSE2Index.X]) - y_diff = np.diff(points_array[..., StateSE2Index.Y]) + elif points_array.shape[-1] == len(PoseSE2Index): + x_diff = np.diff(points_array[..., PoseSE2Index.X]) + y_diff = np.diff(points_array[..., PoseSE2Index.Y]) else: raise ValueError( f"Invalid points_array shape: {points_array.shape}. Expected last dimension to be {len(Point2DIndex)} or " - f"{len(StateSE2Index)}." + f"{len(PoseSE2Index)}." ) points_diff: npt.NDArray[np.float64] = np.concatenate(([x_diff], [y_diff]), axis=0, dtype=np.float64) progress_diff = np.append(0.0, np.linalg.norm(points_diff, axis=0)) - return np.cumsum(progress_diff, dtype=np.float64) # type: ignore + return np.cumsum(progress_diff, dtype=np.float64) + + +def get_path_progress_3d(points_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Compute the cumulative path progress along a series of 3D points. + + :param points_array: Array of shape (..., 3) representing 3D points. + :raises ValueError: If the input points_array is not valid. + :return: Array of shape (...) representing the cumulative path progress. + """ + if points_array.shape[-1] == len(Point3DIndex): + x_diff = np.diff(points_array[..., Point3DIndex.X]) + y_diff = np.diff(points_array[..., Point3DIndex.Y]) + z_diff = np.diff(points_array[..., Point3DIndex.Z]) + else: + raise ValueError( + f"Invalid points_array shape: {points_array.shape}. Expected last dimension to be {len(Point3DIndex)}." + ) + points_diff: npt.NDArray[np.float64] = np.concatenate(([x_diff], [y_diff], [z_diff]), axis=0, dtype=np.float64) + progress_diff = np.append(0.0, np.linalg.norm(points_diff, axis=0)) + return np.cumsum(progress_diff, dtype=np.float64) def offset_points_perpendicular(points_array: npt.NDArray[np.float64], offset: float) -> npt.NDArray[np.float64]: + """Offset 2D points or SE2 poses perpendicularly by a given offset. + + :param points_array: Array points of shape (..., 2) representing 2D points \ + or shape (..., 3) representing SE2 poses. + :param offset: Offset distance to apply perpendicularly. + :raises ValueError: If the input points_array is not valid. + :return: Array of shape (..., 2) representing the offset points. + """ if points_array.shape[-1] == len(Point2DIndex): xy = points_array[..., Point2DIndex.XY] yaws = get_points_2d_yaws(points_array[..., Point2DIndex.XY]) - elif points_array.shape[-1] == len(StateSE2Index): - xy = points_array[..., StateSE2Index.XY] - yaws = points_array[..., StateSE2Index.YAW] + elif points_array.shape[-1] == len(PoseSE2Index): + xy = points_array[..., PoseSE2Index.XY] + yaws = points_array[..., PoseSE2Index.YAW] else: raise ValueError( f"Invalid points_array shape: {points_array.shape}. Expected last dimension to be {len(Point2DIndex)} or " - f"{len(StateSE2Index)}." + f"{len(PoseSE2Index)}." ) return translate_2d_along_body_frame( points_2d=xy, yaws=yaws, - y_translate=offset, - x_translate=0.0, + y_translate=offset, # type: ignore + x_translate=0.0, # type: ignore ) diff --git a/src/py123d/geometry/utils/rotation_utils.py b/src/py123d/geometry/utils/rotation_utils.py index 2429ffda..6a035f5b 100644 --- a/src/py123d/geometry/utils/rotation_utils.py +++ b/src/py123d/geometry/utils/rotation_utils.py @@ -17,27 +17,27 @@ def batch_matmul(A: npt.NDArray[np.float64], B: npt.NDArray[np.float64]) -> npt. :return: Array of shape (..., M, P) resulting from batch matrix multiplication of A and B. """ assert A.ndim >= 2 and B.ndim >= 2 - assert ( - A.shape[-1] == B.shape[-2] - ), f"Inner dimensions must match for matrix multiplication, got {A.shape} and {B.shape}" + assert A.shape[-1] == B.shape[-2], ( + f"Inner dimensions must match for matrix multiplication, got {A.shape} and {B.shape}" + ) return np.einsum("...ij,...jk->...ik", A, B) def normalize_angle(angle: Union[float, npt.NDArray[np.float64]]) -> Union[float, npt.NDArray[np.float64]]: - """ - Map a angle in range [-π, π] - :param angle: any angle as float or array of floats - :return: normalized angle or array of normalized angles + """Normalizes an angle or array of angles to the range [-pi, pi]. + + :param angle: Angle or array of angles in radians to normalize. + :return: Normalized angle or array of angles. """ return ((angle + np.pi) % (2 * np.pi)) - np.pi def get_rotation_matrices_from_euler_array(euler_angles_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - """ - Convert Euler angles to rotation matrices using Tait-Bryan ZYX convention (yaw-pitch-roll). + """Convert Euler angles to rotation matrices using Tait-Bryan ZYX convention (yaw-pitch-roll). - Convention: Intrinsic rotations in order Z-Y-X (yaw, pitch, roll) - Equivalent to: R = R_z(yaw) @ R_y(pitch) @ R_x(roll) + :param euler_angles_array: Array of Euler angles of shape (..., 3), \ + indexed by :class:`~py123d.geometry.EulerAnglesIndex` + :return: Array of rotation matrices of shape (..., 3, 3) """ assert euler_angles_array.ndim >= 1 and euler_angles_array.shape[-1] == len(EulerAnglesIndex) @@ -128,11 +128,21 @@ def get_euler_array_from_rotation_matrices(rotation_matrices: npt.NDArray[np.flo def get_euler_array_from_rotation_matrix(rotation_matrix: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Convert a rotation matrix to Euler angles using Tait-Bryan ZYX convention (yaw-pitch-roll). + + :param rotation_matrix: Rotation matrix of shape (3, 3) + :return: Euler angles of shape (3,), indexed by :class:`~py123d.geometry.EulerAnglesIndex` + """ assert rotation_matrix.ndim == 2 and rotation_matrix.shape == (3, 3) return get_euler_array_from_rotation_matrices(rotation_matrix[None, ...])[0] def get_quaternion_array_from_rotation_matrices(rotation_matrices: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Convert rotation matrices to quaternions. + + :param rotation_matrices: Rotation matrices of shape (..., 3, 3) + :return: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + """ assert rotation_matrices.ndim >= 2 assert rotation_matrices.shape[-1] == rotation_matrices.shape[-2] == 3 # http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/ @@ -198,11 +208,21 @@ def get_quaternion_array_from_rotation_matrices(rotation_matrices: npt.NDArray[n def get_quaternion_array_from_rotation_matrix(rotation_matrix: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Convert a rotation matrix to a quaternion. + + :param rotation_matrix: Rotation matrix of shape (3, 3) + :return: Quaternion of shape (4,), indexed by :class:`~py123d.geometry.QuaternionIndex`. + """ assert rotation_matrix.ndim == 2 and rotation_matrix.shape == (3, 3) return get_quaternion_array_from_rotation_matrices(rotation_matrix[None, ...])[0] def get_quaternion_array_from_euler_array(euler_angles: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Converts array of euler angles to array of quaternions. + + :param euler_angles: Euler angles of shape (..., 3), indexed by :class:`~py123d.geometry.EulerAnglesIndex` + :return: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + """ assert euler_angles.ndim >= 1 and euler_angles.shape[-1] == len(EulerAnglesIndex) # Store original shape for reshaping later @@ -245,15 +265,25 @@ def get_quaternion_array_from_euler_array(euler_angles: npt.NDArray[np.float64]) if len(original_shape) > 1: quaternions = quaternions.reshape(original_shape + (len(QuaternionIndex),)) - return normalize_quaternion_array(quaternions) + return normalize_quaternion_array(quaternions) # type: ignore def get_rotation_matrix_from_euler_array(euler_angles: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Convert Euler angles to rotation matrix using Tait-Bryan ZYX convention (yaw-pitch-roll). + + :param euler_angles: Euler angles of shape (3,), indexed by :class:`~py123d.geometry.EulerAnglesIndex` + :return: Rotation matrix of shape (3, 3) + """ assert euler_angles.ndim == 1 and euler_angles.shape[0] == len(EulerAnglesIndex) return get_rotation_matrices_from_euler_array(euler_angles[None, ...])[0] def get_rotation_matrices_from_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Convert array of quaternions to array of rotation matrices. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Rotation matrices of shape (..., 3, 3) + """ assert quaternion_array.ndim >= 1 and quaternion_array.shape[-1] == len(QuaternionIndex) # Store original shape for reshaping later @@ -279,12 +309,21 @@ def get_rotation_matrices_from_quaternion_array(quaternion_array: npt.NDArray[np def get_rotation_matrix_from_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - # TODO: Check if this function is necessary or batch-wise function is universally applicable + """Convert a quaternion to a rotation matrix. + + :param quaternion_array: Quaternion of shape (4,), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Rotation matrix of shape (3, 3) + """ assert quaternion_array.ndim == 1 and quaternion_array.shape[0] == len(QuaternionIndex) return get_rotation_matrices_from_quaternion_array(quaternion_array[None, :])[0] def get_euler_array_from_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Converts array of quaternions to array of euler angles. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Euler angles of shape (..., 3), indexed by :class:`~py123d.geometry.EulerAnglesIndex` + """ assert quaternion_array.ndim >= 1 and quaternion_array.shape[-1] == len(QuaternionIndex) norm_quaternion = normalize_quaternion_array(quaternion_array) QW, QX, QY, QZ = ( @@ -311,11 +350,12 @@ def get_euler_array_from_quaternion_array(quaternion_array: npt.NDArray[np.float def conjugate_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - """Computes the conjugate of an array of quaternions. - in the order [qw, qx, qy, qz]. - :param quaternion: Array of quaternions. - :return: Array of conjugated quaternions. + """Computes the conjugate of an array of quaternions, i.e. negating the vector part. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Conjugated quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` """ + assert quaternion_array.ndim >= 1 assert quaternion_array.shape[-1] == len(QuaternionIndex) conjugated_quaternions = np.zeros_like(quaternion_array) @@ -325,10 +365,10 @@ def conjugate_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt def invert_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - """Computes the inverse of an array of quaternions. - in the order [qw, qx, qy, qz]. - :param quaternion: Array of quaternions. - :return: Array of inverted quaternions. + """Computes the inverse of an array of quaternions, i.e. conjugate divided by norm squared. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Inverted quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` """ assert quaternion_array.ndim >= 1 assert quaternion_array.shape[-1] == len(QuaternionIndex) @@ -340,10 +380,10 @@ def invert_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.ND def normalize_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - """Normalizes an array of quaternions. - in the order [qw, qx, qy, qz]. - :param quaternion: Array of quaternions. - :return: Array of normalized quaternions. + """Normalizes an array of quaternions to unit length. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Normalized quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` """ assert quaternion_array.ndim >= 1 assert quaternion_array.shape[-1] == len(QuaternionIndex) @@ -354,11 +394,12 @@ def normalize_quaternion_array(quaternion_array: npt.NDArray[np.float64]) -> npt def multiply_quaternion_arrays(q1: npt.NDArray[np.float64], q2: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - """Multiplies two arrays of quaternions element-wise. - in the order [qw, qx, qy, qz]. - :param q1: First array of quaternions. - :param q2: Second array of quaternions. - :return: Array of resulting quaternions after multiplication. + """Multiplies two arrays of quaternions. + + :param q1: First array of quaternions, indexed by :class:`~py123d.geometry.QuaternionIndex` in the last dim. + :param q2: Second array of quaternions, indexed by :class:`~py123d.geometry.QuaternionIndex` in the last dim. + :return: Array of resulting quaternions after multiplication, \ + indexed by :class:`~py123d.geometry.QuaternionIndex` in the last dim. """ assert q1.ndim >= 1 assert q2.ndim >= 1 @@ -392,9 +433,9 @@ def multiply_quaternion_arrays(q1: npt.NDArray[np.float64], q2: npt.NDArray[np.f def get_q_matrices(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: """Computes the Q matrices for an array of quaternions. - in the order [qw, qx, qy, qz]. - :param quaternion: Array of quaternions. - :return: Array of Q matrices. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Array of Q matrices of shape (..., 4, 4) """ assert quaternion_array.ndim >= 1 assert quaternion_array.shape[-1] == len(QuaternionIndex) @@ -432,9 +473,9 @@ def get_q_matrices(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np. def get_q_bar_matrices(quaternion_array: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: """Computes the Q-bar matrices for an array of quaternions. - in the order [qw, qx, qy, qz]. - :param quaternion: Array of quaternions. - :return: Array of Q-bar matrices. + + :param quaternion_array: Quaternions of shape (..., 4), indexed by :class:`~py123d.geometry.QuaternionIndex` + :return: Array of Q-bar matrices of shape (..., 4, 4) """ assert quaternion_array.ndim >= 1 assert quaternion_array.shape[-1] == len(QuaternionIndex) diff --git a/src/py123d/geometry/vector.py b/src/py123d/geometry/vector.py index f2846c1f..90809c19 100644 --- a/src/py123d/geometry/vector.py +++ b/src/py123d/geometry/vector.py @@ -1,12 +1,9 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import Iterable - import numpy as np import numpy.typing as npt -from py123d.common.utils.mixin import ArrayMixin +from py123d.common.utils.mixin import ArrayMixin, indexed_array_repr from py123d.geometry.geometry_index import Vector2DIndex, Vector3DIndex @@ -15,6 +12,7 @@ class Vector2D(ArrayMixin): Class to represents 2D vectors, in x, y direction. Example: + >>> from py123d.geometry import Vector2D >>> v1 = Vector2D(3.0, 4.0) >>> v2 = Vector2D(1.0, 2.0) >>> v3 = v1 + v2 @@ -26,6 +24,7 @@ class Vector2D(ArrayMixin): 5.0 """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, x: float, y: float): @@ -37,10 +36,10 @@ def __init__(self, x: float, y: float): @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Vector2D: - """Constructs a Vector2D from a numpy array. + """Constructs a :class:`Vector2D` from a numpy array of shape (2,), \ + indexed by :class:`~py123d.geometry.geometry_index.Vector2DIndex`. - :param array: Array of shape (2,) representing the vector components [x, y], indexed by \ - :class:`~py123d.geometry.Vector2DIndex`. + :param array: The array of shape (2,) with the x,y components. :param copy: Whether to copy the input array. Defaults to True. :return: A Vector2D instance. """ @@ -52,27 +51,17 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Vector @property def x(self) -> float: - """The x component of the vector. - - :return: The x component of the vector. - """ + """The x component of the vector.""" return self._array[Vector2DIndex.X] @property def y(self) -> float: - """The y component of the vector. - - :return: The y component of the vector. - """ + """The y component of the vector.""" return self._array[Vector2DIndex.Y] @property def array(self) -> npt.NDArray[np.float64]: - """The array representation of the 2D vector. - - :return: A numpy array of shape (2,) containing the vector components [x, y], indexed by \ - :class:`~py123d.geometry.Vector2DIndex`. - """ + """The vector as array of shape (2,), indexed by :class:`~py123d.geometry.Vector2DIndex`.""" array = np.zeros(len(Vector2DIndex), dtype=np.float64) array[Vector2DIndex.X] = self.x array[Vector2DIndex.Y] = self.y @@ -80,18 +69,12 @@ def array(self) -> npt.NDArray[np.float64]: @property def magnitude(self) -> float: - """Calculates the magnitude (length) of the 2D vector. - - :return: The magnitude of the vector. - """ + """The magnitude (length) of the 2D vector.""" return float(np.linalg.norm(self.array)) @property def vector_2d(self) -> Vector2D: - """The 2D vector itself. Handy for polymorphism. - - :return: A Vector2D instance representing the 2D vector. - """ + """The :class:`Vector2D` itself.""" return self def __add__(self, other: Vector2D) -> Vector2D: @@ -126,21 +109,17 @@ def __truediv__(self, scalar: float) -> Vector2D: """ return Vector2D(self.x / scalar, self.y / scalar) - def __iter__(self) -> Iterable[float]: - """Iterator over vector components.""" - return iter((self.x, self.y)) - - def __hash__(self) -> int: - """Hash method""" - return hash((self.x, self.y)) + def __repr__(self) -> str: + """String representation of :class:`Vector2D`.""" + return indexed_array_repr(self, Vector2DIndex) -@dataclass class Vector3D(ArrayMixin): """ Class to represents 3D vectors, in x, y, z direction. Example: + >>> from py123d.geometry import Vector3D >>> v1 = Vector3D(1.0, 2.0, 3.0) >>> v2 = Vector3D(4.0, 5.0, 6.0) >>> v3 = v1 + v2 @@ -152,6 +131,7 @@ class Vector3D(ArrayMixin): 3.7416573867739413 """ + __slots__ = ("_array",) _array: npt.NDArray[np.float64] def __init__(self, x: float, y: float, z: float): @@ -164,9 +144,10 @@ def __init__(self, x: float, y: float, z: float): @classmethod def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Vector3D: - """Constructs a Vector3D from a numpy array. + """Constructs a :class:`Vector3D` from a numpy array of shape (3,), \ + indexed by :class:`~py123d.geometry.geometry_index.Vector3DIndex`. - :param array: Array of shape (3,), indexed by :class:`~py123d.geometry.geometry_index.Vector3DIndex`. + :param array: The array of shape (3,) with the x,y,z components. :param copy: Whether to copy the input array. Defaults to True. :return: A Vector3D instance. """ @@ -178,52 +159,37 @@ def from_array(cls, array: npt.NDArray[np.float64], copy: bool = True) -> Vector @property def x(self) -> float: - """The x component of the vector. - - :return: The x component of the vector. - """ + """The x component of the vector.""" return self._array[Vector3DIndex.X] @property def y(self) -> float: - """The y component of the vector. - - :return: The y component of the vector. - """ + """The y component of the vector.""" return self._array[Vector3DIndex.Y] @property def z(self) -> float: - """The z component of the vector. - - :return: The z component of the vector. - """ + """The z component of the vector.""" return self._array[Vector3DIndex.Z] @property def array(self) -> npt.NDArray[np.float64]: - """ - Returns the vector components as a numpy array - - :return: A numpy array representing the vector components [x, y, z], indexed by \ - :class:`~py123d.geometry.geometry_index.Vector3DIndex`. - """ + """The vector as array of shape (3,), indexed by :class:`~py123d.geometry.Vector3DIndex`.""" return self._array @property def magnitude(self) -> float: - """Calculates the magnitude (length) of the 3D vector. - - :return: The magnitude of the vector. - """ + """The magnitude (length) of the 3D vector.""" return float(np.linalg.norm(self.array)) @property - def vector_2d(self) -> Vector2D: - """Returns the 2D vector projection (x, y) of the 3D vector. + def vector_3d(self) -> Vector3D: + """The :class:`Vector3D` itself.""" + return self - :return: A Vector2D instance representing the 2D projection. - """ + @property + def vector_2d(self) -> Vector2D: + """The 2D vector projection (x, y) of the 3D vector.""" return Vector2D(self.x, self.y) def __add__(self, other: Vector3D) -> Vector3D: @@ -258,10 +224,6 @@ def __truediv__(self, scalar: float) -> Vector3D: """ return Vector3D(self.x / scalar, self.y / scalar, self.z / scalar) - def __iter__(self) -> Iterable[float]: - """Iterator over vector components.""" - return iter((self.x, self.y, self.z)) - - def __hash__(self) -> int: - """Hash method""" - return hash((self.x, self.y, self.z)) + def __repr__(self) -> str: + """String representation of :class:`Vector3D`.""" + return indexed_array_repr(self, Vector3DIndex) diff --git a/src/py123d/script/builders/scene_builder_builder.py b/src/py123d/script/builders/scene_builder_builder.py index 3c7523a4..d5dc4420 100644 --- a/src/py123d/script/builders/scene_builder_builder.py +++ b/src/py123d/script/builders/scene_builder_builder.py @@ -5,7 +5,7 @@ from hydra.utils import instantiate from omegaconf import DictConfig -from py123d.datatypes.scene.abstract_scene_builder import SceneBuilder +from py123d.api.scene.scene_builder import SceneBuilder logger = logging.getLogger(__name__) diff --git a/src/py123d/script/builders/scene_filter_builder.py b/src/py123d/script/builders/scene_filter_builder.py index 191af8ea..d4ad8d42 100644 --- a/src/py123d/script/builders/scene_filter_builder.py +++ b/src/py123d/script/builders/scene_filter_builder.py @@ -4,7 +4,7 @@ from hydra.utils import instantiate from omegaconf import DictConfig -from py123d.datatypes.scene.scene_filter import SceneFilter +from py123d.api.scene.scene_filter import SceneFilter logger = logging.getLogger(__name__) diff --git a/src/py123d/script/builders/utils/utils_type.py b/src/py123d/script/builders/utils/utils_type.py index 0fe5e9fd..53e700de 100644 --- a/src/py123d/script/builders/utils/utils_type.py +++ b/src/py123d/script/builders/utils/utils_type.py @@ -29,9 +29,9 @@ def validate_type(instantiated_class: Any, desired_type: Type[Any]) -> None: :param instantiated_class: class that was created :param desired_type: type that the created class should have """ - assert isinstance( - instantiated_class, desired_type - ), f"Class to be of type {desired_type}, but is {type(instantiated_class)}!" + assert isinstance(instantiated_class, desired_type), ( + f"Class to be of type {desired_type}, but is {type(instantiated_class)}!" + ) def are_the_same_type(lhs: Any, rhs: Any) -> None: diff --git a/src/py123d/script/config/common/scene_builder/default_scene_builder.yaml b/src/py123d/script/config/common/scene_builder/default_scene_builder.yaml index cf2e553a..650ac14c 100644 --- a/src/py123d/script/config/common/scene_builder/default_scene_builder.yaml +++ b/src/py123d/script/config/common/scene_builder/default_scene_builder.yaml @@ -1,4 +1,4 @@ -_target_: py123d.datatypes.scene.arrow.arrow_scene_builder.ArrowSceneBuilder +_target_: py123d.api.scene.arrow.arrow_scene_builder.ArrowSceneBuilder _convert_: 'all' logs_root: ${dataset_paths.py123d_logs_root} diff --git a/src/py123d/script/config/common/scene_filter/log_scenes.yaml b/src/py123d/script/config/common/scene_filter/all_scenes.yaml similarity index 65% rename from src/py123d/script/config/common/scene_filter/log_scenes.yaml rename to src/py123d/script/config/common/scene_filter/all_scenes.yaml index 726b1d73..7e150d18 100644 --- a/src/py123d/script/config/common/scene_filter/log_scenes.yaml +++ b/src/py123d/script/config/common/scene_filter/all_scenes.yaml @@ -1,4 +1,4 @@ -_target_: py123d.datatypes.scene.scene_filter.SceneFilter +_target_: py123d.api.scene.scene_filter.SceneFilter _convert_: 'all' split_types: null @@ -12,9 +12,9 @@ timestamp_threshold_s: null ego_displacement_minimum_m: null duration_s: null -history_s: 1.0 +history_s: null -camera_types: null +pinhole_camera_types: null max_num_scenes: null -shuffle: false +shuffle: True diff --git a/src/py123d/script/config/common/scene_filter/nuplan_logs.yaml b/src/py123d/script/config/common/scene_filter/nuplan_logs.yaml deleted file mode 100644 index a9dc8275..00000000 --- a/src/py123d/script/config/common/scene_filter/nuplan_logs.yaml +++ /dev/null @@ -1,19 +0,0 @@ -_target_: py123d.datatypes.scene.scene_filter.SceneFilter -_convert_: 'all' - -split_types: null -split_names: - - "nuplan_train" - - "nuplan_val" - - "nuplan_test" -log_names: null - - -locations: null -scene_uuids: null -timestamp_threshold_s: null -ego_displacement_minimum_m: null -max_num_scenes: null - -duration_s: null -history_s: null diff --git a/src/py123d/script/config/common/scene_filter/nuplan_mini_logs.yaml b/src/py123d/script/config/common/scene_filter/nuplan_mini_logs.yaml deleted file mode 100644 index 9fd43e0d..00000000 --- a/src/py123d/script/config/common/scene_filter/nuplan_mini_logs.yaml +++ /dev/null @@ -1,19 +0,0 @@ -_target_: py123d.datatypes.scene.scene_filter.SceneFilter -_convert_: 'all' - -split_types: null -split_names: - - "nuplan_mini_train" - - "nuplan_mini_val" - - "nuplan_mini_test" -log_names: null - - -locations: null -scene_uuids: null -timestamp_threshold_s: null -ego_displacement_minimum_m: null -max_num_scenes: null - -duration_s: null -history_s: null diff --git a/src/py123d/script/config/common/scene_filter/viser_scenes.yaml b/src/py123d/script/config/common/scene_filter/viser_scenes.yaml index 1b619af1..8ed7eb3a 100644 --- a/src/py123d/script/config/common/scene_filter/viser_scenes.yaml +++ b/src/py123d/script/config/common/scene_filter/viser_scenes.yaml @@ -1,4 +1,4 @@ -_target_: py123d.datatypes.scene.scene_filter.SceneFilter +_target_: py123d.api.scene.scene_filter.SceneFilter _convert_: 'all' split_types: null @@ -9,7 +9,6 @@ log_names: null locations: null scene_uuids: null timestamp_threshold_s: null -ego_displacement_minimum_m: null duration_s: null history_s: null diff --git a/src/py123d/script/config/common/worker/single_machine_thread_pool.yaml b/src/py123d/script/config/common/worker/single_machine_thread_pool.yaml index 1344c762..308a00a3 100644 --- a/src/py123d/script/config/common/worker/single_machine_thread_pool.yaml +++ b/src/py123d/script/config/common/worker/single_machine_thread_pool.yaml @@ -1,4 +1,4 @@ _target_: py123d.common.multithreading.worker_parallel.SingleMachineParallelExecutor _convert_: 'all' -use_process_pool: True # If true, use ProcessPoolExecutor as the backend, otherwise uses ThreadPoolExecutor -max_workers: 16 # Number of CPU workers (threads/processes) to use per node, "null" means all available +use_process_pool: False # If true, use ProcessPoolExecutor as the backend, otherwise uses ThreadPoolExecutor +max_workers: null # Number of CPU workers (threads/processes) to use per node, "null" means all available diff --git a/src/py123d/script/config/conversion/datasets/av2_sensor_dataset.yaml b/src/py123d/script/config/conversion/datasets/av2_sensor_dataset.yaml index ff8a2433..92e10e79 100644 --- a/src/py123d/script/config/conversion/datasets/av2_sensor_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/av2_sensor_dataset.yaml @@ -2,7 +2,7 @@ av2_sensor_dataset: _target_: py123d.conversion.datasets.av2.av2_sensor_converter.AV2SensorConverter _convert_: 'all' - splits: ["av2-sensor_train"] + splits: ["av2-sensor_val"] # ["av2-sensor_train", "av2-sensor_val", "av2-sensor_test"] av2_data_root: ${dataset_paths.av2_data_root} dataset_converter_config: @@ -14,6 +14,7 @@ av2_sensor_dataset: # Map include_map: true + remap_map_ids: true # Ego include_ego: true @@ -23,12 +24,11 @@ av2_sensor_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" # "path", "binary", "mp4" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" # "path", "path_merged", "binary" - + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_traffic_lights: false include_scenario_tags: false diff --git a/src/py123d/script/config/conversion/datasets/kitti360_dataset.yaml b/src/py123d/script/config/conversion/datasets/kitti360_dataset.yaml index 544eb879..6e62f2c6 100644 --- a/src/py123d/script/config/conversion/datasets/kitti360_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/kitti360_dataset.yaml @@ -22,6 +22,7 @@ kitti360_dataset: # Map include_map: true + remap_map_ids: true # Ego include_ego: true @@ -31,15 +32,15 @@ kitti360_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # Fisheye Cameras include_fisheye_mei_cameras: true - fisheye_mei_camera_store_option: "binary" + fisheye_mei_camera_store_option: ${fisheye_mei_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "path" + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_traffic_lights: false diff --git a/src/py123d/script/config/conversion/datasets/nuplan_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuplan_dataset.yaml index 1241c2ae..5f40b985 100644 --- a/src/py123d/script/config/conversion/datasets/nuplan_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuplan_dataset.yaml @@ -16,6 +16,7 @@ nuplan_dataset: # Map include_map: true + remap_map_ids: false # Ego include_ego: true @@ -28,11 +29,11 @@ nuplan_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" # "path", "binary", "mp4" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" # "path", "path_merged", "binary" + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Scenario tag / Route include_scenario_tags: true diff --git a/src/py123d/script/config/conversion/datasets/nuplan_mini_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuplan_mini_dataset.yaml index fb7818b2..c57c17ee 100644 --- a/src/py123d/script/config/conversion/datasets/nuplan_mini_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuplan_mini_dataset.yaml @@ -16,6 +16,7 @@ nuplan_mini_dataset: # Map include_map: true + remap_map_ids: false # Ego include_ego: true @@ -28,12 +29,11 @@ nuplan_mini_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" # "path", "binary", "mp4" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" # "path", "path_merged", "binary" - + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Scenario tag / Route include_scenario_tags: true include_route: true diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml index 52ce8207..3e368d05 100644 --- a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml @@ -17,6 +17,7 @@ nuscenes_dataset: # Map include_map: true + remap_map_ids: true # Ego include_ego: true @@ -26,11 +27,11 @@ nuscenes_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_fisheye_mei_cameras: false diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml index 9bc7f019..ac7b6a4a 100644 --- a/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml @@ -1,4 +1,4 @@ -nuscenes_dataset: +nuscenes_mini_dataset: _target_: py123d.conversion.datasets.nuscenes.nuscenes_converter.NuScenesConverter _convert_: 'all' @@ -17,6 +17,7 @@ nuscenes_dataset: # Map include_map: true + remap_map_ids: true # Ego include_ego: true @@ -26,11 +27,11 @@ nuscenes_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_fisheye_mei_cameras: false diff --git a/src/py123d/script/config/conversion/datasets/pandaset_dataset.yaml b/src/py123d/script/config/conversion/datasets/pandaset_dataset.yaml index 30747b42..aa71dd85 100644 --- a/src/py123d/script/config/conversion/datasets/pandaset_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/pandaset_dataset.yaml @@ -20,11 +20,11 @@ pandaset_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs include_lidars: true - lidar_store_option: "binary" + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_map: false diff --git a/src/py123d/script/config/conversion/datasets/wopd_dataset.yaml b/src/py123d/script/config/conversion/datasets/wopd_dataset.yaml index 14a643d8..5cc6a590 100644 --- a/src/py123d/script/config/conversion/datasets/wopd_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/wopd_dataset.yaml @@ -18,6 +18,7 @@ wopd_dataset: # Map include_map: true + remap_map_ids: false # Ego include_ego: true @@ -27,11 +28,11 @@ wopd_dataset: # Pinhole Cameras include_pinhole_cameras: true - pinhole_camera_store_option: "binary" # "path", "binary", "mp4" + pinhole_camera_store_option: ${pinhole_camera_store_option} # "path", "jpeg_binary", "png_binary", "mp4" # LiDARs - include_lidars: false - lidar_store_option: "binary" # "path", "path_merged", "binary" + include_lidars: true + lidar_store_option: ${lidar_store_option} # "path", "path_merged", "laz_binary", "draco_binary" # Not available: include_traffic_lights: false diff --git a/src/py123d/script/config/conversion/default_conversion.yaml b/src/py123d/script/config/conversion/default_conversion.yaml index 48e55dcc..fe3be70a 100644 --- a/src/py123d/script/config/conversion/default_conversion.yaml +++ b/src/py123d/script/config/conversion/default_conversion.yaml @@ -16,7 +16,7 @@ defaults: - log_writer: arrow_log_writer - map_writer: gpkg_map_writer - datasets: - - kitti360_dataset + - av2_sensor_dataset - _self_ @@ -24,3 +24,7 @@ terminate_on_exception: True force_map_conversion: True force_log_conversion: True + +pinhole_camera_store_option: "jpeg_binary" # "path", "jpeg_binary", "png_binary", "mp4" +fisheye_mei_camera_store_option: "jpeg_binary" # "path", "jpeg_binary", "png_binary", "mp4" +lidar_store_option: "laz_binary" # "path", "path_merged", "laz_binary", "draco_binary" diff --git a/src/py123d/script/config/conversion/log_writer/arrow_log_writer.yaml b/src/py123d/script/config/conversion/log_writer/arrow_log_writer.yaml index 61a4ead6..6caedc37 100644 --- a/src/py123d/script/config/conversion/log_writer/arrow_log_writer.yaml +++ b/src/py123d/script/config/conversion/log_writer/arrow_log_writer.yaml @@ -3,5 +3,6 @@ _convert_: 'all' logs_root: ${dataset_paths.py123d_logs_root} +sensors_root: ${dataset_paths.py123d_sensors_root} ipc_compression: null # Compression method for ipc files. Options: None, 'lz4', 'zstd' ipc_compression_level: null # Compression level for ipc files. Options: None, or depending on compression method diff --git a/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml b/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml index 86bf8e0b..6bfb4877 100644 --- a/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml +++ b/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml @@ -2,4 +2,3 @@ _target_: py123d.conversion.map_writer.gpkg_map_writer.GPKGMapWriter _convert_: 'all' maps_root: ${dataset_paths.py123d_maps_root} -remap_ids: true diff --git a/src/py123d/script/config/viser/default_viser.yaml b/src/py123d/script/config/viser/default_viser.yaml index b9ff6f9b..18798ffd 100644 --- a/src/py123d/script/config/viser/default_viser.yaml +++ b/src/py123d/script/config/viser/default_viser.yaml @@ -12,6 +12,7 @@ defaults: - default_common - default_dataset_paths - override scene_filter: viser_scenes + - override worker: single_machine_thread_pool - _self_ viser_config: @@ -56,7 +57,7 @@ viser_config: # -> GUI camera_gui_visible: true - camera_gui_types: [] + # camera_gui_types: null # Loads front camera by default camera_gui_image_scale: 0.25 # Fisheye MEI Cameras diff --git a/src/py123d/script/run_conversion.py b/src/py123d/script/run_conversion.py index 52ed5cd6..08cd75aa 100644 --- a/src/py123d/script/run_conversion.py +++ b/src/py123d/script/run_conversion.py @@ -1,5 +1,6 @@ import gc import logging +import traceback from functools import partial from typing import Dict, List @@ -13,7 +14,7 @@ from py123d.script.utils.dataset_path_utils import setup_dataset_paths logging.basicConfig(level=logging.INFO) -logger = logging.getLogger() +logger = logging.getLogger(__name__) CONFIG_PATH = "config/conversion" CONFIG_NAME = "default_conversion" @@ -32,7 +33,6 @@ def main(cfg: DictConfig) -> None: dataset_converters: List[AbstractDatasetConverter] = build_dataset_converters(cfg.datasets) for dataset_converter in dataset_converters: - worker = build_worker(cfg) logger.info(f"Processing dataset: {dataset_converter.__class__.__name__}") @@ -55,28 +55,34 @@ def main(cfg: DictConfig) -> None: def _convert_maps(args: List[Dict[str, int]], cfg: DictConfig, dataset_converter: AbstractDatasetConverter) -> List: - + setup_dataset_paths(cfg.dataset_paths) map_writer = build_map_writer(cfg.map_writer) for arg in args: - dataset_converter.convert_map(arg["map_index"], map_writer) + try: + dataset_converter.convert_map(arg["map_index"], map_writer) + except Exception as e: + logger.error(f"Error converting map index {arg['map_index']}: {e}") + logger.error(traceback.format_exc()) # noqa: F821 + map_writer.close() + gc.collect() + if cfg.terminate_on_failure: + raise e return [] -def _convert_logs(args: List[Dict[str, int]], cfg: DictConfig, dataset_converter: AbstractDatasetConverter) -> None: - +def _convert_logs(args: List[Dict[str, int]], cfg: DictConfig, dataset_converter: AbstractDatasetConverter) -> List: setup_dataset_paths(cfg.dataset_paths) - - def _internal_convert_log(args: Dict[str, int], dataset_converter_: AbstractDatasetConverter) -> int: - # for i2 in tqdm(range(300), leave=False) - log_writer = build_log_writer(cfg.log_writer) - for arg in args: - dataset_converter_.convert_log(arg["log_index"], log_writer) - del log_writer - gc.collect() - - # for arg in : - _internal_convert_log(args, dataset_converter) - gc.collect() + log_writer = build_log_writer(cfg.log_writer) + for arg in args: + try: + dataset_converter.convert_log(arg["log_index"], log_writer) + except Exception as e: + logger.error(f"Error converting log index {arg['log_index']}: {e}") + logger.error(traceback.format_exc()) # noqa: F821 + log_writer.close() + gc.collect() + if cfg.terminate_on_failure: + raise e return [] diff --git a/src/py123d/script/run_viser.py b/src/py123d/script/run_viser.py index df14382c..0d76b33d 100644 --- a/src/py123d/script/run_viser.py +++ b/src/py123d/script/run_viser.py @@ -18,7 +18,6 @@ @hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) def main(cfg: DictConfig) -> None: - # Initialize dataset paths setup_dataset_paths(cfg.dataset_paths) diff --git a/src/py123d/script/utils/dataset_path_utils.py b/src/py123d/script/utils/dataset_path_utils.py index 393c05f4..d557246c 100644 --- a/src/py123d/script/utils/dataset_path_utils.py +++ b/src/py123d/script/utils/dataset_path_utils.py @@ -15,7 +15,7 @@ def setup_dataset_paths(cfg: DictConfig) -> None: :return: None """ - global _global_dataset_paths + global _global_dataset_paths # noqa: PLW0603 if _global_dataset_paths is None: # Make it immutable @@ -23,8 +23,6 @@ def setup_dataset_paths(cfg: DictConfig) -> None: OmegaConf.set_readonly(cfg, True) # Prevents any modifications _global_dataset_paths = cfg - return None - def get_dataset_paths() -> DictConfig: """Get the global dataset paths from anywhere in your code. diff --git a/src/py123d/visualization/color/color.py b/src/py123d/visualization/color/color.py index 5ba3c866..1d9329d7 100644 --- a/src/py123d/visualization/color/color.py +++ b/src/py123d/visualization/color/color.py @@ -8,31 +8,40 @@ @dataclass(frozen=True) class Color: + """Class representing a color in hexadecimal format.""" hex: str @classmethod def from_rgb(cls, rgb: Tuple[int, int, int]) -> Color: + """Create a Color instance from an RGB tuple.""" r, g, b = rgb return cls(f"#{r:02x}{g:02x}{b:02x}") @property def rgb(self) -> Tuple[int, int, int]: + """The RGB representation of the color.""" return ImageColor.getcolor(self.hex, "RGB") @property - def rgba(self) -> Tuple[int, int, int]: + def rgba(self) -> Tuple[int, int, int, int]: + """The RGBA representation of the color.""" return ImageColor.getcolor(self.hex, "RGBA") @property def rgb_norm(self) -> Tuple[float, float, float]: - return tuple([c / 255 for c in self.rgb]) + """The normalized RGB representation of the color.""" + r, g, b = self.rgb + return (r / 255, g / 255, b / 255) @property - def rgba_norm(self) -> Tuple[float, float, float]: - return tuple([c / 255 for c in self.rgba]) + def rgba_norm(self) -> Tuple[float, float, float, float]: + """The normalized RGBA representation of the color.""" + r, g, b, a = self.rgba + return (r / 255, g / 255, b / 255, a / 255) def set_brightness(self, factor: float) -> Color: + """Return a new Color with adjusted brightness.""" r, g, b = self.rgb return Color.from_rgb( ( @@ -43,8 +52,14 @@ def set_brightness(self, factor: float) -> Color: ) def __str__(self) -> str: + """Return the string representation of the color.""" return self.hex + def __repr__(self) -> str: + """Return the official string representation of the color.""" + r, g, b = self.rgb + return f"Color(hex='\x1b[48;2;{r};{g};{b}m{self.hex}\x1b[0m')" + BLACK: Color = Color("#000000") WHITE: Color = Color("#FFFFFF") @@ -73,7 +88,7 @@ def __str__(self) -> str: 9: Color("#17becf"), # cyan } -NEW_TAB_10: Dict[int, str] = { +NEW_TAB_10: Dict[int, Color] = { 0: Color("#4e79a7"), # blue 1: Color("#f28e2b"), # orange 2: Color("#e15759"), # red diff --git a/src/py123d/visualization/color/default.py b/src/py123d/visualization/color/default.py index ed5ccd25..37dea1f6 100644 --- a/src/py123d/visualization/color/default.py +++ b/src/py123d/visualization/color/default.py @@ -2,7 +2,7 @@ from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel from py123d.datatypes.detections.traffic_light_detections import TrafficLightStatus -from py123d.datatypes.maps.map_datatypes import MapLayer +from py123d.datatypes.map_objects.map_layer_types import MapLayer from py123d.visualization.color.color import ( BLACK, DARKER_GREY, @@ -46,9 +46,9 @@ zorder=1, ), MapLayer.CROSSWALK: PlotConfig( - fill_color=Color("#c69fbb"), + fill_color=Color("#d0b9ca"), fill_color_alpha=1.0, - line_color=Color("#c69fbb"), + line_color=Color("#d0b9ca"), line_color_alpha=0.0, line_width=1.0, line_style="-", @@ -81,11 +81,32 @@ line_style="-", zorder=1, ), + MapLayer.STOP_ZONE: PlotConfig( + fill_color=TAB_10[3], + fill_color_alpha=1.0, + line_color=TAB_10[3], + line_color_alpha=0.0, + line_width=1.0, + line_style="-", + zorder=1, + ), } + BOX_DETECTION_CONFIG: Dict[DefaultBoxDetectionLabel, PlotConfig] = { + # Vehicles + DefaultBoxDetectionLabel.EGO: PlotConfig( + fill_color=Color("#DE7061"), + fill_color_alpha=1.0, + line_color=BLACK, + line_color_alpha=1.0, + line_width=1.0, + line_style="-", + marker_style=HEADING_MARKER_STYLE, + zorder=4, + ), DefaultBoxDetectionLabel.VEHICLE: PlotConfig( - fill_color=ELLIS_5[4], + fill_color=Color("#699CDB"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, @@ -95,19 +116,20 @@ marker_size=1.0, zorder=3, ), - DefaultBoxDetectionLabel.PEDESTRIAN: PlotConfig( - fill_color=NEW_TAB_10[6], + DefaultBoxDetectionLabel.TRAIN: PlotConfig( + fill_color=Color("#76b7b2"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, line_width=1.0, line_style="-", - marker_style=None, + marker_style=HEADING_MARKER_STYLE, marker_size=1.0, - zorder=2, + zorder=3, ), + # VRUs DefaultBoxDetectionLabel.BICYCLE: PlotConfig( - fill_color=ELLIS_5[3], + fill_color=Color("#4e79a7"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, @@ -117,28 +139,31 @@ marker_size=1.0, zorder=2, ), - DefaultBoxDetectionLabel.TRAFFIC_CONE: PlotConfig( - fill_color=NEW_TAB_10[5], + DefaultBoxDetectionLabel.PERSON: PlotConfig( + fill_color=Color("#b07aa1"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, line_width=1.0, line_style="-", marker_style=None, + marker_size=1.0, zorder=2, ), - DefaultBoxDetectionLabel.BARRIER: PlotConfig( - fill_color=NEW_TAB_10[5], + DefaultBoxDetectionLabel.ANIMAL: PlotConfig( + fill_color=Color("#9467bd"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, line_width=1.0, line_style="-", marker_style=None, + marker_size=1.0, zorder=2, ), - DefaultBoxDetectionLabel.CZONE_SIGN: PlotConfig( - fill_color=NEW_TAB_10[5], + # Traffic Control + DefaultBoxDetectionLabel.TRAFFIC_SIGN: PlotConfig( + fill_color=Color("#E38C47"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, @@ -147,8 +172,8 @@ marker_style=None, zorder=2, ), - DefaultBoxDetectionLabel.GENERIC_OBJECT: PlotConfig( - fill_color=NEW_TAB_10[5], + DefaultBoxDetectionLabel.TRAFFIC_CONE: PlotConfig( + fill_color=Color("#E38C47"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, @@ -157,8 +182,8 @@ marker_style=None, zorder=2, ), - DefaultBoxDetectionLabel.SIGN: PlotConfig( - fill_color=NEW_TAB_10[8], + DefaultBoxDetectionLabel.TRAFFIC_LIGHT: PlotConfig( + fill_color=Color("#E38C47"), fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, @@ -167,15 +192,26 @@ marker_style=None, zorder=2, ), - DefaultBoxDetectionLabel.EGO: PlotConfig( - fill_color=ELLIS_5[0], + # Other Obstacles + DefaultBoxDetectionLabel.BARRIER: PlotConfig( + fill_color=NEW_TAB_10[5], fill_color_alpha=1.0, line_color=BLACK, line_color_alpha=1.0, line_width=1.0, line_style="-", - marker_style=HEADING_MARKER_STYLE, - zorder=4, + marker_style=None, + zorder=2, + ), + DefaultBoxDetectionLabel.GENERIC_OBJECT: PlotConfig( + fill_color=NEW_TAB_10[5], + fill_color_alpha=1.0, + line_color=BLACK, + line_color_alpha=1.0, + line_width=1.0, + line_style="-", + marker_style=None, + zorder=2, ), } @@ -199,6 +235,7 @@ line_style="--", zorder=3, ) + ROUTE_CONFIG: PlotConfig = PlotConfig( fill_color=Color("#f2c6c0ff"), fill_color_alpha=1.0, diff --git a/src/py123d/visualization/matplotlib/camera.py b/src/py123d/visualization/matplotlib/camera.py index 5155aae4..078fdc98 100644 --- a/src/py123d/visualization/matplotlib/camera.py +++ b/src/py123d/visualization/matplotlib/camera.py @@ -1,71 +1,54 @@ -# from typing import List, Optional, Tuple - from typing import List, Optional, Tuple import cv2 import matplotlib.pyplot as plt import numpy as np import numpy.typing as npt - -# from PIL import ImageColor from pyquaternion import Quaternion -from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel -from py123d.datatypes.detections.box_detections import BoxDetectionSE3, BoxDetectionWrapper -from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeIntrinsics -from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 +from py123d.conversion.registry import DefaultBoxDetectionLabel, LiDARIndex +from py123d.datatypes.detections import BoxDetectionSE3, BoxDetectionWrapper +from py123d.datatypes.sensors import LiDAR, PinholeCamera, PinholeIntrinsics +from py123d.datatypes.vehicle_state import EgoStateSE3 from py123d.geometry import BoundingBoxSE3Index, Corners3DIndex -from py123d.geometry.transform.transform_se3 import convert_absolute_to_relative_se3_array +from py123d.geometry.transform import convert_absolute_to_relative_se3_array from py123d.visualization.color.default import BOX_DETECTION_CONFIG +from py123d.visualization.matplotlib.lidar import get_lidar_pc_color -# from navsim.common.dataclasses import Annotations, Camera, Lidar -# from navsim.common.enums import BoundingBoxIndex, LidarIndex -# from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types -# from navsim.visualization.config import AGENT_CONFIG -# from navsim.visualization.lidar import filter_lidar_pc, get_lidar_pc_color +def add_pinhole_camera_ax(ax: plt.Axes, pinhole_camera: PinholeCamera) -> plt.Axes: + """Add pinhole camera image to matplotlib axis -def add_camera_ax(ax: plt.Axes, camera: PinholeCamera) -> plt.Axes: - """ - Adds camera image to matplotlib ax object - :param ax: matplotlib ax object - :param camera: navsim camera dataclass - :return: ax object with image + :param ax: matplotlib axis + :param pinhole_camera: pinhole camera object + :return: matplotlib axis with image """ - ax.imshow(camera.image) + ax.imshow(pinhole_camera.image) return ax -# FIXME: -# def add_lidar_to_camera_ax(ax: plt.Axes, camera: Camera, lidar: Lidar) -> plt.Axes: -# """ -# Adds camera image with lidar point cloud on matplotlib ax object -# :param ax: matplotlib ax object -# :param camera: navsim camera dataclass -# :param lidar: navsim lidar dataclass -# :return: ax object with image -# """ +def add_lidar_to_camera_ax(ax: plt.Axes, camera: PinholeCamera, lidar: LiDAR) -> plt.Axes: + """Add lidar point cloud to camera image on matplotlib axis -# image, lidar_pc = camera.image.copy(), lidar.lidar_pc.copy() -# image_height, image_width = image.shape[:2] + :param ax: matplotlib axis + :param camera: pinhole camera object + :param lidar: lidar object + :return: matplotlib axis with lidar points overlaid on camera image + """ -# lidar_pc = filter_lidar_pc(lidar_pc) -# lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc)) + image, lidar_pc = camera.image.copy(), lidar.point_cloud.copy() + lidar_index = lidar.metadata.lidar_index -# pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images( -# lidar_pc, -# camera.sensor2lidar_rotation, -# camera.sensor2lidar_translation, -# camera.intrinsics, -# img_shape=(image_height, image_width), -# ) + # lidar_pc = filter_lidar_pc(lidar_pc) + lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc, lidar_index, feature="distance")) + pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images(lidar_pc, lidar_index, camera) -# for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): -# color = (int(color[0]), int(color[1]), int(color[2])) -# cv2.circle(image, (int(x), int(y)), 5, color, -1) + for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): + color = (int(color[0]), int(color[1]), int(color[2])) + cv2.circle(image, (int(x), int(y)), 5, color, -1) -# ax.imshow(image) -# return ax + ax.imshow(image) + return ax def add_box_detections_to_camera_ax( @@ -73,22 +56,29 @@ def add_box_detections_to_camera_ax( camera: PinholeCamera, box_detections: BoxDetectionWrapper, ego_state_se3: EgoStateSE3, - return_image: bool = False, ) -> plt.Axes: + """Add box detections to camera image on matplotlib axis + + :param ax: matplotlib axis + :param camera: pinhole camera object + :param box_detections: box detection wrapper object + :param ego_state_se3: ego state object + :return: matplotlib axis with box detections overlaid on camera image + """ box_detection_array = np.zeros((len(box_detections.box_detections), len(BoundingBoxSE3Index)), dtype=np.float64) default_labels = np.array( [detection.metadata.default_label for detection in box_detections.box_detections], dtype=object ) for idx, box_detection in enumerate(box_detections.box_detections): - assert isinstance( - box_detection, BoxDetectionSE3 - ), f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + assert isinstance(box_detection, BoxDetectionSE3), ( + f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + ) box_detection_array[idx] = box_detection.bounding_box_se3.array # FIXME - box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] = convert_absolute_to_relative_se3_array( - ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] + box_detection_array[..., BoundingBoxSE3Index.SE3] = convert_absolute_to_relative_se3_array( + ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.SE3] ) # box_detection_array[..., BoundingBoxSE3Index.XYZ] -= ego_state_se3.rear_axle_se3.point_3d.array detection_positions, detection_extents, detection_yaws = _transform_annotations_to_camera( @@ -112,27 +102,16 @@ def add_box_detections_to_camera_ax( box_corners, default_labels = box_corners[valid_corners], default_labels[valid_corners] image = _plot_rect_3d_on_img(camera.image.copy(), box_corners, default_labels) - if return_image: - # ax.imshow(image) - return ax, image - ax.imshow(image) return ax -def _transform_annotations_to_camera( - boxes: npt.NDArray[np.float32], - # sensor2lidar_rotation: npt.NDArray[np.float32], - # sensor2lidar_translation: npt.NDArray[np.float32], - extrinsic: npt.NDArray[np.float32], -) -> npt.NDArray[np.float32]: - """ - Helper function to transform bounding boxes into camera frame - TODO: Refactor - :param boxes: array representation of bounding boxes - :param sensor2lidar_rotation: camera rotation - :param sensor2lidar_translation: camera translation - :return: bounding boxes in camera coordinates +def _transform_annotations_to_camera(boxes: npt.NDArray, extrinsic: npt.NDArray) -> npt.NDArray: + """Transforms the box annotations from sensor frame to camera frame. + + :param boxes: array of bounding box parameters. + :param extrinsic: The (4x4) transformation matrix from ego to camera frame. + :return: transformed bounding box parameters in camera frame. """ sensor2lidar_rotation = extrinsic[:3, :3] sensor2lidar_translation = extrinsic[:3, 3] @@ -165,16 +144,7 @@ def _transform_annotations_to_camera( def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np.float32], axis: int = 0): - """ - Rotate 3D points by angles according to axis. - TODO: Refactor - :param points: array of points - :param angles: array of angles - :param axis: axis to perform rotation, defaults to 0 - :raises value: _description_ - :raises ValueError: if axis invalid - :return: rotated points - """ + """Rotate points in 3D along specific axis.""" rot_sin = np.sin(angles) rot_cos = np.cos(angles) ones = np.ones_like(rot_cos) @@ -187,7 +157,7 @@ def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np np.stack([rot_sin, zeros, rot_cos]), ] ) - elif axis == 2 or axis == -1: + elif axis in [2, -1]: rot_mat_T = np.stack( [ np.stack([rot_cos, -rot_sin, zeros]), @@ -209,20 +179,21 @@ def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np def _plot_rect_3d_on_img( - image: npt.NDArray[np.float32], + image: npt.NDArray[np.uint8], box_corners: npt.NDArray[np.float32], labels: List[DefaultBoxDetectionLabel], thickness: int = 3, ) -> npt.NDArray[np.uint8]: - """ - Plot the boundary lines of 3D rectangular on 2D images. + """Plot 3D bounding boxes on image + TODO: refactor - :param image: The numpy array of image. - :param box_corners: Coordinates of the corners of 3D, shape of [N, 8, 2]. - :param box_labels: labels of boxes for coloring - :param thickness: pixel width of liens, defaults to 3 - :return: image with 3D bounding boxes + :param image: The image to plot on + :param box_corners: The corners of the boxes to plot + :param labels: The labels of the boxes to plot + :param thickness: The thickness of the lines, defaults to 3 + :return: The image with 3D bounding boxes plotted """ + line_indices = ( (0, 1), (0, 3), @@ -258,8 +229,8 @@ def _transform_points_to_image( image_shape: Optional[Tuple[int, int]] = None, eps: float = 1e-3, ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: - """ - Transforms points in camera frame to image pixel coordinates + """Transforms points in camera frame to image pixel coordinates + TODO: refactor :param points: points in camera frame :param intrinsic: camera intrinsics @@ -292,50 +263,49 @@ def _transform_points_to_image( return pc_img, cur_pc_in_fov -# def _transform_pcs_to_images( -# lidar_pc: npt.NDArray[np.float32], -# sensor2lidar_rotation: npt.NDArray[np.float32], -# sensor2lidar_translation: npt.NDArray[np.float32], -# intrinsic: npt.NDArray[np.float32], -# img_shape: Optional[Tuple[int, int]] = None, -# eps: float = 1e-3, -# ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: -# """ -# Transforms points in camera frame to image pixel coordinates -# TODO: refactor -# :param lidar_pc: lidar point cloud -# :param sensor2lidar_rotation: camera rotation -# :param sensor2lidar_translation: camera translation -# :param intrinsic: camera intrinsics -# :param img_shape: image shape in pixels, defaults to None -# :param eps: threshold for lidar pc height, defaults to 1e-3 -# :return: lidar pc in pixel coordinates, mask of values in frame -# """ -# pc_xyz = lidar_pc[LidarIndex.POSITION, :].T - -# lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) -# lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T -# lidar2cam_rt = np.eye(4) -# lidar2cam_rt[:3, :3] = lidar2cam_r.T -# lidar2cam_rt[3, :3] = -lidar2cam_t - -# viewpad = np.eye(4) -# viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic -# lidar2img_rt = viewpad @ lidar2cam_rt.T - -# cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) -# cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T -# cur_pc_cam = cur_pc_cam.T -# cur_pc_in_fov = cur_pc_cam[:, 2] > eps -# cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) - -# if img_shape is not None: -# img_h, img_w = img_shape -# cur_pc_in_fov = ( -# cur_pc_in_fov -# & (cur_pc_cam[:, 0] < (img_w - 1)) -# & (cur_pc_cam[:, 0] > 0) -# & (cur_pc_cam[:, 1] < (img_h - 1)) -# & (cur_pc_cam[:, 1] > 0) -# ) -# return cur_pc_cam, cur_pc_in_fov +def _transform_pcs_to_images( + lidar_pc: npt.NDArray[np.float32], + lidar_index: LiDARIndex, + camera: PinholeCamera, + eps: float = 1e-3, +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + """Transforms lidar point cloud to image pixel coordinates + + TODO: refactor + :param lidar_pc: lidar point cloud + :param lidar_index: lidar index + :param camera: pinhole camera + :param eps: lower threshold of points, defaults to 1e-3 + :return: points in pixel coordinates, mask of values in frame + """ + + pc_xyz = lidar_pc[..., lidar_index.XYZ] + + lidar2cam_r = np.linalg.inv(camera.extrinsic.rotation_matrix) + lidar2cam_t = camera.extrinsic.point_3d @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + + camera_matrix = camera.metadata.intrinsics.camera_matrix + viewpad = np.eye(4) + viewpad[: camera_matrix.shape[0], : camera_matrix.shape[1]] = camera_matrix + lidar2img_rt = viewpad @ lidar2cam_rt.T + img_shape = camera.image.shape[:2] + + cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) + cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T + cur_pc_cam = cur_pc_cam.T + cur_pc_in_fov = cur_pc_cam[:, 2] > eps + cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) + + if img_shape is not None: + img_h, img_w = img_shape + cur_pc_in_fov = ( + cur_pc_in_fov + & (cur_pc_cam[:, 0] < (img_w - 1)) + & (cur_pc_cam[:, 0] > 0) + & (cur_pc_cam[:, 1] < (img_h - 1)) + & (cur_pc_cam[:, 1] > 0) + ) + return cur_pc_cam, cur_pc_in_fov diff --git a/src/py123d/visualization/matplotlib/lidar.py b/src/py123d/visualization/matplotlib/lidar.py index 29ffeea8..b51ab89a 100644 --- a/src/py123d/visualization/matplotlib/lidar.py +++ b/src/py123d/visualization/matplotlib/lidar.py @@ -1,70 +1,38 @@ -# TODO: implement -# from typing import Any, List - -# import matplotlib -# import matplotlib.pyplot as plt -# import numpy as np -# import numpy.typing as npt - -# from navsim.common.enums import LidarIndex -# from navsim.visualization.config import LIDAR_CONFIG - - -# def filter_lidar_pc(lidar_pc: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: -# """ -# Filter lidar point cloud according to global configuration -# :param lidar_pc: numpy array of shape (6,n) -# :return: filtered point cloud -# """ - -# pc = lidar_pc.T -# mask = ( -# np.ones((len(pc)), dtype=bool) -# & (pc[:, LidarIndex.X] > LIDAR_CONFIG["x_lim"][0]) -# & (pc[:, LidarIndex.X] < LIDAR_CONFIG["x_lim"][1]) -# & (pc[:, LidarIndex.Y] > LIDAR_CONFIG["y_lim"][0]) -# & (pc[:, LidarIndex.Y] < LIDAR_CONFIG["y_lim"][1]) -# & (pc[:, LidarIndex.Z] > LIDAR_CONFIG["z_lim"][0]) -# & (pc[:, LidarIndex.Z] < LIDAR_CONFIG["z_lim"][1]) -# ) -# pc = pc[mask] -# return pc.T - - -# def get_lidar_pc_color(lidar_pc: npt.NDArray[np.float32], as_hex: bool = False) -> List[Any]: -# """ -# Compute color map of lidar point cloud according to global configuration -# :param lidar_pc: numpy array of shape (6,n) -# :param as_hex: whether to return hex values, defaults to False -# :return: list of RGB or hex values -# """ - -# pc = lidar_pc.T -# if LIDAR_CONFIG["color_element"] == "none": -# colors_rgb = np.zeros((len(pc), 3), dtype=np.uin8) - -# else: -# if LIDAR_CONFIG["color_element"] == "distance": -# color_intensities = np.linalg.norm(pc[:, LidarIndex.POSITION], axis=-1) -# else: -# color_element_map = { -# "x": LidarIndex.X, -# "y": LidarIndex.Y, -# "z": LidarIndex.Z, -# "intensity": LidarIndex.INTENSITY, -# "ring": LidarIndex.RING, -# "id": LidarIndex.ID, -# } -# color_intensities = pc[:, color_element_map[LIDAR_CONFIG["color_element"]]] - -# min, max = color_intensities.min(), color_intensities.max() -# norm_intensities = [(value - min) / (max - min) for value in color_intensities] -# colormap = plt.get_cmap("viridis") -# colors_rgb = np.array([colormap(value) for value in norm_intensities]) -# colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8) - -# assert len(colors_rgb) == len(pc) -# if as_hex: -# return [matplotlib.colors.to_hex(tuple(c / 255.0 for c in rgb)) for rgb in colors_rgb] - -# return [tuple(value) for value in colors_rgb] +from typing import Literal + +import matplotlib.pyplot as plt +import numpy as np +import numpy.typing as npt + +from py123d.conversion.registry.lidar_index_registry import LiDARIndex + + +def get_lidar_pc_color( + lidar_pc: npt.NDArray[np.float32], + lidar_index: LiDARIndex, + feature: Literal["none", "distance", "intensity"], +) -> npt.NDArray[np.uint8]: + """ + Compute color map of lidar point cloud according to global configuration + :param lidar_pc: numpy array of shape (6,n) + :param as_hex: whether to return hex values, defaults to False + :return: list of RGB or hex values + """ + + lidar_xyz = lidar_pc[:, lidar_index.XYZ] + if feature == "none": + colors_rgb = np.zeros((len(lidar_xyz), 3), dtype=np.uin8) + else: + if feature == "distance": + color_intensities = np.linalg.norm(lidar_xyz, axis=-1) + elif feature == "intensity": + assert lidar_index.INTENSITY is not None, "LiDARIndex.INTENSITY is not defined" + color_intensities = lidar_pc[:, lidar_index.INTENSITY] + + min, max = color_intensities.min(), color_intensities.max() + norm_intensities = [(value - min) / (max - min) for value in color_intensities] + colormap = plt.get_cmap("viridis") + colors_rgb = np.array([colormap(value) for value in norm_intensities]) + colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8) + + return colors_rgb diff --git a/src/py123d/visualization/matplotlib/observation.py b/src/py123d/visualization/matplotlib/observation.py index e0a826f2..6d3ca8ff 100644 --- a/src/py123d/visualization/matplotlib/observation.py +++ b/src/py123d/visualization/matplotlib/observation.py @@ -1,18 +1,17 @@ +import traceback from typing import List, Optional, Union import matplotlib.pyplot as plt import numpy as np import shapely.geometry as geom -from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel +from py123d.api import MapAPI, SceneAPI from py123d.datatypes.detections.box_detections import BoxDetectionWrapper from py123d.datatypes.detections.traffic_light_detections import TrafficLightDetectionWrapper -from py123d.datatypes.maps.abstract_map import AbstractMap -from py123d.datatypes.maps.abstract_map_objects import AbstractLane -from py123d.datatypes.maps.map_datatypes import MapLayer -from py123d.datatypes.scene.abstract_scene import AbstractScene +from py123d.datatypes.map_objects.map_layer_types import MapLayer +from py123d.datatypes.map_objects.map_objects import Lane from py123d.datatypes.vehicle_state.ego_state import EgoStateSE2, EgoStateSE3 -from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, Point2D, StateSE2Index, Vector2D +from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, Point2D, PoseSE2Index, Vector2D from py123d.geometry.transform.transform_se2 import translate_se2_along_body_frame from py123d.visualization.color.config import PlotConfig from py123d.visualization.color.default import ( @@ -31,9 +30,31 @@ ) +def add_scene_on_ax(ax: plt.Axes, scene: SceneAPI, iteration: int = 0, radius: float = 80) -> plt.Axes: + ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) + box_detections = scene.get_box_detections_at_iteration(iteration) + traffic_light_detections = scene.get_traffic_light_detections_at_iteration(iteration) + map_api = scene.get_map_api() + + assert ego_vehicle_state is not None, "Ego vehicle state is required to plot the scene." + point_2d = ego_vehicle_state.bounding_box_se2.center_se2.pose_se2.point_2d + if map_api is not None: + add_default_map_on_ax(ax, map_api, point_2d, radius=radius) + if traffic_light_detections is not None: + add_traffic_lights_to_ax(ax, traffic_light_detections, map_api) + + add_box_detections_to_ax(ax, box_detections) + add_ego_vehicle_to_ax(ax, ego_vehicle_state) + + ax.set_xlim(point_2d.x - radius, point_2d.x + radius) + ax.set_ylim(point_2d.y - radius, point_2d.y + radius) + ax.set_aspect("equal", adjustable="box") + return ax + + def add_default_map_on_ax( ax: plt.Axes, - map_api: AbstractMap, + map_api: MapAPI, point_2d: Point2D, radius: float, route_lane_group_ids: Optional[List[int]] = None, @@ -59,7 +80,12 @@ def add_default_map_on_ax( if route_lane_group_ids is not None and int(map_object.object_id) in route_lane_group_ids: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, ROUTE_CONFIG) else: - add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) + add_shapely_polygon_to_ax( + ax, + map_object.shapely_polygon, + MAP_SURFACE_CONFIG[layer], + label=layer.serialize(), + ) if layer in [ MapLayer.GENERIC_DRIVABLE, MapLayer.CARPARK, @@ -67,13 +93,21 @@ def add_default_map_on_ax( MapLayer.INTERSECTION, MapLayer.WALKWAY, ]: - add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) + add_shapely_polygon_to_ax( + ax, + map_object.shapely_polygon, + MAP_SURFACE_CONFIG[layer], + label=layer.serialize(), + ) if layer in [MapLayer.LANE]: - map_object: AbstractLane - add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG) + map_object: Lane + add_shapely_linestring_to_ax( + ax, + map_object.centerline.linestring, + CENTERLINE_CONFIG, + label=layer.serialize(), + ) except Exception: - import traceback - print(f"Error adding map object of type {layer.name} and id {map_object.object_id}") traceback.print_exc() @@ -82,55 +116,20 @@ def add_default_map_on_ax( def add_box_detections_to_ax(ax: plt.Axes, box_detections: BoxDetectionWrapper) -> None: for box_detection in box_detections: - # TODO: Optionally, continue on boxes outside of plot. - # if box_detection.metadata.detection_type == DetectionType.GENERIC_OBJECT: - # continue plot_config = BOX_DETECTION_CONFIG[box_detection.metadata.default_label] - add_bounding_box_to_ax(ax, box_detection.bounding_box, plot_config) - - -def add_box_future_detections_to_ax(ax: plt.Axes, scene: AbstractScene, iteration: int) -> None: - - # TODO: Refactor this function - initial_agents = scene.get_box_detections_at_iteration(iteration) - agents_poses = { - agent.metadata.track_token: [agent.center_se3] - for agent in initial_agents - if agent.metadata.default_label == DefaultBoxDetectionLabel.VEHICLE - } - frequency = 1 - for iteration in range(iteration + frequency, scene.number_of_iterations, frequency): - agents = scene.get_box_detections_at_iteration(iteration) - for agent in agents: - if agent.metadata.track_token in agents_poses: - agents_poses[agent.metadata.track_token].append(agent.center_se3) - - for track_token, poses in agents_poses.items(): - if len(poses) < 2: - continue - poses = np.array([pose.point_2d.array for pose in poses]) - num_poses = poses.shape[0] - alphas = 1 - np.linspace(0.2, 1.0, num_poses) # Start low, end high - for i in range(num_poses - 1): - ax.plot( - poses[i : i + 2, 0], - poses[i : i + 2, 1], - color=BOX_DETECTION_CONFIG[DefaultBoxDetectionLabel.VEHICLE].fill_color.hex, - alpha=alphas[i + 1], - linewidth=BOX_DETECTION_CONFIG[DefaultBoxDetectionLabel.VEHICLE].line_width * 5, - zorder=BOX_DETECTION_CONFIG[DefaultBoxDetectionLabel.VEHICLE].zorder, - ) + add_bounding_box_to_ax(ax, box_detection.bounding_box_se2, plot_config) def add_ego_vehicle_to_ax(ax: plt.Axes, ego_vehicle_state: Union[EgoStateSE3, EgoStateSE2]) -> None: - add_bounding_box_to_ax(ax, ego_vehicle_state.bounding_box, EGO_VEHICLE_CONFIG) + add_bounding_box_to_ax(ax, ego_vehicle_state.bounding_box_se2, EGO_VEHICLE_CONFIG) def add_traffic_lights_to_ax( - ax: plt.Axes, traffic_light_detections: TrafficLightDetectionWrapper, map_api: AbstractMap + ax: plt.Axes, traffic_light_detections: TrafficLightDetectionWrapper, map_api: MapAPI ) -> None: for traffic_light_detection in traffic_light_detections: - lane: AbstractLane = map_api.get_map_object(str(traffic_light_detection.lane_id), MapLayer.LANE) + lane = map_api.get_map_object(traffic_light_detection.lane_id, MapLayer.LANE) + assert isinstance(lane, Lane), f"Lane with id {traffic_light_detection.lane_id} not found." if lane is not None: add_shapely_linestring_to_ax( ax, @@ -146,7 +145,6 @@ def add_bounding_box_to_ax( bounding_box: Union[BoundingBoxSE2, BoundingBoxSE3], plot_config: PlotConfig, ) -> None: - add_shapely_polygon_to_ax(ax, bounding_box.shapely_polygon, plot_config) if plot_config.marker_style is not None: @@ -158,7 +156,7 @@ def add_bounding_box_to_ax( arrow[1] = translate_se2_along_body_frame( center_se2, Vector2D(bounding_box.length / 2.0 + 0.5, 0.0), - ).array[StateSE2Index.XY] + ).array[PoseSE2Index.XY] ax.plot( arrow[:, 0], arrow[:, 1], @@ -169,9 +167,10 @@ def add_bounding_box_to_ax( linestyle=plot_config.line_style, ) elif plot_config.marker_style == "^": - marker_size = min(plot_config.marker_size, min(bounding_box.length, bounding_box.width)) + min_extent = min(bounding_box.length, bounding_box.width) + marker_size = min(plot_config.marker_size, min_extent) marker_polygon = get_pose_triangle(marker_size) - global_marker_polygon = shapely_geometry_local_coords(marker_polygon, bounding_box.center) + global_marker_polygon = shapely_geometry_local_coords(marker_polygon, bounding_box.center_se2) add_shapely_polygon_to_ax(ax, global_marker_polygon, plot_config, disable_smoothing=True) else: raise ValueError(f"Unknown marker style: {plot_config.marker_style}") diff --git a/src/py123d/visualization/matplotlib/plots.py b/src/py123d/visualization/matplotlib/plots.py index 01100f01..aa3e23d7 100644 --- a/src/py123d/visualization/matplotlib/plots.py +++ b/src/py123d/visualization/matplotlib/plots.py @@ -1,54 +1,22 @@ from pathlib import Path from typing import Optional, Tuple, Union -import matplotlib.animation as animation import matplotlib.pyplot as plt +from matplotlib import animation from tqdm import tqdm -from py123d.datatypes.scene.abstract_scene import AbstractScene -from py123d.visualization.matplotlib.observation import ( - add_box_detections_to_ax, - add_default_map_on_ax, - add_ego_vehicle_to_ax, - add_traffic_lights_to_ax, -) +from py123d.api.scene.scene_api import SceneAPI +from py123d.visualization.matplotlib.observation import add_scene_on_ax -def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes: - - ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) - box_detections = scene.get_box_detections_at_iteration(iteration) - traffic_light_detections = scene.get_traffic_light_detections_at_iteration(iteration) - route_lane_group_ids = scene.get_route_lane_group_ids(iteration) - map_api = scene.get_map_api() - - point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d - if map_api is not None: - add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=route_lane_group_ids) - if traffic_light_detections is not None: - add_traffic_lights_to_ax(ax, traffic_light_detections, map_api) - - add_box_detections_to_ax(ax, box_detections) - add_ego_vehicle_to_ax(ax, ego_vehicle_state) - - ax.set_xlim(point_2d.x - radius, point_2d.x + radius) - ax.set_ylim(point_2d.y - radius, point_2d.y + radius) - - ax.set_aspect("equal", adjustable="box") - return ax - - -def plot_scene_at_iteration( - scene: AbstractScene, iteration: int = 0, radius: float = 80 -) -> Tuple[plt.Figure, plt.Axes]: - +def plot_scene_at_iteration(scene: SceneAPI, iteration: int = 0, radius: float = 80) -> Tuple[plt.Figure, plt.Axes]: fig, ax = plt.subplots(figsize=(10, 10)) - _plot_scene_on_ax(ax, scene, iteration, radius) + add_scene_on_ax(ax, scene, iteration, radius) return fig, ax def render_scene_animation( - scene: AbstractScene, + scene: SceneAPI, output_path: Union[str, Path], start_idx: int = 0, end_idx: Optional[int] = None, @@ -56,7 +24,7 @@ def render_scene_animation( fps: float = 20.0, dpi: int = 300, format: str = "mp4", - radius: float = 100, + radius: float = 80, ) -> None: assert format in ["mp4", "gif"], "Format must be either 'mp4' or 'gif'." output_path = Path(output_path) @@ -70,7 +38,7 @@ def render_scene_animation( def update(i): ax.clear() - _plot_scene_on_ax(ax, scene, i, radius) + add_scene_on_ax(ax, scene, i, radius) plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05) pbar.update(1) @@ -78,5 +46,5 @@ def update(i): pbar = tqdm(total=len(frames), desc=f"Rendering {scene.log_name} as {format}") ani = animation.FuncAnimation(fig, update, frames=frames, repeat=False) - ani.save(output_path / f"{scene.log_name}_{scene.uuid}.{format}", writer="ffmpeg", fps=fps, dpi=dpi) + ani.save(output_path / f"{scene.log_name}_{scene.scene_uuid}.{format}", writer="ffmpeg", fps=fps, dpi=dpi) plt.close(fig) diff --git a/src/py123d/visualization/matplotlib/utils.py b/src/py123d/visualization/matplotlib/utils.py index 81c60260..b888d67b 100644 --- a/src/py123d/visualization/matplotlib/utils.py +++ b/src/py123d/visualization/matplotlib/utils.py @@ -1,13 +1,13 @@ -from typing import Union +from typing import Optional, Union -import matplotlib.patches as patches import matplotlib.pyplot as plt import numpy as np -import shapely.affinity as affinity import shapely.geometry as geom +from matplotlib import patches from matplotlib.path import Path +from shapely import affinity -from py123d.geometry import StateSE2, StateSE3 +from py123d.geometry import PoseSE2, PoseSE3 from py123d.visualization.color.config import PlotConfig @@ -16,9 +16,10 @@ def add_shapely_polygon_to_ax( polygon: geom.Polygon, plot_config: PlotConfig, disable_smoothing: bool = False, + label: Optional[str] = None, ) -> plt.Axes: - """ - Adds shapely polygon to birds-eye-view visualization with proper hole handling + """Adds shapely polygon to birds-eye-view visualization with proper hole handling + :param ax: matplotlib ax object :param polygon: shapely Polygon :param plot_config: dictionary containing plot parameters @@ -59,6 +60,7 @@ def create_polygon_path(polygon): edgecolor=plot_config.line_color.hex, linewidth=plot_config.line_width, zorder=plot_config.zorder, + label=label, ) ax.add_patch(patch) @@ -77,12 +79,13 @@ def add_shapely_linestring_to_ax( ax: plt.Axes, linestring: geom.LineString, plot_config: PlotConfig, + label: Optional[str] = None, ) -> plt.Axes: - """ - Adds shapely linestring (polyline) to birds-eye-view visualization + """Adds shapely linestring (polyline) to birds-eye-view visualization + :param ax: matplotlib ax object :param linestring: shapely LineString - :param config: dictionary containing plot parameters + :param plot_config: dictionary containing plot parameters :return: ax with plot """ @@ -95,6 +98,7 @@ def add_shapely_linestring_to_ax( linewidth=plot_config.line_width, linestyle=plot_config.line_style, zorder=plot_config.zorder, + label=label, ) return ax @@ -113,7 +117,7 @@ def get_pose_triangle(size: float) -> geom.Polygon: def shapely_geometry_local_coords( - geometry: geom.base.BaseGeometry, origin: Union[StateSE2, StateSE3] + geometry: geom.base.BaseGeometry, origin: Union[PoseSE2, PoseSE3] ) -> geom.base.BaseGeometry: """Helper for transforming shapely geometry in coord-frame""" # TODO: move somewhere else for general use diff --git a/src/py123d/visualization/viser/elements/detection_elements.py b/src/py123d/visualization/viser/elements/detection_elements.py index be08021b..c77679c4 100644 --- a/src/py123d/visualization/viser/elements/detection_elements.py +++ b/src/py123d/visualization/viser/elements/detection_elements.py @@ -5,10 +5,10 @@ import trimesh import viser +from py123d.api.scene.scene_api import SceneAPI from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel -from py123d.datatypes.scene.abstract_scene import AbstractScene from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.geometry.geometry_index import BoundingBoxSE3Index, Corners3DIndex, StateSE3Index +from py123d.geometry.geometry_index import BoundingBoxSE3Index, Corners3DIndex, PoseSE3Index from py123d.geometry.utils.bounding_box_utils import ( bbse3_array_to_corners_array, corners_array_to_3d_mesh, @@ -19,7 +19,7 @@ def add_box_detections_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, scene_interation: int, initial_ego_state: EgoStateSE3, viser_server: viser.ViserServer, @@ -47,15 +47,15 @@ def add_box_detections_to_viser_server( ) # viser_server.scene.add_batched_axes( # "frames", - # batched_wxyzs=se3_array[:-1, StateSE3Index.QUATERNION], - # batched_positions=se3_array[:-1, StateSE3Index.XYZ], + # batched_wxyzs=se3_array[:-1, PoseSE3Index.QUATERNION], + # batched_positions=se3_array[:-1, PoseSE3Index.XYZ], # ) # ego_rear_axle_se3 = scene.get_ego_state_at_iteration(scene_interation).rear_axle_se3.array - # ego_rear_axle_se3[StateSE3Index.XYZ] -= initial_ego_state.center_se3.array[StateSE3Index.XYZ] + # ego_rear_axle_se3[PoseSE3Index.XYZ] -= initial_ego_state.center_se3.array[PoseSE3Index.XYZ] # viser_server.scene.add_frame( # "ego_rear_axle", - # position=ego_rear_axle_se3[StateSE3Index.XYZ], - # wxyz=ego_rear_axle_se3[StateSE3Index.QUATERNION], + # position=ego_rear_axle_se3[PoseSE3Index.XYZ], + # wxyz=ego_rear_axle_se3[PoseSE3Index.QUATERNION], # ) visible_handle_keys.append("lines") @@ -67,8 +67,7 @@ def add_box_detections_to_viser_server( box_detection_handles[key].visible = False -def _get_bounding_box_meshes(scene: AbstractScene, iteration: int, initial_ego_state: EgoStateSE3) -> trimesh.Trimesh: - +def _get_bounding_box_meshes(scene: SceneAPI, iteration: int, initial_ego_state: EgoStateSE3) -> trimesh.Trimesh: ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) box_detections = scene.get_box_detections_at_iteration(iteration) @@ -78,7 +77,7 @@ def _get_bounding_box_meshes(scene: AbstractScene, iteration: int, initial_ego_s # create meshes for all boxes box_se3_array = np.array([box.array for box in boxes]) - box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[StateSE3Index.XYZ] + box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[PoseSE3Index.XYZ] box_corners_array = bbse3_array_to_corners_array(box_se3_array) box_vertices, box_faces = corners_array_to_3d_mesh(box_corners_array) @@ -111,7 +110,7 @@ def _get_bounding_box_meshes(scene: AbstractScene, iteration: int, initial_ego_s # # Create lines for all boxes # box_se3_array = np.array([box.array for box in boxes]) -# box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[StateSE3Index.XYZ] +# box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[PoseSE3Index.XYZ] # box_corners_array = bbse3_array_to_corners_array(box_se3_array) # box_outlines = corners_array_to_edge_lines(box_corners_array) @@ -127,9 +126,8 @@ def _get_bounding_box_meshes(scene: AbstractScene, iteration: int, initial_ego_s def _get_bounding_box_outlines( - scene: AbstractScene, iteration: int, initial_ego_state: EgoStateSE3 + scene: SceneAPI, iteration: int, initial_ego_state: EgoStateSE3 ) -> npt.NDArray[np.float64]: - ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) box_detections = scene.get_box_detections_at_iteration(iteration) @@ -139,7 +137,7 @@ def _get_bounding_box_outlines( # Create lines for all boxes box_se3_array = np.array([box.array for box in boxes]) - box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[StateSE3Index.XYZ] + box_se3_array[..., BoundingBoxSE3Index.XYZ] -= initial_ego_state.center_se3.array[PoseSE3Index.XYZ] box_corners_array = bbse3_array_to_corners_array(box_se3_array) box_outlines = corners_array_to_edge_lines(box_corners_array) diff --git a/src/py123d/visualization/viser/elements/map_elements.py b/src/py123d/visualization/viser/elements/map_elements.py index 6e16726e..c211b982 100644 --- a/src/py123d/visualization/viser/elements/map_elements.py +++ b/src/py123d/visualization/viser/elements/map_elements.py @@ -4,9 +4,9 @@ import trimesh import viser -from py123d.datatypes.maps.abstract_map import MapLayer -from py123d.datatypes.maps.abstract_map_objects import AbstractSurfaceMapObject -from py123d.datatypes.scene.abstract_scene import AbstractScene +from py123d.api import SceneAPI +from py123d.datatypes.map_objects.base_map_objects import BaseMapSurfaceObject +from py123d.datatypes.map_objects.map_layer_types import MapLayer from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 from py123d.geometry import Point3D, Point3DIndex from py123d.visualization.color.default import MAP_SURFACE_CONFIG @@ -16,17 +16,16 @@ def add_map_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, iteration: int, initial_ego_state: EgoStateSE3, viser_server: viser.ViserServer, viser_config: ViserConfig, map_handles: Dict[MapLayer, viser.GlbHandle], ) -> None: - global last_query_position + global last_query_position # noqa: PLW0603 if viser_config.map_visible: - map_trimesh_dict: Optional[Dict[MapLayer, trimesh.Trimesh]] = None if len(map_handles) == 0 or viser_config._force_map_update: @@ -66,19 +65,18 @@ def add_map_to_viser_server( def _get_map_trimesh_dict( - scene: AbstractScene, + scene: SceneAPI, initial_ego_state: EgoStateSE3, current_ego_state: Optional[EgoStateSE3], viser_config: ViserConfig, ) -> Dict[MapLayer, trimesh.Trimesh]: - # Dictionary to hold the output trimesh meshes. output_trimesh_dict: Dict[MapLayer, trimesh.Trimesh] = {} # Unpack scene center for translation of map objects. - scene_center: Point3D = initial_ego_state.center.point_3d + scene_center: Point3D = initial_ego_state.center_se3.point_3d scene_center_array = scene_center.array - scene_query_position = current_ego_state.center.point_3d + scene_query_position = current_ego_state.center_se3.point_3d # Load map objects within a certain radius around the scene center. map_layers = [ @@ -91,7 +89,7 @@ def _get_map_trimesh_dict( ] map_api = scene.get_map_api() if map_api is not None: - map_objects_dict = map_api.get_proximal_map_objects( + map_objects_dict = map_api.get_map_objects_in_radius( scene_query_position.point_2d, radius=viser_config.map_radius, layers=map_layers, @@ -101,7 +99,7 @@ def _get_map_trimesh_dict( for map_layer in map_objects_dict.keys(): surface_meshes = [] for map_surface in map_objects_dict[map_layer]: - map_surface: AbstractSurfaceMapObject + map_surface: BaseMapSurfaceObject trimesh_mesh = map_surface.trimesh_mesh trimesh_mesh.vertices -= scene_center_array diff --git a/src/py123d/visualization/viser/elements/render_elements.py b/src/py123d/visualization/viser/elements/render_elements.py index f807033e..47784cb2 100644 --- a/src/py123d/visualization/viser/elements/render_elements.py +++ b/src/py123d/visualization/viser/elements/render_elements.py @@ -1,24 +1,21 @@ import numpy as np +from py123d.api.scene.scene_api import SceneAPI from py123d.conversion.utils.sensor_utils.camera_conventions import convert_camera_convention -from py123d.datatypes.scene.abstract_scene import AbstractScene from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.geometry.geometry_index import StateSE3Index -from py123d.geometry.rotation import EulerAngles -from py123d.geometry.se import StateSE3 +from py123d.geometry import EulerAngles, PoseSE3, PoseSE3Index, Vector3D from py123d.geometry.transform.transform_se3 import translate_se3_along_body_frame -from py123d.geometry.vector import Vector3D def get_ego_3rd_person_view_position( - scene: AbstractScene, + scene: SceneAPI, iteration: int, initial_ego_state: EgoStateSE3, -) -> StateSE3: - scene_center_array = initial_ego_state.center.point_3d.array +) -> PoseSE3: + scene_center_array = initial_ego_state.center_se3.point_3d.array ego_pose = scene.get_ego_state_at_iteration(iteration).rear_axle_se3.array - ego_pose[StateSE3Index.XYZ] -= scene_center_array - ego_pose_se3 = StateSE3.from_array(ego_pose) + ego_pose[PoseSE3Index.XYZ] -= scene_center_array + ego_pose_se3 = PoseSE3.from_array(ego_pose) ego_pose_se3 = translate_se3_along_body_frame(ego_pose_se3, Vector3D(-15.0, 0.0, 15)) # adjust the pitch to -10 degrees. @@ -36,18 +33,18 @@ def get_ego_3rd_person_view_position( def get_ego_bev_view_position( - scene: AbstractScene, + scene: SceneAPI, iteration: int, initial_ego_state: EgoStateSE3, -) -> StateSE3: - scene_center_array = initial_ego_state.center.point_3d.array - ego_center = scene.get_ego_state_at_iteration(iteration).center.array - ego_center[StateSE3Index.XYZ] -= scene_center_array - ego_center_planar = StateSE3.from_array(ego_center) +) -> PoseSE3: + scene_center_array = initial_ego_state.center_se3.point_3d.array + ego_center = scene.get_ego_state_at_iteration(iteration).center_se3.array + ego_center[PoseSE3Index.XYZ] -= scene_center_array + ego_center_planar = PoseSE3.from_array(ego_center) planar_euler_angles = EulerAngles(0.0, 0.0, ego_center_planar.euler_angles.yaw) quaternion = planar_euler_angles.quaternion - ego_center_planar._array[StateSE3Index.QUATERNION] = quaternion.array + ego_center_planar._array[PoseSE3Index.QUATERNION] = quaternion.array ego_center_planar = translate_se3_along_body_frame(ego_center_planar, Vector3D(0.0, 0.0, 50)) ego_center_planar = _pitch_se3_by_degrees(ego_center_planar, 90.0) @@ -59,14 +56,13 @@ def get_ego_bev_view_position( ) -def _pitch_se3_by_degrees(state_se3: StateSE3, degrees: float) -> StateSE3: +def _pitch_se3_by_degrees(pose_se3: PoseSE3, degrees: float) -> PoseSE3: + quaternion = EulerAngles(0.0, np.deg2rad(degrees), pose_se3.yaw).quaternion - quaternion = EulerAngles(0.0, np.deg2rad(degrees), state_se3.yaw).quaternion - - return StateSE3( - x=state_se3.x, - y=state_se3.y, - z=state_se3.z, + return PoseSE3( + x=pose_se3.x, + y=pose_se3.y, + z=pose_se3.z, qw=quaternion.qw, qx=quaternion.qx, qy=quaternion.qy, diff --git a/src/py123d/visualization/viser/elements/sensor_elements.py b/src/py123d/visualization/viser/elements/sensor_elements.py index 2cc6bacb..3751da6a 100644 --- a/src/py123d/visualization/viser/elements/sensor_elements.py +++ b/src/py123d/visualization/viser/elements/sensor_elements.py @@ -6,12 +6,17 @@ import numpy.typing as npt import viser -from py123d.datatypes.scene.abstract_scene import AbstractScene -from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICamera, FisheyeMEICameraMetadata, FisheyeMEICameraType -from py123d.datatypes.sensors.lidar import LiDARType -from py123d.datatypes.sensors.pinhole_camera import PinholeCamera, PinholeCameraType +from py123d.api.scene.scene_api import SceneAPI +from py123d.datatypes.sensors import ( + FisheyeMEICamera, + FisheyeMEICameraMetadata, + FisheyeMEICameraType, + LiDARType, + PinholeCamera, + PinholeCameraType, +) from py123d.datatypes.vehicle_state.ego_state import EgoStateSE3 -from py123d.geometry import StateSE3Index +from py123d.geometry import PoseSE3Index from py123d.geometry.transform.transform_se3 import ( convert_relative_to_absolute_points_3d_array, convert_relative_to_absolute_se3_array, @@ -21,18 +26,17 @@ def add_camera_frustums_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, scene_interation: int, initial_ego_state: EgoStateSE3, viser_server: viser.ViserServer, viser_config: ViserConfig, camera_frustum_handles: Dict[PinholeCameraType, viser.CameraFrustumHandle], ) -> None: - if viser_config.camera_frustum_visible: - scene_center_array = initial_ego_state.center.point_3d.array + scene_center_array = initial_ego_state.center_se3.point_3d.array ego_pose = scene.get_ego_state_at_iteration(scene_interation).rear_axle_se3.array - ego_pose[StateSE3Index.XYZ] -= scene_center_array + ego_pose[PoseSE3Index.XYZ] -= scene_center_array def _add_camera_frustums_to_viser_server(camera_type: PinholeCameraType) -> None: camera = scene.get_pinhole_camera_at_iteration(scene_interation, camera_type) @@ -57,8 +61,6 @@ def _add_camera_frustums_to_viser_server(camera_type: PinholeCameraType) -> None wxyz=camera_quaternion, ) - return None - # NOTE; In order to speed up adding camera frustums, we use multithreading and resize the images. with concurrent.futures.ThreadPoolExecutor(max_workers=len(viser_config.camera_frustum_types)) as executor: future_to_camera = { @@ -76,7 +78,7 @@ def _add_camera_frustums_to_viser_server(camera_type: PinholeCameraType) -> None def add_fisheye_frustums_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, scene_interation: int, initial_ego_state: EgoStateSE3, viser_server: viser.ViserServer, @@ -84,9 +86,9 @@ def add_fisheye_frustums_to_viser_server( fisheye_frustum_handles: Dict[FisheyeMEICameraType, viser.CameraFrustumHandle], ) -> None: if viser_config.fisheye_frustum_visible: - scene_center_array = initial_ego_state.center.point_3d.array + scene_center_array = initial_ego_state.center_se3.point_3d.array ego_pose = scene.get_ego_state_at_iteration(scene_interation).rear_axle_se3.array - ego_pose[StateSE3Index.XYZ] -= scene_center_array + ego_pose[PoseSE3Index.XYZ] -= scene_center_array def _add_fisheye_frustums_to_viser_server(fisheye_camera_type: FisheyeMEICameraType) -> None: camera = scene.get_fisheye_mei_camera_at_iteration(scene_interation, fisheye_camera_type) @@ -113,8 +115,6 @@ def _add_fisheye_frustums_to_viser_server(fisheye_camera_type: FisheyeMEICameraT wxyz=fcam_quaternion, ) - return None - # NOTE; In order to speed up adding camera frustums, we use multithreading and resize the images. with concurrent.futures.ThreadPoolExecutor( max_workers=len(viser_config.fisheye_mei_camera_frustum_types) @@ -130,7 +130,7 @@ def _add_fisheye_frustums_to_viser_server(fisheye_camera_type: FisheyeMEICameraT def add_camera_gui_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, scene_interation: int, viser_server: viser.ViserServer, viser_config: ViserConfig, @@ -153,7 +153,7 @@ def add_camera_gui_to_viser_server( def add_lidar_pc_to_viser_server( - scene: AbstractScene, + scene: SceneAPI, scene_interation: int, initial_ego_state: EgoStateSE3, viser_server: viser.ViserServer, @@ -161,10 +161,9 @@ def add_lidar_pc_to_viser_server( lidar_pc_handles: Dict[LiDARType, Optional[viser.PointCloudHandle]], ) -> None: if viser_config.lidar_visible: - - scene_center_array = initial_ego_state.center.point_3d.array + scene_center_array = initial_ego_state.center_se3.point_3d.array ego_pose = scene.get_ego_state_at_iteration(scene_interation).rear_axle_se3.array - ego_pose[StateSE3Index.XYZ] -= scene_center_array + ego_pose[PoseSE3Index.XYZ] -= scene_center_array def _load_lidar_points(lidar_type: LiDARType) -> npt.NDArray[np.float32]: lidar = scene.get_lidar_at_iteration(scene_interation, lidar_type) @@ -202,8 +201,8 @@ def _load_lidar_points(lidar_type: LiDARType) -> npt.NDArray[np.float32]: # viser_server.scene.add_frame( # "lidar_frame", - # position=lidar_extrinsic[StateSE3Index.XYZ], - # wxyz=lidar_extrinsic[StateSE3Index.QUATERNION], + # position=lidar_extrinsic[PoseSE3Index.XYZ], + # wxyz=lidar_extrinsic[PoseSE3Index.QUATERNION], # ) if lidar_pc_handles[LiDARType.LIDAR_MERGED] is not None: @@ -226,13 +225,13 @@ def _get_camera_values( ego_pose: npt.NDArray[np.float64], resize_factor: Optional[float] = None, ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.uint8]]: - assert ego_pose.ndim == 1 and len(ego_pose) == len(StateSE3Index) + assert ego_pose.ndim == 1 and len(ego_pose) == len(PoseSE3Index) rel_camera_pose = camera.extrinsic.array abs_camera_pose = convert_relative_to_absolute_se3_array(origin=ego_pose, se3_array=rel_camera_pose) - camera_position = abs_camera_pose[StateSE3Index.XYZ] - camera_rotation = abs_camera_pose[StateSE3Index.QUATERNION] + camera_position = abs_camera_pose[PoseSE3Index.XYZ] + camera_rotation = abs_camera_pose[PoseSE3Index.QUATERNION] camera_image = _rescale_image(camera.image, resize_factor) return camera_position, camera_rotation, camera_image @@ -243,13 +242,13 @@ def _get_fisheye_camera_values( ego_pose: npt.NDArray[np.float64], resize_factor: Optional[float] = None, ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64], npt.NDArray[np.uint8]]: - assert ego_pose.ndim == 1 and len(ego_pose) == len(StateSE3Index) + assert ego_pose.ndim == 1 and len(ego_pose) == len(PoseSE3Index) rel_camera_pose = camera.extrinsic.array abs_camera_pose = convert_relative_to_absolute_se3_array(origin=ego_pose, se3_array=rel_camera_pose) - camera_position = abs_camera_pose[StateSE3Index.XYZ] - camera_rotation = abs_camera_pose[StateSE3Index.QUATERNION] + camera_position = abs_camera_pose[PoseSE3Index.XYZ] + camera_rotation = abs_camera_pose[PoseSE3Index.QUATERNION] camera_image = _rescale_image(camera.image, resize_factor) return camera_position, camera_rotation, camera_image @@ -264,9 +263,6 @@ def _rescale_image(image: npt.NDArray[np.uint8], scale: float) -> npt.NDArray[np return downscaled_image -import numpy as np - - def calculate_fov(metadata: FisheyeMEICameraMetadata) -> tuple[float, float]: """ Calculate horizontal and vertical FOV in degrees. diff --git a/src/py123d/visualization/viser/viser_config.py b/src/py123d/visualization/viser/viser_config.py index be22f7f5..2907a45d 100644 --- a/src/py123d/visualization/viser/viser_config.py +++ b/src/py123d/visualization/viser/viser_config.py @@ -33,7 +33,6 @@ @dataclass class ViserConfig: - # Server server_host: str = "localhost" server_port: int = 8080 @@ -79,7 +78,7 @@ class ViserConfig: # -> Frustum fisheye_frustum_visible: bool = True fisheye_mei_camera_frustum_visible: bool = True - fisheye_mei_camera_frustum_types: List[PinholeCameraType] = field( + fisheye_mei_camera_frustum_types: List[FisheyeMEICameraType] = field( default_factory=lambda: [fcam for fcam in FisheyeMEICameraType] ) fisheye_frustum_scale: float = 1.0 @@ -98,7 +97,6 @@ def __post_init__(self): def _resolve_enum_arguments( serial_enum_cls: SerialIntEnum, input: Optional[List[Union[int, str, SerialIntEnum]]] ) -> List[SerialIntEnum]: - if input is None: return None assert isinstance(input, list), f"input must be a list of {serial_enum_cls.__name__}" @@ -109,7 +107,7 @@ def _resolve_enum_arguments( self.camera_frustum_types, ) self.camera_gui_types = _resolve_enum_arguments( - FisheyeMEICameraType, + PinholeCameraType, self.camera_gui_types, ) self.fisheye_mei_camera_frustum_types = _resolve_enum_arguments( diff --git a/src/py123d/visualization/viser/viser_viewer.py b/src/py123d/visualization/viser/viser_viewer.py index 1fb5559e..5b3dc13c 100644 --- a/src/py123d/visualization/viser/viser_viewer.py +++ b/src/py123d/visualization/viser/viser_viewer.py @@ -8,8 +8,8 @@ from tqdm import tqdm from viser.theme import TitlebarButton, TitlebarConfig, TitlebarImage -from py123d.datatypes.maps.map_datatypes import MapLayer -from py123d.datatypes.scene.abstract_scene import AbstractScene +from py123d.api.scene.scene_api import SceneAPI +from py123d.datatypes.map_objects.map_layer_types import MapLayer from py123d.datatypes.sensors.fisheye_mei_camera import FisheyeMEICameraType from py123d.datatypes.sensors.lidar import LiDARType from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType @@ -43,24 +43,24 @@ def _build_viser_server(viser_config: ViserConfig) -> viser.ViserServer: TitlebarButton( text="Getting Started", icon=None, - href="https://danieldauner.github.io/py123d", + href="https://autonomousvision.github.io/py123d", ), TitlebarButton( text="Github", icon="GitHub", - href="https://github.com/DanielDauner/py123d", + href="https://github.com/autonomousvision/py123d", ), TitlebarButton( text="Documentation", icon="Description", - href="https://danieldauner.github.io/py123d", + href="https://autonomousvision.github.io/py123d", ), ) image = TitlebarImage( - image_url_light="https://danieldauner.github.io/py123d/_static/logo_black.png", - image_url_dark="https://danieldauner.github.io/py123d/_static/logo_white.png", + image_url_light="https://autonomousvision.github.io/py123d/_static/123D_logo_transparent_black.svg", + image_url_dark="https://autonomousvision.github.io/py123d/_static/123D_logo_transparent_white.svg", image_alt="123D", - href="https://danieldauner.github.io/py123d/", + href="https://autonomousvision.github.io/py123d/", ) titlebar_theme = TitlebarConfig(buttons=buttons, image=image) @@ -79,7 +79,7 @@ def _build_viser_server(viser_config: ViserConfig) -> viser.ViserServer: class ViserViewer: def __init__( self, - scenes: List[AbstractScene], + scenes: List[SceneAPI], viser_config: ViserConfig = ViserConfig(), scene_index: int = 0, ) -> None: @@ -99,14 +99,16 @@ def next(self) -> None: self._scene_index = (self._scene_index + 1) % len(self._scenes) self.set_scene(self._scenes[self._scene_index]) - def set_scene(self, scene: AbstractScene) -> None: + def set_scene(self, scene: SceneAPI) -> None: num_frames = scene.number_of_iterations - initial_ego_state: EgoStateSE3 = scene.get_ego_state_at_iteration(0) + initial_ego_state = scene.get_ego_state_at_iteration(0) + assert initial_ego_state is not None and isinstance(initial_ego_state, EgoStateSE3) + server_playing = True server_rendering = False with self._viser_server.gui.add_folder("Playback"): - gui_info = self._viser_server.gui.add_markdown(content=_get_scene_info_markdown(scene)) + self._viser_server.gui.add_markdown(content=_get_scene_info_markdown(scene)) gui_timestep = self._viser_server.gui.add_slider( "Timestep", @@ -316,7 +318,7 @@ def _(event: viser.GuiEvent) -> None: elif format == "mp4": iio.imwrite(buffer, images, extension=".mp4", fps=20) content = buffer.getvalue() - scene_name = f"{scene.log_metadata.split}_{scene.uuid}" + scene_name = f"{scene.log_metadata.split}_{scene.scene_uuid}" client.send_file_download(f"{scene_name}.{format}", content, save_immediately=True) server_rendering = False @@ -382,6 +384,8 @@ def _(event: viser.GuiEvent) -> None: while server_playing: if gui_playing.value and not server_rendering: gui_timestep.value = (gui_timestep.value + 1) % num_frames + else: + time.sleep(0.1) # update config self._viser_config.playback_speed = gui_speed.value @@ -390,10 +394,11 @@ def _(event: viser.GuiEvent) -> None: self.next() -def _get_scene_info_markdown(scene: AbstractScene) -> str: +def _get_scene_info_markdown(scene: SceneAPI) -> str: markdown = f""" - Dataset: {scene.log_metadata.split} - - Location: {scene.log_metadata.location if scene.log_metadata.location else 'N/A'} - - UUID: {scene.uuid} + - Location: {scene.log_metadata.location if scene.log_metadata.location else "N/A"} + - Log: {scene.log_metadata.log_name} + - UUID: {scene.scene_uuid} """ return markdown diff --git a/tests/unit/api/__init__.py b/tests/unit/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/api/api/__init__.py b/tests/unit/api/api/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/api/api/test_scene_api.py b/tests/unit/api/api/test_scene_api.py new file mode 100644 index 00000000..91ad1e26 --- /dev/null +++ b/tests/unit/api/api/test_scene_api.py @@ -0,0 +1,207 @@ +from typing import Optional +from unittest.mock import Mock + +import pytest + +from py123d.api import MapAPI, SceneAPI, SceneMetadata +from py123d.datatypes.detections import BoxDetectionWrapper, TrafficLightDetectionWrapper +from py123d.datatypes.metadata import LogMetadata, MapMetadata +from py123d.datatypes.sensors import ( + FisheyeMEICamera, + FisheyeMEICameraType, + LiDAR, + LiDARType, + PinholeCamera, + PinholeCameraType, +) +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import EgoStateSE3, VehicleParameters + + +class ConcreteSceneAPI(SceneAPI): + """Concrete implementation for testing purposes.""" + + def __init__(self): + self._log_metadata = Mock(spec=LogMetadata) + self._scene_metadata = Mock(spec=SceneMetadata) + self._map_api = Mock(spec=MapAPI) + + def get_log_metadata(self) -> LogMetadata: + """Inherited, see super class.""" + return self._log_metadata + + def get_scene_metadata(self) -> SceneMetadata: + """Inherited, see super class.""" + return self._scene_metadata + + def get_map_api(self) -> Optional[MapAPI]: + """Inherited, see super class.""" + return self._map_api + + def get_timepoint_at_iteration(self, iteration: int) -> TimePoint: + """Inherited, see super class.""" + return Mock(spec=TimePoint) + + def get_ego_state_at_iteration(self, iteration: int) -> Optional[EgoStateSE3]: + """Inherited, see super class.""" + return Mock(spec=EgoStateSE3) + + def get_box_detections_at_iteration(self, iteration: int) -> Optional[BoxDetectionWrapper]: + """Inherited, see super class.""" + return Mock(spec=BoxDetectionWrapper) + + def get_traffic_light_detections_at_iteration(self, iteration: int) -> Optional[TrafficLightDetectionWrapper]: + """Inherited, see super class.""" + return Mock(spec=TrafficLightDetectionWrapper) + + def get_route_lane_group_ids(self, iteration: int) -> Optional[list]: + """Inherited, see super class.""" + return [1, 2, 3] + + def get_pinhole_camera_at_iteration( + self, iteration: int, camera_type: PinholeCameraType + ) -> Optional[PinholeCamera]: + """Inherited, see super class.""" + return Mock(spec=PinholeCamera) + + def get_fisheye_mei_camera_at_iteration( + self, iteration: int, camera_type: FisheyeMEICameraType + ) -> Optional[FisheyeMEICamera]: + """Inherited, see super class.""" + return Mock(spec=FisheyeMEICamera) + + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> Optional[LiDAR]: + """Inherited, see super class.""" + return Mock(spec=LiDAR) + + +@pytest.fixture +def scene_api(): + """Fixture providing a concrete SceneAPI instance.""" + api = ConcreteSceneAPI() + api._log_metadata.dataset = "test_dataset" + api._log_metadata.split = "test_split" + api._log_metadata.location = "test_location" + api._log_metadata.log_name = "test_log" + api._log_metadata.version = "1.0.0" + api._log_metadata.map_metadata = Mock(spec=MapMetadata) + api._log_metadata.vehicle_parameters = Mock(spec=VehicleParameters) + api._log_metadata.pinhole_camera_metadata = {PinholeCameraType.PCAM_B0: Mock()} + api._log_metadata.fisheye_mei_camera_metadata = {FisheyeMEICameraType.FCAM_L: Mock()} + api._log_metadata.lidar_metadata = {LiDARType.LIDAR_TOP: Mock()} + api._scene_metadata.initial_uuid = "test-uuid-123" + api._scene_metadata.number_of_iterations = 100 + api._scene_metadata.number_of_history_iterations = 10 + return api + + +class TestSceneAPIProperties: + """Test property accessors of SceneAPI.""" + + def test_log_metadata(self, scene_api): + """Test log_metadata property.""" + assert scene_api.log_metadata == scene_api._log_metadata + + def test_scene_metadata(self, scene_api): + """Test scene_metadata property.""" + assert scene_api.scene_metadata == scene_api._scene_metadata + + def test_map_metadata(self, scene_api): + """Test map_metadata property.""" + assert scene_api.map_metadata == scene_api._log_metadata.map_metadata + + def test_map_api(self, scene_api): + """Test map_api property.""" + assert scene_api.map_api == scene_api._map_api + + def test_dataset(self, scene_api): + """Test dataset property.""" + assert scene_api.dataset == "test_dataset" + + def test_split(self, scene_api): + """Test split property.""" + assert scene_api.split == "test_split" + + def test_location(self, scene_api): + """Test location property.""" + assert scene_api.location == "test_location" + + def test_log_name(self, scene_api): + """Test log_name property.""" + assert scene_api.log_name == "test_log" + + def test_version(self, scene_api): + """Test version property.""" + assert scene_api.version == "1.0.0" + + def test_scene_uuid(self, scene_api): + """Test scene_uuid property.""" + assert scene_api.scene_uuid == "test-uuid-123" + + def test_number_of_iterations(self, scene_api): + """Test number_of_iterations property.""" + assert scene_api.number_of_iterations == 100 + + def test_number_of_history_iterations(self, scene_api): + """Test number_of_history_iterations property.""" + assert scene_api.number_of_history_iterations == 10 + + def test_vehicle_parameters(self, scene_api): + """Test vehicle_parameters property.""" + assert scene_api.vehicle_parameters == scene_api._log_metadata.vehicle_parameters + + def test_available_pinhole_camera_types(self, scene_api): + """Test available_pinhole_camera_types property.""" + assert scene_api.available_pinhole_camera_types == [PinholeCameraType.PCAM_B0] + + def test_available_fisheye_mei_camera_types(self, scene_api): + """Test available_fisheye_mei_camera_types property.""" + assert scene_api.available_fisheye_mei_camera_types == [FisheyeMEICameraType.FCAM_L] + + def test_available_lidar_types(self, scene_api): + """Test available_lidar_types property.""" + assert scene_api.available_lidar_types == [LiDARType.LIDAR_TOP] + + +class TestSceneAPIMethods: + """Test abstract method implementations.""" + + def test_get_timepoint_at_iteration(self, scene_api): + """Test get_timepoint_at_iteration method.""" + result = scene_api.get_timepoint_at_iteration(0) + assert isinstance(result, Mock) + + def test_get_ego_state_at_iteration(self, scene_api): + """Test get_ego_state_at_iteration method.""" + result = scene_api.get_ego_state_at_iteration(0) + assert result is not None + + def test_get_box_detections_at_iteration(self, scene_api): + """Test get_box_detections_at_iteration method.""" + result = scene_api.get_box_detections_at_iteration(0) + assert result is not None + + def test_get_traffic_light_detections_at_iteration(self, scene_api): + """Test get_traffic_light_detections_at_iteration method.""" + result = scene_api.get_traffic_light_detections_at_iteration(0) + assert result is not None + + def test_get_route_lane_group_ids(self, scene_api): + """Test get_route_lane_group_ids method.""" + result = scene_api.get_route_lane_group_ids(0) + assert result == [1, 2, 3] + + def test_get_pinhole_camera_at_iteration(self, scene_api): + """Test get_pinhole_camera_at_iteration method.""" + result = scene_api.get_pinhole_camera_at_iteration(0, PinholeCameraType.PCAM_B0) + assert result is not None + + def test_get_fisheye_mei_camera_at_iteration(self, scene_api): + """Test get_fisheye_mei_camera_at_iteration method.""" + result = scene_api.get_fisheye_mei_camera_at_iteration(0, FisheyeMEICameraType.FCAM_L) + assert result is not None + + def test_get_lidar_at_iteration(self, scene_api): + """Test get_lidar_at_iteration method.""" + result = scene_api.get_lidar_at_iteration(0, LiDARType.LIDAR_TOP) + assert result is not None diff --git a/tests/unit/api/scene/__init__.py b/tests/unit/api/scene/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/conversion/registry/__init__.py b/tests/unit/conversion/registry/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/conversion/registry/test_box_detection_label_registry.py b/tests/unit/conversion/registry/test_box_detection_label_registry.py new file mode 100644 index 00000000..393e3527 --- /dev/null +++ b/tests/unit/conversion/registry/test_box_detection_label_registry.py @@ -0,0 +1,42 @@ +from py123d.conversion.registry.box_detection_label_registry import BOX_DETECTION_LABEL_REGISTRY, BoxDetectionLabel + + +class TestBoxDetectionLabelRegistry: + def test_correct_type(self): + """Test that all registered box detection labels are of correct type.""" + for label_class in BOX_DETECTION_LABEL_REGISTRY.values(): + assert issubclass(label_class, BoxDetectionLabel) + + def test_initialize_all_labels(self): + """Test that all registered box detection labels can be initialized.""" + for label_enum_class in BOX_DETECTION_LABEL_REGISTRY.values(): + label_enum_class: BoxDetectionLabel + for integer in range(len(label_enum_class)): + label_a = label_enum_class.from_int(integer) + label_b = label_enum_class(integer) + assert isinstance(label_a, label_enum_class) + assert isinstance(label_b, label_enum_class) + + def test_serialize_deserialize(self): + """Test that all registered box detection labels can be serialized and deserialized.""" + for label_enum_class in BOX_DETECTION_LABEL_REGISTRY.values(): + label_enum_class: BoxDetectionLabel + for integer in range(len(label_enum_class)): + label = label_enum_class.from_int(integer) + serialized_lower = label.serialize(lower=True) + serialized_upper = label.serialize(lower=False) + deserialized_lower = label_enum_class.deserialize(serialized_lower) + deserialized_upper = label_enum_class.deserialize(serialized_upper) + assert label == deserialized_lower + assert label == deserialized_upper + + def test_to_default(self): + """Test that all registered box detection labels can be converted to DefaultBoxDetectionLabel.""" + from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel + + for label_enum_class in BOX_DETECTION_LABEL_REGISTRY.values(): + label_enum_class: BoxDetectionLabel + for integer in range(len(label_enum_class)): + label = label_enum_class.from_int(integer) + default_label = label.to_default() + assert isinstance(default_label, DefaultBoxDetectionLabel) diff --git a/tests/unit/conversion/registry/test_lidar_registry.py b/tests/unit/conversion/registry/test_lidar_registry.py new file mode 100644 index 00000000..c93bc538 --- /dev/null +++ b/tests/unit/conversion/registry/test_lidar_registry.py @@ -0,0 +1,38 @@ +from enum import IntEnum + +import numpy as np + +from py123d.conversion.registry.lidar_index_registry import LIDAR_INDEX_REGISTRY, LiDARIndex + + +class TestLiDARRegistry: + def test_registered_types(self): + """Test that all registered LiDAR types are of correct type.""" + for lidar_class in LIDAR_INDEX_REGISTRY.values(): + assert issubclass(lidar_class, LiDARIndex) + + def test_initialize_all_types(self): + """Test that all registered LiDAR types can be initialized.""" + for lidar_enum_class in LIDAR_INDEX_REGISTRY.values(): + lidar_enum_class: LiDARIndex + for integer in range(len(lidar_enum_class)): + lidar_pc_index = lidar_enum_class(integer) + assert isinstance(lidar_pc_index, LiDARIndex) + assert isinstance(lidar_pc_index, IntEnum) + assert isinstance(lidar_pc_index, int) + + def test_xy_slice(self): + """Test that all registered LiDAR types have correct xy slice.""" + for lidar_enum_class in LIDAR_INDEX_REGISTRY.values(): + lidar_enum_class: LiDARIndex + dummy_lidar_pc = np.zeros((42, len(lidar_enum_class)), dtype=np.float32) + lidar_pc_xy_slice = dummy_lidar_pc[..., lidar_enum_class.XY] + assert lidar_pc_xy_slice.shape[-1] == 2 + + def test_xyz_slice(self): + """Test that all registered LiDAR types have correct xyz slice.""" + for lidar_enum_class in LIDAR_INDEX_REGISTRY.values(): + lidar_enum_class: LiDARIndex + dummy_lidar_pc = np.zeros((42, len(lidar_enum_class)), dtype=np.float32) + lidar_pc_xyz_slice = dummy_lidar_pc[..., lidar_enum_class.XYZ] + assert lidar_pc_xyz_slice.shape[-1] == 3 diff --git a/tests/unit/datatypes/detections/test_box_detections.py b/tests/unit/datatypes/detections/test_box_detections.py new file mode 100644 index 00000000..891bed1c --- /dev/null +++ b/tests/unit/datatypes/detections/test_box_detections.py @@ -0,0 +1,357 @@ +import pytest + +from py123d.conversion.registry.box_detection_label_registry import BoxDetectionLabel, DefaultBoxDetectionLabel +from py123d.datatypes.detections import ( + BoxDetectionMetadata, + BoxDetectionSE2, + BoxDetectionSE3, + BoxDetectionWrapper, +) +from py123d.datatypes.time.time_point import TimePoint +from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, PoseSE2, PoseSE3, Vector2D, Vector3D + + +class DummyBoxDetectionLabel(BoxDetectionLabel): + CAR = 1 + PEDESTRIAN = 2 + BICYCLE = 3 + + def to_default(self): + mapping = { + DummyBoxDetectionLabel.CAR: DefaultBoxDetectionLabel.VEHICLE, + DummyBoxDetectionLabel.PEDESTRIAN: DefaultBoxDetectionLabel.PERSON, + DummyBoxDetectionLabel.BICYCLE: DefaultBoxDetectionLabel.BICYCLE, + } + return mapping[self] + + +sample_metadata_args = { + "label": DummyBoxDetectionLabel.CAR, + "track_token": "sample_token", + "num_lidar_points": 10, + "timepoint": TimePoint.from_s(0.0), +} + + +class TestBoxDetectionMetadata: + def test_initialization(self): + metadata = BoxDetectionMetadata(**sample_metadata_args) + assert isinstance(metadata, BoxDetectionMetadata) + assert metadata.label == DummyBoxDetectionLabel.CAR + assert metadata.track_token == "sample_token" + assert metadata.num_lidar_points == 10 + assert isinstance(metadata.timepoint, TimePoint) + + def test_default_label(self): + metadata = BoxDetectionMetadata(**sample_metadata_args) + label = metadata.label + default_label = metadata.default_label + assert label == DummyBoxDetectionLabel.CAR + assert label.to_default() == DefaultBoxDetectionLabel.VEHICLE + assert default_label == DefaultBoxDetectionLabel.VEHICLE + + def test_default_label_with_default_label(self): + sample_args = sample_metadata_args.copy() + sample_args["label"] = DefaultBoxDetectionLabel.PERSON + metadata = BoxDetectionMetadata(**sample_args) + label = metadata.label + default_label = metadata.default_label + assert label == DefaultBoxDetectionLabel.PERSON + assert default_label == DefaultBoxDetectionLabel.PERSON + + def test_optional_args(self): + sample_args = { + "label": DummyBoxDetectionLabel.BICYCLE, + "track_token": "another_token", + } + metadata = BoxDetectionMetadata(**sample_args) + assert isinstance(metadata, BoxDetectionMetadata) + assert metadata.label == DummyBoxDetectionLabel.BICYCLE + assert metadata.track_token == "another_token" + assert metadata.num_lidar_points is None + assert metadata.timepoint is None + + def test_missing_args(self): + sample_args = { + "label": DummyBoxDetectionLabel.CAR, + } + with pytest.raises(TypeError): + BoxDetectionMetadata(**sample_args) + + sample_args = { + "track_token": "token_only", + } + with pytest.raises(TypeError): + BoxDetectionMetadata(**sample_args) + + sample_args = { + "timepoint": TimePoint.from_s(0.0), + } + with pytest.raises(TypeError): + BoxDetectionMetadata(**sample_args) + + +class TestBoxDetectionSE2: + def setup_method(self): + self.metadata = BoxDetectionMetadata(**sample_metadata_args) + self.bounding_box_se2 = BoundingBoxSE2( + center_se2=PoseSE2(x=0.0, y=0.0, yaw=0.0), + length=4.0, + width=2.0, + ) + self.velocity = None + + def test_initialization(self): + box_detection = BoxDetectionSE2( + metadata=self.metadata, + bounding_box_se2=self.bounding_box_se2, + velocity_2d=self.velocity, + ) + assert isinstance(box_detection, BoxDetectionSE2) + assert box_detection.metadata == self.metadata + assert box_detection.bounding_box_se2 == self.bounding_box_se2 + assert box_detection.velocity_2d is None + + def test_properties(self): + box_detection = BoxDetectionSE2( + metadata=self.metadata, + bounding_box_se2=self.bounding_box_se2, + velocity_2d=self.velocity, + ) + assert box_detection.shapely_polygon == self.bounding_box_se2.shapely_polygon + assert box_detection.center_se2 == self.bounding_box_se2.center_se2 + assert box_detection.bounding_box_se2 == self.bounding_box_se2 + + def test_optional_velocity(self): + box_detection_no_velo = BoxDetectionSE2( + metadata=self.metadata, + bounding_box_se2=self.bounding_box_se2, + ) + assert isinstance(box_detection_no_velo, BoxDetectionSE2) + assert box_detection_no_velo.velocity_2d is None + + box_detection_velo = BoxDetectionSE2( + metadata=self.metadata, + bounding_box_se2=self.bounding_box_se2, + velocity_2d=Vector2D(x=1.0, y=0.0), + ) + assert isinstance(box_detection_velo, BoxDetectionSE2) + assert box_detection_velo.velocity_2d == Vector2D(x=1.0, y=0.0) + + +class TestBoxBoxDetectionSE3: + def setup_method(self): + self.metadata = BoxDetectionMetadata(**sample_metadata_args) + self.bounding_box_se3 = BoundingBoxSE3( + center_se3=PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0), + length=4.0, + width=2.0, + height=1.5, + ) + self.velocity = Vector3D(x=1.0, y=0.0, z=0.0) + + def test_initialization(self): + box_detection = BoxDetectionSE3( + metadata=self.metadata, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=self.velocity, + ) + assert isinstance(box_detection, BoxDetectionSE3) + assert box_detection.metadata == self.metadata + assert box_detection.bounding_box_se3 == self.bounding_box_se3 + assert box_detection.velocity_3d == self.velocity + + def test_properties(self): + box_detection = BoxDetectionSE3( + metadata=self.metadata, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=self.velocity, + ) + assert box_detection.shapely_polygon == self.bounding_box_se3.shapely_polygon + assert box_detection.center_se3 == self.bounding_box_se3.center_se3 + assert box_detection.center_se2 == self.bounding_box_se3.center_se2 + assert box_detection.bounding_box_se3 == self.bounding_box_se3 + assert box_detection.bounding_box_se2 == self.bounding_box_se3.bounding_box_se2 + assert box_detection.velocity_3d == self.velocity + assert box_detection.velocity_2d == self.velocity.vector_2d + + def test_box_detection_se2_conversion(self): + box_detection = BoxDetectionSE3( + metadata=self.metadata, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=Vector3D(x=1.0, y=0.0, z=0.0), + ) + box_detection_se2 = box_detection.box_detection_se2 + assert isinstance(box_detection_se2, BoxDetectionSE2) + assert box_detection_se2.metadata == self.metadata + assert box_detection_se2.bounding_box_se2 == self.bounding_box_se3.bounding_box_se2 + assert box_detection_se2.velocity_2d == Vector2D(x=1.0, y=0.0) + + def test_box_detection_se3_conversion(self): + box_detection_se2 = BoxDetectionSE2( + metadata=self.metadata, + bounding_box_se2=self.bounding_box_se3.bounding_box_se2, + velocity_2d=Vector2D(x=1.0, y=0.0), + ) + box_detection_se3 = BoxDetectionSE3( + metadata=box_detection_se2.metadata, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=Vector3D(x=1.0, y=0.0, z=0.0), + ) + assert isinstance(box_detection_se3, BoxDetectionSE3) + assert box_detection_se3.metadata == box_detection_se2.metadata + assert box_detection_se3.bounding_box_se3 == self.bounding_box_se3 + assert box_detection_se3.velocity_2d == Vector2D(x=1.0, y=0.0) + + box_detection_se3_converted = box_detection_se3.box_detection_se2 + assert isinstance(box_detection_se3_converted, BoxDetectionSE2) + assert box_detection_se3_converted.metadata == box_detection_se2.metadata + assert box_detection_se3_converted.bounding_box_se2 == box_detection_se2.bounding_box_se2 + assert box_detection_se3_converted.velocity_2d == box_detection_se2.velocity_2d + + def test_optional_velocity(self): + box_detection_no_velo = BoxDetectionSE3( + metadata=self.metadata, + bounding_box_se3=self.bounding_box_se3, + ) + assert isinstance(box_detection_no_velo, BoxDetectionSE3) + assert box_detection_no_velo.velocity_3d is None + + box_detection_velo = BoxDetectionSE3( + metadata=self.metadata, + bounding_box_se3=self.bounding_box_se3, + velocity_3d=Vector3D(x=1.0, y=0.0, z=0.0), + ) + assert isinstance(box_detection_velo, BoxDetectionSE3) + assert box_detection_velo.velocity_3d == Vector3D(x=1.0, y=0.0, z=0.0) + + +class TestBoxDetectionWrapper: + def setup_method(self): + self.metadata1 = BoxDetectionMetadata( + label=DummyBoxDetectionLabel.CAR, + track_token="token1", + num_lidar_points=10, + timepoint=TimePoint.from_s(0.0), + ) + self.metadata2 = BoxDetectionMetadata( + label=DummyBoxDetectionLabel.PEDESTRIAN, + track_token="token2", + num_lidar_points=5, + timepoint=TimePoint.from_s(0.0), + ) + self.metadata3 = BoxDetectionMetadata( + label=DummyBoxDetectionLabel.BICYCLE, + track_token="token3", + num_lidar_points=8, + timepoint=TimePoint.from_s(0.0), + ) + + self.box_detection1 = BoxDetectionSE2( + metadata=self.metadata1, + bounding_box_se2=BoundingBoxSE2( + center_se2=PoseSE2(x=0.0, y=0.0, yaw=0.0), + length=4.0, + width=2.0, + ), + velocity_2d=Vector2D(x=1.0, y=0.0), + ) + self.box_detection2 = BoxDetectionSE2( + metadata=self.metadata2, + bounding_box_se2=BoundingBoxSE2( + center_se2=PoseSE2(x=5.0, y=5.0, yaw=0.0), + length=1.0, + width=0.5, + ), + velocity_2d=Vector2D(x=0.5, y=0.5), + ) + self.box_detection3 = BoxDetectionSE3( + metadata=self.metadata3, + bounding_box_se3=BoundingBoxSE3( + center_se3=PoseSE3(x=10.0, y=10.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0), + length=2.0, + width=1.0, + height=1.5, + ), + velocity_3d=Vector3D(x=0.0, y=1.0, z=0.0), + ) + + def test_initialization(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + assert isinstance(wrapper, BoxDetectionWrapper) + assert len(wrapper.box_detections) == 2 + + def test_empty_initialization(self): + wrapper = BoxDetectionWrapper(box_detections=[]) + assert isinstance(wrapper, BoxDetectionWrapper) + assert len(wrapper.box_detections) == 0 + + def test_getitem(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + assert wrapper[0] == self.box_detection1 + assert wrapper[1] == self.box_detection2 + + def test_getitem_out_of_range(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1]) + with pytest.raises(IndexError): + _ = wrapper[1] + + def test_len(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2, self.box_detection3]) + assert len(wrapper) == 3 + + def test_len_empty(self): + wrapper = BoxDetectionWrapper(box_detections=[]) + assert len(wrapper) == 0 + + def test_iter(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + detections = list(wrapper) + assert len(detections) == 2 + assert detections[0] == self.box_detection1 + assert detections[1] == self.box_detection2 + + def test_get_detection_by_track_token_found(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2, self.box_detection3]) + detection = wrapper.get_detection_by_track_token("token2") + assert detection is not None + assert detection == self.box_detection2 + assert detection.metadata.track_token == "token2" + + def test_get_detection_by_track_token_not_found(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + detection = wrapper.get_detection_by_track_token("nonexistent_token") + assert detection is None + + def test_get_detection_by_track_token_empty_wrapper(self): + wrapper = BoxDetectionWrapper(box_detections=[]) + detection = wrapper.get_detection_by_track_token("token1") + assert detection is None + + def test_occupancy_map(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + occupancy_map = wrapper.occupancy_map_2d + assert occupancy_map is not None + assert len(occupancy_map.geometries) == 2 + assert len(occupancy_map.ids) == 2 + assert "token1" in occupancy_map.ids + assert "token2" in occupancy_map.ids + + def test_occupancy_map_cached(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection2]) + occupancy_map1 = wrapper.occupancy_map_2d + occupancy_map2 = wrapper.occupancy_map_2d + assert occupancy_map1 is occupancy_map2 + + def test_occupancy_map_empty(self): + wrapper = BoxDetectionWrapper(box_detections=[]) + occupancy_map = wrapper.occupancy_map_2d + assert occupancy_map is not None + assert len(occupancy_map.geometries) == 0 + assert len(occupancy_map.ids) == 0 + + def test_mixed_detection_types(self): + wrapper = BoxDetectionWrapper(box_detections=[self.box_detection1, self.box_detection3]) + assert len(wrapper) == 2 + assert isinstance(wrapper[0], BoxDetectionSE2) + assert isinstance(wrapper[1], BoxDetectionSE3) diff --git a/tests/unit/datatypes/detections/test_traffic_lights.py b/tests/unit/datatypes/detections/test_traffic_lights.py new file mode 100644 index 00000000..958bfe51 --- /dev/null +++ b/tests/unit/datatypes/detections/test_traffic_lights.py @@ -0,0 +1,83 @@ +from py123d.datatypes.detections import TrafficLightDetection, TrafficLightDetectionWrapper, TrafficLightStatus +from py123d.datatypes.time.time_point import TimePoint + + +class TestTrafficLightStatus: + def test_status_values(self): + """Test that TrafficLightStatus enum has correct values.""" + assert TrafficLightStatus.GREEN.value == 0 + assert TrafficLightStatus.YELLOW.value == 1 + assert TrafficLightStatus.RED.value == 2 + assert TrafficLightStatus.OFF.value == 3 + assert TrafficLightStatus.UNKNOWN.value == 4 + + +class TestTrafficLightDetection: + def test_creation_with_required_fields(self): + """Test that TrafficLightDetection can be created with required fields.""" + detection = TrafficLightDetection(lane_id=1, status=TrafficLightStatus.GREEN) + assert detection.lane_id == 1 + assert detection.status == TrafficLightStatus.GREEN + assert detection.timepoint is None + + def test_creation_with_timepoint(self): + """Test that TrafficLightDetection can be created with timepoint.""" + timepoint = TimePoint.from_s(0) + detection = TrafficLightDetection( + lane_id=2, + status=TrafficLightStatus.RED, + timepoint=timepoint, + ) + assert detection.lane_id == 2 + assert detection.status == TrafficLightStatus.RED + assert detection.timepoint == timepoint + + +class TestTrafficLightDetectionWrapper: + def setup_method(self): + self.detection1 = TrafficLightDetection(lane_id=1, status=TrafficLightStatus.GREEN) + self.detection2 = TrafficLightDetection(lane_id=2, status=TrafficLightStatus.RED) + self.detection3 = TrafficLightDetection(lane_id=3, status=TrafficLightStatus.YELLOW) + self.wrapper = TrafficLightDetectionWrapper( + traffic_light_detections=[self.detection1, self.detection2, self.detection3] + ) + + def test_getitem(self): + """Test __getitem__ method of TrafficLightDetectionWrapper.""" + assert self.wrapper[0] == self.detection1 + assert self.wrapper[1] == self.detection2 + assert self.wrapper[2] == self.detection3 + + def test_len(self): + """Test __len__ method of TrafficLightDetectionWrapper.""" + assert len(self.wrapper) == 3 + + def test_iter(self): + """Test __iter__ method of TrafficLightDetectionWrapper.""" + detections = list(self.wrapper) + assert detections == [self.detection1, self.detection2, self.detection3] + + def test_get_detection_by_lane_id_found(self): + """Test get_detection_by_lane_id method of TrafficLightDetectionWrapper.""" + result = self.wrapper.get_detection_by_lane_id(2) + assert result == self.detection2 + assert result.status == TrafficLightStatus.RED + + def test_get_detection_by_lane_id_not_found(self): + """Test get_detection_by_lane_id method of TrafficLightDetectionWrapper when not found.""" + result = self.wrapper.get_detection_by_lane_id(99) + assert result is None + + def test_get_detection_by_lane_id_first_match(self): + """Test get_detection_by_lane_id method returns first match.""" + duplicate = TrafficLightDetection(lane_id=1, status=TrafficLightStatus.OFF) + wrapper = TrafficLightDetectionWrapper(traffic_light_detections=[self.detection1, duplicate]) + result = wrapper.get_detection_by_lane_id(1) + assert result == self.detection1 + + def test_empty_wrapper(self): + """Test behavior of an empty TrafficLightDetectionWrapper.""" + empty_wrapper = TrafficLightDetectionWrapper(traffic_light_detections=[]) + assert len(empty_wrapper) == 0 + assert list(empty_wrapper) == [] + assert empty_wrapper.get_detection_by_lane_id(1) is None diff --git a/tests/unit/datatypes/map_objects/__init__.py b/tests/unit/datatypes/map_objects/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/datatypes/map_objects/mock_map_api.py b/tests/unit/datatypes/map_objects/mock_map_api.py new file mode 100644 index 00000000..9158df1f --- /dev/null +++ b/tests/unit/datatypes/map_objects/mock_map_api.py @@ -0,0 +1,118 @@ +from typing import Dict, Iterable, List, Optional, Union + +import shapely + +from py123d.api import MapAPI +from py123d.datatypes.map_objects import BaseMapObject, MapLayer +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + GenericDrivable, + Intersection, + Lane, + LaneGroup, + RoadEdge, + RoadLine, + StopZone, + Walkway, +) +from py123d.datatypes.metadata import MapMetadata +from py123d.geometry import Point2D + + +class MockMapAPI(MapAPI): + def __init__( + self, + lanes: List[Lane] = [], + lane_groups: List[LaneGroup] = [], + intersections: List[Intersection] = [], + crosswalks: List[Intersection] = [], + carparks: List[Carpark] = [], + walkways: List[Walkway] = [], + generic_drivables: List[GenericDrivable] = [], + stop_zones: List[StopZone] = [], + road_edges: List[RoadEdge] = [], + road_lines: List[RoadLine] = [], + add_map_api_links: bool = False, + ): + self._layers: Dict[MapLayer, List[BaseMapObject]] = { + MapLayer.LANE: lanes, + MapLayer.LANE_GROUP: lane_groups, + MapLayer.INTERSECTION: intersections, + MapLayer.CROSSWALK: crosswalks, + MapLayer.WALKWAY: walkways, + MapLayer.CARPARK: carparks, + MapLayer.GENERIC_DRIVABLE: generic_drivables, + MapLayer.STOP_ZONE: stop_zones, + MapLayer.ROAD_EDGE: road_edges, + MapLayer.ROAD_LINE: road_lines, + } + + for layer, layer_objects in self._layers.items(): + if layer in [ + MapLayer.LANE, + MapLayer.LANE_GROUP, + MapLayer.INTERSECTION, + ]: + for obj in layer_objects: + if add_map_api_links: + obj._map_api = self # type: ignore + else: + obj._map_api = None # type: ignore + + def get_map_metadata(self) -> MapMetadata: + return MapMetadata( + dataset="test", + split="test_split", + log_name="test_log_name", + location="test_location", + map_has_z=True, + map_is_local=True, + ) + + def get_available_map_layers(self) -> List[MapLayer]: + return list(self._layers.keys()) + + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[BaseMapObject]: + target_layer = self._layers.get(layer, []) + map_object: Optional[BaseMapObject] = None + for obj in target_layer: + if obj.object_id == object_id: + map_object = obj + break + return map_object + + def get_map_objects_in_radius( + self, point: Point2D, radius: float, layers: List[MapLayer] + ) -> Dict[MapLayer, List[BaseMapObject]]: + return {} + + def query( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]]: + return {} + + def query_object_ids( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + predicate: Optional[str] = None, + sort: bool = False, + distance: Optional[float] = None, + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: + return {} + + def query_nearest( + self, + geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], + layers: List[MapLayer], + return_all: bool = True, + max_distance: Optional[float] = None, + return_distance: bool = False, + exclusive: bool = False, + ) -> Dict[MapLayer, Union[List[BaseMapObject], Dict[int, List[BaseMapObject]]]]: + return {} diff --git a/tests/unit/datatypes/map_objects/test_base_map_objects.py b/tests/unit/datatypes/map_objects/test_base_map_objects.py new file mode 100644 index 00000000..8e22dea8 --- /dev/null +++ b/tests/unit/datatypes/map_objects/test_base_map_objects.py @@ -0,0 +1,208 @@ +import numpy as np +import pytest +import shapely.geometry as geom + +from py123d.datatypes.map_objects.base_map_objects import BaseMapLineObject, BaseMapObject, BaseMapSurfaceObject +from py123d.datatypes.map_objects.map_layer_types import MapLayer +from py123d.geometry import Polyline2D, Polyline3D + + +class ConcreteMapObject(BaseMapObject): + """Concrete implementation for testing BaseMapObject.""" + + def __init__(self, object_id, layer_type=MapLayer.GENERIC_DRIVABLE): + super().__init__(object_id) + self._layer = layer_type + + @property + def layer(self) -> MapLayer: + return self._layer + + +class ConcreteMapSurfaceObject(BaseMapSurfaceObject): + """Concrete implementation for testing BaseMapSurfaceObject.""" + + def __init__(self, object_id, outline=None, shapely_polygon=None, layer_type=MapLayer.GENERIC_DRIVABLE): + super().__init__(object_id, outline, shapely_polygon) + self._layer = layer_type + + @property + def layer(self) -> MapLayer: + return self._layer + + +class ConcreteMapLineObject(BaseMapLineObject): + """Concrete implementation for testing BaseMapLineObject.""" + + def __init__(self, object_id, polyline, layer_type=MapLayer.GENERIC_DRIVABLE): + super().__init__(object_id, polyline) + self._layer = layer_type + + @property + def layer(self) -> MapLayer: + return self._layer + + +class TestBaseMapObject: + """Test cases for BaseMapObject class.""" + + def test_init_with_string_id(self): + """Test initialization with string object ID.""" + obj = ConcreteMapObject("test_id_123") + assert obj.object_id == "test_id_123" + + def test_init_with_int_id(self): + """Test initialization with integer object ID.""" + obj = ConcreteMapObject(42) + assert obj.object_id == 42 + + def test_object_id_property(self): + """Test object_id property.""" + obj = ConcreteMapObject("unique_id") + assert obj.object_id == "unique_id" + + def test_layer_property(self): + """Test layer property.""" + obj = ConcreteMapObject("id1", MapLayer.GENERIC_DRIVABLE) + assert obj.layer == MapLayer.GENERIC_DRIVABLE + + def test_abstract_instantiation_fails(self): + """Test that instantiating BaseMapObject directly raises TypeError.""" + with pytest.raises(TypeError): + BaseMapObject("test_id") + + +class TestBaseMapSurfaceObject: + """Test cases for BaseMapSurfaceObject class.""" + + def test_init_with_polyline2d(self): + coords = np.array([[0, 0], [1, 0], [1, 1], [0, 1], [0, 0]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_1", outline=polyline) + assert obj.object_id == "surf_1" + assert isinstance(obj.outline, Polyline2D) + + def test_init_with_polyline3d(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 0]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_2", outline=polyline) + assert isinstance(obj.outline, Polyline3D) + + def test_init_with_shapely_polygon(self): + polygon = geom.Polygon([(0, 0), (1, 0), (1, 1), (0, 1)]) + obj = ConcreteMapSurfaceObject("surf_3", shapely_polygon=polygon) + assert obj.shapely_polygon.equals(polygon) + + def test_init_without_outline_or_polygon_raises_error(self): + with pytest.raises(ValueError): + ConcreteMapSurfaceObject("surf_4") + + def test_outline_property(self): + coords = np.array([[0, 0], [1, 0], [1, 1], [0, 0]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_5", outline=polyline) + assert obj.outline is polyline + + def test_outline_2d_from_2d_polyline(self): + coords = np.array([[0, 0], [1, 0], [1, 1], [0, 0]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_6", outline=polyline) + assert isinstance(obj.outline_2d, Polyline2D) + + def test_outline_2d_from_3d_polyline(self): + coords = np.array([[0, 0, 5], [1, 0, 5], [1, 1, 5], [0, 0, 5]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_7", outline=polyline) + outline_2d = obj.outline_2d + assert isinstance(outline_2d, Polyline2D) + + def test_outline_3d_from_3d_polyline(self): + coords = np.array([[0, 0, 2], [1, 0, 2], [1, 1, 2], [0, 0, 2]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_8", outline=polyline) + assert isinstance(obj.outline_3d, Polyline3D) + + def test_outline_3d_from_2d_polyline(self): + coords = np.array([[0, 0], [1, 0], [1, 1], [0, 0]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_9", outline=polyline) + outline_3d = obj.outline_3d + assert isinstance(outline_3d, Polyline3D) + + def test_shapely_polygon_property(self): + polygon = geom.Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]) + obj = ConcreteMapSurfaceObject("surf_10", shapely_polygon=polygon) + assert obj.shapely_polygon.equals(polygon) + + def test_trimesh_mesh_property(self): + coords = np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 0, 0]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapSurfaceObject("surf_11", outline=polyline) + mesh = obj.trimesh_mesh + assert mesh is not None + assert len(mesh.vertices) > 0 + assert len(mesh.faces) > 0 + + +class TestBaseMapLineObject: + """Test cases for BaseMapLineObject class.""" + + def test_init_with_polyline2d(self): + coords = np.array([[0, 0], [1, 1], [2, 2]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject("line_1", polyline) + assert obj.object_id == "line_1" + assert isinstance(obj.polyline, Polyline2D) + + def test_init_with_polyline3d(self): + coords = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]]) + polyline = Polyline3D(coords) + obj = ConcreteMapLineObject("line_2", polyline) + assert isinstance(obj.polyline, Polyline3D) + + def test_polyline_property(self): + coords = np.array([[0, 0], [1, 1], [2, 2]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject("line_3", polyline) + assert obj.polyline == polyline + + def test_polyline_2d_from_2d_polyline(self): + coords = np.array([[0, 0], [1, 1], [2, 2]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject("line_4", polyline) + assert isinstance(obj.polyline_2d, Polyline2D) + assert obj.polyline_2d is polyline + + def test_polyline_2d_from_3d_polyline(self): + coords = np.array([[0, 0, 5], [1, 1, 5], [2, 2, 5]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapLineObject("line_5", polyline) + polyline_2d = obj.polyline_2d + assert isinstance(polyline_2d, Polyline2D) + + def test_polyline_3d_from_3d_polyline(self): + coords = np.array([[0, 0, 3], [1, 1, 3], [2, 2, 3]]) + polyline = Polyline3D.from_array(coords) + obj = ConcreteMapLineObject("line_6", polyline) + assert isinstance(obj.polyline_3d, Polyline3D) + assert obj.polyline_3d is polyline + + def test_polyline_3d_from_2d_polyline(self): + coords = np.array([[0, 0], [1, 1], [2, 2]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject("line_7", polyline) + polyline_3d = obj.polyline_3d + assert isinstance(polyline_3d, Polyline3D) + + def test_shapely_linestring_property(self): + coords = np.array([[0, 0], [1, 1], [2, 2]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject("line_8", polyline) + linestring = obj.shapely_linestring + assert isinstance(linestring, geom.LineString) + + def test_object_id_with_integer(self): + coords = np.array([[0, 0], [1, 1]]) + polyline = Polyline2D.from_array(coords) + obj = ConcreteMapLineObject(999, polyline) + assert obj.object_id == 999 diff --git a/tests/unit/datatypes/map_objects/test_map_objects.py b/tests/unit/datatypes/map_objects/test_map_objects.py new file mode 100644 index 00000000..252500cd --- /dev/null +++ b/tests/unit/datatypes/map_objects/test_map_objects.py @@ -0,0 +1,907 @@ +from typing import List, Tuple + +import numpy as np +import pytest +import shapely +import trimesh + +from py123d.datatypes.map_objects import Intersection, Lane, LaneGroup, MapLayer +from py123d.datatypes.map_objects.map_layer_types import RoadEdgeType, RoadLineType +from py123d.datatypes.map_objects.map_objects import ( + Carpark, + Crosswalk, + GenericDrivable, + RoadEdge, + RoadLine, + StopZone, + Walkway, +) +from py123d.geometry.polyline import Polyline2D, Polyline3D + +from .mock_map_api import MockMapAPI + + +def _get_linked_map_object_setup() -> Tuple[List[Lane], List[LaneGroup], List[Intersection]]: + """Helper function to create linked map objects for testing.""" + + Z = 0.0 + + # Lanes: + lanes: List[Lane] = [] + + # Middle Lane 0, group 0 + middle_left_boundary = np.array([[0.0, 1.0, Z], [50.0, 1.0, Z]]) + middle_right_boundary = np.array([[0.0, -1.0, Z], [50.0, -1.0, Z]]) + middle_centerline = np.mean(np.array([middle_right_boundary, middle_left_boundary]), axis=0) + lanes.append( + Lane( + object_id=0, + lane_group_id=0, + left_boundary=Polyline3D.from_array(middle_left_boundary), + right_boundary=Polyline3D.from_array(middle_right_boundary), + centerline=Polyline3D.from_array(middle_centerline), + left_lane_id=1, + right_lane_id=2, + predecessor_ids=[3], + successor_ids=[4], + speed_limit_mps=0.0, + ) + ) + + # Left Lane 1, group 0 + left_left_boundary = np.array([[0.0, 2.0, Z], [50.0, 2.0, Z]]) + left_right_boundary = middle_left_boundary.copy() + left_centerline = np.mean(np.array([left_right_boundary, left_left_boundary]), axis=0) + lanes.append( + Lane( + object_id=1, + lane_group_id=0, + left_boundary=Polyline3D.from_array(left_left_boundary), + right_boundary=Polyline3D.from_array(left_right_boundary), + centerline=Polyline3D.from_array(left_centerline), + left_lane_id=None, + right_lane_id=0, + predecessor_ids=[], + successor_ids=[], + speed_limit_mps=0.0, + ) + ) + + # Right Lane 2, group 0 + right_right_boundary = np.array([[0.0, -2.0, Z], [50.0, -2.0, Z]]) + right_left_boundary = middle_right_boundary.copy() + right_centerline = np.mean(np.array([right_right_boundary, right_left_boundary]), axis=0) + lanes.append( + Lane( + object_id=2, + lane_group_id=0, + left_boundary=Polyline3D.from_array(right_left_boundary), + right_boundary=Polyline3D.from_array(right_right_boundary), + centerline=Polyline3D.from_array(right_centerline), + left_lane_id=0, + right_lane_id=None, + predecessor_ids=[], + successor_ids=[], + speed_limit_mps=0.0, + ) + ) + + # Predecessor lane 3, group 1 + predecessor_left_boundary = np.array([[-50.0, 1.0, Z], [0.0, 1.0, Z]]) + predecessor_right_boundary = np.array([[-50.0, -1.0, Z], [0.0, -1.0, Z]]) + predecessor_centerline = np.mean(np.array([predecessor_right_boundary, predecessor_left_boundary]), axis=0) + lanes.append( + Lane( + object_id=3, + lane_group_id=1, + left_boundary=Polyline3D.from_array(predecessor_left_boundary), + right_boundary=Polyline3D.from_array(predecessor_right_boundary), + centerline=Polyline3D.from_array(predecessor_centerline), + left_lane_id=None, + right_lane_id=None, + predecessor_ids=[], + successor_ids=[0], + speed_limit_mps=0.0, + ) + ) + + # Successor lane 4, group 2 + successor_left_boundary = np.array([[50.0, 1.0, Z], [100.0, 1.0, Z]]) + successor_right_boundary = np.array([[50.0, -1.0, Z], [100.0, -1.0, Z]]) + successor_centerline = np.mean(np.array([successor_right_boundary, successor_left_boundary]), axis=0) + lanes.append( + Lane( + object_id=4, + lane_group_id=2, + left_boundary=Polyline3D.from_array(successor_left_boundary), + right_boundary=Polyline3D.from_array(successor_right_boundary), + centerline=Polyline3D.from_array(successor_centerline), + left_lane_id=None, + right_lane_id=None, + predecessor_ids=[0], + successor_ids=[], + speed_limit_mps=0.0, + ) + ) + + # Lane Groups: + lane_groups = [] + + # Middle lane group 0, lanes 0,1,2 + middle_lane_group = LaneGroup( + object_id=0, + lane_ids=[0, 1, 2], + left_boundary=Polyline3D.from_array(left_left_boundary), + right_boundary=Polyline3D.from_array(left_right_boundary), + intersection_id=None, + predecessor_ids=[1], + successor_ids=[2], + ) + lane_groups.append(middle_lane_group) + + # Predecessor lane group 1, lane 3, intersection 0 + predecessor_lane_group = LaneGroup( + object_id=1, + lane_ids=[3], + left_boundary=Polyline3D.from_array(predecessor_left_boundary), + right_boundary=Polyline3D.from_array(predecessor_right_boundary), + intersection_id=0, + predecessor_ids=[], + successor_ids=[0], + ) + lane_groups.append(predecessor_lane_group) + + # Successor lane group 2, lane 4, intersection 1 + successor_lane_group = LaneGroup( + object_id=2, + lane_ids=[4], + left_boundary=Polyline3D.from_array(successor_left_boundary), + right_boundary=Polyline3D.from_array(successor_right_boundary), + intersection_id=1, + predecessor_ids=[0], + successor_ids=[], + ) + lane_groups.append(successor_lane_group) + + # Intersections: + intersections = [] + + # Intersection 0, includes lane groups 1 + intersection_predecessor = Intersection( + object_id=0, + lane_group_ids=[1], + outline=predecessor_lane_group.outline, + ) + intersections.append(intersection_predecessor) + + intersection_successor = Intersection( + object_id=1, + lane_group_ids=[2], + outline=successor_lane_group.outline, + ) + intersections.append(intersection_successor) + + return lanes, lane_groups, intersections + + +class TestLane: + def setup_method(self) -> None: + lanes, lane_groups, intersections = _get_linked_map_object_setup() + self.lanes = lanes + self.lane_groups = lane_groups + self.intersections = intersections + + def test_set_up(self): + """Test that the setup function creates the correct number of map objects.""" + assert len(self.lanes) == 5 + assert len(self.lane_groups) == 3 + assert len(self.intersections) == 2 + + def test_properties(self): + """Test that the properties of the Lane objects are correct.""" + lane0 = self.lanes[0] + assert lane0.layer == MapLayer.LANE + assert lane0.lane_group_id == 0 + assert isinstance(lane0.left_boundary, Polyline3D) + assert isinstance(lane0.right_boundary, Polyline3D) + assert isinstance(lane0.centerline, Polyline3D) + + assert lane0.left_lane_id == 1 + assert lane0.right_lane_id == 2 + assert lane0.predecessor_ids == [3] + assert lane0.successor_ids == [4] + assert lane0.speed_limit_mps == 0.0 + assert isinstance(lane0.trimesh_mesh, trimesh.base.Trimesh) + + def test_base_properties(self): + """Test that the base_surface property of the Lane objects is correct.""" + lane0 = self.lanes[0] + assert lane0.object_id == 0 + assert isinstance(lane0.outline, Polyline3D) + assert isinstance(lane0.outline_2d, Polyline2D) + assert isinstance(lane0.outline_3d, Polyline3D) + assert isinstance(lane0.shapely_polygon, shapely.Polygon) + + def test_left_links(self): + """Test that the left neighboring lanes are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_left_neighbor(lane: Lane): + assert lane is not None + assert lane.left_lane is None + assert lane.left_lane_id is None + + # Middle Lane 0 + lane0: Lane = map_api.get_map_object(0, MapLayer.LANE) + assert lane0 is not None + assert lane0.left_lane is not None + assert isinstance(lane0.left_lane, Lane) + assert lane0.left_lane.object_id == 1 + assert lane0.left_lane.object_id == lane0.left_lane_id + + # Left Lane 1 + lane1: Lane = map_api.get_map_object(1, MapLayer.LANE) + _no_left_neighbor(lane1) + + # Right Lane 2 + lane2: Lane = map_api.get_map_object(2, MapLayer.LANE) + assert lane2 is not None + assert lane2.left_lane is not None + assert isinstance(lane2.left_lane, Lane) + assert lane2.left_lane.object_id == 0 + assert lane2.left_lane.object_id == lane2.left_lane_id + + # Predecessor Lane 3 + lane3: Lane = map_api.get_map_object(3, MapLayer.LANE) + _no_left_neighbor(lane3) + + # Successor Lane 4 + lane4: Lane = map_api.get_map_object(4, MapLayer.LANE) + _no_left_neighbor(lane4) + + def test_right_links(self): + """Test that the right neighboring lanes are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_right_neighbor(lane: Lane): + assert lane is not None + assert lane.right_lane is None + assert lane.right_lane_id is None + + # Middle Lane 0 + lane0: Lane = map_api.get_map_object(0, MapLayer.LANE) + assert lane0 is not None + assert lane0.right_lane is not None + assert isinstance(lane0.right_lane, Lane) + assert lane0.right_lane.object_id == 2 + assert lane0.right_lane.object_id == lane0.right_lane_id + + # Left Lane 1 + lane1: Lane = map_api.get_map_object(1, MapLayer.LANE) + assert lane1 is not None + assert lane1.right_lane is not None + assert isinstance(lane1.right_lane, Lane) + assert lane1.right_lane.object_id == 0 + assert lane1.right_lane.object_id == lane1.right_lane_id + + # Right Lane 2 + lane2: Lane = map_api.get_map_object(2, MapLayer.LANE) + _no_right_neighbor(lane2) + + # Predecessor Lane 3 + lane3: Lane = map_api.get_map_object(3, MapLayer.LANE) + _no_right_neighbor(lane3) + + # Successor Lane 4 + lane4: Lane = map_api.get_map_object(4, MapLayer.LANE) + _no_right_neighbor(lane4) + + def test_predecessor_links(self): + """Test that the predecessor lanes are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_predecessors(lane: Lane): + assert lane is not None + assert lane.predecessors == [] + assert lane.predecessor_ids == [] + + # Middle Lane 0 + lane0: Lane = map_api.get_map_object(0, MapLayer.LANE) + assert lane0 is not None + assert lane0.predecessors is not None + assert len(lane0.predecessors) == 1 + assert isinstance(lane0.predecessors[0], Lane) + assert lane0.predecessors[0].object_id == 3 + assert lane0.predecessor_ids == [3] + + # Left Lane 1 + lane1: Lane = map_api.get_map_object(1, MapLayer.LANE) + _no_predecessors(lane1) + + # Right Lane 2 + lane2: Lane = map_api.get_map_object(2, MapLayer.LANE) + _no_predecessors(lane2) + + # Predecessor Lane 3 + lane3: Lane = map_api.get_map_object(3, MapLayer.LANE) + _no_predecessors(lane3) + + # Successor Lane 4 + lane4: Lane = map_api.get_map_object(4, MapLayer.LANE) + assert lane4 is not None + assert lane4.predecessors is not None + assert len(lane4.predecessors) == 1 + assert isinstance(lane4.predecessors[0], Lane) + assert lane4.predecessors[0].object_id == 0 + assert lane4.predecessor_ids == [0] + + def test_successor_links(self): + """Test that the successor lanes are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_successors(lane: Lane): + assert lane is not None + assert lane.successors == [] + assert lane.successor_ids == [] + + # Middle Lane 0 + lane0: Lane = map_api.get_map_object(0, MapLayer.LANE) + assert lane0 is not None + assert lane0.successors is not None + assert len(lane0.successors) == 1 + assert isinstance(lane0.successors[0], Lane) + assert lane0.successors[0].object_id == 4 + assert lane0.successor_ids == [4] + + # Left Lane 1 + lane1: Lane = map_api.get_map_object(1, MapLayer.LANE) + _no_successors(lane1) + + # Right Lane 2 + lane2: Lane = map_api.get_map_object(2, MapLayer.LANE) + _no_successors(lane2) + + # Predecessor Lane 3 + lane3: Lane = map_api.get_map_object(3, MapLayer.LANE) + assert lane3 is not None + assert lane3.successors is not None + assert len(lane3.successors) == 1 + assert isinstance(lane3.successors[0], Lane) + assert lane3.successors[0].object_id == 0 + assert lane3.successor_ids == [0] + + # Successor Lane 4 + lane4: Lane = map_api.get_map_object(4, MapLayer.LANE) + _no_successors(lane4) + + def test_no_links(self): + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=False, + ) + for lane in self.lanes: + lane_from_api: Lane = map_api.get_map_object(lane.object_id, MapLayer.LANE) + assert lane_from_api is not None + assert lane_from_api.left_lane is None + assert lane_from_api.right_lane is None + assert lane_from_api.predecessors is None + assert lane_from_api.successors is None + + def test_lane_group_links(self): + """Test that the lane group links are correct.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + for lane in self.lanes: + lane_from_api: Lane = map_api.get_map_object(lane.object_id, MapLayer.LANE) + assert lane_from_api is not None + assert lane_from_api.lane_group is not None + assert isinstance(lane_from_api.lane_group, LaneGroup) + assert lane_from_api.lane_group.object_id == lane_from_api.lane_group_id + + +class TestLaneGroup: + def setup_method(self): + lanes, lane_groups, intersections = _get_linked_map_object_setup() + self.lanes = lanes + self.lane_groups = lane_groups + self.intersections = intersections + + def test_properties(self): + """Test that the properties of the LaneGroup objects are correct.""" + lane_group0 = self.lane_groups[0] + assert lane_group0.layer == MapLayer.LANE_GROUP + assert lane_group0.lane_ids == [0, 1, 2] + assert isinstance(lane_group0.left_boundary, Polyline3D) + assert isinstance(lane_group0.right_boundary, Polyline3D) + assert lane_group0.intersection_id is None + assert lane_group0.predecessor_ids == [1] + assert lane_group0.successor_ids == [2] + assert isinstance(lane_group0.trimesh_mesh, trimesh.base.Trimesh) + + def test_base_properties(self): + """Test that the base surface properties of the LaneGroup objects are correct.""" + lane_group0 = self.lane_groups[0] + assert lane_group0.object_id == 0 + assert isinstance(lane_group0.outline, Polyline3D) + assert isinstance(lane_group0.outline_2d, Polyline2D) + assert isinstance(lane_group0.outline_3d, Polyline3D) + assert isinstance(lane_group0.shapely_polygon, shapely.Polygon) + + def test_lane_links(self): + """Test that the lanes are correctly linked to the lane group.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + # Lane group 0 contains lanes 0, 1, 2 + lane_group0: LaneGroup = map_api.get_map_object(0, MapLayer.LANE_GROUP) + assert lane_group0 is not None + assert lane_group0.lanes is not None + assert len(lane_group0.lanes) == 3 + for i, lane in enumerate(lane_group0.lanes): + assert isinstance(lane, Lane) + assert lane.object_id == i + + # Lane group 1 contains lane 3 + lane_group1: LaneGroup = map_api.get_map_object(1, MapLayer.LANE_GROUP) + assert lane_group1 is not None + assert lane_group1.lanes is not None + assert len(lane_group1.lanes) == 1 + assert isinstance(lane_group1.lanes[0], Lane) + assert lane_group1.lanes[0].object_id == 3 + + # Lane group 2 contains lane 4 + lane_group2: LaneGroup = map_api.get_map_object(2, MapLayer.LANE_GROUP) + assert lane_group2 is not None + assert lane_group2.lanes is not None + assert len(lane_group2.lanes) == 1 + assert isinstance(lane_group2.lanes[0], Lane) + assert lane_group2.lanes[0].object_id == 4 + + def test_predecessor_links(self): + """Test that the predecessor lane groups are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_predecessors(lane_group: LaneGroup): + assert lane_group is not None + assert lane_group.predecessors == [] + assert lane_group.predecessor_ids == [] + + # Lane group 0 has predecessor lane group 1 + lane_group0: LaneGroup = map_api.get_map_object(0, MapLayer.LANE_GROUP) + assert lane_group0 is not None + assert lane_group0.predecessors is not None + assert len(lane_group0.predecessors) == 1 + assert isinstance(lane_group0.predecessors[0], LaneGroup) + assert lane_group0.predecessors[0].object_id == 1 + assert lane_group0.predecessor_ids == [1] + + # Lane group 1 has no predecessors + lane_group1: LaneGroup = map_api.get_map_object(1, MapLayer.LANE_GROUP) + _no_predecessors(lane_group1) + + # Lane group 2 has predecessor lane group 0 + lane_group2: LaneGroup = map_api.get_map_object(2, MapLayer.LANE_GROUP) + assert lane_group2 is not None + assert lane_group2.predecessors is not None + assert len(lane_group2.predecessors) == 1 + assert isinstance(lane_group2.predecessors[0], LaneGroup) + assert lane_group2.predecessors[0].object_id == 0 + assert lane_group2.predecessor_ids == [0] + + def test_successor_links(self): + """Test that the successor lane groups are correctly linked.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + def _no_successors(lane_group: LaneGroup): + assert lane_group is not None + assert lane_group.successors == [] + assert lane_group.successor_ids == [] + + # Lane group 0 has successor lane group 2 + lane_group0: LaneGroup = map_api.get_map_object(0, MapLayer.LANE_GROUP) + assert lane_group0 is not None + assert lane_group0.successors is not None + assert len(lane_group0.successors) == 1 + assert isinstance(lane_group0.successors[0], LaneGroup) + assert lane_group0.successors[0].object_id == 2 + assert lane_group0.successor_ids == [2] + + # Lane group 1 has successor lane group 0 + lane_group1: LaneGroup = map_api.get_map_object(1, MapLayer.LANE_GROUP) + assert lane_group1 is not None + assert lane_group1.successors is not None + assert len(lane_group1.successors) == 1 + assert isinstance(lane_group1.successors[0], LaneGroup) + assert lane_group1.successors[0].object_id == 0 + assert lane_group1.successor_ids == [0] + + # Lane group 2 has no successors + lane_group2: LaneGroup = map_api.get_map_object(2, MapLayer.LANE_GROUP) + _no_successors(lane_group2) + + def test_intersection_links(self): + """Test that the intersection links are correct.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + # Lane group 0 has no intersection + lane_group0: LaneGroup = map_api.get_map_object(0, MapLayer.LANE_GROUP) + assert lane_group0 is not None + assert lane_group0.intersection_id is None + assert lane_group0.intersection is None + + # Lane group 1 has intersection 0 + lane_group1: LaneGroup = map_api.get_map_object(1, MapLayer.LANE_GROUP) + assert lane_group1 is not None + assert lane_group1.intersection_id == 0 + assert lane_group1.intersection is not None + assert isinstance(lane_group1.intersection, Intersection) + assert lane_group1.intersection.object_id == 0 + + # Lane group 2 has intersection 1 + lane_group2: LaneGroup = map_api.get_map_object(2, MapLayer.LANE_GROUP) + assert lane_group2 is not None + assert lane_group2.intersection_id == 1 + assert lane_group2.intersection is not None + assert isinstance(lane_group2.intersection, Intersection) + assert lane_group2.intersection.object_id == 1 + + def test_no_links(self): + """Test that when map_api is not provided, no links are available.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=False, + ) + for lane_group in self.lane_groups: + lg_from_api: LaneGroup = map_api.get_map_object(lane_group.object_id, MapLayer.LANE_GROUP) + assert lg_from_api is not None + assert lg_from_api.lanes is None + assert lg_from_api.predecessors is None + assert lg_from_api.successors is None + assert lg_from_api.intersection is None + + +class TestIntersection: + def setup_method(self): + lanes, lane_groups, intersections = _get_linked_map_object_setup() + self.lanes = lanes + self.lane_groups = lane_groups + self.intersections = intersections + + def test_properties(self): + """Test that the properties of the Intersection objects are correct.""" + intersection0 = self.intersections[0] + assert intersection0.layer == MapLayer.INTERSECTION + assert intersection0.lane_group_ids == [1] + assert isinstance(intersection0.outline, Polyline3D) + + def test_base_properties(self): + """Test that the base surface properties of the Intersection objects are correct.""" + intersection0 = self.intersections[0] + assert intersection0.object_id == 0 + assert isinstance(intersection0.outline, Polyline3D) + assert isinstance(intersection0.outline_2d, Polyline2D) + assert isinstance(intersection0.outline_3d, Polyline3D) + assert isinstance(intersection0.shapely_polygon, shapely.Polygon) + + def test_lane_group_links(self): + """Test that the lane groups are correctly linked to the intersection.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=True, + ) + + # Intersection 0 contains lane group 1 + intersection0: Intersection = map_api.get_map_object(0, MapLayer.INTERSECTION) + assert intersection0 is not None + assert intersection0.lane_groups is not None + assert len(intersection0.lane_groups) == 1 + assert isinstance(intersection0.lane_groups[0], LaneGroup) + assert intersection0.lane_groups[0].object_id == 1 + + # Intersection 1 contains lane group 2 + intersection1: Intersection = map_api.get_map_object(1, MapLayer.INTERSECTION) + assert intersection1 is not None + assert intersection1.lane_groups is not None + assert len(intersection1.lane_groups) == 1 + assert isinstance(intersection1.lane_groups[0], LaneGroup) + assert intersection1.lane_groups[0].object_id == 2 + + def test_no_links(self): + """Test that when map_api is not provided, no links are available.""" + map_api = MockMapAPI( + lanes=self.lanes, + lane_groups=self.lane_groups, + intersections=self.intersections, + add_map_api_links=False, + ) + for intersection in self.intersections: + int_from_api: Intersection = map_api.get_map_object(intersection.object_id, MapLayer.INTERSECTION) + assert int_from_api is not None + assert int_from_api.lane_groups is None + + +class TestCrosswalk: + def test_properties(self): + """Test that the properties of the Crosswalk object are correct.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) + ) + crosswalk = Crosswalk(object_id=0, outline=outline) + assert crosswalk.layer == MapLayer.CROSSWALK + assert crosswalk.object_id == 0 + assert isinstance(crosswalk.outline, Polyline3D) + assert isinstance(crosswalk.shapely_polygon, shapely.Polygon) + + def test_init_with_shapely_polygon(self): + """Test initialization with shapely polygon.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 1.0)]) + crosswalk = Crosswalk(object_id=0, shapely_polygon=shapely_polygon) + assert crosswalk.object_id == 0 + assert isinstance(crosswalk.shapely_polygon, shapely.Polygon) + assert isinstance(crosswalk.outline_2d, Polyline2D) + + def test_init_with_polyline2d(self): + """Test initialization with Polyline2D outline.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [0.0, 0.0]])) + crosswalk = Crosswalk(object_id=0, outline=outline) + assert isinstance(crosswalk.outline_2d, Polyline2D) + assert isinstance(crosswalk.shapely_polygon, shapely.Polygon) + + def test_base_surface_properties(self): + """Test base surface object properties.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) + ) + crosswalk = Crosswalk(object_id=0, outline=outline) + assert isinstance(crosswalk.outline_3d, Polyline3D) + assert crosswalk.shapely_polygon.is_valid + + +class TestCarpark: + def test_properties(self): + """Test that the properties of the Carpark object are correct.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [2.0, 2.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 0.0]]) + ) + carpark = Carpark(object_id=1, outline=outline) + assert carpark.layer == MapLayer.CARPARK + assert carpark.object_id == 1 + assert isinstance(carpark.outline, Polyline3D) + assert isinstance(carpark.shapely_polygon, shapely.Polygon) + + def test_init_with_shapely_polygon(self): + """Test initialization with shapely polygon.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (2.0, 0.0), (2.0, 2.0), (0.0, 2.0)]) + carpark = Carpark(object_id=1, shapely_polygon=shapely_polygon) + assert carpark.object_id == 1 + assert isinstance(carpark.shapely_polygon, shapely.Polygon) + assert isinstance(carpark.outline_2d, Polyline2D) + + def test_init_with_polyline2d(self): + """Test initialization with Polyline2D outline.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [2.0, 0.0], [2.0, 2.0], [0.0, 2.0], [0.0, 0.0]])) + carpark = Carpark(object_id=1, outline=outline) + assert isinstance(carpark.outline_2d, Polyline2D) + assert isinstance(carpark.shapely_polygon, shapely.Polygon) + + def test_polygon_area(self): + """Test that the polygon area is calculated correctly.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [2.0, 2.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 0.0]]) + ) + carpark = Carpark(object_id=1, outline=outline) + assert carpark.shapely_polygon.area == pytest.approx(4.0) + + +class TestWalkway: + def test_properties(self): + """Test that the properties of the Walkway object are correct.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [3.0, 0.0], [3.0, 1.0], [0.0, 1.0], [0.0, 0.0]])) + walkway = Walkway(object_id=2, outline=outline) + assert walkway.layer == MapLayer.WALKWAY + assert walkway.object_id == 2 + assert isinstance(walkway.outline_2d, Polyline2D) + assert isinstance(walkway.shapely_polygon, shapely.Polygon) + + def test_init_with_shapely_polygon(self): + """Test initialization with shapely polygon.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (3.0, 0.0), (3.0, 1.0), (0.0, 1.0)]) + walkway = Walkway(object_id=2, shapely_polygon=shapely_polygon) + assert walkway.object_id == 2 + assert isinstance(walkway.shapely_polygon, shapely.Polygon) + + def test_init_with_polyline3d(self): + """Test initialization with Polyline3D outline.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [3.0, 0.0, 0.0], [3.0, 1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]) + ) + walkway = Walkway(object_id=2, outline=outline) + assert isinstance(walkway.outline_3d, Polyline3D) + assert isinstance(walkway.shapely_polygon, shapely.Polygon) + + def test_polygon_bounds(self): + """Test that polygon bounds are correct.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [3.0, 0.0], [3.0, 1.0], [0.0, 1.0], [0.0, 0.0]])) + walkway = Walkway(object_id=2, outline=outline) + bounds = walkway.shapely_polygon.bounds + assert bounds == (0.0, 0.0, 3.0, 1.0) + + +class TestGenericDrivable: + def test_properties(self): + """Test that the properties of the GenericDrivable object are correct.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [5.0, 3.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]]) + ) + generic_drivable = GenericDrivable(object_id=3, outline=outline) + assert generic_drivable.layer == MapLayer.GENERIC_DRIVABLE + assert generic_drivable.object_id == 3 + assert isinstance(generic_drivable.outline, Polyline3D) + assert isinstance(generic_drivable.shapely_polygon, shapely.Polygon) + + def test_init_with_shapely_polygon(self): + """Test initialization with shapely polygon.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (5.0, 0.0), (5.0, 3.0), (0.0, 3.0)]) + generic_drivable = GenericDrivable(object_id=3, shapely_polygon=shapely_polygon) + assert generic_drivable.object_id == 3 + assert isinstance(generic_drivable.shapely_polygon, shapely.Polygon) + + def test_init_with_polyline2d(self): + """Test initialization with Polyline2D outline.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [5.0, 0.0], [5.0, 3.0], [0.0, 3.0], [0.0, 0.0]])) + generic_drivable = GenericDrivable(object_id=3, outline=outline) + assert isinstance(generic_drivable.outline_2d, Polyline2D) + assert isinstance(generic_drivable.shapely_polygon, shapely.Polygon) + + def test_polygon_area(self): + """Test that the polygon area is calculated correctly.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [5.0, 3.0, 0.0], [0.0, 3.0, 0.0], [0.0, 0.0, 0.0]]) + ) + generic_drivable = GenericDrivable(object_id=3, outline=outline) + assert generic_drivable.shapely_polygon.area == pytest.approx(15.0) + + +class TestStopZone: + def test_properties(self): + """Test that the properties of the StopZone object are correct.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 0.5), (0.0, 0.5)]) + stop_zone = StopZone(object_id=4, shapely_polygon=shapely_polygon) + assert stop_zone.layer == MapLayer.STOP_ZONE + assert stop_zone.object_id == 4 + assert isinstance(stop_zone.shapely_polygon, shapely.Polygon) + assert isinstance(stop_zone.outline_2d, Polyline2D) + + def test_init_with_polyline3d(self): + """Test initialization with Polyline3D outline.""" + outline = Polyline3D.from_array( + np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.5, 0.0], [0.0, 0.5, 0.0], [0.0, 0.0, 0.0]]) + ) + stop_zone = StopZone(object_id=4, outline=outline) + assert isinstance(stop_zone.outline, Polyline3D) + assert isinstance(stop_zone.shapely_polygon, shapely.Polygon) + + def test_init_with_polyline2d(self): + """Test initialization with Polyline2D outline.""" + outline = Polyline2D.from_array(np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 0.5], [0.0, 0.5], [0.0, 0.0]])) + stop_zone = StopZone(object_id=4, outline=outline) + assert isinstance(stop_zone.outline_2d, Polyline2D) + assert isinstance(stop_zone.shapely_polygon, shapely.Polygon) + + def test_polygon_area(self): + """Test that the polygon area is calculated correctly.""" + shapely_polygon = shapely.Polygon([(0.0, 0.0), (1.0, 0.0), (1.0, 0.5), (0.0, 0.5)]) + stop_zone = StopZone(object_id=4, shapely_polygon=shapely_polygon) + assert stop_zone.shapely_polygon.area == pytest.approx(0.5) + + +class TestRoadEdge: + def test_properties(self): + """Test that the properties of the RoadEdge object are correct.""" + polyline = Polyline3D.from_array(np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0]])) + road_edge = RoadEdge(object_id=5, road_edge_type=1, polyline=polyline) + assert road_edge.layer == MapLayer.ROAD_EDGE + assert road_edge.object_id == 5 + assert road_edge.road_edge_type == 1 + assert isinstance(road_edge.polyline, Polyline3D) + + def test_init_with_polyline2d(self): + """Test initialization with Polyline2D.""" + polyline = Polyline2D.from_array(np.array([[0.0, 0.0], [10.0, 0.0], [20.0, 0.0]])) + road_edge = RoadEdge(object_id=5, road_edge_type=1, polyline=polyline) + assert isinstance(road_edge.polyline, Polyline2D) + assert road_edge.road_edge_type == 1 + + def test_polyline_length(self): + """Test that the polyline has correct number of points.""" + polyline = Polyline3D.from_array(np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [20.0, 0.0, 0.0]])) + road_edge = RoadEdge(object_id=5, road_edge_type=1, polyline=polyline) + assert len(road_edge.polyline.array) == 3 + + def test_different_road_edge_types(self): + """Test different road edge types.""" + polyline = Polyline3D.from_array(np.array([[0.0, 0.0, 0.0], [10.0, 0.0, 0.0]])) + for edge_type in RoadEdgeType: + road_edge = RoadEdge(object_id=5, road_edge_type=edge_type, polyline=polyline) + assert road_edge.road_edge_type == edge_type + + +class TestRoadLine: + def test_properties(self): + """Test that the properties of the RoadLine object are correct.""" + polyline = Polyline2D.from_array(np.array([[0.0, 1.0], [10.0, 1.0], [20.0, 1.0]])) + road_line = RoadLine(object_id=6, road_line_type=2, polyline=polyline) + assert road_line.layer == MapLayer.ROAD_LINE + assert road_line.object_id == 6 + assert road_line.road_line_type == 2 + assert isinstance(road_line.polyline, Polyline2D) + + def test_init_with_polyline3d(self): + """Test initialization with Polyline3D.""" + polyline = Polyline3D.from_array(np.array([[0.0, 1.0, 0.0], [10.0, 1.0, 0.0], [20.0, 1.0, 0.0]])) + road_line = RoadLine(object_id=6, road_line_type=2, polyline=polyline) + assert isinstance(road_line.polyline, Polyline3D) + assert road_line.road_line_type == 2 + + def test_polyline_length(self): + """Test that the polyline has correct number of points.""" + polyline = Polyline2D.from_array(np.array([[0.0, 1.0], [10.0, 1.0], [20.0, 1.0], [30.0, 1.0]])) + road_line = RoadLine(object_id=6, road_line_type=2, polyline=polyline) + assert len(road_line.polyline.array) == 4 + + def test_different_road_line_types(self): + """Test different road line types.""" + polyline = Polyline2D.from_array(np.array([[0.0, 1.0], [10.0, 1.0]])) + for line_type in RoadLineType: + road_line = RoadLine(object_id=6, road_line_type=line_type, polyline=polyline) + assert road_line.road_line_type == line_type diff --git a/tests/unit/datatypes/metadata/__init__.py b/tests/unit/datatypes/metadata/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/datatypes/metadata/test_log_metadata.py b/tests/unit/datatypes/metadata/test_log_metadata.py new file mode 100644 index 00000000..8d8348b8 --- /dev/null +++ b/tests/unit/datatypes/metadata/test_log_metadata.py @@ -0,0 +1,152 @@ +from unittest.mock import MagicMock, patch + +import pytest + +from py123d.datatypes.metadata.log_metadata import LogMetadata +from py123d.datatypes.metadata.map_metadata import MapMetadata +from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters + + +class TestLogMetadata: + def test_init_minimal(self): + """Test LogMetadata initialization with minimal required fields.""" + log_metadata = LogMetadata( + dataset="test_dataset", split="train", log_name="log_001", location="test_location", timestep_seconds=0.1 + ) + assert log_metadata.dataset == "test_dataset" + assert log_metadata.split == "train" + assert log_metadata.log_name == "log_001" + assert log_metadata.location == "test_location" + assert log_metadata.timestep_seconds == 0.1 + assert log_metadata.vehicle_parameters is None + assert log_metadata.box_detection_label_class is None + assert log_metadata.pinhole_camera_metadata == {} + assert log_metadata.fisheye_mei_camera_metadata == {} + assert log_metadata.lidar_metadata == {} + assert log_metadata.map_metadata is None + + def test_to_dict_minimal(self): + """Test to_dict with minimal fields.""" + log_metadata = LogMetadata( + dataset="test_dataset", split="train", log_name="log_001", location="test_location", timestep_seconds=0.1 + ) + result = log_metadata.to_dict() + assert result["dataset"] == "test_dataset" + assert result["split"] == "train" + assert result["log_name"] == "log_001" + assert result["location"] == "test_location" + assert result["timestep_seconds"] == 0.1 + assert result["vehicle_parameters"] is None + assert result["box_detection_label_class"] is None + assert result["pinhole_camera_metadata"] == {} + + def test_from_dict_minimal(self): + """Test from_dict with minimal fields.""" + data_dict = { + "dataset": "test_dataset", + "split": "train", + "log_name": "log_001", + "location": "test_location", + "timestep_seconds": 0.1, + "vehicle_parameters": None, + "box_detection_label_class": None, + "map_metadata": None, + "version": "1.0.0", + } + log_metadata = LogMetadata.from_dict(data_dict) + assert log_metadata.dataset == "test_dataset" + assert log_metadata.split == "train" + assert log_metadata.vehicle_parameters is None + + @patch.object(VehicleParameters, "from_dict") + def test_from_dict_with_vehicle_parameters(self, mock_vehicle_params): + """Test from_dict with vehicle parameters.""" + mock_vehicle = MagicMock() + mock_vehicle_params.return_value = mock_vehicle + + data_dict = { + "dataset": "test_dataset", + "split": "train", + "log_name": "log_001", + "location": "test_location", + "timestep_seconds": 0.1, + "vehicle_parameters": {"some": "data"}, + "box_detection_label_class": None, + "map_metadata": None, + "version": "1.0.0", + } + log_metadata = LogMetadata.from_dict(data_dict) + mock_vehicle_params.assert_called_once_with({"some": "data"}) + assert log_metadata.vehicle_parameters == mock_vehicle + + @patch("py123d.datatypes.metadata.log_metadata.BOX_DETECTION_LABEL_REGISTRY", {"TestLabel": MagicMock}) + def test_from_dict_with_box_detection_label(self): + """Test from_dict with box detection label class.""" + data_dict = { + "dataset": "test_dataset", + "split": "train", + "log_name": "log_001", + "location": "test_location", + "timestep_seconds": 0.1, + "vehicle_parameters": None, + "box_detection_label_class": "TestLabel", + "map_metadata": None, + "version": "1.0.0", + } + log_metadata = LogMetadata.from_dict(data_dict) + assert log_metadata.box_detection_label_class is not None + + def test_from_dict_with_invalid_box_detection_label(self): + """Test from_dict with invalid box detection label class.""" + data_dict = { + "dataset": "test_dataset", + "split": "train", + "log_name": "log_001", + "location": "test_location", + "timestep_seconds": 0.1, + "vehicle_parameters": None, + "box_detection_label_class": "InvalidLabel", + "map_metadata": None, + "version": "1.0.0", + } + with pytest.raises(ValueError): + LogMetadata.from_dict(data_dict) + + @patch.object(MapMetadata, "from_dict") + def test_from_dict_with_map_metadata(self, mock_map_metadata): + """Test from_dict with map metadata.""" + mock_map = MagicMock() + mock_map_metadata.return_value = mock_map + + data_dict = { + "dataset": "test_dataset", + "split": "train", + "log_name": "log_001", + "location": "test_location", + "timestep_seconds": 0.1, + "vehicle_parameters": None, + "box_detection_label_class": None, + "map_metadata": {"some": "data"}, + "version": "1.0.0", + } + log_metadata = LogMetadata.from_dict(data_dict) + mock_map_metadata.assert_called_once_with({"some": "data"}) + assert log_metadata.map_metadata == mock_map + + def test_roundtrip_serialization(self): + """Test that to_dict and from_dict are inverses.""" + original = LogMetadata( + dataset="test_dataset", + split="train", + log_name="log_001", + location="test_location", + timestep_seconds=0.1, + ) + data_dict = original.to_dict() + reconstructed = LogMetadata.from_dict(data_dict) + + assert original.dataset == reconstructed.dataset + assert original.split == reconstructed.split + assert original.log_name == reconstructed.log_name + assert original.location == reconstructed.location + assert original.timestep_seconds == reconstructed.timestep_seconds diff --git a/tests/unit/datatypes/metadata/test_map_metadata.py b/tests/unit/datatypes/metadata/test_map_metadata.py new file mode 100644 index 00000000..2af85186 --- /dev/null +++ b/tests/unit/datatypes/metadata/test_map_metadata.py @@ -0,0 +1,104 @@ +from py123d.datatypes.metadata.map_metadata import MapMetadata + + +class TestMapMetadata: + def test_map_metadata_initialization(self): + """Test that MapMetadata can be initialized with required fields.""" + metadata = MapMetadata( + dataset="test_dataset", + split="train", + log_name="log_001", + location="test_location", + map_has_z=True, + map_is_local=False, + ) + + assert metadata.dataset == "test_dataset" + assert metadata.split == "train" + assert metadata.log_name == "log_001" + assert metadata.location == "test_location" + assert metadata.map_has_z is True + assert metadata.map_is_local is False + assert metadata.version is not None + + def test_map_metadata_to_dict(self): + """Test conversion of MapMetadata to dictionary.""" + metadata = MapMetadata( + dataset="test_dataset", + split="val", + log_name="log_002", + location="test_location", + map_has_z=False, + map_is_local=True, + ) + + result = metadata.to_dict() + + assert isinstance(result, dict) + assert result["dataset"] == "test_dataset" + assert result["split"] == "val" + assert result["log_name"] == "log_002" + assert result["location"] == "test_location" + assert result["map_has_z"] is False + assert result["map_is_local"] is True + assert "version" in result + + def test_map_metadata_from_dict(self): + """Test creation of MapMetadata from dictionary.""" + data = { + "dataset": "test_dataset", + "split": "test", + "log_name": "log_003", + "location": "test_location", + "map_has_z": True, + "map_is_local": False, + "version": "1.0.0", + } + + metadata = MapMetadata.from_dict(data) + + assert metadata.dataset == "test_dataset" + assert metadata.split == "test" + assert metadata.log_name == "log_003" + assert metadata.location == "test_location" + assert metadata.map_has_z is True + assert metadata.map_is_local is False + assert metadata.version == "1.0.0" + + def test_map_metadata_with_none_values(self): + """Test MapMetadata with None values for optional fields.""" + metadata = MapMetadata( + dataset="test_dataset", + split=None, + log_name=None, + location="test_location", + map_has_z=True, + map_is_local=True, + ) + + assert metadata.split is None + assert metadata.log_name is None + assert metadata.dataset == "test_dataset" + + def test_map_metadata_roundtrip(self): + """Test that converting to dict and back preserves data.""" + original = MapMetadata( + dataset="roundtrip_dataset", + split="train", + log_name="log_roundtrip", + location="location_test", + map_has_z=False, + map_is_local=True, + version="2.0.0", + ) + + data_dict = original.to_dict() + restored = MapMetadata.from_dict(data_dict) + + assert restored.dataset == original.dataset + assert restored.split == original.split + assert restored.log_name == original.log_name + assert restored.location == original.location + assert restored.map_has_z == original.map_has_z + assert restored.map_is_local == original.map_is_local + assert restored.version == original.version diff --git a/tests/unit/datatypes/sensors/test_fisheye_mei_camera.py b/tests/unit/datatypes/sensors/test_fisheye_mei_camera.py new file mode 100644 index 00000000..0731b56d --- /dev/null +++ b/tests/unit/datatypes/sensors/test_fisheye_mei_camera.py @@ -0,0 +1,368 @@ +import numpy as np +import pytest + +from py123d.datatypes.sensors.fisheye_mei_camera import ( + FisheyeMEICamera, + FisheyeMEICameraMetadata, + FisheyeMEICameraType, + FisheyeMEIDistortion, + FisheyeMEIDistortionIndex, + FisheyeMEIProjection, + FisheyeMEIProjectionIndex, +) +from py123d.geometry import PoseSE3 + + +class TestFisheyeMEICameraType: + def test_camera_type_values(self): + """Test that camera type enum has expected values.""" + assert FisheyeMEICameraType.FCAM_L.value == 0 + assert FisheyeMEICameraType.FCAM_R.value == 1 + + def test_camera_type_from_int(self): + """Test creating camera type from integer values.""" + assert FisheyeMEICameraType(0) == FisheyeMEICameraType.FCAM_L + assert FisheyeMEICameraType(1) == FisheyeMEICameraType.FCAM_R + + def test_camera_type_members(self): + """Test that all expected members exist.""" + members = list(FisheyeMEICameraType) + assert len(members) == 2 + assert FisheyeMEICameraType.FCAM_L in members + assert FisheyeMEICameraType.FCAM_R in members + + def test_camera_type_comparison(self): + """Test comparison between camera types.""" + assert FisheyeMEICameraType.FCAM_L != FisheyeMEICameraType.FCAM_R + assert FisheyeMEICameraType.FCAM_L == FisheyeMEICameraType.FCAM_L + + +class TestFisheyeMEIDistortion: + def test_distortion_initialization(self): + """Test distortion parameter initialization.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + assert distortion.k1 == 0.1 + assert distortion.k2 == 0.2 + assert distortion.p1 == 0.3 + assert distortion.p2 == 0.4 + + def test_distortion_from_array(self): + """Test creating distortion from array.""" + array = np.array([0.1, 0.2, 0.3, 0.4]) + distortion = FisheyeMEIDistortion.from_array(array) + assert distortion.k1 == 0.1 + assert distortion.k2 == 0.2 + assert distortion.p1 == 0.3 + assert distortion.p2 == 0.4 + + def test_distortion_from_array_copy(self): + """Test that from_array copies data by default.""" + array = np.array([0.1, 0.2, 0.3, 0.4]) + distortion = FisheyeMEIDistortion.from_array(array, copy=True) + array[0] = 999.0 + assert distortion.k1 == 0.1 + + def test_distortion_from_array_no_copy(self): + """Test that from_array can avoid copying.""" + array = np.array([0.1, 0.2, 0.3, 0.4]) + distortion = FisheyeMEIDistortion.from_array(array, copy=False) + array[0] = 999.0 + assert distortion.k1 == 999.0 + + def test_distortion_array_property(self): + """Test array property returns correct values.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + array = distortion.array + assert len(array) == 4 + np.testing.assert_array_equal(array, [0.1, 0.2, 0.3, 0.4]) + + def test_distortion_index_mapping(self): + """Test that distortion indices map correctly.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + assert distortion.array[FisheyeMEIDistortionIndex.K1] == 0.1 + assert distortion.array[FisheyeMEIDistortionIndex.K2] == 0.2 + assert distortion.array[FisheyeMEIDistortionIndex.P1] == 0.3 + assert distortion.array[FisheyeMEIDistortionIndex.P2] == 0.4 + + +class TestFisheyeMEIProjection: + def test_projection_initialization(self): + """Test projection parameter initialization.""" + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + assert projection.gamma1 == 1.0 + assert projection.gamma2 == 2.0 + assert projection.u0 == 3.0 + assert projection.v0 == 4.0 + + def test_projection_from_array(self): + """Test creating projection from array.""" + array = np.array([1.0, 2.0, 3.0, 4.0]) + projection = FisheyeMEIProjection.from_array(array) + assert projection.gamma1 == 1.0 + assert projection.gamma2 == 2.0 + assert projection.u0 == 3.0 + assert projection.v0 == 4.0 + + def test_projection_from_array_copy(self): + """Test that from_array copies data by default.""" + array = np.array([1.0, 2.0, 3.0, 4.0]) + projection = FisheyeMEIProjection.from_array(array, copy=True) + array[0] = 999.0 + assert projection.gamma1 == 1.0 + + def test_projection_from_array_no_copy(self): + """Test that from_array can avoid copying.""" + array = np.array([1.0, 2.0, 3.0, 4.0]) + projection = FisheyeMEIProjection.from_array(array, copy=False) + array[0] = 999.0 + assert projection.gamma1 == 999.0 + + def test_projection_array_property(self): + """Test array property returns correct values.""" + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + array = projection.array + assert len(array) == 4 + np.testing.assert_array_equal(array, [1.0, 2.0, 3.0, 4.0]) + + def test_projection_index_mapping(self): + """Test that projection indices map correctly.""" + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + assert projection.array[FisheyeMEIProjectionIndex.GAMMA1] == 1.0 + assert projection.array[FisheyeMEIProjectionIndex.GAMMA2] == 2.0 + assert projection.array[FisheyeMEIProjectionIndex.U0] == 3.0 + assert projection.array[FisheyeMEIProjectionIndex.V0] == 4.0 + + +class TestFisheyeMEICameraMetadata: + def test_metadata_initialization(self): + """Test metadata initialization with all parameters.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=distortion, + projection=projection, + width=1920, + height=1080, + ) + assert metadata.camera_type == FisheyeMEICameraType.FCAM_L + assert metadata.mirror_parameter == 0.5 + assert metadata.distortion == distortion + assert metadata.projection == projection + assert metadata.aspect_ratio == 1920 / 1080 + + def test_metadata_initialization_with_none(self): + """Test metadata initialization with None distortion and projection.""" + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_R, + mirror_parameter=None, + distortion=None, + projection=None, + width=640, + height=480, + ) + assert metadata.camera_type == FisheyeMEICameraType.FCAM_R + assert metadata.mirror_parameter is None + assert metadata.distortion is None + assert metadata.projection is None + assert metadata.aspect_ratio == 640 / 480 + + def test_metadata_to_dict(self): + """Test converting metadata to dictionary.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=distortion, + projection=projection, + width=1920, + height=1080, + ) + result = metadata.to_dict() + assert result["camera_type"] == 0 + assert result["mirror_parameter"] == 0.5 + assert result["distortion"] == [0.1, 0.2, 0.3, 0.4] + assert result["projection"] == [1.0, 2.0, 3.0, 4.0] + assert result["width"] == 1920 + assert result["height"] == 1080 + + def test_metadata_to_dict_with_none(self): + """Test converting metadata with None values to dictionary.""" + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_R, + mirror_parameter=None, + distortion=None, + projection=None, + width=640, + height=480, + ) + result = metadata.to_dict() + assert result["camera_type"] == 1 + assert result["mirror_parameter"] is None + assert result["distortion"] is None + assert result["projection"] is None + assert result["width"] == 640 + assert result["height"] == 480 + + def test_metadata_from_dict(self): + """Test creating metadata from dictionary.""" + data = { + "camera_type": 0, + "mirror_parameter": 0.5, + "distortion": [0.1, 0.2, 0.3, 0.4], + "projection": [1.0, 2.0, 3.0, 4.0], + "width": 1920, + "height": 1080, + } + metadata = FisheyeMEICameraMetadata.from_dict(data) + assert metadata.camera_type == FisheyeMEICameraType.FCAM_L + assert metadata.mirror_parameter == 0.5 + assert metadata.distortion.k1 == 0.1 + assert metadata.projection.gamma1 == 1.0 + assert metadata.aspect_ratio == 1920 / 1080 + + def test_metadata_from_dict_with_none(self): + """Test creating metadata from dictionary with None values.""" + data = { + "camera_type": 1, + "mirror_parameter": None, + "distortion": None, + "projection": None, + "width": 640, + "height": 480, + } + metadata = FisheyeMEICameraMetadata.from_dict(data) + assert metadata.camera_type == FisheyeMEICameraType.FCAM_R + assert metadata.mirror_parameter is None + assert metadata.distortion is None + assert metadata.projection is None + + def test_metadata_roundtrip(self): + """Test that to_dict and from_dict are inverses.""" + distortion = FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4) + projection = FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0) + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=distortion, + projection=projection, + width=1920, + height=1080, + ) + data_dict = metadata.to_dict() + metadata_restored = FisheyeMEICameraMetadata.from_dict(data_dict) + assert metadata.camera_type == metadata_restored.camera_type + assert metadata.mirror_parameter == metadata_restored.mirror_parameter + np.testing.assert_array_equal(metadata.distortion.array, metadata_restored.distortion.array) + np.testing.assert_array_equal(metadata.projection.array, metadata_restored.projection.array) + assert metadata.aspect_ratio == metadata_restored.aspect_ratio + + def test_aspect_ratio_calculation(self): + """Test aspect ratio calculation.""" + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=None, + projection=None, + width=1920, + height=1080, + ) + assert metadata.aspect_ratio == pytest.approx(16 / 9, abs=1e-05) + + +class TestFisheyeMEICamera: + def test_camera_initialization(self): + """Test FisheyeMEICamera initialization.""" + + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=FisheyeMEIDistortion(k1=0.1, k2=0.2, p1=0.3, p2=0.4), + projection=FisheyeMEIProjection(gamma1=1.0, gamma2=2.0, u0=3.0, v0=4.0), + width=1920, + height=1080, + ) + image = np.zeros((1080, 1920), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = FisheyeMEICamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.metadata == metadata + np.testing.assert_array_equal(camera.image, image) + assert camera.extrinsic == extrinsic + + def test_camera_metadata_property(self): + """Test that metadata property returns correct metadata.""" + + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_R, + mirror_parameter=0.8, + distortion=None, + projection=None, + width=640, + height=480, + ) + image = np.ones((480, 640), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = FisheyeMEICamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.metadata is metadata + assert camera.metadata.camera_type == FisheyeMEICameraType.FCAM_R + + def test_camera_image_property(self): + """Test that image property returns correct image.""" + + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=None, + projection=None, + width=640, + height=480, + ) + image = np.random.randint(0, 255, (480, 640), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = FisheyeMEICamera(metadata=metadata, image=image, extrinsic=extrinsic) + + np.testing.assert_array_equal(camera.image, image) + assert camera.image.dtype == np.uint8 + + def test_camera_extrinsic_property(self): + """Test that extrinsic property returns correct pose.""" + + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=None, + projection=None, + width=640, + height=480, + ) + image = np.zeros((480, 640), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = FisheyeMEICamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.extrinsic is extrinsic + + def test_camera_with_color_image(self): + """Test camera with color (3-channel) image.""" + + metadata = FisheyeMEICameraMetadata( + camera_type=FisheyeMEICameraType.FCAM_L, + mirror_parameter=0.5, + distortion=None, + projection=None, + width=640, + height=480, + ) + image = np.zeros((480, 640, 3), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = FisheyeMEICamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.image.shape == (480, 640, 3) diff --git a/tests/unit/datatypes/sensors/test_lidar.py b/tests/unit/datatypes/sensors/test_lidar.py new file mode 100644 index 00000000..211b9e42 --- /dev/null +++ b/tests/unit/datatypes/sensors/test_lidar.py @@ -0,0 +1,279 @@ +import numpy as np +import pytest + +from py123d.conversion.registry.lidar_index_registry import LIDAR_INDEX_REGISTRY +from py123d.datatypes.sensors.lidar import LiDAR, LiDARMetadata, LiDARType +from py123d.geometry import PoseSE3 + + +class TestLiDARType: + def test_lidar_type_enum_values(self): + """Test that LiDARType enum has correct values.""" + assert LiDARType.LIDAR_UNKNOWN.value == 0 + assert LiDARType.LIDAR_MERGED.value == 1 + assert LiDARType.LIDAR_TOP.value == 2 + assert LiDARType.LIDAR_FRONT.value == 3 + assert LiDARType.LIDAR_SIDE_LEFT.value == 4 + assert LiDARType.LIDAR_SIDE_RIGHT.value == 5 + assert LiDARType.LIDAR_BACK.value == 6 + assert LiDARType.LIDAR_DOWN.value == 7 + + def test_lidar_type_enum_names(self): + """Test that LiDARType enum members have correct names.""" + assert LiDARType.LIDAR_UNKNOWN.name == "LIDAR_UNKNOWN" + assert LiDARType.LIDAR_MERGED.name == "LIDAR_MERGED" + assert LiDARType.LIDAR_TOP.name == "LIDAR_TOP" + assert LiDARType.LIDAR_FRONT.name == "LIDAR_FRONT" + assert LiDARType.LIDAR_SIDE_LEFT.name == "LIDAR_SIDE_LEFT" + assert LiDARType.LIDAR_SIDE_RIGHT.name == "LIDAR_SIDE_RIGHT" + assert LiDARType.LIDAR_BACK.name == "LIDAR_BACK" + assert LiDARType.LIDAR_DOWN.name == "LIDAR_DOWN" + + def test_lidar_type_from_value(self): + """Test that LiDARType can be created from integer values.""" + assert LiDARType(0) == LiDARType.LIDAR_UNKNOWN + assert LiDARType(1) == LiDARType.LIDAR_MERGED + assert LiDARType(2) == LiDARType.LIDAR_TOP + assert LiDARType(3) == LiDARType.LIDAR_FRONT + assert LiDARType(4) == LiDARType.LIDAR_SIDE_LEFT + assert LiDARType(5) == LiDARType.LIDAR_SIDE_RIGHT + assert LiDARType(6) == LiDARType.LIDAR_BACK + assert LiDARType(7) == LiDARType.LIDAR_DOWN + + def test_lidar_type_unique_values(self): + """Test that all LiDARType enum values are unique.""" + values = [member.value for member in LiDARType] + assert len(values) == len(set(values)) + + def test_lidar_type_count(self): + """Test that LiDARType has expected number of members.""" + assert len(LiDARType) == 8 + + +class TestLiDARMetadata: + def setup_method(self): + """Set up test fixtures.""" + + # Get a lidar index class from registry (assuming at least one exists) + self.lidar_index_class = next(iter(LIDAR_INDEX_REGISTRY.values())) + self.lidar_type = LiDARType.LIDAR_TOP + self.extrinsic = PoseSE3.from_list([1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0]) + + def test_lidar_metadata_creation_with_extrinsic(self): + """Test creating LiDARMetadata with extrinsic.""" + metadata = LiDARMetadata( + lidar_type=self.lidar_type, + lidar_index=self.lidar_index_class, + extrinsic=self.extrinsic, + ) + assert metadata.lidar_type == self.lidar_type + assert metadata.lidar_index == self.lidar_index_class + assert metadata.extrinsic is not None + + def test_lidar_metadata_creation_without_extrinsic(self): + """Test creating LiDARMetadata without extrinsic.""" + metadata = LiDARMetadata(lidar_type=self.lidar_type, lidar_index=self.lidar_index_class) + assert metadata.lidar_type == self.lidar_type + assert metadata.lidar_index == self.lidar_index_class + assert metadata.extrinsic is None + + def test_lidar_metadata_to_dict_with_extrinsic(self): + """Test serializing LiDARMetadata to dict with extrinsic.""" + metadata = LiDARMetadata( + lidar_type=self.lidar_type, + lidar_index=self.lidar_index_class, + extrinsic=self.extrinsic, + ) + data_dict = metadata.to_dict() + assert data_dict["lidar_type"] == self.lidar_type.name + assert data_dict["lidar_index"] == self.lidar_index_class.__name__ + assert data_dict["extrinsic"] is not None + assert isinstance(data_dict["extrinsic"], list) + + def test_lidar_metadata_to_dict_without_extrinsic(self): + """Test serializing LiDARMetadata to dict without extrinsic.""" + metadata = LiDARMetadata(lidar_type=self.lidar_type, lidar_index=self.lidar_index_class) + data_dict = metadata.to_dict() + assert data_dict["lidar_type"] == self.lidar_type.name + assert data_dict["lidar_index"] == self.lidar_index_class.__name__ + assert data_dict["extrinsic"] is None + + def test_lidar_metadata_from_dict_with_extrinsic(self): + """Test deserializing LiDARMetadata from dict with extrinsic.""" + data_dict = { + "lidar_type": self.lidar_type.name, + "lidar_index": self.lidar_index_class.__name__, + "extrinsic": [1.0, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0], + } + metadata = LiDARMetadata.from_dict(data_dict) + assert metadata.lidar_type == self.lidar_type + assert metadata.lidar_index == self.lidar_index_class + assert metadata.extrinsic is not None + + def test_lidar_metadata_from_dict_without_extrinsic(self): + """Test deserializing LiDARMetadata from dict without extrinsic.""" + data_dict = { + "lidar_type": self.lidar_type.name, + "lidar_index": self.lidar_index_class.__name__, + "extrinsic": None, + } + metadata = LiDARMetadata.from_dict(data_dict) + assert metadata.lidar_type == self.lidar_type + assert metadata.lidar_index == self.lidar_index_class + assert metadata.extrinsic is None + + def test_lidar_metadata_roundtrip_with_extrinsic(self): + """Test roundtrip serialization/deserialization with extrinsic.""" + metadata = LiDARMetadata( + lidar_type=self.lidar_type, + lidar_index=self.lidar_index_class, + extrinsic=self.extrinsic, + ) + data_dict = metadata.to_dict() + restored_metadata = LiDARMetadata.from_dict(data_dict) + assert restored_metadata.lidar_type == metadata.lidar_type + assert restored_metadata.lidar_index == metadata.lidar_index + + def test_lidar_metadata_roundtrip_without_extrinsic(self): + """Test roundtrip serialization/deserialization without extrinsic.""" + metadata = LiDARMetadata(lidar_type=self.lidar_type, lidar_index=self.lidar_index_class) + data_dict = metadata.to_dict() + restored_metadata = LiDARMetadata.from_dict(data_dict) + assert restored_metadata.lidar_type == metadata.lidar_type + assert restored_metadata.lidar_index == metadata.lidar_index + assert restored_metadata.extrinsic is None + + def test_lidar_metadata_from_dict_unknown_index_raises_error(self): + """Test that unknown lidar index raises ValueError.""" + data_dict = {"lidar_type": self.lidar_type.name, "lidar_index": "UnknownLiDARIndex", "extrinsic": None} + with pytest.raises(ValueError): + LiDARMetadata.from_dict(data_dict) + + +class TestLiDAR: + def setup_method(self): + """Set up test fixtures.""" + # Get a lidar index class from registry + + self.lidars = {} + self.extrinsic = PoseSE3.from_list([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0]) + + for lidar_index_name, lidar_index_class in LIDAR_INDEX_REGISTRY.items(): + metadata = LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=lidar_index_class, + extrinsic=self.extrinsic, + ) + point_cloud = np.random.rand(100, len(lidar_index_class)).astype(np.float32) + self.lidars[lidar_index_name] = LiDAR(metadata=metadata, point_cloud=point_cloud) + + def test_lidar_xyz_property(self): + """Test xyz property returns correct shape and values.""" + for lidar in self.lidars.values(): + xyz = lidar.xyz + assert xyz.shape[0] == lidar.point_cloud.shape[0] + assert xyz.shape[1] == 3 + + def test_lidar_xy_property(self): + """Test xy property returns correct shape and values.""" + for lidar in self.lidars.values(): + xy = lidar.xy + assert xy.shape[0] == lidar.point_cloud.shape[0] + assert xy.shape[1] == 2 + + def test_lidar_intensity_property_when_available(self): + """Test intensity property when INTENSITY attribute exists.""" + for lidar in self.lidars.values(): + intensity = lidar.intensity + if hasattr(lidar.metadata.lidar_index, "INTENSITY"): + assert intensity is not None + assert intensity.shape[0] == lidar.point_cloud.shape[0] + else: + assert intensity is None + + def test_lidar_intensity_property_when_not_available(self): + """Test intensity property returns None when not available.""" + for lidar in self.lidars.values(): + if not hasattr(lidar.metadata.lidar_index, "INTENSITY"): + assert lidar.intensity is None + + def test_lidar_range_property_when_available(self): + """Test range property when RANGE attribute exists.""" + for lidar in self.lidars.values(): + range_values = lidar.range + if hasattr(lidar.metadata.lidar_index, "RANGE"): + assert range_values is not None + assert range_values.shape[0] == lidar.point_cloud.shape[0] + else: + assert range_values is None + + def test_lidar_range_property_when_not_available(self): + """Test range property returns None when not available.""" + for lidar in self.lidars.values(): + if not hasattr(lidar.metadata.lidar_index, "RANGE"): + assert lidar.range is None + + def test_lidar_elongation_property_when_available(self): + """Test elongation property when ELONGATION attribute exists.""" + for lidar in self.lidars.values(): + elongation = lidar.elongation + if hasattr(lidar.metadata.lidar_index, "ELONGATION"): + assert elongation is not None + assert elongation.shape[0] == lidar.point_cloud.shape[0] + else: + assert elongation is None + + def test_lidar_elongation_property_when_not_available(self): + """Test elongation property returns None when not available.""" + for lidar in self.lidars.values(): + if not hasattr(lidar.metadata.lidar_index, "ELONGATION"): + assert lidar.elongation is None + + def test_lidar_ring_property_when_available(self): + """Test ring property when RING attribute exists.""" + for lidar in self.lidars.values(): + ring = lidar.ring + if hasattr(lidar.metadata.lidar_index, "RING"): + assert ring is not None + assert ring.shape[0] == lidar.point_cloud.shape[0] + else: + assert ring is None + + def test_lidar_ring_property_when_not_available(self): + """Test ring property returns None when not available.""" + for lidar in self.lidars.values(): + if not hasattr(lidar.metadata.lidar_index, "RING"): + assert lidar.ring is None + + def test_lidar_with_empty_point_cloud(self): + """Test LiDAR with empty point cloud.""" + for lidar_index_class in LIDAR_INDEX_REGISTRY.values(): + metadata = LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=lidar_index_class, + extrinsic=self.extrinsic, + ) + empty_point_cloud = np.empty((0, len(lidar_index_class)), dtype=np.float32) + lidar = LiDAR(metadata=metadata, point_cloud=empty_point_cloud) + assert lidar.xyz.shape == (0, 3) + assert lidar.xy.shape == (0, 2) + + def test_lidar_with_single_point(self): + """Test LiDAR with single point.""" + for lidar_index_class in LIDAR_INDEX_REGISTRY.values(): + metadata = LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=lidar_index_class, + extrinsic=self.extrinsic, + ) + single_point_cloud = np.random.rand(1, len(lidar_index_class)).astype(np.float32) + lidar = LiDAR(metadata=metadata, point_cloud=single_point_cloud) + assert lidar.xyz.shape == (1, 3) + assert lidar.xy.shape == (1, 2) + + def test_lidar_point_cloud_dtype(self): + """Test that point cloud maintains float32 dtype.""" + for lidar in self.lidars.values(): + assert lidar.point_cloud.dtype == np.float32 + assert lidar.xyz.dtype == np.float32 + assert lidar.xy.dtype == np.float32 diff --git a/tests/unit/datatypes/sensors/test_pinhole_camera.py b/tests/unit/datatypes/sensors/test_pinhole_camera.py new file mode 100644 index 00000000..f05ab970 --- /dev/null +++ b/tests/unit/datatypes/sensors/test_pinhole_camera.py @@ -0,0 +1,497 @@ +import numpy as np +import pytest + +from py123d.datatypes.sensors.pinhole_camera import ( + PinholeCamera, + PinholeCameraMetadata, + PinholeCameraType, + PinholeDistortion, + PinholeDistortionIndex, + PinholeIntrinsics, +) +from py123d.geometry import PoseSE3 + + +class TestPinholeCameraType: + def test_camera_type_values(self): + """Test that camera type enum has expected values.""" + assert PinholeCameraType.PCAM_F0 == PinholeCameraType.PCAM_F0 + assert PinholeCameraType.PCAM_B0 == PinholeCameraType.PCAM_B0 + assert PinholeCameraType.PCAM_L0 == PinholeCameraType.PCAM_L0 + assert PinholeCameraType.PCAM_L1 == PinholeCameraType.PCAM_L1 + assert PinholeCameraType.PCAM_L2 == PinholeCameraType.PCAM_L2 + assert PinholeCameraType.PCAM_R0 == PinholeCameraType.PCAM_R0 + assert PinholeCameraType.PCAM_R1 == PinholeCameraType.PCAM_R1 + assert PinholeCameraType.PCAM_R2 == PinholeCameraType.PCAM_R2 + assert PinholeCameraType.PCAM_STEREO_L == PinholeCameraType.PCAM_STEREO_L + assert PinholeCameraType.PCAM_STEREO_R == PinholeCameraType.PCAM_STEREO_R + + def test_camera_type_from_int(self): + """Test creating camera type from integer.""" + assert PinholeCameraType(0) == PinholeCameraType.PCAM_F0 + assert PinholeCameraType(5) == PinholeCameraType.PCAM_R0 + assert PinholeCameraType(9) == PinholeCameraType.PCAM_STEREO_R + + def test_camera_type_count(self): + """Test that all camera types are defined.""" + camera_types = list(PinholeCameraType) + assert len(camera_types) == 10 + + def test_camera_type_unique_values(self): + """Test that all camera type values are unique.""" + values = [ct.value for ct in PinholeCameraType] + assert len(values) == len(set(values)) + + +class TestPinholeIntrinsics: + def test_intrinsics_creation(self): + """Test creating PinholeIntrinsics instance.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0, skew=0.0) + + assert intrinsics.fx == 500.0 + assert intrinsics.fy == 500.0 + assert intrinsics.cx == 320.0 + assert intrinsics.cy == 240.0 + assert intrinsics.skew == 0.0 + + def test_intrinsics_default_skew(self): + """Test that skew defaults to 0.0 when not provided.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + + assert intrinsics.skew == 0.0 + + def test_intrinsics_from_array(self): + """Test creating intrinsics from array.""" + array = np.array([500.0, 500.0, 320.0, 240.0, 0.0], dtype=np.float64) + intrinsics = PinholeIntrinsics.from_array(array) + + assert intrinsics.fx == 500.0 + assert intrinsics.fy == 500.0 + assert intrinsics.cx == 320.0 + assert intrinsics.cy == 240.0 + assert intrinsics.skew == 0.0 + + def test_intrinsics_from_array_copy(self): + """Test that from_array creates a copy by default.""" + array = np.array([500.0, 500.0, 320.0, 240.0, 0.0], dtype=np.float64) + intrinsics = PinholeIntrinsics.from_array(array, copy=True) + + # Modify original array + array[0] = 1000.0 + + # Intrinsics should still have original value + assert intrinsics.fx == 500.0 + + def test_intrinsics_from_array_no_copy(self): + """Test that from_array can avoid copying.""" + array = np.array([500.0, 500.0, 320.0, 240.0, 0.0], dtype=np.float64) + intrinsics = PinholeIntrinsics.from_array(array, copy=False) + + # Modify original array + array[0] = 1000.0 + + # Intrinsics should reflect the change + assert intrinsics.fx == 1000.0 + + def test_intrinsics_from_camera_matrix(self): + """Test creating intrinsics from 3x3 camera matrix.""" + K = np.array([[500.0, 0.5, 320.0], [0.0, 500.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64) + + intrinsics = PinholeIntrinsics.from_camera_matrix(K) + + assert intrinsics.fx == 500.0 + assert intrinsics.fy == 500.0 + assert intrinsics.cx == 320.0 + assert intrinsics.cy == 240.0 + assert intrinsics.skew == 0.5 + + def test_intrinsics_camera_matrix_property(self): + """Test getting camera matrix from intrinsics.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=600.0, cx=320.0, cy=240.0, skew=0.5) + K = intrinsics.camera_matrix + + expected_K = np.array([[500.0, 0.5, 320.0], [0.0, 600.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64) + + np.testing.assert_array_almost_equal(K, expected_K) + + def test_intrinsics_camera_matrix_roundtrip(self): + """Test converting to camera matrix and back.""" + original_K = np.array([[500.0, 0.5, 320.0], [0.0, 600.0, 240.0], [0.0, 0.0, 1.0]], dtype=np.float64) + + intrinsics = PinholeIntrinsics.from_camera_matrix(original_K) + restored_K = intrinsics.camera_matrix + + np.testing.assert_array_almost_equal(restored_K, original_K) + + def test_intrinsics_array_property(self): + """Test accessing the underlying array.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0, skew=0.5) + array = intrinsics.array + + assert isinstance(array, np.ndarray) + assert array.shape == (5,) + np.testing.assert_array_almost_equal(array, [500.0, 500.0, 320.0, 240.0, 0.5]) + + def test_intrinsics_from_list(self): + """Test creating intrinsics from list via from_list method.""" + intrinsics_list = [500.0, 500.0, 320.0, 240.0, 0.0] + intrinsics = PinholeIntrinsics.from_list(intrinsics_list) + + assert intrinsics.fx == 500.0 + assert intrinsics.fy == 500.0 + assert intrinsics.cx == 320.0 + assert intrinsics.cy == 240.0 + assert intrinsics.skew == 0.0 + + def test_intrinsics_tolist(self): + """Test converting intrinsics to list.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0, skew=0.5) + intrinsics_list = intrinsics.tolist() + + assert isinstance(intrinsics_list, list) + assert len(intrinsics_list) == 5 + assert intrinsics_list[0] == pytest.approx(500.0) + assert intrinsics_list[1] == pytest.approx(500.0) + assert intrinsics_list[2] == pytest.approx(320.0) + assert intrinsics_list[3] == pytest.approx(240.0) + assert intrinsics_list[4] == pytest.approx(0.5) + + def test_intrinsics_different_fx_fy(self): + """Test intrinsics with different focal lengths.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=600.0, cx=320.0, cy=240.0) + + assert intrinsics.fx == 500.0 + assert intrinsics.fy == 600.0 + assert intrinsics.fx != intrinsics.fy + + def test_intrinsics_non_centered_principal_point(self): + """Test intrinsics with non-centered principal point.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=100.0, cy=100.0) + + assert intrinsics.cx == 100.0 + assert intrinsics.cy == 100.0 + + +class TestPinholeDistortion: + def test_distortion_creation(self): + """Test creating PinholeDistortion instance.""" + distortion = PinholeDistortion(k1=0.1, k2=0.01, p1=0.001, p2=0.001, k3=0.001) + + assert distortion.k1 == 0.1 + assert distortion.k2 == 0.01 + assert distortion.p1 == 0.001 + assert distortion.p2 == 0.001 + assert distortion.k3 == 0.001 + + def test_distortion_from_array(self): + """Test creating distortion from array.""" + array = np.array([0.1, 0.01, 0.001, 0.001, 0.001], dtype=np.float64) + distortion = PinholeDistortion.from_array(array) + + assert distortion.k1 == 0.1 + assert distortion.k2 == 0.01 + assert distortion.p1 == 0.001 + assert distortion.p2 == 0.001 + assert distortion.k3 == 0.001 + + def test_distortion_from_array_copy(self): + """Test that from_array creates a copy by default.""" + array = np.array([0.1, 0.01, 0.001, 0.001, 0.001], dtype=np.float64) + distortion = PinholeDistortion.from_array(array, copy=True) + + # Modify original array + array[0] = 0.5 + + # Distortion should still have original value + assert distortion.k1 == 0.1 + + def test_distortion_from_array_no_copy(self): + """Test that from_array can avoid copying.""" + array = np.array([0.1, 0.01, 0.001, 0.001, 0.001], dtype=np.float64) + distortion = PinholeDistortion.from_array(array, copy=False) + + # Modify original array + array[0] = 0.5 + + # Distortion should reflect the change + assert distortion.k1 == 0.5 + + def test_distortion_array_property(self): + """Test accessing the underlying array.""" + distortion = PinholeDistortion(k1=0.1, k2=0.01, p1=0.001, p2=0.001, k3=0.001) + array = distortion.array + + assert isinstance(array, np.ndarray) + assert array.shape == (len(PinholeDistortionIndex),) + np.testing.assert_array_almost_equal(array, [0.1, 0.01, 0.001, 0.001, 0.001]) + + def test_distortion_zero_values(self): + """Test distortion with zero values.""" + distortion = PinholeDistortion(k1=0.0, k2=0.0, p1=0.0, p2=0.0, k3=0.0) + + assert distortion.k1 == 0.0 + assert distortion.k2 == 0.0 + assert distortion.p1 == 0.0 + assert distortion.p2 == 0.0 + assert distortion.k3 == 0.0 + + def test_distortion_negative_values(self): + """Test distortion with negative values.""" + distortion = PinholeDistortion(k1=-0.1, k2=-0.01, p1=-0.001, p2=-0.001, k3=-0.001) + + assert distortion.k1 == -0.1 + assert distortion.k2 == -0.01 + assert distortion.p1 == -0.001 + assert distortion.p2 == -0.001 + assert distortion.k3 == -0.001 + + def test_distortion_from_list(self): + """Test creating distortion from list via from_list method.""" + distortion_list = [0.1, 0.01, 0.001, 0.001, 0.001] + distortion = PinholeDistortion.from_list(distortion_list) + + assert distortion.k1 == 0.1 + assert distortion.k2 == 0.01 + assert distortion.p1 == 0.001 + assert distortion.p2 == 0.001 + assert distortion.k3 == 0.001 + + def test_distortion_tolist(self): + """Test converting distortion to list.""" + distortion = PinholeDistortion(k1=0.1, k2=0.01, p1=0.001, p2=0.001, k3=0.001) + distortion_list = distortion.tolist() + + assert isinstance(distortion_list, list) + assert len(distortion_list) == 5 + assert distortion_list[0] == pytest.approx(0.1) + assert distortion_list[1] == pytest.approx(0.01) + assert distortion_list[2] == pytest.approx(0.001) + assert distortion_list[3] == pytest.approx(0.001) + assert distortion_list[4] == pytest.approx(0.001) + + +class TestPinholeMetadata: + def test_metadata_from_dict_with_none_intrinsics(self): + """Test creating metadata from dict with None intrinsics.""" + data_dict = { + "camera_type": 1, + "intrinsics": None, + "distortion": [0.1, 0.01, 0.001, 0.001, 0.001], + "width": 800, + "height": 600, + } + + metadata = PinholeCameraMetadata.from_dict(data_dict) + + assert metadata.camera_type == PinholeCameraType.PCAM_B0 + assert metadata.intrinsics is None + assert metadata.distortion is not None + assert metadata.width == 800 + assert metadata.height == 600 + + def test_metadata_from_dict_with_none_distortion(self): + """Test creating metadata from dict with None distortion.""" + data_dict = { + "camera_type": 2, + "intrinsics": [600.0, 600.0, 400.0, 300.0, 0.0], + "distortion": None, + "width": 800, + "height": 600, + } + + metadata = PinholeCameraMetadata.from_dict(data_dict) + + assert metadata.camera_type == PinholeCameraType.PCAM_L0 + assert metadata.intrinsics is not None + assert metadata.distortion is None + + def test_metadata_different_aspect_ratios(self): + """Test metadata with different aspect ratios.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + + # 16:9 aspect ratio + metadata_16_9 = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=1920, height=1080 + ) + assert metadata_16_9.aspect_ratio == pytest.approx(16 / 9) + + # 4:3 aspect ratio + metadata_4_3 = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + assert metadata_4_3.aspect_ratio == pytest.approx(4 / 3) + + def test_metadata_fov_with_different_focal_lengths(self): + """Test FOV calculation with different focal lengths.""" + intrinsics_narrow = PinholeIntrinsics(fx=1000.0, fy=1000.0, cx=320.0, cy=240.0) + intrinsics_wide = PinholeIntrinsics(fx=250.0, fy=250.0, cx=320.0, cy=240.0) + + metadata_narrow = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics_narrow, distortion=None, width=640, height=480 + ) + metadata_wide = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics_wide, distortion=None, width=640, height=480 + ) + + # Wider focal length should result in larger FOV + assert metadata_wide.fov_x > metadata_narrow.fov_x + assert metadata_wide.fov_y > metadata_narrow.fov_y + + def test_metadata_to_dict_preserves_types(self): + """Test that to_dict preserves correct types.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + distortion = PinholeDistortion(k1=0.1, k2=0.01, p1=0.001, p2=0.001, k3=0.001) + + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_R1, intrinsics=intrinsics, distortion=distortion, width=1280, height=720 + ) + + data_dict = metadata.to_dict() + + assert isinstance(data_dict["camera_type"], int) + assert isinstance(data_dict["width"], int) + assert isinstance(data_dict["height"], int) + assert isinstance(data_dict["intrinsics"], list) + assert isinstance(data_dict["distortion"], list) + + def test_metadata_all_camera_types(self): + """Test metadata creation with all camera types.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + + for camera_type in PinholeCameraType: + metadata = PinholeCameraMetadata( + camera_type=camera_type, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + assert metadata.camera_type == camera_type + + def test_metadata_square_image(self): + """Test metadata with square image dimensions.""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=256.0, cy=256.0) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=512, height=512 + ) + + assert metadata.aspect_ratio == 1.0 + assert metadata.fov_x == pytest.approx(metadata.fov_y) + + def test_metadata_non_square_pixels(self): + """Test metadata with non-square pixels (different fx and fy).""" + intrinsics = PinholeIntrinsics(fx=500.0, fy=600.0, cx=320.0, cy=240.0) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + + expected_fov_x = 2 * np.arctan(640 / (2 * 500.0)) + expected_fov_y = 2 * np.arctan(480 / (2 * 600.0)) + + assert metadata.fov_x == pytest.approx(expected_fov_x) + assert metadata.fov_y == pytest.approx(expected_fov_y) + assert metadata.fov_x != pytest.approx(metadata.fov_y) + + +class TestPinholeCamera: + def test_pinhole_camera_creation(self): + """Test creating PinholeCamera instance.""" + + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + image = np.zeros((480, 640, 3), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.metadata == metadata + assert np.array_equal(camera.image, image) + assert camera.extrinsic == extrinsic + + def test_pinhole_camera_with_color_image(self): + """Test PinholeCamera with color image.""" + + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.image.shape == (480, 640, 3) + assert camera.image.dtype == np.uint8 + + def test_pinhole_camera_with_grayscale_image(self): + """Test PinholeCamera with grayscale image.""" + + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_L0, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + image = np.random.randint(0, 255, (480, 640), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.image.shape == (480, 640) + + def test_pinhole_camera_with_distortion(self): + """Test PinholeCamera with distortion parameters.""" + + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + distortion = PinholeDistortion(k1=0.1, k2=0.01, p1=0.001, p2=0.001, k3=0.001) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, + intrinsics=intrinsics, + distortion=distortion, + width=640, + height=480, + ) + image = np.zeros((480, 640, 3), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.metadata.distortion is not None + assert camera.metadata.distortion.k1 == 0.1 + + def test_pinhole_camera_different_types(self): + """Test PinholeCamera with different camera types.""" + + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=320.0, cy=240.0) + image = np.zeros((480, 640, 3), dtype=np.uint8) + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + for camera_type in [ + PinholeCameraType.PCAM_F0, + PinholeCameraType.PCAM_B0, + PinholeCameraType.PCAM_STEREO_L, + PinholeCameraType.PCAM_STEREO_R, + ]: + metadata = PinholeCameraMetadata( + camera_type=camera_type, intrinsics=intrinsics, distortion=None, width=640, height=480 + ) + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + assert camera.metadata.camera_type == camera_type + + def test_pinhole_camera_with_different_resolutions(self): + """Test PinholeCamera with different image resolutions.""" + + resolutions = [(640, 480), (1920, 1080), (1280, 720), (800, 600)] + extrinsic = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + + for width, height in resolutions: + intrinsics = PinholeIntrinsics(fx=500.0, fy=500.0, cx=width / 2, cy=height / 2) + metadata = PinholeCameraMetadata( + camera_type=PinholeCameraType.PCAM_F0, + intrinsics=intrinsics, + distortion=None, + width=width, + height=height, + ) + image = np.zeros((height, width, 3), dtype=np.uint8) + camera = PinholeCamera(metadata=metadata, image=image, extrinsic=extrinsic) + + assert camera.metadata.width == width + assert camera.metadata.height == height + assert camera.image.shape[:2] == (height, width) diff --git a/tests/unit/datatypes/time/test_time.py b/tests/unit/datatypes/time/test_time.py new file mode 100644 index 00000000..69ab6166 --- /dev/null +++ b/tests/unit/datatypes/time/test_time.py @@ -0,0 +1,67 @@ +import pytest + +from py123d.datatypes.time.time_point import TimePoint + + +class TestTimePoint: + def test_from_ns(self): + """Test constructing TimePoint from nanoseconds.""" + tp = TimePoint.from_ns(1000000) + assert tp.time_ns == 1000000 + assert tp.time_us == 1000 + + def test_from_us(self): + """Test constructing TimePoint from microseconds.""" + tp = TimePoint.from_us(1000) + assert tp.time_us == 1000 + assert tp.time_ns == 1000000 + + def test_from_ms(self): + """Test constructing TimePoint from milliseconds.""" + tp = TimePoint.from_ms(1.5) + assert tp.time_ms == 1.5 + assert tp.time_us == 1500 + + def test_from_s(self): + """Test constructing TimePoint from seconds.""" + tp = TimePoint.from_s(2.5) + assert tp.time_s == 2.5 + assert tp.time_us == 2500000 + + def test_time_ns_property(self): + """Test accessing time value in nanoseconds.""" + tp = TimePoint.from_us(1000) + assert tp.time_ns == 1000000 + + def test_time_us_property(self): + """Test accessing time value in microseconds.""" + tp = TimePoint.from_us(1000) + assert tp.time_us == 1000 + + def test_time_ms_property(self): + """Test accessing time value in milliseconds.""" + tp = TimePoint.from_us(1500) + assert tp.time_ms == 1.5 + + def test_time_s_property(self): + """Test accessing time value in seconds.""" + tp = TimePoint.from_us(2500000) + assert tp.time_s == 2.5 + + def test_from_ns_integer_assertion(self): + """Test that from_ns raises AssertionError for non-integer input.""" + with pytest.raises(AssertionError): + TimePoint.from_ns(1000.5) + + def test_from_us_integer_assertion(self): + """Test that from_us raises AssertionError for non-integer input.""" + with pytest.raises(AssertionError): + TimePoint.from_us(1000.5) + + def test_conversion_chain(self): + """Test conversions between different time units.""" + original_us = 123456 + tp = TimePoint.from_us(original_us) + assert TimePoint.from_ns(tp.time_ns).time_us == original_us + assert TimePoint.from_ms(tp.time_ms).time_us == original_us + assert TimePoint.from_s(tp.time_s).time_us == original_us diff --git a/tests/unit/datatypes/vehicle_state/test_dynamic_state.py b/tests/unit/datatypes/vehicle_state/test_dynamic_state.py new file mode 100644 index 00000000..b89950d1 --- /dev/null +++ b/tests/unit/datatypes/vehicle_state/test_dynamic_state.py @@ -0,0 +1,103 @@ +import numpy as np + +from py123d.datatypes.vehicle_state.dynamic_state import ( + DynamicStateSE2, + DynamicStateSE2Index, + DynamicStateSE3, + DynamicStateSE3Index, +) +from py123d.geometry import Vector2D, Vector3D + + +class TestDynamicStateSE3: + def test_init(self): + velocity = Vector3D(1.0, 2.0, 3.0) + acceleration = Vector3D(4.0, 5.0, 6.0) + angular_velocity = Vector3D(7.0, 8.0, 9.0) + + state = DynamicStateSE3(velocity, acceleration, angular_velocity) + + assert np.allclose(state.velocity_3d.array, velocity.array) + assert np.allclose(state.acceleration_3d.array, acceleration.array) + assert np.allclose(state.angular_velocity.array, angular_velocity.array) + + def test_from_array(self): + array = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]) + state = DynamicStateSE3.from_array(array) + + assert np.allclose(state.array, array) + assert state.array is not array # Default copy=True + assert len(state.array) == len(DynamicStateSE3Index) + + def test_from_array_no_copy(self): + array = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]) + state = DynamicStateSE3.from_array(array, copy=False) + + assert state.array is array + + def test_velocity_properties(self): + velocity = Vector3D(1.0, 2.0, 3.0) + state = DynamicStateSE3(velocity, Vector3D(0, 0, 0), Vector3D(0, 0, 0)) + + assert np.allclose(state.velocity_3d.array, [1.0, 2.0, 3.0]) + assert np.allclose(state.velocity_2d.array, [1.0, 2.0]) + + def test_acceleration_properties(self): + acceleration = Vector3D(4.0, 5.0, 6.0) + state = DynamicStateSE3(Vector3D(0, 0, 0), acceleration, Vector3D(0, 0, 0)) + + assert np.allclose(state.acceleration_3d.array, [4.0, 5.0, 6.0]) + assert np.allclose(state.acceleration_2d.array, [4.0, 5.0]) + + def test_dynamic_state_se2_projection(self): + velocity = Vector3D(1.0, 2.0, 3.0) + acceleration = Vector3D(4.0, 5.0, 6.0) + angular_velocity = Vector3D(7.0, 8.0, 9.0) + + state_se3 = DynamicStateSE3(velocity, acceleration, angular_velocity) + state_se2 = state_se3.dynamic_state_se2 + + assert np.allclose(state_se2.velocity_2d.array, [1.0, 2.0]) + assert np.allclose(state_se2.acceleration_2d.array, [4.0, 5.0]) + assert np.isclose(state_se2.angular_velocity, 9.0) + + +class TestDynamicStateSE2: + def test_init(self): + velocity = Vector2D(1.0, 2.0) + acceleration = Vector2D(3.0, 4.0) + angular_velocity = 5.0 + + state = DynamicStateSE2(velocity, acceleration, angular_velocity) + + assert np.allclose(state.velocity_2d.array, velocity.array) + assert np.allclose(state.acceleration_2d.array, acceleration.array) + assert np.isclose(state.angular_velocity, angular_velocity) + assert len(state.array) == len(DynamicStateSE2Index) + + def test_from_array(self): + array = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + state = DynamicStateSE2.from_array(array) + + assert np.allclose(state.array, array) + + def test_velocity_properties(self): + velocity = Vector2D(1.0, 2.0) + state = DynamicStateSE2(velocity, Vector2D(0, 0), 0.0) + assert np.allclose(state.velocity_2d.array, [1.0, 2.0]) + + def test_acceleration_properties(self): + acceleration = Vector2D(3.0, 4.0) + state = DynamicStateSE2(Vector2D(0, 0), acceleration, 0.0) + assert np.allclose(state.acceleration_2d.array, [3.0, 4.0]) + + def test_angular_velocity_property(self): + state = DynamicStateSE2(Vector2D(0, 0), Vector2D(0, 0), 5.0) + + assert np.isclose(state.angular_velocity, 5.0) + + def test_array_property(self): + array = np.array([1.0, 2.0, 3.0, 4.0, 5.0]) + state = DynamicStateSE2.from_array(array) + + assert np.array_equal(state.array, array) diff --git a/tests/unit/datatypes/vehicle_state/test_ego_state.py b/tests/unit/datatypes/vehicle_state/test_ego_state.py new file mode 100644 index 00000000..d07de432 --- /dev/null +++ b/tests/unit/datatypes/vehicle_state/test_ego_state.py @@ -0,0 +1,297 @@ +import pytest + +from py123d.conversion.registry.box_detection_label_registry import DefaultBoxDetectionLabel +from py123d.datatypes.time import TimePoint +from py123d.datatypes.vehicle_state import ( + DynamicStateSE2, + DynamicStateSE3, + EgoStateSE2, + EgoStateSE3, + VehicleParameters, +) +from py123d.datatypes.vehicle_state.ego_state import EGO_TRACK_TOKEN +from py123d.geometry import PoseSE2, PoseSE3, Vector2D, Vector3D +from py123d.geometry.bounding_box import BoundingBoxSE2 + + +class TestEgoStateSE2: + def setup_method(self): + """Set up test fixtures for EgoStateSE2 tests.""" + self.rear_axle_pose = PoseSE2(x=0.0, y=0.0, yaw=0.0) + self.vehicle_params = VehicleParameters( + vehicle_name="test_vehicle", + length=4.5, + width=2.0, + height=1.5, + wheel_base=2.7, + rear_axle_to_center_longitudinal=1.35, + rear_axle_to_center_vertical=0.5, + ) + self.dynamic_state = DynamicStateSE2( + velocity=Vector2D(1.0, 0.0), + acceleration=Vector2D(0.1, 0.0), + angular_velocity=0.1, + ) + self.timepoint = TimePoint.from_us(1000000) + self.tire_steering_angle = 0.2 + + def test_init(self): + """Test EgoStateSE2 initialization.""" + ego_state = EgoStateSE2( + rear_axle_se2=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se2=self.dynamic_state, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se2 == self.rear_axle_pose + assert ego_state.vehicle_parameters == self.vehicle_params + assert ego_state.dynamic_state_se2 == self.dynamic_state + assert ego_state.timepoint == self.timepoint + assert ego_state.tire_steering_angle == self.tire_steering_angle + + def test_from_rear_axle(self): + """Test creating EgoStateSE2 from rear axle.""" + ego_state = EgoStateSE2.from_rear_axle( + rear_axle_se2=self.rear_axle_pose, + dynamic_state_se2=self.dynamic_state, + vehicle_parameters=self.vehicle_params, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se2 == self.rear_axle_pose + assert ego_state.vehicle_parameters == self.vehicle_params + + def test_from_center(self): + """Test creating EgoStateSE2 from center pose.""" + center_pose = PoseSE2(x=1.35, y=0.0, yaw=0.0) + ego_state = EgoStateSE2.from_center( + center_se2=center_pose, + dynamic_state_se2=self.dynamic_state, + vehicle_parameters=self.vehicle_params, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se2 is not None + assert ego_state.vehicle_parameters == self.vehicle_params + + def test_rear_axle_property(self): + """Test rear_axle property.""" + ego_state = EgoStateSE2(rear_axle_se2=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + assert ego_state.rear_axle_se2 == self.rear_axle_pose + + def test_center_property(self): + """Test center property calculation.""" + ego_state = EgoStateSE2(rear_axle_se2=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + center = ego_state.center_se2 + assert center is not None + assert center.x == pytest.approx(self.vehicle_params.rear_axle_to_center_longitudinal) + assert center.y == pytest.approx(0.0) + assert center.yaw == pytest.approx(0.0) + + def test_bounding_box_property(self): + """Test bounding box properties.""" + ego_state = EgoStateSE2(rear_axle_se2=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + bbox = ego_state.bounding_box_se2 + bbox_center = BoundingBoxSE2(ego_state.center_se2, self.vehicle_params.length, self.vehicle_params.width) + assert bbox is not None + assert bbox.length == self.vehicle_params.length + assert bbox.width == self.vehicle_params.width + assert ego_state.bounding_box_se2 == bbox_center + + def test_box_detection_property(self): + """Test box detection properties.""" + ego_state = EgoStateSE2( + rear_axle_se2=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se2=self.dynamic_state, + timepoint=self.timepoint, + ) + + box_det = ego_state.box_detection_se2 + assert box_det is not None + assert box_det.metadata.label == DefaultBoxDetectionLabel.EGO + assert box_det.metadata.track_token == EGO_TRACK_TOKEN + assert box_det.metadata.timepoint == self.timepoint + + def test_optional_parameters_none(self): + """Test EgoStateSE2 with optional parameters as None.""" + ego_state = EgoStateSE2( + rear_axle_se2=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se2=None, + timepoint=None, + tire_steering_angle=None, + ) + + assert ego_state.dynamic_state_se2 is None + assert ego_state.timepoint is None + assert ego_state.tire_steering_angle is None + + def test_default_tire_steering_angle(self): + """Test default tire steering angle value.""" + ego_state = EgoStateSE2(rear_axle_se2=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + assert ego_state.tire_steering_angle == 0.0 + + +class TestEgoStateSE3: + def setup_method(self): + """Set up test fixtures for EgoStateSE3 tests.""" + + self.rear_axle_pose = PoseSE3(x=0.0, y=0.0, z=0.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + self.vehicle_params = VehicleParameters( + vehicle_name="test_vehicle", + length=4.5, + width=2.0, + height=1.5, + wheel_base=2.7, + rear_axle_to_center_longitudinal=1.35, + rear_axle_to_center_vertical=0.5, + ) + self.dynamic_state = DynamicStateSE3( + velocity=Vector3D(1.0, 0.0, 0.0), + acceleration=Vector3D(0.1, 0.0, 0.0), + angular_velocity=Vector3D(0.0, 0.0, 0.1), + ) + self.timepoint = TimePoint.from_us(1000000) + self.tire_steering_angle = 0.2 + + def test_init(self): + """Test EgoStateSE3 initialization.""" + ego_state = EgoStateSE3( + rear_axle_se3=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=self.dynamic_state, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se3 == self.rear_axle_pose + assert ego_state.vehicle_parameters == self.vehicle_params + assert ego_state.dynamic_state_se3 == self.dynamic_state + assert ego_state.timepoint == self.timepoint + assert ego_state.tire_steering_angle == self.tire_steering_angle + + def test_from_rear_axle(self): + """Test creating EgoStateSE3 from rear axle.""" + ego_state = EgoStateSE3.from_rear_axle( + rear_axle_se3=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=self.dynamic_state, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se3 == self.rear_axle_pose + assert ego_state.vehicle_parameters == self.vehicle_params + + def test_from_center(self): + """Test creating EgoStateSE3 from center pose.""" + center_pose = PoseSE3(x=1.35, y=0.0, z=0.5, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + ego_state = EgoStateSE3.from_center( + center_se3=center_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=self.dynamic_state, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + assert ego_state.rear_axle_se3 is not None + assert ego_state.vehicle_parameters == self.vehicle_params + + def test_rear_axle_properties(self): + """Test rear_axle properties.""" + ego_state = EgoStateSE3(rear_axle_se3=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + assert ego_state.rear_axle_se3 == self.rear_axle_pose + assert ego_state.rear_axle_se2 is not None + + def test_center_properties(self): + """Test center property calculation.""" + ego_state = EgoStateSE3(rear_axle_se3=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + center = ego_state.center_se3 + assert center is not None + assert center.x == pytest.approx(self.vehicle_params.rear_axle_to_center_longitudinal) + assert center.y == pytest.approx(0.0) + + center_se2 = ego_state.center_se2 + assert center_se2 is not None + + def test_bounding_box_properties(self): + """Test bounding box properties.""" + ego_state = EgoStateSE3(rear_axle_se3=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + bbox_se3 = ego_state.bounding_box_se3 + assert bbox_se3 is not None + assert bbox_se3.length == self.vehicle_params.length + assert bbox_se3.width == self.vehicle_params.width + assert bbox_se3.height == self.vehicle_params.height + + bbox_se2 = ego_state.bounding_box_se2 + assert bbox_se2 is not None + assert ego_state.bounding_box_se3 == bbox_se3 + + def test_box_detection_properties(self): + """Test box detection properties.""" + ego_state = EgoStateSE3( + rear_axle_se3=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=self.dynamic_state, + timepoint=self.timepoint, + ) + + box_det_se3 = ego_state.box_detection_se3 + assert box_det_se3 is not None + assert box_det_se3.metadata.label == DefaultBoxDetectionLabel.EGO + assert box_det_se3.metadata.track_token == EGO_TRACK_TOKEN + assert box_det_se3.metadata.timepoint == self.timepoint + + box_det_se2 = ego_state.box_detection_se2 + assert box_det_se2 is not None + assert box_det_se2.metadata.label == DefaultBoxDetectionLabel.EGO + assert box_det_se2.metadata.track_token == EGO_TRACK_TOKEN + assert box_det_se2.metadata.timepoint == self.timepoint + + def test_ego_state_se2_projection(self): + """Test projection to EgoStateSE2.""" + ego_state = EgoStateSE3( + rear_axle_se3=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=self.dynamic_state, + timepoint=self.timepoint, + tire_steering_angle=self.tire_steering_angle, + ) + + ego_state_se2 = ego_state.ego_state_se2 + assert ego_state_se2 is not None + assert isinstance(ego_state_se2, EgoStateSE2) + assert ego_state_se2.vehicle_parameters == self.vehicle_params + assert ego_state_se2.timepoint == self.timepoint + assert ego_state_se2.tire_steering_angle == self.tire_steering_angle + + def test_optional_parameters_none(self): + """Test EgoStateSE3 with optional parameters as None.""" + ego_state = EgoStateSE3( + rear_axle_se3=self.rear_axle_pose, + vehicle_parameters=self.vehicle_params, + dynamic_state_se3=None, + timepoint=None, + tire_steering_angle=None, + ) + + assert ego_state.dynamic_state_se3 is None + assert ego_state.timepoint is None + assert ego_state.tire_steering_angle is None + + def test_default_tire_steering_angle(self): + """Test default tire steering angle value.""" + ego_state = EgoStateSE3(rear_axle_se3=self.rear_axle_pose, vehicle_parameters=self.vehicle_params) + + assert ego_state.tire_steering_angle == 0.0 diff --git a/tests/unit/datatypes/vehicle_state/test_vehicle_parameters.py b/tests/unit/datatypes/vehicle_state/test_vehicle_parameters.py new file mode 100644 index 00000000..f04b5cff --- /dev/null +++ b/tests/unit/datatypes/vehicle_state/test_vehicle_parameters.py @@ -0,0 +1,77 @@ +from py123d.datatypes.vehicle_state.vehicle_parameters import VehicleParameters + + +class TestVehicleParameters: + def setup_method(self): + self.params = VehicleParameters( + vehicle_name="test_vehicle", + width=2.0, + length=5.0, + height=1.8, + wheel_base=3.0, + rear_axle_to_center_vertical=0.5, + rear_axle_to_center_longitudinal=1.5, + ) + + def test_initialization(self): + """Test that VehicleParameters initializes correctly.""" + assert self.params.vehicle_name == "test_vehicle" + assert self.params.width == 2.0 + assert self.params.length == 5.0 + assert self.params.height == 1.8 + assert self.params.wheel_base == 3.0 + assert self.params.rear_axle_to_center_vertical == 0.5 + assert self.params.rear_axle_to_center_longitudinal == 1.5 + + def test_half_width(self): + """Test half_width property.""" + assert self.params.half_width == 1.0 + + def test_half_length(self): + """Test half_length property.""" + assert self.params.half_length == 2.5 + + def test_half_height(self): + """Test half_height property.""" + assert self.params.half_height == 0.9 + + def test_to_dict(self): + """Test to_dict method.""" + result = self.params.to_dict() + expected = { + "vehicle_name": "test_vehicle", + "width": 2.0, + "length": 5.0, + "height": 1.8, + "wheel_base": 3.0, + "rear_axle_to_center_vertical": 0.5, + "rear_axle_to_center_longitudinal": 1.5, + } + assert result == expected + + def test_from_dict(self): + """Test from_dict method.""" + data = { + "vehicle_name": "from_dict_vehicle", + "width": 1.5, + "length": 4.0, + "height": 1.6, + "wheel_base": 2.5, + "rear_axle_to_center_vertical": 0.4, + "rear_axle_to_center_longitudinal": 1.2, + } + params = VehicleParameters.from_dict(data) + assert params.vehicle_name == "from_dict_vehicle" + assert params.width == 1.5 + assert params.length == 4.0 + assert params.height == 1.6 + assert params.wheel_base == 2.5 + assert params.rear_axle_to_center_vertical == 0.4 + assert params.rear_axle_to_center_longitudinal == 1.2 + + def test_from_dict_to_dict_round_trip(self): + """Test that from_dict and to_dict are inverses.""" + original_dict = self.params.to_dict() + recreated_params = VehicleParameters.from_dict(original_dict) + recreated_dict = recreated_params.to_dict() + assert original_dict == recreated_dict diff --git a/tests/unit/geometry/test_bounding_box.py b/tests/unit/geometry/test_bounding_box.py index 34c24fc6..c58ae8f7 100644 --- a/tests/unit/geometry/test_bounding_box.py +++ b/tests/unit/geometry/test_bounding_box.py @@ -1,10 +1,9 @@ -import unittest - import numpy as np +import pytest import shapely.geometry as geom from py123d.common.utils.mixin import ArrayMixin -from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, Point2D, Point3D, StateSE2, StateSE3 +from py123d.geometry import BoundingBoxSE2, BoundingBoxSE3, Point2D, Point3D, PoseSE2, PoseSE3 from py123d.geometry.geometry_index import ( BoundingBoxSE2Index, BoundingBoxSE3Index, @@ -14,12 +13,12 @@ ) -class TestBoundingBoxSE2(unittest.TestCase): +class TestBoundingBoxSE2: """Unit tests for BoundingBoxSE2 class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" - self.center = StateSE2(1.0, 2.0, 0.5) + self.center = PoseSE2(1.0, 2.0, 0.5) self.length = 4.0 self.width = 2.0 self.bbox = BoundingBoxSE2(self.center, self.length, self.width) @@ -27,9 +26,9 @@ def setUp(self): def test_init(self): """Test BoundingBoxSE2 initialization.""" bbox = BoundingBoxSE2(self.center, self.length, self.width) - self.assertEqual(bbox.length, self.length) - self.assertEqual(bbox.width, self.width) - np.testing.assert_array_equal(bbox.center.array, self.center.array) + assert bbox.length == self.length + assert bbox.width == self.width + np.testing.assert_array_equal(bbox.center_se2.array, self.center.array) def test_from_array(self): """Test BoundingBoxSE2.from_array method.""" @@ -44,14 +43,13 @@ def test_from_array_copy(self): bbox_no_copy = BoundingBoxSE2.from_array(array, copy=False) array[0] = 999.0 - self.assertNotEqual(bbox_copy.array[0], 999.0) - self.assertEqual(bbox_no_copy.array[0], 999.0) + assert bbox_copy.array[0] != 999.0 + assert bbox_no_copy.array[0] == 999.0 def test_properties(self): """Test BoundingBoxSE2 properties.""" - self.assertEqual(self.bbox.length, self.length) - self.assertEqual(self.bbox.width, self.width) - np.testing.assert_array_equal(self.bbox.center.array, self.center.array) + assert self.bbox.length == self.length + assert self.bbox.width == self.width np.testing.assert_array_equal(self.bbox.center_se2.array, self.center.array) def test_array_property(self): @@ -61,68 +59,68 @@ def test_array_property(self): def test_array_mixin(self): """Test that BoundingBoxSE2 is an instance of ArrayMixin.""" - self.assertIsInstance(self.bbox, ArrayMixin) + assert isinstance(self.bbox, ArrayMixin) expected = np.array([1.0, 2.0, 0.5, 4.0, 2.0], dtype=np.float16) output_array = np.array(self.bbox, dtype=np.float16) np.testing.assert_array_equal(output_array, expected) - self.assertEqual(output_array.dtype, np.float16) - self.assertEqual(output_array.shape, (len(BoundingBoxSE2Index),)) + assert output_array.dtype == np.float16 + assert output_array.shape == (len(BoundingBoxSE2Index),) def test_bounding_box_se2_property(self): """Test bounding_box_se2 property returns self.""" - self.assertIs(self.bbox.bounding_box_se2, self.bbox) + assert self.bbox.bounding_box_se2 is self.bbox def test_corners_array(self): """Test corners_array property.""" corners = self.bbox.corners_array - self.assertEqual(corners.shape, (len(Corners2DIndex), len(Point2DIndex))) - self.assertIsInstance(corners, np.ndarray) + assert corners.shape == (len(Corners2DIndex), len(Point2DIndex)) + assert isinstance(corners, np.ndarray) def test_corners_dict(self): """Test corners_dict property.""" corners_dict = self.bbox.corners_dict - self.assertEqual(len(corners_dict), len(Corners2DIndex)) + assert len(corners_dict) == len(Corners2DIndex) for index in Corners2DIndex: - self.assertIn(index, corners_dict) - self.assertIsInstance(corners_dict[index], Point2D) + assert index in corners_dict + assert isinstance(corners_dict[index], Point2D) def test_shapely_polygon(self): """Test shapely_polygon property.""" polygon = self.bbox.shapely_polygon - self.assertIsInstance(polygon, geom.Polygon) - self.assertAlmostEqual(polygon.area, self.length * self.width) + assert isinstance(polygon, geom.Polygon) + assert polygon.area == pytest.approx(self.length * self.width) def test_array_assertions(self): """Test array assertions in from_array.""" # Test 2D array - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): BoundingBoxSE2.from_array(np.array([[1, 2, 3, 4, 5]])) # Test wrong size - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): BoundingBoxSE2.from_array(np.array([1, 2, 3, 4])) -class TestBoundingBoxSE3(unittest.TestCase): +class TestBoundingBoxSE3: """Unit tests for BoundingBoxSE3 class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.array = np.array([1.0, 2.0, 3.0, 0.98185617, 0.06407135, 0.09115755, 0.1534393, 4.0, 2.0, 1.5]) - self.center = StateSE3(1.0, 2.0, 3.0, 0.98185617, 0.06407135, 0.09115755, 0.1534393) + self.center_se3 = PoseSE3(1.0, 2.0, 3.0, 0.98185617, 0.06407135, 0.09115755, 0.1534393) self.length = 4.0 self.width = 2.0 self.height = 1.5 - self.bbox = BoundingBoxSE3(self.center, self.length, self.width, self.height) + self.bbox = BoundingBoxSE3(self.center_se3, self.length, self.width, self.height) def test_init(self): """Test BoundingBoxSE3 initialization.""" - bbox = BoundingBoxSE3(self.center, self.length, self.width, self.height) - self.assertEqual(bbox.length, self.length) - self.assertEqual(bbox.width, self.width) - self.assertEqual(bbox.height, self.height) - np.testing.assert_array_equal(bbox.center.array, self.center.array) + bbox = BoundingBoxSE3(self.center_se3, self.length, self.width, self.height) + assert bbox.length == self.length + assert bbox.width == self.width + assert bbox.height == self.height + np.testing.assert_array_equal(bbox.center_se3.array, self.center_se3.array) def test_from_array(self): """Test BoundingBoxSE3.from_array method.""" @@ -137,16 +135,15 @@ def test_from_array_copy(self): bbox_no_copy = BoundingBoxSE3.from_array(array, copy=False) array[0] = 999.0 - self.assertNotEqual(bbox_copy.array[0], 999.0) - self.assertEqual(bbox_no_copy.array[0], 999.0) + assert bbox_copy.array[0] != 999.0 + assert bbox_no_copy.array[0] == 999.0 def test_properties(self): """Test BoundingBoxSE3 properties.""" - self.assertEqual(self.bbox.length, self.length) - self.assertEqual(self.bbox.width, self.width) - self.assertEqual(self.bbox.height, self.height) - np.testing.assert_array_equal(self.bbox.center.array, self.center.array) - np.testing.assert_array_equal(self.bbox.center_se3.array, self.center.array) + assert self.bbox.length == self.length + assert self.bbox.width == self.width + assert self.bbox.height == self.height + np.testing.assert_array_equal(self.bbox.center_se3.array, self.center_se3.array) def test_array_property(self): """Test array property.""" @@ -155,66 +152,62 @@ def test_array_property(self): def test_array_mixin(self): """Test that BoundingBoxSE3 is an instance of ArrayMixin.""" - self.assertIsInstance(self.bbox, ArrayMixin) + assert isinstance(self.bbox, ArrayMixin) expected = np.array(self.array, dtype=np.float16) output_array = np.array(self.bbox, dtype=np.float16) np.testing.assert_array_equal(output_array, expected) - self.assertEqual(output_array.dtype, np.float16) - self.assertEqual(output_array.shape, (len(BoundingBoxSE3Index),)) + assert output_array.dtype == np.float16 + assert output_array.shape == (len(BoundingBoxSE3Index),) def test_bounding_box_se2_property(self): """Test bounding_box_se2 property.""" bbox_2d = self.bbox.bounding_box_se2 - self.assertIsInstance(bbox_2d, BoundingBoxSE2) - self.assertEqual(bbox_2d.length, self.length) - self.assertEqual(bbox_2d.width, self.width) - self.assertEqual(bbox_2d.center.x, self.center.x) - self.assertEqual(bbox_2d.center.y, self.center.y) - self.assertEqual(bbox_2d.center.yaw, self.center.euler_angles.yaw) + assert isinstance(bbox_2d, BoundingBoxSE2) + assert bbox_2d.length == self.length + assert bbox_2d.width == self.width + assert bbox_2d.center_se2.x == self.center_se3.x + assert bbox_2d.center_se2.y == self.center_se3.y + assert bbox_2d.center_se2.yaw == self.center_se3.euler_angles.yaw def test_corners_array(self): """Test corners_array property.""" corners = self.bbox.corners_array - self.assertEqual(corners.shape, (8, 3)) - self.assertIsInstance(corners, np.ndarray) + assert corners.shape == (8, 3) + assert isinstance(corners, np.ndarray) def test_corners_dict(self): """Test corners_dict property.""" corners_dict = self.bbox.corners_dict - self.assertEqual(len(corners_dict), 8) + assert len(corners_dict) == 8 for index in Corners3DIndex: - self.assertIn(index, corners_dict) - self.assertIsInstance(corners_dict[index], Point3D) + assert index in corners_dict + assert isinstance(corners_dict[index], Point3D) def test_shapely_polygon(self): """Test shapely_polygon property.""" polygon = self.bbox.shapely_polygon - self.assertIsInstance(polygon, geom.Polygon) - self.assertAlmostEqual(polygon.area, self.length * self.width) + assert isinstance(polygon, geom.Polygon) + assert polygon.area == pytest.approx(self.length * self.width) def test_array_assertions(self): """Test array assertions in from_array.""" # Test 2D array - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): BoundingBoxSE3.from_array(np.array([[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]])) # Test wrong size, less than required - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): BoundingBoxSE3.from_array(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])) # Test wrong size, greater than required - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): BoundingBoxSE3.from_array(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])) def test_zero_dimensions(self): """Test bounding box with zero dimensions.""" - center = StateSE3(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) + center = PoseSE3(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0) bbox = BoundingBoxSE3(center, 0.0, 0.0, 0.0) - self.assertEqual(bbox.length, 0.0) - self.assertEqual(bbox.width, 0.0) - self.assertEqual(bbox.height, 0.0) - - -if __name__ == "__main__": - unittest.main() + assert bbox.length == 0.0 + assert bbox.width == 0.0 + assert bbox.height == 0.0 diff --git a/tests/unit/geometry/test_occupancy_map.py b/tests/unit/geometry/test_occupancy_map.py index 2390300d..a50190b4 100644 --- a/tests/unit/geometry/test_occupancy_map.py +++ b/tests/unit/geometry/test_occupancy_map.py @@ -1,15 +1,14 @@ -import unittest - import numpy as np +import pytest import shapely.geometry as geom from py123d.geometry import OccupancyMap2D -class TestOccupancyMap2D(unittest.TestCase): +class TestOccupancyMap2D: """Unit tests for OccupancyMap2D class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures with various geometries.""" self.square1 = geom.Polygon([(0, 0), (2, 0), (2, 2), (0, 2)]) self.square2 = geom.Polygon([(3, 3), (5, 3), (5, 5), (3, 5)]) @@ -24,35 +23,35 @@ def test_init_with_default_ids(self): """Test initialization with default string IDs.""" occ_map = OccupancyMap2D(self.geometries) - self.assertEqual(len(occ_map), 4) - self.assertEqual(occ_map.ids, ["0", "1", "2", "3"]) - self.assertEqual(len(occ_map.geometries), 4) + assert len(occ_map) == 4 + assert occ_map.ids == ["0", "1", "2", "3"] + assert len(occ_map.geometries) == 4 def test_init_with_string_ids(self): """Test initialization with custom string IDs.""" occ_map = OccupancyMap2D(self.geometries, ids=self.string_ids) - self.assertEqual(len(occ_map), 4) - self.assertEqual(occ_map.ids, self.string_ids) - self.assertEqual(occ_map["square1"], self.square1) + assert len(occ_map) == 4 + assert occ_map.ids == self.string_ids + assert occ_map["square1"] == self.square1 def test_init_with_int_ids(self): """Test initialization with integer IDs.""" occ_map = OccupancyMap2D(self.geometries, ids=self.int_ids) - self.assertEqual(len(occ_map), 4) - self.assertEqual(occ_map.ids, self.int_ids) - self.assertEqual(occ_map[1], self.square1) + assert len(occ_map) == 4 + assert occ_map.ids == self.int_ids + assert occ_map[1] == self.square1 def test_init_with_mismatched_ids_length(self): """Test that initialization fails with mismatched IDs length.""" - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): OccupancyMap2D(self.geometries, ids=["id1", "id2"]) def test_init_with_custom_node_capacity(self): """Test initialization with custom node capacity.""" occ_map = OccupancyMap2D(self.geometries, node_capacity=5) - self.assertEqual(occ_map._node_capacity, 5) + assert occ_map._node_capacity == 5 def test_from_dict_constructor(self): """Test construction from dictionary.""" @@ -60,54 +59,54 @@ def test_from_dict_constructor(self): occ_map = OccupancyMap2D.from_dict(geometry_dict) - self.assertEqual(len(occ_map), 3) - self.assertEqual(set(occ_map.ids), set(["square", "circle", "line"])) - self.assertEqual(occ_map["square"], self.square1) + assert len(occ_map) == 3 + assert set(occ_map.ids) == set(["square", "circle", "line"]) + assert occ_map["square"] == self.square1 def test_getitem_string_id(self): """Test geometry retrieval by string ID.""" occ_map = OccupancyMap2D(self.geometries, ids=self.string_ids) - self.assertEqual(occ_map["square1"], self.square1) - self.assertEqual(occ_map["circle"], self.circle) + assert occ_map["square1"] == self.square1 + assert occ_map["circle"] == self.circle def test_getitem_int_id(self): """Test geometry retrieval by integer ID.""" occ_map = OccupancyMap2D(self.geometries, ids=self.int_ids) - self.assertEqual(occ_map[1], self.square1) - self.assertEqual(occ_map[3], self.circle) + assert occ_map[1] == self.square1 + assert occ_map[3] == self.circle def test_getitem_invalid_id(self): """Test that invalid ID raises KeyError.""" occ_map = OccupancyMap2D(self.geometries, ids=self.string_ids) - with self.assertRaises(KeyError): + with pytest.raises(KeyError): _ = occ_map["nonexistent"] def test_len(self): """Test length property.""" occ_map = OccupancyMap2D(self.geometries) - self.assertEqual(len(occ_map), 4) + assert len(occ_map) == 4 empty_map = OccupancyMap2D([]) - self.assertEqual(len(empty_map), 0) + assert len(empty_map) == 0 def test_ids_property(self): """Test IDs property getter.""" occ_map = OccupancyMap2D(self.geometries, ids=self.string_ids) - self.assertEqual(occ_map.ids, self.string_ids) + assert occ_map.ids == self.string_ids def test_geometries_property(self): """Test geometries property getter.""" occ_map = OccupancyMap2D(self.geometries) - self.assertEqual(list(occ_map.geometries), self.geometries) + assert list(occ_map.geometries) == self.geometries def test_id_to_idx_property(self): """Test id_to_idx property.""" occ_map = OccupancyMap2D(self.geometries, ids=self.string_ids) expected_mapping = {"square1": 0, "square2": 1, "circle": 2, "line": 3} - self.assertEqual(occ_map.id_to_idx, expected_mapping) + assert occ_map.id_to_idx == expected_mapping def test_intersects_with_overlapping_geometry(self): """Test intersects method with overlapping geometry.""" @@ -118,10 +117,10 @@ def test_intersects_with_overlapping_geometry(self): intersecting_ids = occ_map.intersects(query_geom) # NOTE: square2 does not intersect with the query geometry, the rest does. - self.assertIn("square1", intersecting_ids) - self.assertIn("circle", intersecting_ids) - self.assertIn("line", intersecting_ids) - self.assertEqual(len(intersecting_ids), 3) + assert "square1" in intersecting_ids + assert "circle" in intersecting_ids + assert "line" in intersecting_ids + assert len(intersecting_ids) == 3 def test_intersects_with_non_overlapping_geometry(self): """Test intersects method with non-overlapping geometry.""" @@ -131,7 +130,7 @@ def test_intersects_with_non_overlapping_geometry(self): query_geom = geom.Polygon([(10, 10), (12, 10), (12, 12), (10, 12)]) intersecting_ids = occ_map.intersects(query_geom) - self.assertEqual(len(intersecting_ids), 0) + assert len(intersecting_ids) == 0 def test_query_with_intersects_predicate(self): """Test query method with intersects predicate.""" @@ -139,11 +138,11 @@ def test_query_with_intersects_predicate(self): query_geom = geom.Point(1, 1) indices = occ_map.query(query_geom, predicate="intersects") - self.assertIsInstance(indices, np.ndarray) - self.assertIn(occ_map.id_to_idx["square1"], indices) - self.assertIn(occ_map.id_to_idx["circle"], indices) - self.assertIn(occ_map.id_to_idx["line"], indices) - self.assertNotIn(occ_map.id_to_idx["square2"], indices) + assert isinstance(indices, np.ndarray) + assert occ_map.id_to_idx["square1"] in indices + assert occ_map.id_to_idx["circle"] in indices + assert occ_map.id_to_idx["line"] in indices + assert occ_map.id_to_idx["square2"] not in indices def test_query_with_contains_predicate(self): """Test query method with contains predicate.""" @@ -152,11 +151,11 @@ def test_query_with_contains_predicate(self): query_geom = geom.Point(4, 4) indices = occ_map.query(query_geom, predicate="within") - self.assertIsInstance(indices, np.ndarray) - self.assertIn(occ_map.id_to_idx["square2"], indices) - self.assertNotIn(occ_map.id_to_idx["square1"], indices) - self.assertNotIn(occ_map.id_to_idx["circle"], indices) - self.assertNotIn(occ_map.id_to_idx["line"], indices) + assert isinstance(indices, np.ndarray) + assert occ_map.id_to_idx["square2"] in indices + assert occ_map.id_to_idx["square1"] not in indices + assert occ_map.id_to_idx["circle"] not in indices + assert occ_map.id_to_idx["line"] not in indices def test_query_with_distance(self): """Test query method with distance parameter.""" @@ -165,11 +164,11 @@ def test_query_with_distance(self): query_geom = geom.Point(4, 4) indices = occ_map.query(query_geom, predicate="dwithin", distance=3.0) - self.assertIsInstance(indices, np.ndarray) - self.assertIn(occ_map.id_to_idx["square2"], indices) - self.assertIn(occ_map.id_to_idx["square1"], indices) - self.assertNotIn(occ_map.id_to_idx["circle"], indices) - self.assertNotIn(occ_map.id_to_idx["line"], indices) + assert isinstance(indices, np.ndarray) + assert occ_map.id_to_idx["square2"] in indices + assert occ_map.id_to_idx["square1"] in indices + assert occ_map.id_to_idx["circle"] not in indices + assert occ_map.id_to_idx["line"] not in indices def test_query_nearest_basic(self): """Test query_nearest method basic functionality.""" @@ -178,7 +177,7 @@ def test_query_nearest_basic(self): query_geom = geom.Point(4, 4) nearest_indices = occ_map.query_nearest(query_geom) - self.assertIsInstance(nearest_indices, np.ndarray) + assert isinstance(nearest_indices, np.ndarray) def test_query_nearest_with_distance(self): """Test query_nearest method with return_distance=True.""" @@ -187,11 +186,11 @@ def test_query_nearest_with_distance(self): query_geom = geom.Point(1, 1) result = occ_map.query_nearest(query_geom, return_distance=True) - self.assertIsInstance(result, tuple) - self.assertEqual(len(result), 2) + assert isinstance(result, tuple) + assert len(result) == 2 indices, distances = result - self.assertIsInstance(indices, np.ndarray) - self.assertIsInstance(distances, np.ndarray) + assert isinstance(indices, np.ndarray) + assert isinstance(distances, np.ndarray) def test_query_nearest_with_max_distance(self): """Test query_nearest method with max_distance.""" @@ -200,12 +199,12 @@ def test_query_nearest_with_max_distance(self): query_geom = geom.Point(10, 10) nearest_indices = occ_map.query_nearest(query_geom, max_distance=1.0) - self.assertIsInstance(nearest_indices, np.ndarray) - self.assertEqual(len(nearest_indices), 0) + assert isinstance(nearest_indices, np.ndarray) + assert len(nearest_indices) == 0 nearest_indices = occ_map.query_nearest(query_geom, max_distance=10.0) - self.assertIsInstance(nearest_indices, np.ndarray) - self.assertTrue(len(nearest_indices) > 0) + assert isinstance(nearest_indices, np.ndarray) + assert len(nearest_indices) > 0 def test_contains_vectorized_single_point(self): """Test contains_vectorized with a single point.""" @@ -214,9 +213,9 @@ def test_contains_vectorized_single_point(self): points = np.array([[1.0, 1.0]]) # Point inside square1 and circle result = occ_map.contains_vectorized(points) - self.assertEqual(result.shape, (4, 1)) - self.assertIsInstance(result, np.ndarray) - self.assertEqual(result.dtype, bool) + assert result.shape == (4, 1) + assert isinstance(result, np.ndarray) + assert result.dtype == bool def test_contains_vectorized_multiple_points(self): """Test contains_vectorized with multiple points.""" @@ -231,28 +230,28 @@ def test_contains_vectorized_multiple_points(self): ) result = occ_map.contains_vectorized(points) - self.assertEqual(result.shape, (4, 3)) - self.assertIsInstance(result, np.ndarray) - self.assertEqual(result.dtype, bool) + assert result.shape == (4, 3) + assert isinstance(result, np.ndarray) + assert result.dtype == bool # Check specific containment results # Point [1.0, 1.0] should be in square1 (index 0) and circle (index 2) - self.assertTrue(result[0, 0]) # square1 contains point 0 - self.assertFalse(result[1, 0]) # square2 does not contain point 0 - self.assertTrue(result[2, 0]) # circle contains point 0 - self.assertFalse(result[3, 0]) # line does not contain point 0 + assert result[0, 0] # square1 contains point 0 + assert not result[1, 0] # square2 does not contain point 0 + assert result[2, 0] # circle contains point 0 + assert not result[3, 0] # line does not contain point 0 # Point [4.0, 4.0] should be in square2 (index 1) only - self.assertFalse(result[0, 1]) # square1 does not contain point 1 - self.assertTrue(result[1, 1]) # square2 contains point 1 - self.assertFalse(result[2, 1]) # circle does not contain point 1 - self.assertFalse(result[3, 1]) # line does not contain point 1 + assert not result[0, 1] # square1 does not contain point 1 + assert result[1, 1] # square2 contains point 1 + assert not result[2, 1] # circle does not contain point 1 + assert not result[3, 1] # line does not contain point 1 # Point [10.0, 10.0] should not be in any geometry - self.assertFalse(result[0, 2]) # square1 does not contain point 2 - self.assertFalse(result[1, 2]) # square2 does not contain point 2 - self.assertFalse(result[2, 2]) # circle does not contain point 2 - self.assertFalse(result[3, 2]) # line does not contain point 2 + assert not result[0, 2] # square1 does not contain point 2 + assert not result[1, 2] # square2 does not contain point 2 + assert not result[2, 2] # circle does not contain point 2 + assert not result[3, 2] # line does not contain point 2 def test_contains_vectorized_empty_points(self): """Test contains_vectorized with empty points array.""" @@ -261,25 +260,20 @@ def test_contains_vectorized_empty_points(self): points = np.empty((0, 2)) result = occ_map.contains_vectorized(points) - self.assertEqual(result.shape, (4, 0)) + assert result.shape == (4, 0) def test_empty_occupancy_map(self): """Test behavior with empty geometry list.""" occ_map = OccupancyMap2D([]) - self.assertEqual(len(occ_map), 0) - self.assertEqual(occ_map.ids, []) - self.assertEqual(len(occ_map.geometries), 0) + assert len(occ_map) == 0 + assert occ_map.ids == [] + assert len(occ_map.geometries) == 0 def test_single_geometry_map(self): """Test behavior with single geometry.""" occ_map = OccupancyMap2D([self.square1], ids=["single"]) - self.assertEqual(len(occ_map), 1) - self.assertEqual(occ_map.ids, ["single"]) - self.assertEqual(occ_map["single"], self.square1) - - -if __name__ == "__main__": - - unittest.main() + assert len(occ_map) == 1 + assert occ_map.ids == ["single"] + assert occ_map["single"] == self.square1 diff --git a/tests/unit/geometry/test_point.py b/tests/unit/geometry/test_point.py index 15932540..d293f8b5 100644 --- a/tests/unit/geometry/test_point.py +++ b/tests/unit/geometry/test_point.py @@ -1,16 +1,16 @@ -import unittest from unittest.mock import MagicMock, patch import numpy as np +import pytest from py123d.geometry import Point2D, Point3D from py123d.geometry.geometry_index import Point2DIndex, Point3DIndex -class TestPoint2D(unittest.TestCase): +class TestPoint2D: """Unit tests for Point2D class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.x_coord = 3.5 self.y_coord = 4.2 @@ -22,54 +22,54 @@ def setUp(self): def test_init(self): """Test Point2D initialization.""" point = Point2D(1.0, 2.0) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 2.0) + assert point.x == 1.0 + assert point.y == 2.0 def test_from_array_valid(self): """Test from_array class method with valid input.""" # Mock Point2DIndex enum values point = Point2D.from_array(self.test_array) - self.assertEqual(point.x, self.x_coord) - self.assertEqual(point.y, self.y_coord) + assert point.x == self.x_coord + assert point.y == self.y_coord def test_from_array_invalid_dimensions(self): """Test from_array with invalid array dimensions.""" # 2D array should raise assertion error array_2d = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point2D.from_array(array_2d) # 3D array should raise assertion error array_3d = np.array([[[1.0]]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point2D.from_array(array_3d) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" array_wrong_length = np.array([1.0, 2.0, 3.0], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point2D.from_array(array_wrong_length) # Empty array empty_array = np.array([], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point2D.from_array(empty_array) def test_array_property(self): """Test the array property.""" expected_array = np.array([self.x_coord, self.y_coord], dtype=np.float64) np.testing.assert_array_equal(self.point.array, expected_array) - self.assertEqual(self.point.array.dtype, np.float64) - self.assertEqual(self.point.array.shape, (2,)) + assert self.point.array.dtype == np.float64 + assert self.point.array.shape == (2,) def test_array_like(self): """Test the __array__ behavior.""" expected_array = np.array([self.x_coord, self.y_coord], dtype=np.float32) output_array = np.array(self.point, dtype=np.float32) np.testing.assert_array_equal(output_array, expected_array) - self.assertEqual(output_array.dtype, np.float32) - self.assertEqual(output_array.shape, (2,)) + assert output_array.dtype == np.float32 + assert output_array.shape == (2,) def test_shapely_point_property(self): """Test the shapely_point property.""" @@ -80,29 +80,23 @@ def test_shapely_point_property(self): result = self.point.shapely_point mock_point.assert_called_once_with(self.x_coord, self.y_coord) - self.assertEqual(result, mock_point_instance) + assert result == mock_point_instance def test_iter(self): """Test the __iter__ method.""" coords = list(self.point) - self.assertEqual(coords, [self.x_coord, self.y_coord]) + assert coords == [self.x_coord, self.y_coord] # Test that it's actually iterable x, y = self.point - self.assertEqual(x, self.x_coord) - self.assertEqual(y, self.y_coord) + assert x == self.x_coord + assert y == self.y_coord - def test_hash(self): - """Test the __hash__ method.""" - point_dict = {self.point: "test"} - self.assertIn(self.point, point_dict) - self.assertEqual(point_dict[self.point], "test") - -class TestPoint3D(unittest.TestCase): +class TestPoint3D: """Unit tests for Point3D class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.x_coord = 3.5 self.y_coord = 4.2 @@ -116,56 +110,56 @@ def setUp(self): def test_init(self): """Test Point3D initialization.""" point = Point3D(1.0, 2.0, 3.0) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 2.0) - self.assertEqual(point.z, 3.0) + assert point.x == 1.0 + assert point.y == 2.0 + assert point.z == 3.0 def test_from_array_valid(self): """Test from_array class method with valid input.""" # Mock Point3DIndex enum values point = Point3D.from_array(self.test_array) - self.assertEqual(point.x, self.x_coord) - self.assertEqual(point.y, self.y_coord) - self.assertEqual(point.z, self.z_coord) + assert point.x == self.x_coord + assert point.y == self.y_coord + assert point.z == self.z_coord def test_from_array_invalid_dimensions(self): """Test from_array with invalid array dimensions.""" # 2D array should raise assertion error array_2d = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point3D.from_array(array_2d) # 3D array should raise assertion error array_3d = np.array([[[1.0]]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point3D.from_array(array_3d) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" array_wrong_length = np.array([1.0, 2.0], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point3D.from_array(array_wrong_length) # Empty array empty_array = np.array([], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Point3D.from_array(empty_array) def test_array_property(self): """Test the array property.""" expected_array = np.array([self.x_coord, self.y_coord, self.z_coord], dtype=np.float64) np.testing.assert_array_equal(self.point.array, expected_array) - self.assertEqual(self.point.array.dtype, np.float64) - self.assertEqual(self.point.array.shape, (3,)) + assert self.point.array.dtype == np.float64 + assert self.point.array.shape == (3,) def test_array_like(self): """Test the __array__ behavior.""" expected_array = np.array([self.x_coord, self.y_coord, self.z_coord], dtype=np.float32) output_array = np.array(self.point, dtype=np.float32) np.testing.assert_array_equal(output_array, expected_array) - self.assertEqual(output_array.dtype, np.float32) - self.assertEqual(output_array.shape, (3,)) + assert output_array.dtype == np.float32 + assert output_array.shape == (3,) def test_shapely_point_property(self): """Test the shapely_point property.""" @@ -176,25 +170,15 @@ def test_shapely_point_property(self): result = self.point.shapely_point mock_point.assert_called_once_with(self.x_coord, self.y_coord, self.z_coord) - self.assertEqual(result, mock_point_instance) + assert result == mock_point_instance def test_iter(self): """Test the __iter__ method.""" coords = list(self.point) - self.assertEqual(coords, [self.x_coord, self.y_coord, self.z_coord]) + assert coords == [self.x_coord, self.y_coord, self.z_coord] # Test that it's actually iterable x, y, z = self.point - self.assertEqual(x, self.x_coord) - self.assertEqual(y, self.y_coord) - self.assertEqual(z, self.z_coord) - - def test_hash(self): - """Test the __hash__ method.""" - point_dict = {self.point: "test"} - self.assertIn(self.point, point_dict) - self.assertEqual(point_dict[self.point], "test") - - -if __name__ == "__main__": - unittest.main() + assert x == self.x_coord + assert y == self.y_coord + assert z == self.z_coord diff --git a/tests/unit/geometry/test_polyline.py b/tests/unit/geometry/test_polyline.py index a2614410..f5fe56db 100644 --- a/tests/unit/geometry/test_polyline.py +++ b/tests/unit/geometry/test_polyline.py @@ -1,12 +1,11 @@ -import unittest - import numpy as np +import pytest import shapely.geometry as geom -from py123d.geometry import Point2D, Point3D, Polyline2D, Polyline3D, PolylineSE2, StateSE2 +from py123d.geometry import Point2D, Point3D, Polyline2D, Polyline3D, PolylineSE2, PoseSE2 -class TestPolyline2D(unittest.TestCase): +class TestPolyline2D: """Test class for Polyline2D.""" def test_from_linestring(self): @@ -14,36 +13,36 @@ def test_from_linestring(self): coords = [(0.0, 0.0), (1.0, 1.0), (2.0, 0.0)] linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) - self.assertIsInstance(polyline, Polyline2D) - self.assertTrue(polyline.linestring.equals(linestring)) + assert isinstance(polyline, Polyline2D) + assert polyline.linestring.equals(linestring) def test_from_linestring_with_z(self): """Test creating Polyline2D from LineString with Z coordinates.""" coords = [(0.0, 0.0, 1.0), (1.0, 1.0, 2.0), (2.0, 0.0, 3.0)] linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) - self.assertIsInstance(polyline, Polyline2D) - self.assertFalse(polyline.linestring.has_z) + assert isinstance(polyline, Polyline2D) + assert not polyline.linestring.has_z def test_from_array_2d(self): """Test creating Polyline2D from 2D array.""" array = np.array([[0.0, 0.0], [1.0, 1.0], [2.0, 0.0]], dtype=np.float32) polyline = Polyline2D.from_array(array) - self.assertIsInstance(polyline, Polyline2D) + assert isinstance(polyline, Polyline2D) np.testing.assert_array_almost_equal(polyline.array, array) def test_from_array_3d(self): """Test creating Polyline2D from 3D array.""" array = np.array([[0.0, 0.0, 1.0], [1.0, 1.0, 2.0], [2.0, 0.0, 3.0]], dtype=np.float32) polyline = Polyline2D.from_array(array) - self.assertIsInstance(polyline, Polyline2D) + assert isinstance(polyline, Polyline2D) expected = array[:, :2] np.testing.assert_array_almost_equal(polyline.array, expected) def test_from_array_invalid_shape(self): """Test creating Polyline2D from invalid array shape.""" array = np.array([[0.0], [1.0], [2.0]], dtype=np.float32) - with self.assertRaises(ValueError): + with pytest.raises(ValueError): Polyline2D.from_array(array) def test_array_property(self): @@ -52,8 +51,8 @@ def test_array_property(self): linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) array = polyline.array - self.assertEqual(array.shape, (3, 2)) - self.assertEqual(array.dtype, np.float64) + assert array.shape == (3, 2) + assert array.dtype == np.float64 np.testing.assert_array_almost_equal(array, coords) def test_length_property(self): @@ -61,7 +60,7 @@ def test_length_property(self): coords = [(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)] linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) - self.assertEqual(polyline.length, 2.0) + assert polyline.length == 2.0 def test_interpolate_single_distance(self): """Test interpolation with single distance.""" @@ -69,9 +68,9 @@ def test_interpolate_single_distance(self): linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) point = polyline.interpolate(1.0) - self.assertIsInstance(point, Point2D) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 0.0) + assert isinstance(point, Point2D) + assert point.x == 1.0 + assert point.y == 0.0 def test_interpolate_multiple_distances(self): """Test interpolation with multiple distances.""" @@ -79,8 +78,8 @@ def test_interpolate_multiple_distances(self): linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) points = polyline.interpolate(np.array([0.0, 1.0, 2.0])) - self.assertIsInstance(points, np.ndarray) - self.assertEqual(points.shape, (3, 2)) + assert isinstance(points, np.ndarray) + assert points.shape == (3, 2) expected = np.array([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]]) np.testing.assert_array_almost_equal(points, expected) @@ -90,9 +89,9 @@ def test_interpolate_normalized(self): linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) point = polyline.interpolate(0.5, normalized=True) - self.assertIsInstance(point, Point2D) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 0.0) + assert isinstance(point, Point2D) + assert point.x == 1.0 + assert point.y == 0.0 def test_project_point2d(self): """Test projecting Point2D onto polyline.""" @@ -101,16 +100,16 @@ def test_project_point2d(self): polyline = Polyline2D.from_linestring(linestring) point = Point2D(1.0, 1.0) distance = polyline.project(point) - self.assertEqual(distance, 1.0) + assert distance == 1.0 def test_project_statese2(self): """Test projecting StateSE2 onto polyline.""" coords = [(0.0, 0.0), (2.0, 0.0)] linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) - state = StateSE2(1.0, 1.0, 0.0) + state = PoseSE2(1.0, 1.0, 0.0) distance = polyline.project(state) - self.assertEqual(distance, 1.0) + assert distance == 1.0 def test_polyline_se2_property(self): """Test polyline_se2 property.""" @@ -118,10 +117,10 @@ def test_polyline_se2_property(self): linestring = geom.LineString(coords) polyline = Polyline2D.from_linestring(linestring) polyline_se2 = polyline.polyline_se2 - self.assertIsInstance(polyline_se2, PolylineSE2) + assert isinstance(polyline_se2, PolylineSE2) -class TestPolylineSE2(unittest.TestCase): +class TestPolylineSE2: """Test class for PolylineSE2.""" def test_from_linestring(self): @@ -129,67 +128,60 @@ def test_from_linestring(self): coords = [(0.0, 0.0), (1.0, 0.0), (2.0, 0.0)] linestring = geom.LineString(coords) polyline = PolylineSE2.from_linestring(linestring) - self.assertIsInstance(polyline, PolylineSE2) - self.assertEqual(polyline.array.shape, (3, 3)) + assert isinstance(polyline, PolylineSE2) + assert polyline.array.shape == (3, 3) def test_from_array_2d(self): """Test creating PolylineSE2 from 2D array.""" array = np.array([[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]], dtype=np.float32) polyline = PolylineSE2.from_array(array) - self.assertIsInstance(polyline, PolylineSE2) - self.assertEqual(polyline.array.shape, (3, 3)) + assert isinstance(polyline, PolylineSE2) + assert polyline.array.shape == (3, 3) def test_from_array_se2(self): """Test creating PolylineSE2 from SE2 array.""" array = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float32) polyline = PolylineSE2.from_array(array) - self.assertIsInstance(polyline, PolylineSE2) + assert isinstance(polyline, PolylineSE2) np.testing.assert_array_almost_equal(polyline.array, array) def test_from_array_invalid_shape(self): """Test creating PolylineSE2 from invalid array shape.""" array = np.array([[0.0], [1.0], [2.0]], dtype=np.float32) - with self.assertRaises(ValueError): + with pytest.raises(ValueError): PolylineSE2.from_array(array) - def test_from_discrete_se2(self): - """Test creating PolylineSE2 from discrete SE2 states.""" - states = [StateSE2(0.0, 0.0, 0.0), StateSE2(1.0, 0.0, 0.0), StateSE2(2.0, 0.0, 0.0)] - polyline = PolylineSE2.from_discrete_se2(states) - self.assertIsInstance(polyline, PolylineSE2) - self.assertEqual(polyline.array.shape, (3, 3)) - def test_length_property(self): """Test length property.""" array = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float64) polyline = PolylineSE2.from_array(array) - self.assertEqual(polyline.length, 2.0) + assert polyline.length == 2.0 def test_interpolate_single_distance(self): """Test interpolation with single distance.""" array = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float64) polyline = PolylineSE2.from_array(array) state = polyline.interpolate(1.0) - self.assertIsInstance(state, StateSE2) - self.assertEqual(state.x, 1.0) - self.assertEqual(state.y, 0.0) + assert isinstance(state, PoseSE2) + assert state.x == 1.0 + assert state.y == 0.0 def test_interpolate_multiple_distances(self): """Test interpolation with multiple distances.""" array = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float64) polyline = PolylineSE2.from_array(array) states = polyline.interpolate(np.array([0.0, 1.0, 2.0])) - self.assertIsInstance(states, np.ndarray) - self.assertEqual(states.shape, (3, 3)) + assert isinstance(states, np.ndarray) + assert states.shape == (3, 3) def test_interpolate_normalized(self): """Test normalized interpolation.""" array = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float64) polyline = PolylineSE2.from_array(array) state = polyline.interpolate(0.5, normalized=True) - self.assertIsInstance(state, StateSE2) - self.assertEqual(state.x, 1.0) - self.assertEqual(state.y, 0.0) + assert isinstance(state, PoseSE2) + assert state.x == 1.0 + assert state.y == 0.0 def test_project_point2d(self): """Test projecting Point2D onto SE2 polyline.""" @@ -197,18 +189,18 @@ def test_project_point2d(self): polyline = PolylineSE2.from_array(array) point = Point2D(1.0, 1.0) distance = polyline.project(point) - self.assertEqual(distance, 1.0) + assert distance == 1.0 def test_project_statese2(self): """Test projecting StateSE2 onto SE2 polyline.""" array = np.array([[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]], dtype=np.float64) polyline = PolylineSE2.from_array(array) - state = StateSE2(1.0, 1.0, 0.0) + state = PoseSE2(1.0, 1.0, 0.0) distance = polyline.project(state) - self.assertEqual(distance, 1.0) + assert distance == 1.0 -class TestPolyline3D(unittest.TestCase): +class TestPolyline3D: """Test class for Polyline3D.""" def test_from_linestring_with_z(self): @@ -216,28 +208,32 @@ def test_from_linestring_with_z(self): coords = [(0.0, 0.0, 1.0), (1.0, 1.0, 2.0), (2.0, 0.0, 3.0)] linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) - self.assertIsInstance(polyline, Polyline3D) - self.assertTrue(polyline.linestring.has_z) + assert isinstance(polyline, Polyline3D) + assert polyline.linestring.has_z def test_from_linestring_without_z(self): """Test creating Polyline3D from LineString without Z coordinates.""" coords = [(0.0, 0.0), (1.0, 1.0), (2.0, 0.0)] linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) - self.assertIsInstance(polyline, Polyline3D) - self.assertTrue(polyline.linestring.has_z) + assert isinstance(polyline, Polyline3D) + assert polyline.linestring.has_z def test_from_array(self): """Test creating Polyline3D from 3D array.""" array = np.array([[0.0, 0.0, 1.0], [1.0, 1.0, 2.0], [2.0, 0.0, 3.0]], dtype=np.float64) polyline = Polyline3D.from_array(array) - self.assertIsInstance(polyline, Polyline3D) + assert isinstance(polyline, Polyline3D) np.testing.assert_array_almost_equal(polyline.array, array) def test_from_array_invalid_shape(self): """Test creating Polyline3D from invalid array shape.""" - array = np.array([[0.0, 0.0], [1.0, 1.0]], dtype=np.float64) - with self.assertRaises(AssertionError): + array = np.array([[0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]], dtype=np.float64) + with pytest.raises(ValueError): + Polyline3D.from_array(array) + + array = np.array([[0.0], [1.0]], dtype=np.float64) + with pytest.raises(ValueError): Polyline3D.from_array(array) def test_array_property(self): @@ -246,8 +242,8 @@ def test_array_property(self): linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) array = polyline.array - self.assertEqual(array.shape, (3, 3)) - self.assertEqual(array.dtype, np.float64) + assert array.shape == (3, 3) + assert array.dtype == np.float64 np.testing.assert_array_almost_equal(array, coords) def test_polyline_2d_property(self): @@ -256,8 +252,8 @@ def test_polyline_2d_property(self): linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) polyline_2d = polyline.polyline_2d - self.assertIsInstance(polyline_2d, Polyline2D) - self.assertFalse(polyline_2d.linestring.has_z) + assert isinstance(polyline_2d, Polyline2D) + assert not polyline_2d.linestring.has_z def test_polyline_se2_property(self): """Test polyline_se2 property.""" @@ -265,25 +261,35 @@ def test_polyline_se2_property(self): linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) polyline_se2 = polyline.polyline_se2 - self.assertIsInstance(polyline_se2, PolylineSE2) + assert isinstance(polyline_se2, PolylineSE2) def test_length_property(self): """Test length property.""" coords = [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (2.0, 0.0, 0.0)] linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) - self.assertEqual(polyline.length, 2.0) + assert polyline.length == 2.0 + + coords = [(0.0, 0.0, 0.0), (0.0, 0.0, 1.0), (0.0, 0.0, 2.0)] + linestring = geom.LineString(coords) + polyline = Polyline3D.from_linestring(linestring) + assert polyline.length == 2.0 + + coords = [(0.0, 0.0, 0.0), (1.0, 1.0, 1.0), (2.0, 2.0, 2.0)] + linestring = geom.LineString(coords) + polyline = Polyline3D.from_linestring(linestring) + assert polyline.length == 2 * np.sqrt(3) def test_interpolate_single_distance(self): """Test interpolation with single distance.""" coords = [(0.0, 0.0, 0.0), (2.0, 0.0, 2.0)] linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) - point = polyline.interpolate(1.0) - self.assertIsInstance(point, Point3D) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 0.0) - self.assertEqual(point.z, 1.0) + point = polyline.interpolate(np.sqrt(2)) + assert isinstance(point, Point3D) + assert point.x == 1.0 + assert point.y == 0.0 + assert point.z == 1.0 def test_interpolate_multiple_distances(self): """Test interpolation with multiple distances.""" @@ -291,8 +297,8 @@ def test_interpolate_multiple_distances(self): linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) points = polyline.interpolate(np.array([0.0, 1.0, 2.0])) - self.assertIsInstance(points, np.ndarray) - self.assertEqual(points.shape, (3, 3)) + assert isinstance(points, np.ndarray) + assert points.shape == (3, 3) def test_interpolate_normalized(self): """Test normalized interpolation.""" @@ -300,10 +306,10 @@ def test_interpolate_normalized(self): linestring = geom.LineString(coords) polyline = Polyline3D.from_linestring(linestring) point = polyline.interpolate(0.5, normalized=True) - self.assertIsInstance(point, Point3D) - self.assertEqual(point.x, 1.0) - self.assertEqual(point.y, 0.0) - self.assertEqual(point.z, 1.0) + assert isinstance(point, Point3D) + assert point.x == 1.0 + assert point.y == 0.0 + assert point.z == 1.0 def test_project_point2d(self): """Test projecting Point2D onto 3D polyline.""" @@ -312,7 +318,7 @@ def test_project_point2d(self): polyline = Polyline3D.from_linestring(linestring) point = Point2D(1.0, 1.0) distance = polyline.project(point) - self.assertEqual(distance, 1.0) + assert distance == 1.0 def test_project_point3d(self): """Test projecting Point3D onto 3D polyline.""" @@ -321,4 +327,4 @@ def test_project_point3d(self): polyline = Polyline3D.from_linestring(linestring) point = Point3D(1.0, 1.0, 1.0) distance = polyline.project(point) - self.assertEqual(distance, 1.0) + assert distance == 1.0 diff --git a/tests/unit/geometry/test_pose.py b/tests/unit/geometry/test_pose.py new file mode 100644 index 00000000..1936cbc2 --- /dev/null +++ b/tests/unit/geometry/test_pose.py @@ -0,0 +1,422 @@ +import numpy as np +import pytest + +from py123d.geometry import Point2D, PoseSE2, PoseSE3 +from py123d.geometry.geometry_index import PoseSE2Index +from py123d.geometry.pose import EulerPoseSE3 + + +class TestPoseSE2: + def test_init(self): + """Test basic initialization with explicit x, y, yaw values.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.yaw == 0.5 + + def test_from_array(self): + """Test creation from numpy array.""" + array = np.array([1.0, 2.0, 0.5]) + pose = PoseSE2.from_array(array) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.yaw == 0.5 + + def test_from_array_copy(self): + """Test that copy=True creates independent pose from array.""" + array = np.array([1.0, 2.0, 0.5]) + pose = PoseSE2.from_array(array, copy=True) + array[0] = 99.0 + assert pose.x == 1.0 + + def test_from_array_no_copy(self): + """Test that copy=False links pose to original array.""" + array = np.array([1.0, 2.0, 0.5]) + pose = PoseSE2.from_array(array, copy=False) + array[0] = 99.0 + assert pose.x == 99.0 + + def test_properties(self): + """Test access to individual pose component properties.""" + pose = PoseSE2(x=3.0, y=4.0, yaw=np.pi / 4) + assert pose.x == 3.0 + assert pose.y == 4.0 + assert pytest.approx(pose.yaw) == np.pi / 4 + + def test_array_property(self): + """Test that the array property returns correct numpy array.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + array = pose.array + assert array.shape == (3,) + assert array[PoseSE2Index.X] == 1.0 + assert array[PoseSE2Index.Y] == 2.0 + assert array[PoseSE2Index.YAW] == 0.5 + + def test_point_2d(self): + """Test extraction of 2D position as Point2D.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + point = pose.point_2d + assert isinstance(point, Point2D) + assert point.x == 1.0 + assert point.y == 2.0 + + def test_rotation_matrix(self): + """Test extraction of 2x2 rotation matrix.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.0) + rot_mat = pose.rotation_matrix + expected = np.array([[1.0, 0.0], [0.0, 1.0]]) + np.testing.assert_allclose(rot_mat, expected) + + def test_rotation_matrix_pi_half(self): + """Test extraction of 2x2 rotation matrix for 90 degree rotation.""" + pose = PoseSE2(x=0.0, y=0.0, yaw=np.pi / 2) + rot_mat = pose.rotation_matrix + expected = np.array([[0.0, -1.0], [1.0, 0.0]]) + np.testing.assert_allclose(rot_mat, expected, atol=1e-10) + + def test_transformation_matrix(self): + """Test extraction of 3x3 transformation matrix.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.0) + trans_mat = pose.transformation_matrix + assert trans_mat.shape == (3, 3) + expected = np.array([[1.0, 0.0, 1.0], [0.0, 1.0, 2.0], [0.0, 0.0, 0.0]]) + np.testing.assert_allclose(trans_mat, expected) + + def test_shapely_point(self): + """Test extraction of Shapely Point representation.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + shapely_point = pose.shapely_point + assert shapely_point.x == 1.0 + assert shapely_point.y == 2.0 + + def test_pose_se2_property(self): + """Test that pose_se2 property returns self.""" + pose = PoseSE2(x=1.0, y=2.0, yaw=0.5) + assert pose.pose_se2 is pose + + def test_equality(self): + """Test equality comparison of PoseSE2 instances.""" + pose1 = PoseSE2(x=1.0, y=2.0, yaw=0.5) + pose2 = PoseSE2(x=1.0, y=2.0, yaw=0.5) + assert pose1 == pose2 + + def test_inequality(self): + """Test inequality comparison of PoseSE2 instances.""" + pose1 = PoseSE2(x=1.0, y=2.0, yaw=0.5) + pose2 = PoseSE2(x=1.0, y=2.0, yaw=0.6) + assert pose1 != pose2 + + +class TestPoseSE3: + def test_init(self): + """Test basic initialization with explicit x, y, z, and quaternion values.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.qw == 1.0 + assert pose.qx == 0.0 + assert pose.qy == 0.0 + assert pose.qz == 0.0 + + def test_from_array(self): + """Test creation from numpy array.""" + array = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + pose = PoseSE3.from_array(array) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.qw == 1.0 + assert pose.qx == 0.0 + assert pose.qy == 0.0 + assert pose.qz == 0.0 + + def test_from_array_copy(self): + """Test that copy=True creates independent pose from array.""" + array = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + pose = PoseSE3.from_array(array, copy=True) + array[0] = 99.0 + assert pose.x == 1.0 + + def test_from_array_no_copy(self): + """Test that copy=False links pose to original array.""" + array = np.array([1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + pose = PoseSE3.from_array(array, copy=False) + array[0] = 99.0 + assert pose.x == 99.0 + + def test_from_transformation_matrix(self): + """Test creation from 4x4 transformation matrix.""" + trans_mat = np.eye(4) + trans_mat[:3, 3] = [1.0, 2.0, 3.0] + pose = PoseSE3.from_transformation_matrix(trans_mat) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.qw == 1.0 + assert pose.qx == 0.0 + assert pose.qy == 0.0 + assert pose.qz == 0.0 + + def test_properties(self): + """Test access to individual pose component properties.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.qw == 1.0 + assert pose.qx == 0.0 + assert pose.qy == 0.0 + assert pose.qz == 0.0 + + def test_array_property(self): + """Test that the array property returns the correct numpy array representation.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + array = pose.array + assert array.shape == (7,) + np.testing.assert_allclose(array, [1.0, 2.0, 3.0, 1.0, 0.0, 0.0, 0.0]) + + def test_pose_se2(self): + """Test extraction of 2D pose from 3D pose.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + pose_2d = pose.pose_se2 + assert isinstance(pose_2d, PoseSE2) + assert pose_2d.x == 1.0 + assert pose_2d.y == 2.0 + assert pytest.approx(pose_2d.yaw) == 0.0 + + def test_point_3d(self): + """Test extraction of 3D point from 3D pose.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + point = pose.point_3d + assert point.x == 1.0 + assert point.y == 2.0 + assert point.z == 3.0 + + def test_point_2d(self): + """Test extraction of 2D point from 3D pose.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + point = pose.point_2d + assert isinstance(point, Point2D) + assert point.x == 1.0 + assert point.y == 2.0 + + def test_shapely_point(self): + """Test extraction of Shapely Point representation.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + shapely_point = pose.shapely_point + assert shapely_point.x == 1.0 + assert shapely_point.y == 2.0 + assert shapely_point.z == 3.0 + + def test_rotation_matrix(self): + """Test extraction of 3x3 rotation matrix.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + rot_mat = pose.rotation_matrix + expected = np.eye(3) + np.testing.assert_allclose(rot_mat, expected) + + def test_transformation_matrix(self): + """Test extraction of 4x4 transformation matrix.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + trans_mat = pose.transformation_matrix + assert trans_mat.shape == (4, 4) + expected = np.eye(4) + expected[:3, 3] = [1.0, 2.0, 3.0] + np.testing.assert_allclose(trans_mat, expected) + + def test_transformation_matrix_roundtrip(self): + """Test round-trip conversion between pose and transformation matrix.""" + pose1 = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + trans_mat = pose1.transformation_matrix + pose2 = PoseSE3.from_transformation_matrix(trans_mat) + assert pose1 == pose2 + + def test_euler_angles(self): + """Test extraction of Euler angles from quaternion.""" + pose = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + assert pytest.approx(pose.roll) == 0.0 + assert pytest.approx(pose.pitch) == 0.0 + assert pytest.approx(pose.yaw) == 0.0 + + def test_equality(self): + """Test equality comparison of PoseSE3 instances.""" + pose1 = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + pose2 = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + assert pose1 == pose2 + + def test_inequality(self): + """Test inequality comparison of PoseSE3 instances.""" + pose1 = PoseSE3(x=1.0, y=2.0, z=3.0, qw=1.0, qx=0.0, qy=0.0, qz=0.0) + pose2 = PoseSE3(x=1.0, y=2.0, z=3.0, qw=0.9, qx=0.1, qy=0.0, qz=0.0) + assert pose1 != pose2 + + +class TestEulerPoseSE3: + def test_init(self): + """Test initialization of EulerPoseSE3 with position and orientation.""" + + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.roll == 0.1 + assert pose.pitch == 0.2 + assert pose.yaw == 0.3 + + def test_from_array(self): + """Test creation of EulerPoseSE3 from numpy array.""" + + array = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + pose = EulerPoseSE3.from_array(array) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.roll == 0.1 + assert pose.pitch == 0.2 + assert pose.yaw == 0.3 + + def test_from_array_copy(self): + """Test that copy=True creates independent pose from array.""" + + array = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + pose = EulerPoseSE3.from_array(array, copy=True) + array[0] = 99.0 + assert pose.x == 1.0 + + def test_from_array_no_copy(self): + """Test that copy=False links pose to original array.""" + + array = np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + pose = EulerPoseSE3.from_array(array, copy=False) + array[0] = 99.0 + assert pose.x == 99.0 + + def test_from_transformation_matrix(self): + """Test creation of EulerPoseSE3 from 4x4 transformation matrix.""" + trans_mat = np.eye(4) + trans_mat[:3, 3] = [1.0, 2.0, 3.0] + pose = EulerPoseSE3.from_transformation_matrix(trans_mat) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pytest.approx(pose.roll) == 0.0 + assert pytest.approx(pose.pitch) == 0.0 + assert pytest.approx(pose.yaw) == 0.0 + + def test_properties(self): + """Test access to individual pose component properties.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + assert pose.x == 1.0 + assert pose.y == 2.0 + assert pose.z == 3.0 + assert pose.roll == 0.1 + assert pose.pitch == 0.2 + assert pose.yaw == 0.3 + + def test_array_property(self): + """Test that the array property returns the correct numpy array representation.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + array = pose.array + assert array.shape == (6,) + np.testing.assert_allclose(array, [1.0, 2.0, 3.0, 0.1, 0.2, 0.3]) + + def test_pose_se2(self): + """Test extraction of 2D pose from 3D Euler pose.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + pose_2d = pose.pose_se2 + assert isinstance(pose_2d, PoseSE2) + assert pose_2d.x == 1.0 + assert pose_2d.y == 2.0 + assert pose_2d.yaw == 0.3 + + def test_point_3d(self): + """Test extraction of 3D point from 3D Euler pose.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + point = pose.point_3d + assert point.x == 1.0 + assert point.y == 2.0 + assert point.z == 3.0 + + def test_point_2d(self): + """Test extraction of 2D point from 3D Euler pose.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + point = pose.point_2d + assert isinstance(point, Point2D) + assert point.x == 1.0 + assert point.y == 2.0 + + def test_shapely_point(self): + """Test extraction of Shapely Point representation.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + shapely_point = pose.shapely_point + assert shapely_point.x == 1.0 + assert shapely_point.y == 2.0 + assert shapely_point.z == 3.0 + + def test_rotation_matrix(self): + """Test the rotation matrix property of EulerPoseSE3.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0) + rot_mat = pose.rotation_matrix + expected = np.eye(3) + np.testing.assert_allclose(rot_mat, expected) + + def test_transformation_matrix(self): + """Test the transformation matrix property of EulerPoseSE3.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0) + trans_mat = pose.transformation_matrix + assert trans_mat.shape == (4, 4) + expected = np.eye(4) + expected[:3, 3] = [1.0, 2.0, 3.0] + np.testing.assert_allclose(trans_mat, expected) + + def test_transformation_matrix_roundtrip(self): + """Test round-trip conversion between EulerPoseSE3 and transformation matrix.""" + pose1 = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0) + trans_mat = pose1.transformation_matrix + pose2 = EulerPoseSE3.from_transformation_matrix(trans_mat) + np.testing.assert_allclose(pose1.array, pose2.array) + + def test_euler_angles(self): + """Test the euler_angles property of EulerPoseSE3.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + euler = pose.euler_angles + assert pytest.approx(euler.roll) == 0.1 + assert pytest.approx(euler.pitch) == 0.2 + assert pytest.approx(euler.yaw) == 0.3 + + def test_pose_se3_conversion(self): + """Test conversion from EulerPoseSE3 to PoseSE3.""" + euler_pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0) + quat_pose = euler_pose.pose_se3 + assert isinstance(quat_pose, PoseSE3) + assert quat_pose.x == 1.0 + assert quat_pose.y == 2.0 + assert quat_pose.z == 3.0 + assert pytest.approx(quat_pose.qw) == 1.0 + assert pytest.approx(quat_pose.qx) == 0.0 + assert pytest.approx(quat_pose.qy) == 0.0 + assert pytest.approx(quat_pose.qz) == 0.0 + + def test_quaternion(self): + """Test the quaternion property of EulerPoseSE3.""" + pose = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.0, pitch=0.0, yaw=0.0) + quat = pose.quaternion + assert pytest.approx(quat.qw) == 1.0 + assert pytest.approx(quat.qx) == 0.0 + assert pytest.approx(quat.qy) == 0.0 + assert pytest.approx(quat.qz) == 0.0 + + def test_equality(self): + """Test equality comparison of EulerPoseSE3 instances.""" + + pose1 = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + pose2 = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + assert pose1 == pose2 + + def test_inequality(self): + """Test inequality comparison of EulerPoseSE3 instances.""" + + pose1 = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.3) + pose2 = EulerPoseSE3(x=1.0, y=2.0, z=3.0, roll=0.1, pitch=0.2, yaw=0.4) + assert pose1 != pose2 diff --git a/tests/unit/geometry/test_rotation.py b/tests/unit/geometry/test_rotation.py index 66f51c49..a7fc2e09 100644 --- a/tests/unit/geometry/test_rotation.py +++ b/tests/unit/geometry/test_rotation.py @@ -1,15 +1,14 @@ -import unittest - import numpy as np +import pytest from py123d.geometry.geometry_index import EulerAnglesIndex, QuaternionIndex from py123d.geometry.rotation import EulerAngles, Quaternion -class TestEulerAngles(unittest.TestCase): +class TestEulerAngles: """Unit tests for EulerAngles class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.roll = 0.1 self.pitch = 0.2 @@ -23,23 +22,23 @@ def setUp(self): def test_init(self): """Test EulerAngles initialization.""" euler = EulerAngles(roll=0.1, pitch=0.2, yaw=0.3) - self.assertEqual(euler.roll, 0.1) - self.assertEqual(euler.pitch, 0.2) - self.assertEqual(euler.yaw, 0.3) + assert euler.roll == 0.1 + assert euler.pitch == 0.2 + assert euler.yaw == 0.3 def test_from_array_valid(self): """Test from_array class method with valid input.""" euler = EulerAngles.from_array(self.test_array) - self.assertIsInstance(euler, EulerAngles) - self.assertAlmostEqual(euler.roll, self.roll) - self.assertAlmostEqual(euler.pitch, self.pitch) - self.assertAlmostEqual(euler.yaw, self.yaw) + assert isinstance(euler, EulerAngles) + assert euler.roll == pytest.approx(self.roll) + assert euler.pitch == pytest.approx(self.pitch) + assert euler.yaw == pytest.approx(self.yaw) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): EulerAngles.from_array(np.array([1, 2])) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): EulerAngles.from_array(np.array([[1, 2, 3]])) def test_from_array_copy(self): @@ -49,51 +48,37 @@ def test_from_array_copy(self): euler_no_copy = EulerAngles.from_array(original_array, copy=False) original_array[0] = 999.0 - self.assertNotEqual(euler_copy.roll, 999.0) - self.assertEqual(euler_no_copy.roll, 999.0) + assert euler_copy.roll != 999.0 + assert euler_no_copy.roll == 999.0 def test_from_rotation_matrix(self): """Test from_rotation_matrix class method.""" identity_matrix = np.eye(3) euler = EulerAngles.from_rotation_matrix(identity_matrix) - self.assertAlmostEqual(euler.roll, 0.0, places=10) - self.assertAlmostEqual(euler.pitch, 0.0, places=10) - self.assertAlmostEqual(euler.yaw, 0.0, places=10) + assert euler.roll == pytest.approx(0.0, abs=1e-10) + assert euler.pitch == pytest.approx(0.0, abs=1e-10) + assert euler.yaw == pytest.approx(0.0, abs=1e-10) def test_from_rotation_matrix_invalid(self): """Test from_rotation_matrix with invalid input.""" - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): EulerAngles.from_rotation_matrix(np.array([[1, 2]])) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): EulerAngles.from_rotation_matrix(np.array([1, 2, 3])) def test_array_property(self): """Test array property.""" array = self.euler_angles.array - self.assertEqual(array.shape, (3,)) - self.assertEqual(array[EulerAnglesIndex.ROLL], self.roll) - self.assertEqual(array[EulerAnglesIndex.PITCH], self.pitch) - self.assertEqual(array[EulerAnglesIndex.YAW], self.yaw) - - def test_iterator(self): - """Test iterator functionality.""" - values = list(self.euler_angles) - self.assertEqual(values, [self.roll, self.pitch, self.yaw]) - - def test_hash(self): - """Test hash functionality.""" - euler1 = EulerAngles(0.1, 0.2, 0.3) - euler2 = EulerAngles(0.1, 0.2, 0.3) - euler3 = EulerAngles(0.1, 0.2, 0.4) + assert array.shape == (3,) + assert array[EulerAnglesIndex.ROLL] == self.roll + assert array[EulerAnglesIndex.PITCH] == self.pitch + assert array[EulerAnglesIndex.YAW] == self.yaw - self.assertEqual(hash(euler1), hash(euler2)) - self.assertNotEqual(hash(euler1), hash(euler3)) - -class TestQuaternion(unittest.TestCase): +class TestQuaternion: """Unit tests for Quaternion class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.qw = 1.0 self.qx = 0.0 @@ -109,24 +94,24 @@ def setUp(self): def test_init(self): """Test Quaternion initialization.""" quat = Quaternion(1.0, 0.0, 0.0, 0.0) - self.assertEqual(quat.qw, 1.0) - self.assertEqual(quat.qx, 0.0) - self.assertEqual(quat.qy, 0.0) - self.assertEqual(quat.qz, 0.0) + assert quat.qw == 1.0 + assert quat.qx == 0.0 + assert quat.qy == 0.0 + assert quat.qz == 0.0 def test_from_array_valid(self): """Test from_array class method with valid input.""" quat = Quaternion.from_array(self.test_array) - self.assertAlmostEqual(quat.qw, self.qw) - self.assertAlmostEqual(quat.qx, self.qx) - self.assertAlmostEqual(quat.qy, self.qy) - self.assertAlmostEqual(quat.qz, self.qz) + assert quat.qw == pytest.approx(self.qw) + assert quat.qx == pytest.approx(self.qx) + assert quat.qy == pytest.approx(self.qy) + assert quat.qz == pytest.approx(self.qz) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Quaternion.from_array(np.array([1, 2, 3])) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Quaternion.from_array(np.array([[1, 2, 3, 4]])) def test_from_array_copy(self): @@ -136,76 +121,58 @@ def test_from_array_copy(self): quat_no_copy = Quaternion.from_array(original_array, copy=False) original_array[0] = 999.0 - self.assertNotEqual(quat_copy.qw, 999.0) - self.assertEqual(quat_no_copy.qw, 999.0) + assert quat_copy.qw != 999.0 + assert quat_no_copy.qw == 999.0 def test_from_rotation_matrix(self): """Test from_rotation_matrix class method.""" identity_matrix = np.eye(3) quat = Quaternion.from_rotation_matrix(identity_matrix) - self.assertAlmostEqual(quat.qw, 1.0, places=10) - self.assertAlmostEqual(quat.qx, 0.0, places=10) - self.assertAlmostEqual(quat.qy, 0.0, places=10) - self.assertAlmostEqual(quat.qz, 0.0, places=10) + assert quat.qw == pytest.approx(1.0, abs=1e-10) + assert quat.qx == pytest.approx(0.0, abs=1e-10) + assert quat.qy == pytest.approx(0.0, abs=1e-10) + assert quat.qz == pytest.approx(0.0, abs=1e-10) def test_from_rotation_matrix_invalid(self): """Test from_rotation_matrix with invalid input.""" - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Quaternion.from_rotation_matrix(np.array([[1, 2]])) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Quaternion.from_rotation_matrix(np.array([1, 2, 3])) def test_from_euler_angles(self): """Test from_euler_angles class method.""" euler = EulerAngles(0.0, 0.0, 0.0) quat = Quaternion.from_euler_angles(euler) - self.assertAlmostEqual(quat.qw, 1.0, places=10) - self.assertAlmostEqual(quat.qx, 0.0, places=10) - self.assertAlmostEqual(quat.qy, 0.0, places=10) - self.assertAlmostEqual(quat.qz, 0.0, places=10) + assert quat.qw == pytest.approx(1.0, abs=1e-10) + assert quat.qx == pytest.approx(0.0, abs=1e-10) + assert quat.qy == pytest.approx(0.0, abs=1e-10) + assert quat.qz == pytest.approx(0.0, abs=1e-10) def test_array_property(self): """Test array property.""" array = self.quaternion.array - self.assertEqual(array.shape, (4,)) + assert array.shape == (4,) np.testing.assert_array_equal(array, self.test_array) def test_pyquaternion_property(self): """Test pyquaternion property.""" pyquat = self.quaternion.pyquaternion - self.assertEqual(pyquat.w, self.qw) - self.assertEqual(pyquat.x, self.qx) - self.assertEqual(pyquat.y, self.qy) - self.assertEqual(pyquat.z, self.qz) + assert pyquat.w == self.qw + assert pyquat.x == self.qx + assert pyquat.y == self.qy + assert pyquat.z == self.qz def test_euler_angles_property(self): """Test euler_angles property.""" euler = self.quaternion.euler_angles - self.assertIsInstance(euler, EulerAngles) - self.assertAlmostEqual(euler.roll, 0.0, places=10) - self.assertAlmostEqual(euler.pitch, 0.0, places=10) - self.assertAlmostEqual(euler.yaw, 0.0, places=10) + assert isinstance(euler, EulerAngles) + assert euler.roll == pytest.approx(0.0, abs=1e-10) + assert euler.pitch == pytest.approx(0.0, abs=1e-10) + assert euler.yaw == pytest.approx(0.0, abs=1e-10) def test_rotation_matrix_property(self): """Test rotation_matrix property.""" rot_matrix = self.quaternion.rotation_matrix - self.assertEqual(rot_matrix.shape, (3, 3)) + assert rot_matrix.shape == (3, 3) np.testing.assert_array_almost_equal(rot_matrix, np.eye(3)) - - def test_iterator(self): - """Test iterator functionality.""" - values = list(self.quaternion) - self.assertEqual(values, [self.qw, self.qx, self.qy, self.qz]) - - def test_hash(self): - """Test hash functionality.""" - quat1 = Quaternion(1.0, 0.0, 0.0, 0.0) - quat2 = Quaternion(1.0, 0.0, 0.0, 0.0) - quat3 = Quaternion(0.0, 1.0, 0.0, 0.0) - - self.assertEqual(hash(quat1), hash(quat2)) - self.assertNotEqual(hash(quat1), hash(quat3)) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/geometry/test_vector.py b/tests/unit/geometry/test_vector.py index 4f3b159e..60e5e051 100644 --- a/tests/unit/geometry/test_vector.py +++ b/tests/unit/geometry/test_vector.py @@ -1,14 +1,13 @@ -import unittest - import numpy as np +import pytest from py123d.geometry import Vector2D, Vector2DIndex, Vector3D, Vector3DIndex -class TestVector2D(unittest.TestCase): +class TestVector2D: """Unit tests for Vector2D class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.x_coord = 3.5 self.y_coord = 4.2 @@ -20,74 +19,68 @@ def setUp(self): def test_init(self): """Test Vector2D initialization.""" vector = Vector2D(1.0, 2.0) - self.assertEqual(vector.x, 1.0) - self.assertEqual(vector.y, 2.0) + assert vector.x == 1.0 + assert vector.y == 2.0 def test_from_array_valid(self): """Test from_array class method with valid input.""" vector = Vector2D.from_array(self.test_array) - self.assertEqual(vector.x, self.x_coord) - self.assertEqual(vector.y, self.y_coord) + assert vector.x == self.x_coord + assert vector.y == self.y_coord def test_from_array_invalid_dimensions(self): """Test from_array with invalid array dimensions.""" # 2D array should raise assertion error array_2d = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector2D.from_array(array_2d) # 3D array should raise assertion error array_3d = np.array([[[1.0]]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector2D.from_array(array_3d) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" array_wrong_length = np.array([1.0, 2.0, 3.0], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector2D.from_array(array_wrong_length) # Empty array empty_array = np.array([], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector2D.from_array(empty_array) def test_array_property(self): """Test the array property.""" expected_array = np.array([self.x_coord, self.y_coord], dtype=np.float64) np.testing.assert_array_equal(self.vector.array, expected_array) - self.assertEqual(self.vector.array.dtype, np.float64) - self.assertEqual(self.vector.array.shape, (2,)) + assert self.vector.array.dtype == np.float64 + assert self.vector.array.shape == (2,) def test_array_like(self): """Test the __array__ behavior.""" expected_array = np.array([self.x_coord, self.y_coord], dtype=np.float32) output_array = np.array(self.vector, dtype=np.float32) np.testing.assert_array_equal(output_array, expected_array) - self.assertEqual(output_array.dtype, np.float32) - self.assertEqual(output_array.shape, (2,)) + assert output_array.dtype == np.float32 + assert output_array.shape == (2,) def test_iter(self): """Test the __iter__ method.""" coords = list(self.vector) - self.assertEqual(coords, [self.x_coord, self.y_coord]) + assert coords == [self.x_coord, self.y_coord] # Test that it's actually iterable x, y = self.vector - self.assertEqual(x, self.x_coord) - self.assertEqual(y, self.y_coord) - - def test_hash(self): - """Test the __hash__ method.""" - vector_dict = {self.vector: "test"} - self.assertIn(self.vector, vector_dict) - self.assertEqual(vector_dict[self.vector], "test") + assert x == self.x_coord + assert y == self.y_coord -class TestVector3D(unittest.TestCase): +class TestVector3D: """Unit tests for Vector3D class.""" - def setUp(self): + def setup_method(self): """Set up test fixtures.""" self.x_coord = 3.5 self.y_coord = 4.2 @@ -101,72 +94,62 @@ def setUp(self): def test_init(self): """Test Vector3D initialization.""" vector = Vector3D(1.0, 2.0, 3.0) - self.assertEqual(vector.x, 1.0) - self.assertEqual(vector.y, 2.0) - self.assertEqual(vector.z, 3.0) + assert vector.x == 1.0 + assert vector.y == 2.0 + assert vector.z == 3.0 def test_from_array_valid(self): """Test from_array class method with valid input.""" vector = Vector3D.from_array(self.test_array) - self.assertEqual(vector.x, self.x_coord) - self.assertEqual(vector.y, self.y_coord) - self.assertEqual(vector.z, self.z_coord) + assert vector.x == self.x_coord + assert vector.y == self.y_coord + assert vector.z == self.z_coord def test_from_array_invalid_dimensions(self): """Test from_array with invalid array dimensions.""" # 2D array should raise assertion error array_2d = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector3D.from_array(array_2d) # 3D array should raise assertion error array_3d = np.array([[[1.0]]], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector3D.from_array(array_3d) def test_from_array_invalid_shape(self): """Test from_array with invalid array shape.""" array_wrong_length = np.array([1.0, 2.0], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector3D.from_array(array_wrong_length) # Empty array empty_array = np.array([], dtype=np.float64) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): Vector3D.from_array(empty_array) def test_array_property(self): """Test the array property.""" expected_array = np.array([self.x_coord, self.y_coord, self.z_coord], dtype=np.float64) np.testing.assert_array_equal(self.vector.array, expected_array) - self.assertEqual(self.vector.array.dtype, np.float64) - self.assertEqual(self.vector.array.shape, (3,)) + assert self.vector.array.dtype == np.float64 + assert self.vector.array.shape == (3,) def test_array_like(self): """Test the __array__ behavior.""" expected_array = np.array([self.x_coord, self.y_coord, self.z_coord], dtype=np.float32) output_array = np.array(self.vector, dtype=np.float32) np.testing.assert_array_equal(output_array, expected_array) - self.assertEqual(output_array.dtype, np.float32) - self.assertEqual(output_array.shape, (3,)) + assert output_array.dtype == np.float32 + assert output_array.shape == (3,) def test_iter(self): """Test the __iter__ method.""" coords = list(self.vector) - self.assertEqual(coords, [self.x_coord, self.y_coord, self.z_coord]) + assert coords == [self.x_coord, self.y_coord, self.z_coord] # Test that it's actually iterable x, y, z = self.vector - self.assertEqual(x, self.x_coord) - self.assertEqual(y, self.y_coord) - self.assertEqual(z, self.z_coord) - - def test_hash(self): - """Test the __hash__ method.""" - vector_dict = {self.vector: "test"} - self.assertIn(self.vector, vector_dict) - self.assertEqual(vector_dict[self.vector], "test") - - -if __name__ == "__main__": - unittest.main() + assert x == self.x_coord + assert y == self.y_coord + assert z == self.z_coord diff --git a/tests/unit/geometry/transform/test_transform_consistency.py b/tests/unit/geometry/transform/test_transform_consistency.py index a798fabd..2b31c110 100644 --- a/tests/unit/geometry/transform/test_transform_consistency.py +++ b/tests/unit/geometry/transform/test_transform_consistency.py @@ -1,10 +1,8 @@ -import unittest - import numpy as np import numpy.typing as npt -from py123d.geometry import EulerStateSE3, StateSE2, Vector2D, Vector3D -from py123d.geometry.geometry_index import EulerStateSE3Index, Point2DIndex, Point3DIndex, StateSE2Index +from py123d.geometry import EulerPoseSE3, PoseSE2, Vector2D, Vector3D +from py123d.geometry.geometry_index import EulerPoseSE3Index, Point2DIndex, Point3DIndex, PoseSE2Index from py123d.geometry.transform.transform_euler_se3 import ( convert_absolute_to_relative_euler_se3_array, convert_absolute_to_relative_points_3d_array, @@ -15,9 +13,9 @@ translate_euler_se3_along_y, ) from py123d.geometry.transform.transform_se2 import ( - convert_absolute_to_relative_point_2d_array, + convert_absolute_to_relative_points_2d_array, convert_absolute_to_relative_se2_array, - convert_relative_to_absolute_point_2d_array, + convert_relative_to_absolute_points_2d_array, convert_relative_to_absolute_se2_array, translate_se2_along_body_frame, translate_se2_along_x, @@ -27,10 +25,10 @@ from py123d.geometry.utils.rotation_utils import get_rotation_matrices_from_euler_array -class TestTransformConsistency(unittest.TestCase): +class TestTransformConsistency: """Tests to ensure consistency between different transformation functions.""" - def setUp(self): + def setup_method(self): self.decimal = 4 # Decimal places for np.testing.assert_array_almost_equal self.num_consistency_tests = 10 # Number of random test cases for consistency checks @@ -40,19 +38,17 @@ def setUp(self): def _get_random_se2_array(self, size: int) -> npt.NDArray[np.float64]: """Generate a random SE2 pose""" - random_se2_array = np.random.uniform(-self.max_pose_xyz, self.max_pose_xyz, (size, len(StateSE2Index))) - random_se2_array[:, StateSE2Index.YAW] = np.random.uniform(-np.pi, np.pi, size) # yaw angles + random_se2_array = np.random.uniform(-self.max_pose_xyz, self.max_pose_xyz, (size, len(PoseSE2Index))) + random_se2_array[:, PoseSE2Index.YAW] = np.random.uniform(-np.pi, np.pi, size) # yaw angles return random_se2_array def _get_random_se3_array(self, size: int) -> npt.NDArray[np.float64]: """Generate a random SE3 poses""" - random_se3_array = np.zeros((size, len(EulerStateSE3Index)), dtype=np.float64) - random_se3_array[:, EulerStateSE3Index.XYZ] = np.random.uniform( - -self.max_pose_xyz, self.max_pose_xyz, (size, 3) - ) - random_se3_array[:, EulerStateSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) - random_se3_array[:, EulerStateSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) - random_se3_array[:, EulerStateSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array = np.zeros((size, len(EulerPoseSE3Index)), dtype=np.float64) + random_se3_array[:, EulerPoseSE3Index.XYZ] = np.random.uniform(-self.max_pose_xyz, self.max_pose_xyz, (size, 3)) + random_se3_array[:, EulerPoseSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array[:, EulerPoseSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) + random_se3_array[:, EulerPoseSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) return random_se3_array @@ -60,7 +56,7 @@ def test_se2_absolute_relative_conversion_consistency(self) -> None: """Test that converting absolute->relative->absolute returns original poses""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = StateSE2.from_array(self._get_random_se2_array(1)[0]) + reference = PoseSE2.from_array(self._get_random_se2_array(1)[0]) # Generate random absolute poses num_poses = np.random.randint(self.min_random_poses, self.max_random_poses) @@ -76,15 +72,15 @@ def test_se2_points_absolute_relative_conversion_consistency(self) -> None: """Test that converting absolute->relative->absolute returns original points""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = StateSE2.from_array(self._get_random_se2_array(1)[0]) + reference = PoseSE2.from_array(self._get_random_se2_array(1)[0]) # Generate random absolute points num_points = np.random.randint(self.min_random_poses, self.max_random_poses) - absolute_points = self._get_random_se2_array(num_points)[:, StateSE2Index.XY] + absolute_points = self._get_random_se2_array(num_points)[:, PoseSE2Index.XY] # Convert absolute -> relative -> absolute - relative_points = convert_absolute_to_relative_point_2d_array(reference, absolute_points) - recovered_absolute = convert_relative_to_absolute_point_2d_array(reference, relative_points) + relative_points = convert_absolute_to_relative_points_2d_array(reference, absolute_points) + recovered_absolute = convert_relative_to_absolute_points_2d_array(reference, relative_points) np.testing.assert_array_almost_equal(absolute_points, recovered_absolute, decimal=self.decimal) @@ -92,7 +88,7 @@ def test_se2_points_consistency(self) -> None: """Test whether SE2 point and pose conversions are consistent""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = StateSE2.from_array(self._get_random_se2_array(1)[0]) + reference = PoseSE2.from_array(self._get_random_se2_array(1)[0]) # Generate random absolute points num_poses = np.random.randint(self.min_random_poses, self.max_random_poses) @@ -100,24 +96,24 @@ def test_se2_points_consistency(self) -> None: # Convert absolute -> relative -> absolute relative_se2 = convert_absolute_to_relative_se2_array(reference, absolute_se2) - relative_points = convert_absolute_to_relative_point_2d_array( - reference, absolute_se2[..., StateSE2Index.XY] + relative_points = convert_absolute_to_relative_points_2d_array( + reference, absolute_se2[..., PoseSE2Index.XY] ) np.testing.assert_array_almost_equal( - relative_se2[..., StateSE2Index.XY], relative_points, decimal=self.decimal + relative_se2[..., PoseSE2Index.XY], relative_points, decimal=self.decimal ) recovered_absolute_se2 = convert_relative_to_absolute_se2_array(reference, relative_se2) - absolute_points = convert_relative_to_absolute_point_2d_array(reference, relative_points) + absolute_points = convert_relative_to_absolute_points_2d_array(reference, relative_points) np.testing.assert_array_almost_equal( - recovered_absolute_se2[..., StateSE2Index.XY], absolute_points, decimal=self.decimal + recovered_absolute_se2[..., PoseSE2Index.XY], absolute_points, decimal=self.decimal ) def test_se2_translation_consistency(self) -> None: """Test that SE2 translations are consistent between different methods""" for _ in range(self.num_consistency_tests): # Generate random pose - pose = StateSE2.from_array(self._get_random_se2_array(1)[0]) + pose = PoseSE2.from_array(self._get_random_se2_array(1)[0]) # Generate random distances dx = np.random.uniform(-10.0, 10.0) @@ -142,7 +138,7 @@ def test_se3_absolute_relative_conversion_consistency(self) -> None: """Test that converting absolute->relative->absolute returns original poses""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = EulerStateSE3.from_array(self._get_random_se3_array(1)[0]) + reference = EulerPoseSE3.from_array(self._get_random_se3_array(1)[0]) # Generate random absolute poses num_poses = np.random.randint(self.min_random_poses, self.max_random_poses) @@ -153,16 +149,16 @@ def test_se3_absolute_relative_conversion_consistency(self) -> None: recovered_absolute = convert_relative_to_absolute_euler_se3_array(reference, relative_poses) np.testing.assert_array_almost_equal( - absolute_poses[..., EulerStateSE3Index.XYZ], - recovered_absolute[..., EulerStateSE3Index.XYZ], + absolute_poses[..., EulerPoseSE3Index.XYZ], + recovered_absolute[..., EulerPoseSE3Index.XYZ], decimal=self.decimal, ) absolute_rotation_matrices = get_rotation_matrices_from_euler_array( - absolute_poses[..., EulerStateSE3Index.EULER_ANGLES] + absolute_poses[..., EulerPoseSE3Index.EULER_ANGLES] ) recovered_rotation_matrices = get_rotation_matrices_from_euler_array( - recovered_absolute[..., EulerStateSE3Index.EULER_ANGLES] + recovered_absolute[..., EulerPoseSE3Index.EULER_ANGLES] ) np.testing.assert_array_almost_equal( @@ -175,11 +171,11 @@ def test_se3_points_absolute_relative_conversion_consistency(self) -> None: """Test that converting absolute->relative->absolute returns original points""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = EulerStateSE3.from_array(self._get_random_se3_array(1)[0]) + reference = EulerPoseSE3.from_array(self._get_random_se3_array(1)[0]) # Generate random absolute points num_points = np.random.randint(self.min_random_poses, self.max_random_poses) - absolute_points = self._get_random_se3_array(num_points)[:, EulerStateSE3Index.XYZ] + absolute_points = self._get_random_se3_array(num_points)[:, EulerPoseSE3Index.XYZ] # Convert absolute -> relative -> absolute relative_points = convert_absolute_to_relative_points_3d_array(reference, absolute_points) @@ -191,7 +187,7 @@ def test_se3_points_consistency(self) -> None: """Test whether SE3 point and pose conversions are consistent""" for _ in range(self.num_consistency_tests): # Generate random reference pose - reference = EulerStateSE3.from_array(self._get_random_se3_array(1)[0]) + reference = EulerPoseSE3.from_array(self._get_random_se3_array(1)[0]) # Generate random absolute points num_poses = np.random.randint(self.min_random_poses, self.max_random_poses) @@ -200,16 +196,16 @@ def test_se3_points_consistency(self) -> None: # Convert absolute -> relative -> absolute relative_se3 = convert_absolute_to_relative_euler_se3_array(reference, absolute_se3) relative_points = convert_absolute_to_relative_points_3d_array( - reference, absolute_se3[..., EulerStateSE3Index.XYZ] + reference, absolute_se3[..., EulerPoseSE3Index.XYZ] ) np.testing.assert_array_almost_equal( - relative_se3[..., EulerStateSE3Index.XYZ], relative_points, decimal=self.decimal + relative_se3[..., EulerPoseSE3Index.XYZ], relative_points, decimal=self.decimal ) recovered_absolute_se3 = convert_relative_to_absolute_euler_se3_array(reference, relative_se3) absolute_points = convert_relative_to_absolute_points_3d_array(reference, relative_points) np.testing.assert_array_almost_equal( - recovered_absolute_se3[..., EulerStateSE3Index.XYZ], absolute_points, decimal=self.decimal + recovered_absolute_se3[..., EulerPoseSE3Index.XYZ], absolute_points, decimal=self.decimal ) def test_se2_se3_translation_along_body_consistency(self) -> None: @@ -217,8 +213,8 @@ def test_se2_se3_translation_along_body_consistency(self) -> None: for _ in range(self.num_consistency_tests): # Create equivalent SE2 and SE3 poses (SE3 with z=0 and no rotations except yaw) - pose_se2 = StateSE2.from_array(self._get_random_se2_array(1)[0]) - pose_se3 = EulerStateSE3.from_array( + pose_se2 = PoseSE2.from_array(self._get_random_se2_array(1)[0]) + pose_se3 = EulerPoseSE3.from_array( np.array([pose_se2.x, pose_se2.y, 0.0, 0.0, 0.0, pose_se2.yaw], dtype=np.float64) ) @@ -228,13 +224,13 @@ def test_se2_se3_translation_along_body_consistency(self) -> None: translated_se3_x = translate_euler_se3_along_x(pose_se3, dx) np.testing.assert_array_almost_equal( - translated_se2_x.array[StateSE2Index.XY], - translated_se3_x.array[EulerStateSE3Index.XY], + translated_se2_x.array[PoseSE2Index.XY], + translated_se3_x.array[EulerPoseSE3Index.XY], decimal=self.decimal, ) np.testing.assert_almost_equal( - translated_se2_x.array[StateSE2Index.YAW], - translated_se3_x.array[EulerStateSE3Index.YAW], + translated_se2_x.array[PoseSE2Index.YAW], + translated_se3_x.array[EulerPoseSE3Index.YAW], decimal=self.decimal, ) @@ -244,13 +240,13 @@ def test_se2_se3_translation_along_body_consistency(self) -> None: translated_se3_y = translate_euler_se3_along_y(pose_se3, dy) np.testing.assert_array_almost_equal( - translated_se2_y.array[StateSE2Index.XY], - translated_se3_y.array[EulerStateSE3Index.XY], + translated_se2_y.array[PoseSE2Index.XY], + translated_se3_y.array[EulerPoseSE3Index.XY], decimal=self.decimal, ) np.testing.assert_almost_equal( - translated_se2_y.array[StateSE2Index.YAW], - translated_se3_y.array[EulerStateSE3Index.YAW], + translated_se2_y.array[PoseSE2Index.YAW], + translated_se3_y.array[EulerPoseSE3Index.YAW], decimal=self.decimal, ) @@ -260,13 +256,13 @@ def test_se2_se3_translation_along_body_consistency(self) -> None: translated_se2_xy = translate_se2_along_body_frame(pose_se2, Vector2D(dx, dy)) translated_se3_xy = translate_euler_se3_along_body_frame(pose_se3, Vector3D(dx, dy, 0.0)) np.testing.assert_array_almost_equal( - translated_se2_xy.array[StateSE2Index.XY], - translated_se3_xy.array[EulerStateSE3Index.XY], + translated_se2_xy.array[PoseSE2Index.XY], + translated_se3_xy.array[EulerPoseSE3Index.XY], decimal=self.decimal, ) np.testing.assert_almost_equal( - translated_se2_xy.array[StateSE2Index.YAW], - translated_se3_xy.array[EulerStateSE3Index.YAW], + translated_se2_xy.array[PoseSE2Index.YAW], + translated_se3_xy.array[EulerPoseSE3Index.YAW], decimal=self.decimal, ) @@ -278,8 +274,8 @@ def test_se2_se3_point_conversion_consistency(self) -> None: y = np.random.uniform(-10.0, 10.0) yaw = np.random.uniform(-np.pi, np.pi) - reference_se2 = StateSE2.from_array(np.array([x, y, yaw], dtype=np.float64)) - reference_se3 = EulerStateSE3.from_array(np.array([x, y, 0.0, 0.0, 0.0, yaw], dtype=np.float64)) + reference_se2 = PoseSE2.from_array(np.array([x, y, yaw], dtype=np.float64)) + reference_se3 = EulerPoseSE3.from_array(np.array([x, y, 0.0, 0.0, 0.0, yaw], dtype=np.float64)) # Generate 2D points and embed them in 3D with z=0 num_points = np.random.randint(1, 8) @@ -287,8 +283,8 @@ def test_se2_se3_point_conversion_consistency(self) -> None: points_3d = np.column_stack([points_2d, np.zeros(num_points)]) # Convert using SE2 functions - relative_2d = convert_absolute_to_relative_point_2d_array(reference_se2, points_2d) - absolute_2d_recovered = convert_relative_to_absolute_point_2d_array(reference_se2, relative_2d) + relative_2d = convert_absolute_to_relative_points_2d_array(reference_se2, points_2d) + absolute_2d_recovered = convert_relative_to_absolute_points_2d_array(reference_se2, relative_2d) # Convert using SE3 functions relative_3d = convert_absolute_to_relative_points_3d_array(reference_se3, points_3d) @@ -316,15 +312,15 @@ def test_se2_se3_pose_conversion_consistency(self) -> None: y = np.random.uniform(-10.0, 10.0) yaw = np.random.uniform(-np.pi, np.pi) - reference_se2 = StateSE2.from_array(np.array([x, y, yaw], dtype=np.float64)) - reference_se3 = EulerStateSE3.from_array(np.array([x, y, 0.0, 0.0, 0.0, yaw], dtype=np.float64)) + reference_se2 = PoseSE2.from_array(np.array([x, y, yaw], dtype=np.float64)) + reference_se3 = EulerPoseSE3.from_array(np.array([x, y, 0.0, 0.0, 0.0, yaw], dtype=np.float64)) # Generate 2D poses and embed them in 3D with z=0 and zero roll/pitch num_poses = np.random.randint(1, 8) pose_2d = self._get_random_se2_array(num_poses) - pose_3d = np.zeros((num_poses, len(EulerStateSE3Index)), dtype=np.float64) - pose_3d[:, EulerStateSE3Index.XY] = pose_2d[:, StateSE2Index.XY] - pose_3d[:, EulerStateSE3Index.YAW] = pose_2d[:, StateSE2Index.YAW] + pose_3d = np.zeros((num_poses, len(EulerPoseSE3Index)), dtype=np.float64) + pose_3d[:, EulerPoseSE3Index.XY] = pose_2d[:, PoseSE2Index.XY] + pose_3d[:, EulerPoseSE3Index.YAW] = pose_2d[:, PoseSE2Index.YAW] # Convert using SE2 functions relative_se2 = convert_absolute_to_relative_se2_array(reference_se2, pose_2d) @@ -336,20 +332,20 @@ def test_se2_se3_pose_conversion_consistency(self) -> None: # Check that SE2 and SE3 conversions are consistent for the x,y components np.testing.assert_array_almost_equal( - relative_se2[:, StateSE2Index.XY], relative_se3[:, EulerStateSE3Index.XY], decimal=self.decimal + relative_se2[:, PoseSE2Index.XY], relative_se3[:, EulerPoseSE3Index.XY], decimal=self.decimal ) np.testing.assert_array_almost_equal( - absolute_se2_recovered[:, StateSE2Index.XY], - absolute_se3_recovered[:, EulerStateSE3Index.XY], + absolute_se2_recovered[:, PoseSE2Index.XY], + absolute_se3_recovered[:, EulerPoseSE3Index.XY], decimal=self.decimal, ) # Check that SE2 and SE3 conversions are consistent for the yaw component np.testing.assert_array_almost_equal( - relative_se2[:, StateSE2Index.YAW], relative_se3[:, EulerStateSE3Index.YAW], decimal=self.decimal + relative_se2[:, PoseSE2Index.YAW], relative_se3[:, EulerPoseSE3Index.YAW], decimal=self.decimal ) np.testing.assert_array_almost_equal( - absolute_se2_recovered[:, StateSE2Index.YAW], - absolute_se3_recovered[:, EulerStateSE3Index.YAW], + absolute_se2_recovered[:, PoseSE2Index.YAW], + absolute_se3_recovered[:, EulerPoseSE3Index.YAW], decimal=self.decimal, ) @@ -379,7 +375,7 @@ def test_se2_array_translation_consistency(self) -> None: # Translate each pose individually result_individual = np.zeros_like(poses_array) for i in range(num_poses): - pose = StateSE2.from_array(poses_array[i]) + pose = PoseSE2.from_array(poses_array[i]) translated = translate_se2_along_body_frame(pose, translation) result_individual[i] = translated.array @@ -387,57 +383,57 @@ def test_se2_array_translation_consistency(self) -> None: def test_transform_empty_arrays(self) -> None: """Test that transform functions handle empty arrays correctly""" - reference_se2 = StateSE2.from_array(np.array([1.0, 2.0, np.pi / 4], dtype=np.float64)) - reference_se3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], dtype=np.float64)) + reference_se2 = PoseSE2.from_array(np.array([1.0, 2.0, np.pi / 4], dtype=np.float64)) + reference_se3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.1, 0.2, 0.3], dtype=np.float64)) # Test SE2 empty arrays - empty_se2_poses = np.array([], dtype=np.float64).reshape(0, len(StateSE2Index)) + empty_se2_poses = np.array([], dtype=np.float64).reshape(0, len(PoseSE2Index)) empty_2d_points = np.array([], dtype=np.float64).reshape(0, len(Point2DIndex)) result_se2_poses = convert_absolute_to_relative_se2_array(reference_se2, empty_se2_poses) - result_2d_points = convert_absolute_to_relative_point_2d_array(reference_se2, empty_2d_points) + result_2d_points = convert_absolute_to_relative_points_2d_array(reference_se2, empty_2d_points) - self.assertEqual(result_se2_poses.shape, (0, len(StateSE2Index))) - self.assertEqual(result_2d_points.shape, (0, len(Point2DIndex))) + assert result_se2_poses.shape == (0, len(PoseSE2Index)) + assert result_2d_points.shape == (0, len(Point2DIndex)) # Test SE3 empty arrays - empty_se3_poses = np.array([], dtype=np.float64).reshape(0, len(EulerStateSE3Index)) + empty_se3_poses = np.array([], dtype=np.float64).reshape(0, len(EulerPoseSE3Index)) empty_3d_points = np.array([], dtype=np.float64).reshape(0, len(Point3DIndex)) result_se3_poses = convert_absolute_to_relative_euler_se3_array(reference_se3, empty_se3_poses) result_3d_points = convert_absolute_to_relative_points_3d_array(reference_se3, empty_3d_points) - self.assertEqual(result_se3_poses.shape, (0, len(EulerStateSE3Index))) - self.assertEqual(result_3d_points.shape, (0, len(Point3DIndex))) + assert result_se3_poses.shape == (0, len(EulerPoseSE3Index)) + assert result_3d_points.shape == (0, len(Point3DIndex)) def test_transform_identity_operations(self) -> None: """Test that transforms with identity reference frames work correctly""" # Identity SE2 pose - identity_se2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) - identity_se3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + identity_se2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + identity_se3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) for _ in range(self.num_consistency_tests): # Test SE2 identity transforms num_poses = np.random.randint(1, 10) se2_poses = self._get_random_se2_array(num_poses) - se2_points = se2_poses[:, StateSE2Index.XY] + se2_points = se2_poses[:, PoseSE2Index.XY] relative_se2_poses = convert_absolute_to_relative_se2_array(identity_se2, se2_poses) - relative_se2_points = convert_absolute_to_relative_point_2d_array(identity_se2, se2_points) + relative_se2_points = convert_absolute_to_relative_points_2d_array(identity_se2, se2_points) np.testing.assert_array_almost_equal(se2_poses, relative_se2_poses, decimal=self.decimal) np.testing.assert_array_almost_equal(se2_points, relative_se2_points, decimal=self.decimal) # Test SE3 identity transforms se3_poses = self._get_random_se3_array(num_poses) - se3_points = se3_poses[:, EulerStateSE3Index.XYZ] + se3_points = se3_poses[:, EulerPoseSE3Index.XYZ] relative_se3_poses = convert_absolute_to_relative_euler_se3_array(identity_se3, se3_poses) relative_se3_points = convert_absolute_to_relative_points_3d_array(identity_se3, se3_points) np.testing.assert_array_almost_equal( - se3_poses[..., EulerStateSE3Index.EULER_ANGLES], - relative_se3_poses[..., EulerStateSE3Index.EULER_ANGLES], + se3_poses[..., EulerPoseSE3Index.EULER_ANGLES], + relative_se3_poses[..., EulerPoseSE3Index.EULER_ANGLES], decimal=self.decimal, ) np.testing.assert_array_almost_equal(se3_points, relative_se3_points, decimal=self.decimal) @@ -449,16 +445,16 @@ def test_transform_large_rotations(self) -> None: large_yaw_se2 = np.random.uniform(-4 * np.pi, 4 * np.pi) large_euler_se3 = np.random.uniform(-4 * np.pi, 4 * np.pi, 3) - reference_se2 = StateSE2.from_array(np.array([0.0, 0.0, large_yaw_se2], dtype=np.float64)) - reference_se3 = EulerStateSE3.from_array( + reference_se2 = PoseSE2.from_array(np.array([0.0, 0.0, large_yaw_se2], dtype=np.float64)) + reference_se3 = EulerPoseSE3.from_array( np.array([0.0, 0.0, 0.0, large_euler_se3[0], large_euler_se3[1], large_euler_se3[2]], dtype=np.float64) ) # Generate test poses/points test_se2_poses = self._get_random_se2_array(5) test_se3_poses = self._get_random_se3_array(5) - test_2d_points = test_se2_poses[:, StateSE2Index.XY] - test_3d_points = test_se3_poses[:, EulerStateSE3Index.XYZ] + test_2d_points = test_se2_poses[:, PoseSE2Index.XY] + test_3d_points = test_se3_poses[:, EulerPoseSE3Index.XYZ] # Test round-trip conversions should still work relative_se2 = convert_absolute_to_relative_se2_array(reference_se2, test_se2_poses) @@ -467,26 +463,22 @@ def test_transform_large_rotations(self) -> None: relative_se3 = convert_absolute_to_relative_euler_se3_array(reference_se3, test_se3_poses) recovered_se3 = convert_relative_to_absolute_euler_se3_array(reference_se3, relative_se3) - relative_2d_points = convert_absolute_to_relative_point_2d_array(reference_se2, test_2d_points) - recovered_2d_points = convert_relative_to_absolute_point_2d_array(reference_se2, relative_2d_points) + relative_2d_points = convert_absolute_to_relative_points_2d_array(reference_se2, test_2d_points) + recovered_2d_points = convert_relative_to_absolute_points_2d_array(reference_se2, relative_2d_points) relative_3d_points = convert_absolute_to_relative_points_3d_array(reference_se3, test_3d_points) recovered_3d_points = convert_relative_to_absolute_points_3d_array(reference_se3, relative_3d_points) # Check consistency (allowing for angle wrapping) np.testing.assert_array_almost_equal( - test_se2_poses[:, StateSE2Index.XY], - recovered_se2[:, StateSE2Index.XY], + test_se2_poses[:, PoseSE2Index.XY], + recovered_se2[:, PoseSE2Index.XY], decimal=self.decimal, ) np.testing.assert_array_almost_equal( - test_se3_poses[:, EulerStateSE3Index.XYZ], - recovered_se3[:, EulerStateSE3Index.XYZ], + test_se3_poses[:, EulerPoseSE3Index.XYZ], + recovered_se3[:, EulerPoseSE3Index.XYZ], decimal=self.decimal, ) np.testing.assert_array_almost_equal(test_2d_points, recovered_2d_points, decimal=self.decimal) np.testing.assert_array_almost_equal(test_3d_points, recovered_3d_points, decimal=self.decimal) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/geometry/transform/test_transform_euler_se3.py b/tests/unit/geometry/transform/test_transform_euler_se3.py index 8b146fd1..1785acef 100644 --- a/tests/unit/geometry/transform/test_transform_euler_se3.py +++ b/tests/unit/geometry/transform/test_transform_euler_se3.py @@ -1,9 +1,7 @@ -import unittest - import numpy as np import numpy.typing as npt -from py123d.geometry import EulerStateSE3, Vector3D +from py123d.geometry import EulerPoseSE3, Vector3D from py123d.geometry.transform.transform_euler_se3 import ( convert_absolute_to_relative_euler_se3_array, convert_absolute_to_relative_points_3d_array, @@ -16,138 +14,137 @@ ) -class TestTransformEulerSE3(unittest.TestCase): - - def setUp(self): +class TestTransformEulerSE3: + def setup_method(self): self.decimal = 6 # Decimal places for np.testing.assert_array_almost_equal self.num_consistency_tests = 10 # Number of random test cases for consistency checks def test_translate_se3_along_x(self) -> None: """Tests translating a SE3 state along the body frame forward direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 1.0 - result: EulerStateSE3 = translate_euler_se3_along_x(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_x(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_x_negative(self) -> None: """Tests translating a SE3 state along the body frame backward direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = -0.5 - result: EulerStateSE3 = translate_euler_se3_along_x(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.5, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_x(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.5, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_x_with_rotation(self) -> None: """Tests translating a SE3 state along the body frame forward direction with yaw rotation.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64)) distance: float = 2.5 - result: EulerStateSE3 = translate_euler_se3_along_x(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array( + result: EulerPoseSE3 = translate_euler_se3_along_x(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([0.0, 2.5, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_y(self) -> None: """Tests translating a SE3 state along the body frame lateral direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 1.0 - result: EulerStateSE3 = translate_euler_se3_along_y(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_y(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 1.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_y_with_existing_position(self) -> None: """Tests translating a SE3 state along the body frame lateral direction with existing position.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 2.5 - result: EulerStateSE3 = translate_euler_se3_along_y(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 4.5, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_y(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 4.5, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_y_negative(self) -> None: """Tests translating a SE3 state along the body frame lateral direction in the negative direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = -1.0 - result: EulerStateSE3 = translate_euler_se3_along_y(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_y(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_y_with_rotation(self) -> None: """Tests translating a SE3 state along the body frame lateral direction with roll rotation.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, np.pi / 2, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, np.pi / 2, 0.0, 0.0], dtype=np.float64)) distance: float = -1.0 - result: EulerStateSE3 = translate_euler_se3_along_y(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array( + result: EulerPoseSE3 = translate_euler_se3_along_y(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([1.0, 2.0, 2.0, np.pi / 2, 0.0, 0.0], dtype=np.float64) ) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_z(self) -> None: """Tests translating a SE3 state along the body frame vertical direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 1.0 - result: EulerStateSE3 = translate_euler_se3_along_z(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_z(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_z_large_distance(self) -> None: """Tests translating a SE3 state along the body frame vertical direction with a large distance.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 10.0 - result: EulerStateSE3 = translate_euler_se3_along_z(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 15.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_z(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 15.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_z_negative(self) -> None: """Tests translating a SE3 state along the body frame vertical direction in the negative direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) distance: float = -2.0 - result: EulerStateSE3 = translate_euler_se3_along_z(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_z(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_z_with_rotation(self) -> None: """Tests translating a SE3 state along the body frame vertical direction with pitch rotation.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, np.pi / 2, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, np.pi / 2, 0.0], dtype=np.float64)) distance: float = 2.0 - result: EulerStateSE3 = translate_euler_se3_along_z(pose, distance) - expected: EulerStateSE3 = EulerStateSE3.from_array( + result: EulerPoseSE3 = translate_euler_se3_along_z(pose, distance) + expected: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([3.0, 2.0, 3.0, 0.0, np.pi / 2, 0.0], dtype=np.float64) ) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_body_frame(self) -> None: """Tests translating a SE3 state along the body frame forward direction.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) translation: Vector3D = Vector3D.from_array(np.array([1.0, 0.0, 0.0], dtype=np.float64)) - result: EulerStateSE3 = translate_euler_se3_along_body_frame(pose, translation) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_body_frame(pose, translation) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_body_frame_multiple_axes(self) -> None: """Tests translating a SE3 state along the body frame in multiple axes.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) translation: Vector3D = Vector3D.from_array(np.array([0.5, -1.0, 2.0], dtype=np.float64)) - result: EulerStateSE3 = translate_euler_se3_along_body_frame(pose, translation) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.5, 1.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_body_frame(pose, translation) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.5, 1.0, 5.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_body_frame_zero_translation(self) -> None: """Tests translating a SE3 state along the body frame with zero translation.""" - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) translation: Vector3D = Vector3D.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) - result: EulerStateSE3 = translate_euler_se3_along_body_frame(pose, translation) - expected: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) + result: EulerPoseSE3 = translate_euler_se3_along_body_frame(pose, translation) + expected: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 2.0, 3.0, 0.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array) def test_translate_se3_along_body_frame_with_rotation(self) -> None: """Tests translating a SE3 state along the body frame forward direction with yaw rotation.""" # Rotate 90 degrees around z-axis, then translate 1 unit along body x-axis - pose: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64)) + pose: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64)) translation: Vector3D = Vector3D.from_array(np.array([1.0, 0.0, 0.0], dtype=np.float64)) - result: EulerStateSE3 = translate_euler_se3_along_body_frame(pose, translation) + result: EulerPoseSE3 = translate_euler_se3_along_body_frame(pose, translation) # Should move in +Y direction in world frame - expected: EulerStateSE3 = EulerStateSE3.from_array( + expected: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([0.0, 1.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) @@ -169,7 +166,7 @@ def test_translate_se3_along_body_frame_consistency(self) -> None: start_pitch: float = np.random.uniform(-np.pi, np.pi) start_yaw: float = np.random.uniform(-np.pi, np.pi) - original_pose: EulerStateSE3 = EulerStateSE3.from_array( + original_pose: EulerPoseSE3 = EulerPoseSE3.from_array( np.array( [ start_x, @@ -185,37 +182,37 @@ def test_translate_se3_along_body_frame_consistency(self) -> None: # x-axis translation translation_x: Vector3D = Vector3D.from_array(np.array([x_distance, 0.0, 0.0], dtype=np.float64)) - result_body_frame_x: EulerStateSE3 = translate_euler_se3_along_body_frame(original_pose, translation_x) - result_axis_x: EulerStateSE3 = translate_euler_se3_along_x(original_pose, x_distance) + result_body_frame_x: EulerPoseSE3 = translate_euler_se3_along_body_frame(original_pose, translation_x) + result_axis_x: EulerPoseSE3 = translate_euler_se3_along_x(original_pose, x_distance) np.testing.assert_array_almost_equal(result_body_frame_x.array, result_axis_x.array, decimal=self.decimal) # y-axis translation translation_y: Vector3D = Vector3D.from_array(np.array([0.0, y_distance, 0.0], dtype=np.float64)) - result_body_frame_y: EulerStateSE3 = translate_euler_se3_along_body_frame(original_pose, translation_y) - result_axis_y: EulerStateSE3 = translate_euler_se3_along_y(original_pose, y_distance) + result_body_frame_y: EulerPoseSE3 = translate_euler_se3_along_body_frame(original_pose, translation_y) + result_axis_y: EulerPoseSE3 = translate_euler_se3_along_y(original_pose, y_distance) np.testing.assert_array_almost_equal(result_body_frame_y.array, result_axis_y.array, decimal=self.decimal) # z-axis translation translation_z: Vector3D = Vector3D.from_array(np.array([0.0, 0.0, z_distance], dtype=np.float64)) - result_body_frame_z: EulerStateSE3 = translate_euler_se3_along_body_frame(original_pose, translation_z) - result_axis_z: EulerStateSE3 = translate_euler_se3_along_z(original_pose, z_distance) + result_body_frame_z: EulerPoseSE3 = translate_euler_se3_along_body_frame(original_pose, translation_z) + result_axis_z: EulerPoseSE3 = translate_euler_se3_along_z(original_pose, z_distance) np.testing.assert_array_almost_equal(result_body_frame_z.array, result_axis_z.array, decimal=self.decimal) # all axes translation translation_all: Vector3D = Vector3D.from_array( np.array([x_distance, y_distance, z_distance], dtype=np.float64) ) - result_body_frame_all: EulerStateSE3 = translate_euler_se3_along_body_frame(original_pose, translation_all) - intermediate_pose: EulerStateSE3 = translate_euler_se3_along_x(original_pose, x_distance) + result_body_frame_all: EulerPoseSE3 = translate_euler_se3_along_body_frame(original_pose, translation_all) + intermediate_pose: EulerPoseSE3 = translate_euler_se3_along_x(original_pose, x_distance) intermediate_pose = translate_euler_se3_along_y(intermediate_pose, y_distance) - result_axis_all: EulerStateSE3 = translate_euler_se3_along_z(intermediate_pose, z_distance) + result_axis_all: EulerPoseSE3 = translate_euler_se3_along_z(intermediate_pose, z_distance) np.testing.assert_array_almost_equal( result_body_frame_all.array, result_axis_all.array, decimal=self.decimal ) def test_convert_absolute_to_relative_se3_array(self) -> None: """Tests converting absolute SE3 poses to relative SE3 poses.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) absolute_poses: npt.NDArray[np.float64] = np.array( [ [2.0, 2.0, 2.0, 0.0, 0.0, 0.0], @@ -235,7 +232,7 @@ def test_convert_absolute_to_relative_se3_array(self) -> None: def test_convert_absolute_to_relative_se3_array_single_pose(self) -> None: """Tests converting a single absolute SE3 pose to a relative SE3 pose.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) absolute_poses: npt.NDArray[np.float64] = np.array([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_euler_se3_array(reference, absolute_poses) expected: npt.NDArray[np.float64] = np.array([[1.0, 2.0, 3.0, 0.0, 0.0, 0.0]], dtype=np.float64) @@ -243,7 +240,7 @@ def test_convert_absolute_to_relative_se3_array_single_pose(self) -> None: def test_convert_absolute_to_relative_se3_array_with_rotation(self) -> None: """Tests converting absolute SE3 poses to relative SE3 poses with 90 degree yaw rotation.""" - reference: EulerStateSE3 = EulerStateSE3.from_array( + reference: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) absolute_poses: npt.NDArray[np.float64] = np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=np.float64) @@ -253,7 +250,7 @@ def test_convert_absolute_to_relative_se3_array_with_rotation(self) -> None: def test_convert_relative_to_absolute_se3_array(self) -> None: """Tests converting relative SE3 poses to absolute SE3 poses.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) relative_poses: npt.NDArray[np.float64] = np.array( [ [1.0, 1.0, 1.0, 0.0, 0.0, 0.0], @@ -273,7 +270,7 @@ def test_convert_relative_to_absolute_se3_array(self) -> None: def test_convert_relative_to_absolute_se3_array_with_rotation(self) -> None: """Tests converting relative SE3 poses to absolute SE3 poses with 90 degree yaw rotation.""" - reference: EulerStateSE3 = EulerStateSE3.from_array( + reference: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([1.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) relative_poses: npt.NDArray[np.float64] = np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]], dtype=np.float64) @@ -283,7 +280,7 @@ def test_convert_relative_to_absolute_se3_array_with_rotation(self) -> None: def test_convert_absolute_to_relative_points_3d_array(self) -> None: """Tests converting absolute 3D points to relative 3D points.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) absolute_points: npt.NDArray[np.float64] = np.array([[2.0, 2.0, 2.0], [0.0, 1.0, 0.0]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_points_3d_array(reference, absolute_points) expected: npt.NDArray[np.float64] = np.array([[1.0, 1.0, 1.0], [-1.0, 0.0, -1.0]], dtype=np.float64) @@ -291,7 +288,7 @@ def test_convert_absolute_to_relative_points_3d_array(self) -> None: def test_convert_absolute_to_relative_points_3d_array_origin_reference(self) -> None: """Tests converting absolute 3D points to relative 3D points with origin reference.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0], dtype=np.float64)) absolute_points: npt.NDArray[np.float64] = np.array([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_points_3d_array(reference, absolute_points) expected: npt.NDArray[np.float64] = np.array([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]], dtype=np.float64) @@ -299,7 +296,7 @@ def test_convert_absolute_to_relative_points_3d_array_origin_reference(self) -> def test_convert_absolute_to_relative_points_3d_array_with_rotation(self) -> None: """Tests converting absolute 3D points to relative 3D points with 90 degree yaw rotation.""" - reference: EulerStateSE3 = EulerStateSE3.from_array( + reference: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([0.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) absolute_points: npt.NDArray[np.float64] = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 1.0]], dtype=np.float64) @@ -309,7 +306,7 @@ def test_convert_absolute_to_relative_points_3d_array_with_rotation(self) -> Non def test_convert_relative_to_absolute_points_3d_array(self) -> None: """Tests converting relative 3D points to absolute 3D points.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) relative_points: npt.NDArray[np.float64] = np.array([[1.0, 1.0, 1.0], [-1.0, 0.0, -1.0]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_relative_to_absolute_points_3d_array(reference, relative_points) expected: npt.NDArray[np.float64] = np.array([[2.0, 2.0, 2.0], [0.0, 1.0, 0.0]], dtype=np.float64) @@ -317,7 +314,7 @@ def test_convert_relative_to_absolute_points_3d_array(self) -> None: def test_convert_relative_to_absolute_points_3d_array_empty(self) -> None: """Tests converting an empty array of relative 3D points to absolute 3D points.""" - reference: EulerStateSE3 = EulerStateSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) + reference: EulerPoseSE3 = EulerPoseSE3.from_array(np.array([1.0, 1.0, 1.0, 0.0, 0.0, 0.0], dtype=np.float64)) relative_points: npt.NDArray[np.float64] = np.array([], dtype=np.float64).reshape(0, 3) result: npt.NDArray[np.float64] = convert_relative_to_absolute_points_3d_array(reference, relative_points) expected: npt.NDArray[np.float64] = np.array([], dtype=np.float64).reshape(0, 3) @@ -325,7 +322,7 @@ def test_convert_relative_to_absolute_points_3d_array_empty(self) -> None: def test_convert_relative_to_absolute_points_3d_array_with_rotation(self) -> None: """Tests converting relative 3D points to absolute 3D points with 90 degree yaw rotation.""" - reference: EulerStateSE3 = EulerStateSE3.from_array( + reference: EulerPoseSE3 = EulerPoseSE3.from_array( np.array([1.0, 0.0, 0.0, 0.0, 0.0, np.pi / 2], dtype=np.float64) ) relative_points: npt.NDArray[np.float64] = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 1.0]], dtype=np.float64) diff --git a/tests/unit/geometry/transform/test_transform_se2.py b/tests/unit/geometry/transform/test_transform_se2.py index 60af633e..ec8836aa 100644 --- a/tests/unit/geometry/transform/test_transform_se2.py +++ b/tests/unit/geometry/transform/test_transform_se2.py @@ -1,14 +1,14 @@ -import unittest - import numpy as np import numpy.typing as npt -from py123d.geometry import StateSE2, Vector2D -from py123d.geometry.transform.transform_se2 import ( - convert_absolute_to_relative_point_2d_array, +from py123d.geometry import PoseSE2, PoseSE2Index, Vector2D +from py123d.geometry.transform import ( + convert_absolute_to_relative_points_2d_array, convert_absolute_to_relative_se2_array, - convert_relative_to_absolute_point_2d_array, + convert_points_2d_array_between_origins, + convert_relative_to_absolute_points_2d_array, convert_relative_to_absolute_se2_array, + convert_se2_array_between_origins, translate_se2_along_body_frame, translate_se2_along_x, translate_se2_along_y, @@ -16,82 +16,89 @@ ) -class TestTransformSE2(unittest.TestCase): - - def setUp(self): +class TestTransformSE2: + def setup_method(self): self.decimal = 6 # Decimal places for np.testing.assert_array_almost_equal + def _get_random_se2_array(self, num_poses: int) -> npt.NDArray[np.float64]: + """Generates a random SE2 array for testing.""" + x = np.random.uniform(-10.0, 10.0, size=(num_poses,)) + y = np.random.uniform(-10.0, 10.0, size=(num_poses,)) + yaw = np.random.uniform(-np.pi, np.pi, size=(num_poses,)) + se2_array = np.stack((x, y, yaw), axis=-1) + return se2_array + def test_translate_se2_along_x(self) -> None: """Tests translating a SE2 state along the X-axis.""" - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 1.0 - result: StateSE2 = translate_se2_along_x(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([1.0, 0.0, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_x(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([1.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_x_negative(self) -> None: """Tests translating a SE2 state along the X-axis in the negative direction.""" - pose: StateSE2 = StateSE2.from_array(np.array([1.0, 2.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([1.0, 2.0, 0.0], dtype=np.float64)) distance: float = -0.5 - result: StateSE2 = translate_se2_along_x(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([0.5, 2.0, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_x(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([0.5, 2.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_x_with_rotation(self) -> None: """Tests translating a SE2 state along the X-axis with 90 degree rotation.""" - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) distance: float = 1.0 - result: StateSE2 = translate_se2_along_x(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([0.0, 1.0, np.pi / 2], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_x(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([0.0, 1.0, np.pi / 2], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_y(self) -> None: """Tests translating a SE2 state along the Y-axis.""" - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) distance: float = 1.0 - result: StateSE2 = translate_se2_along_y(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([0.0, 1.0, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_y(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([0.0, 1.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_y_negative(self) -> None: """Tests translating a SE2 state along the Y-axis in the negative direction.""" - pose: StateSE2 = StateSE2.from_array(np.array([1.0, 2.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([1.0, 2.0, 0.0], dtype=np.float64)) distance: float = -1.5 - result: StateSE2 = translate_se2_along_y(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([1.0, 0.5, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_y(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([1.0, 0.5, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_y_with_rotation(self) -> None: """Tests translating a SE2 state along the Y-axis with -90 degree rotation.""" - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, -np.pi / 2], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, -np.pi / 2], dtype=np.float64)) distance: float = 2.0 - result: StateSE2 = translate_se2_along_y(pose, distance) - expected: StateSE2 = StateSE2.from_array(np.array([2.0, 0.0, -np.pi / 2], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_y(pose, distance) + expected: PoseSE2 = PoseSE2.from_array(np.array([2.0, 0.0, -np.pi / 2], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_body_frame_forward(self) -> None: """Tests translating a SE2 state along the body frame forward direction, with 90 degree rotation.""" # Move 1 unit forward in the direction of yaw (pi/2 = 90 degrees = +Y direction) - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) vector: Vector2D = Vector2D(1.0, 0.0) - result: StateSE2 = translate_se2_along_body_frame(pose, vector) - expected: StateSE2 = StateSE2.from_array(np.array([0.0, 1.0, np.pi / 2], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_body_frame(pose, vector) + expected: PoseSE2 = PoseSE2.from_array(np.array([0.0, 1.0, np.pi / 2], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_body_frame_backward(self) -> None: """Tests translating a SE2 state along the body frame backward direction.""" - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) vector: Vector2D = Vector2D(-1.0, 0.0) - result: StateSE2 = translate_se2_along_body_frame(pose, vector) - expected: StateSE2 = StateSE2.from_array(np.array([-1.0, 0.0, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_body_frame(pose, vector) + expected: PoseSE2 = PoseSE2.from_array(np.array([-1.0, 0.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_body_frame_diagonal(self) -> None: """Tests translating a SE2 state along the body frame diagonal direction.""" - pose: StateSE2 = StateSE2.from_array(np.array([1.0, 0.0, np.deg2rad(45)], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([1.0, 0.0, np.deg2rad(45)], dtype=np.float64)) vector: Vector2D = Vector2D(1.0, 0.0) - result: StateSE2 = translate_se2_along_body_frame(pose, vector) - expected: StateSE2 = StateSE2.from_array( + result: PoseSE2 = translate_se2_along_body_frame(pose, vector) + expected: PoseSE2 = PoseSE2.from_array( np.array([1.0 + np.sqrt(2.0) / 2, 0.0 + np.sqrt(2.0) / 2, np.deg2rad(45)], dtype=np.float64) ) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) @@ -99,19 +106,19 @@ def test_translate_se2_along_body_frame_diagonal(self) -> None: def test_translate_se2_along_body_frame_lateral(self) -> None: """Tests translating a SE2 state along the body frame lateral direction.""" # Move 1 unit to the right (positive y in body frame) - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) vector: Vector2D = Vector2D(0.0, 1.0) - result: StateSE2 = translate_se2_along_body_frame(pose, vector) - expected: StateSE2 = StateSE2.from_array(np.array([0.0, 1.0, 0.0], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_body_frame(pose, vector) + expected: PoseSE2 = PoseSE2.from_array(np.array([0.0, 1.0, 0.0], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_along_body_frame_lateral_with_rotation(self) -> None: """Tests translating a SE2 state along the body frame lateral direction with 90 degree rotation.""" # Move 1 unit to the right when facing 90 degrees - pose: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) + pose: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) vector: Vector2D = Vector2D(0.0, 1.0) - result: StateSE2 = translate_se2_along_body_frame(pose, vector) - expected: StateSE2 = StateSE2.from_array(np.array([-1.0, 0.0, np.pi / 2], dtype=np.float64)) + result: PoseSE2 = translate_se2_along_body_frame(pose, vector) + expected: PoseSE2 = PoseSE2.from_array(np.array([-1.0, 0.0, np.pi / 2], dtype=np.float64)) np.testing.assert_array_almost_equal(result.array, expected.array, decimal=self.decimal) def test_translate_se2_array_along_body_frame_single_distance(self) -> None: @@ -140,7 +147,7 @@ def test_translate_se2_array_along_body_frame_lateral(self) -> None: def test_convert_absolute_to_relative_se2_array(self) -> None: """Tests converting absolute SE2 poses to relative SE2 poses.""" - origin: StateSE2 = StateSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) + origin: PoseSE2 = PoseSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) absolute_poses: npt.NDArray[np.float64] = np.array([[2.0, 2.0, 0.0], [0.0, 1.0, np.pi / 2]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_se2_array(origin, absolute_poses) expected: npt.NDArray[np.float64] = np.array([[1.0, 1.0, 0.0], [-1.0, 0.0, np.pi / 2]], dtype=np.float64) @@ -148,7 +155,7 @@ def test_convert_absolute_to_relative_se2_array(self) -> None: def test_convert_absolute_to_relative_se2_array_with_rotation(self) -> None: """Tests converting absolute SE2 poses to relative SE2 poses with 90 degree rotation.""" - reference: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) absolute_poses: npt.NDArray[np.float64] = np.array([[1.0, 0.0, np.pi / 2]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_se2_array(reference, absolute_poses) expected: npt.NDArray[np.float64] = np.array([[0.0, -1.0, 0.0]], dtype=np.float64) @@ -156,7 +163,7 @@ def test_convert_absolute_to_relative_se2_array_with_rotation(self) -> None: def test_convert_absolute_to_relative_se2_array_identity(self) -> None: """Tests converting absolute SE2 poses to relative SE2 poses with identity transformation.""" - reference: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, 0.0], dtype=np.float64)) absolute_poses: npt.NDArray[np.float64] = np.array([[1.0, 2.0, np.pi / 4]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_absolute_to_relative_se2_array(reference, absolute_poses) expected: npt.NDArray[np.float64] = np.array([[1.0, 2.0, np.pi / 4]], dtype=np.float64) @@ -164,7 +171,7 @@ def test_convert_absolute_to_relative_se2_array_identity(self) -> None: def test_convert_relative_to_absolute_se2_array(self) -> None: """Tests converting relative SE2 poses to absolute SE2 poses.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) relative_poses: npt.NDArray[np.float64] = np.array([[1.0, 1.0, 0.0], [-1.0, 0.0, np.pi / 2]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_relative_to_absolute_se2_array(reference, relative_poses) expected: npt.NDArray[np.float64] = np.array([[2.0, 2.0, 0.0], [0.0, 1.0, np.pi / 2]], dtype=np.float64) @@ -172,7 +179,7 @@ def test_convert_relative_to_absolute_se2_array(self) -> None: def test_convert_relative_to_absolute_se2_array_with_rotation(self) -> None: """Tests converting relative SE2 poses to absolute SE2 poses with rotation.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 0.0, np.pi / 2], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 0.0, np.pi / 2], dtype=np.float64)) relative_poses: npt.NDArray[np.float64] = np.array([[1.0, 0.0, 0.0]], dtype=np.float64) result: npt.NDArray[np.float64] = convert_relative_to_absolute_se2_array(reference, relative_poses) expected: npt.NDArray[np.float64] = np.array([[1.0, 1.0, np.pi / 2]], dtype=np.float64) @@ -180,40 +187,97 @@ def test_convert_relative_to_absolute_se2_array_with_rotation(self) -> None: def test_convert_absolute_to_relative_point_2d_array(self) -> None: """Tests converting absolute 2D points to relative 2D points.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) absolute_points: npt.NDArray[np.float64] = np.array([[2.0, 2.0], [0.0, 1.0]], dtype=np.float64) - result: npt.NDArray[np.float64] = convert_absolute_to_relative_point_2d_array(reference, absolute_points) + result: npt.NDArray[np.float64] = convert_absolute_to_relative_points_2d_array(reference, absolute_points) expected: npt.NDArray[np.float64] = np.array([[1.0, 1.0], [-1.0, 0.0]], dtype=np.float64) np.testing.assert_array_almost_equal(result, expected, decimal=self.decimal) def test_convert_absolute_to_relative_point_2d_array_with_rotation(self) -> None: """Tests converting absolute 2D points to relative 2D points with 90 degree rotation.""" - reference: StateSE2 = StateSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([0.0, 0.0, np.pi / 2], dtype=np.float64)) absolute_points: npt.NDArray[np.float64] = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.float64) - result: npt.NDArray[np.float64] = convert_absolute_to_relative_point_2d_array(reference, absolute_points) + result: npt.NDArray[np.float64] = convert_absolute_to_relative_points_2d_array(reference, absolute_points) expected: npt.NDArray[np.float64] = np.array([[1.0, 0.0], [0.0, -1.0]], dtype=np.float64) np.testing.assert_array_almost_equal(result, expected, decimal=self.decimal) def test_convert_absolute_to_relative_point_2d_array_empty(self) -> None: """Tests converting an empty array of absolute 2D points to relative 2D points.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) absolute_points: npt.NDArray[np.float64] = np.array([], dtype=np.float64).reshape(0, 2) - result: npt.NDArray[np.float64] = convert_absolute_to_relative_point_2d_array(reference, absolute_points) + result: npt.NDArray[np.float64] = convert_absolute_to_relative_points_2d_array(reference, absolute_points) expected: npt.NDArray[np.float64] = np.array([], dtype=np.float64).reshape(0, 2) np.testing.assert_array_almost_equal(result, expected, decimal=self.decimal) def test_convert_relative_to_absolute_point_2d_array(self) -> None: """Tests converting relative 2D points to absolute 2D points.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 1.0, 0.0], dtype=np.float64)) relative_points: npt.NDArray[np.float64] = np.array([[1.0, 1.0], [-1.0, 0.0]], dtype=np.float64) - result: npt.NDArray[np.float64] = convert_relative_to_absolute_point_2d_array(reference, relative_points) + result: npt.NDArray[np.float64] = convert_relative_to_absolute_points_2d_array(reference, relative_points) expected: npt.NDArray[np.float64] = np.array([[2.0, 2.0], [0.0, 1.0]], dtype=np.float64) np.testing.assert_array_almost_equal(result, expected, decimal=self.decimal) def test_convert_relative_to_absolute_point_2d_array_with_rotation(self) -> None: """Tests converting relative 2D points to absolute 2D points with 90 degree rotation.""" - reference: StateSE2 = StateSE2.from_array(np.array([1.0, 0.0, np.pi / 2], dtype=np.float64)) + reference: PoseSE2 = PoseSE2.from_array(np.array([1.0, 0.0, np.pi / 2], dtype=np.float64)) relative_points: npt.NDArray[np.float64] = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float64) - result: npt.NDArray[np.float64] = convert_relative_to_absolute_point_2d_array(reference, relative_points) + result: npt.NDArray[np.float64] = convert_relative_to_absolute_points_2d_array(reference, relative_points) expected: npt.NDArray[np.float64] = np.array([[1.0, 1.0], [0.0, 0.0]], dtype=np.float64) np.testing.assert_array_almost_equal(result, expected, decimal=self.decimal) + + def test_convert_points_2d_array_between_origins(self): + random_points_2d = np.random.rand(10, 2) + for _ in range(10): + from_se2 = PoseSE2.from_array(self._get_random_se2_array(1)[0]) + to_se2 = PoseSE2.from_array(self._get_random_se2_array(1)[0]) + + identity_se2_array = np.zeros(len(PoseSE2Index), dtype=np.float64) + identity_se2 = PoseSE2.from_array(identity_se2_array) + + # Check if consistent with absolute-relative-absolute conversion + converted_points_quat = convert_points_2d_array_between_origins(from_se2, to_se2, random_points_2d) + abs_from_se2 = convert_relative_to_absolute_points_2d_array(from_se2, random_points_2d) + rel_to_se2 = convert_absolute_to_relative_points_2d_array(to_se2, abs_from_se2) + np.testing.assert_allclose(converted_points_quat, rel_to_se2, atol=1e-6) + + # Check if consistent with absolute conversion to identity origin + absolute_se2 = convert_points_2d_array_between_origins(from_se2, identity_se2, random_points_2d) + np.testing.assert_allclose( + absolute_se2[..., PoseSE2Index.XY], + abs_from_se2[..., PoseSE2Index.XY], + atol=1e-6, + ) + + def test_convert_se2_array_between_origins(self): + for _ in range(10): + random_se2_array = self._get_random_se2_array(np.random.randint(1, 10)) + + from_se2 = PoseSE2.from_array(self._get_random_se2_array(1)[0]) + to_se2 = PoseSE2.from_array(self._get_random_se2_array(1)[0]) + identity_se2_array = np.zeros(len(PoseSE2Index), dtype=np.float64) + identity_se2 = PoseSE2.from_array(identity_se2_array) + + # Check if consistent with absolute-relative-absolute conversion + converted_se2 = convert_se2_array_between_origins(from_se2, to_se2, random_se2_array) + + abs_from_se2 = convert_relative_to_absolute_se2_array(from_se2, random_se2_array) + rel_to_se2 = convert_absolute_to_relative_se2_array(to_se2, abs_from_se2) + + np.testing.assert_allclose( + converted_se2[..., PoseSE2Index.XY], + rel_to_se2[..., PoseSE2Index.XY], + atol=1e-6, + ) + np.testing.assert_allclose( + converted_se2[..., PoseSE2Index.YAW], + rel_to_se2[..., PoseSE2Index.YAW], + atol=1e-6, + ) + + # Check if consistent with absolute conversion to identity origin + absolute_se2 = convert_se2_array_between_origins(from_se2, identity_se2, random_se2_array) + np.testing.assert_allclose( + absolute_se2[..., PoseSE2Index.XY], + abs_from_se2[..., PoseSE2Index.XY], + atol=1e-6, + ) diff --git a/tests/unit/geometry/transform/test_transform_se3.py b/tests/unit/geometry/transform/test_transform_se3.py index 44e86504..f425e066 100644 --- a/tests/unit/geometry/transform/test_transform_se3.py +++ b/tests/unit/geometry/transform/test_transform_se3.py @@ -1,10 +1,8 @@ -import unittest - import numpy as np import numpy.typing as npt import py123d.geometry.transform.transform_euler_se3 as euler_transform_se3 -from py123d.geometry import EulerStateSE3, EulerStateSE3Index, Point3D, StateSE3, StateSE3Index +from py123d.geometry import EulerPoseSE3, EulerPoseSE3Index, Point3D, PoseSE3, PoseSE3Index from py123d.geometry.transform.transform_se3 import ( convert_absolute_to_relative_points_3d_array, convert_absolute_to_relative_se3_array, @@ -24,10 +22,9 @@ ) -class TestTransformSE3(unittest.TestCase): - - def setUp(self): - euler_se3_a = EulerStateSE3( +class TestTransformSE3: + def setup_method(self): + euler_se3_a = EulerPoseSE3( x=1.0, y=2.0, z=3.0, @@ -35,7 +32,7 @@ def setUp(self): pitch=0.0, yaw=0.0, ) - euler_se3_b = EulerStateSE3( + euler_se3_b = EulerPoseSE3( x=1.0, y=-2.0, z=3.0, @@ -43,7 +40,7 @@ def setUp(self): pitch=np.deg2rad(90), yaw=0.0, ) - euler_se3_c = EulerStateSE3( + euler_se3_c = EulerPoseSE3( x=-1.0, y=2.0, z=-3.0, @@ -52,9 +49,9 @@ def setUp(self): yaw=np.deg2rad(90), ) - quat_se3_a: StateSE3 = euler_se3_a.state_se3 - quat_se3_b: StateSE3 = euler_se3_b.state_se3 - quat_se3_c: StateSE3 = euler_se3_c.state_se3 + quat_se3_a: PoseSE3 = euler_se3_a.pose_se3 + quat_se3_b: PoseSE3 = euler_se3_b.pose_se3 + quat_se3_c: PoseSE3 = euler_se3_c.pose_se3 self.euler_se3 = [euler_se3_a, euler_se3_b, euler_se3_c] self.quat_se3 = [quat_se3_a, quat_se3_b, quat_se3_c] @@ -63,13 +60,11 @@ def setUp(self): def _get_random_euler_se3_array(self, size: int) -> npt.NDArray[np.float64]: """Generate a random SE3 poses""" - random_se3_array = np.zeros((size, len(EulerStateSE3Index)), dtype=np.float64) - random_se3_array[:, EulerStateSE3Index.XYZ] = np.random.uniform( - -self.max_pose_xyz, self.max_pose_xyz, (size, 3) - ) - random_se3_array[:, EulerStateSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) - random_se3_array[:, EulerStateSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) - random_se3_array[:, EulerStateSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array = np.zeros((size, len(EulerPoseSE3Index)), dtype=np.float64) + random_se3_array[:, EulerPoseSE3Index.XYZ] = np.random.uniform(-self.max_pose_xyz, self.max_pose_xyz, (size, 3)) + random_se3_array[:, EulerPoseSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array[:, EulerPoseSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) + random_se3_array[:, EulerPoseSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) return random_se3_array @@ -77,10 +72,10 @@ def _convert_euler_se3_array_to_quat_se3_array( self, euler_se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: """Convert an array of SE3 poses from Euler angles to Quaternion representation""" - quat_se3_array = np.zeros((euler_se3_array.shape[0], len(StateSE3Index)), dtype=np.float64) - quat_se3_array[:, StateSE3Index.XYZ] = euler_se3_array[:, EulerStateSE3Index.XYZ] - quat_se3_array[:, StateSE3Index.QUATERNION] = get_quaternion_array_from_euler_array( - euler_se3_array[:, EulerStateSE3Index.EULER_ANGLES] + quat_se3_array = np.zeros((euler_se3_array.shape[0], len(PoseSE3Index)), dtype=np.float64) + quat_se3_array[:, PoseSE3Index.XYZ] = euler_se3_array[:, EulerPoseSE3Index.XYZ] + quat_se3_array[:, PoseSE3Index.QUATERNION] = get_quaternion_array_from_euler_array( + euler_se3_array[:, EulerPoseSE3Index.EULER_ANGLES] ) return quat_se3_array @@ -92,17 +87,16 @@ def _get_random_quat_se3_array(self, size: int) -> npt.NDArray[np.float64]: def test_sanity(self): for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): - for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): - np.testing.assert_allclose( - quat_se3.point_3d.array, - euler_se3.point_3d.array, - atol=1e-6, - ) - np.testing.assert_allclose( - quat_se3.rotation_matrix, - euler_se3.rotation_matrix, - atol=1e-6, - ) + np.testing.assert_allclose( + quat_se3.point_3d.array, + euler_se3.point_3d.array, + atol=1e-6, + ) + np.testing.assert_allclose( + quat_se3.rotation_matrix, + euler_se3.rotation_matrix, + atol=1e-6, + ) def test_random_sanity(self): for _ in range(10): @@ -110,20 +104,19 @@ def test_random_sanity(self): random_quat_se3_array = self._convert_euler_se3_array_to_quat_se3_array(random_euler_se3_array) np.testing.assert_allclose( - random_euler_se3_array[:, EulerStateSE3Index.XYZ], - random_quat_se3_array[:, StateSE3Index.XYZ], + random_euler_se3_array[:, EulerPoseSE3Index.XYZ], + random_quat_se3_array[:, PoseSE3Index.XYZ], atol=1e-6, ) quat_rotation_matrices = get_rotation_matrices_from_quaternion_array( - random_quat_se3_array[:, StateSE3Index.QUATERNION] + random_quat_se3_array[:, PoseSE3Index.QUATERNION] ) euler_rotation_matrices = get_rotation_matrices_from_euler_array( - random_euler_se3_array[:, EulerStateSE3Index.EULER_ANGLES] + random_euler_se3_array[:, EulerPoseSE3Index.EULER_ANGLES] ) np.testing.assert_allclose(euler_rotation_matrices, quat_rotation_matrices, atol=1e-6) def test_convert_absolute_to_relative_points_3d_array(self): - random_points_3d = np.random.rand(10, 3) for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): rel_points_quat = convert_absolute_to_relative_points_3d_array(quat_se3, random_points_3d) @@ -133,7 +126,6 @@ def test_convert_absolute_to_relative_points_3d_array(self): np.testing.assert_allclose(rel_points_quat, rel_points_euler, atol=1e-6) def test_convert_absolute_to_relative_se3_array(self): - for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): random_euler_se3_array = self._get_random_euler_se3_array(np.random.randint(1, 10)) random_quat_se3_array = self._convert_euler_se3_array_to_quat_se3_array(random_euler_se3_array) @@ -143,19 +135,18 @@ def test_convert_absolute_to_relative_se3_array(self): euler_se3, random_euler_se3_array ) np.testing.assert_allclose( - rel_se3_euler[..., EulerStateSE3Index.XYZ], rel_se3_quat[..., StateSE3Index.XYZ], atol=1e-6 + rel_se3_euler[..., EulerPoseSE3Index.XYZ], rel_se3_quat[..., PoseSE3Index.XYZ], atol=1e-6 ) # We compare rotation matrices to avoid issues with quaternion sign ambiguity quat_rotation_matrices = get_rotation_matrices_from_quaternion_array( - rel_se3_quat[..., StateSE3Index.QUATERNION] + rel_se3_quat[..., PoseSE3Index.QUATERNION] ) euler_rotation_matrices = get_rotation_matrices_from_euler_array( - rel_se3_euler[..., EulerStateSE3Index.EULER_ANGLES] + rel_se3_euler[..., EulerPoseSE3Index.EULER_ANGLES] ) np.testing.assert_allclose(quat_rotation_matrices, euler_rotation_matrices, atol=1e-6) def test_convert_relative_to_absolute_points_3d_array(self): - random_points_3d = np.random.rand(10, 3) for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): rel_points_quat = convert_relative_to_absolute_points_3d_array(quat_se3, random_points_3d) @@ -165,7 +156,6 @@ def test_convert_relative_to_absolute_points_3d_array(self): np.testing.assert_allclose(rel_points_quat, rel_points_euler, atol=1e-6) def test_convert_relative_to_absolute_se3_array(self): - for quat_se3, euler_se3 in zip(self.quat_se3, self.euler_se3): random_euler_se3_array = self._get_random_euler_se3_array(np.random.randint(1, 10)) random_quat_se3_array = self._convert_euler_se3_array_to_quat_se3_array(random_euler_se3_array) @@ -175,15 +165,15 @@ def test_convert_relative_to_absolute_se3_array(self): euler_se3, random_euler_se3_array ) np.testing.assert_allclose( - abs_se3_euler[..., EulerStateSE3Index.XYZ], abs_se3_quat[..., StateSE3Index.XYZ], atol=1e-6 + abs_se3_euler[..., EulerPoseSE3Index.XYZ], abs_se3_quat[..., PoseSE3Index.XYZ], atol=1e-6 ) # We compare rotation matrices to avoid issues with quaternion sign ambiguity quat_rotation_matrices = get_rotation_matrices_from_quaternion_array( - abs_se3_quat[..., StateSE3Index.QUATERNION] + abs_se3_quat[..., PoseSE3Index.QUATERNION] ) euler_rotation_matrices = get_rotation_matrices_from_euler_array( - abs_se3_euler[..., EulerStateSE3Index.EULER_ANGLES] + abs_se3_euler[..., EulerPoseSE3Index.EULER_ANGLES] ) np.testing.assert_allclose(quat_rotation_matrices, euler_rotation_matrices, atol=1e-6) # convert_points_3d_array_between_origins(quat_se3, random_quat_se3_array) @@ -192,11 +182,11 @@ def test_convert_se3_array_between_origins(self): for _ in range(10): random_quat_se3_array = self._get_random_quat_se3_array(np.random.randint(1, 10)) - from_se3 = StateSE3.from_array(self._get_random_quat_se3_array(1)[0]) - to_se3 = StateSE3.from_array(self._get_random_quat_se3_array(1)[0]) - identity_se3_array = np.zeros(len(StateSE3Index), dtype=np.float64) - identity_se3_array[StateSE3Index.QW] = 1.0 - identity_se3 = StateSE3.from_array(identity_se3_array) + from_se3 = PoseSE3.from_array(self._get_random_quat_se3_array(1)[0]) + to_se3 = PoseSE3.from_array(self._get_random_quat_se3_array(1)[0]) + identity_se3_array = np.zeros(len(PoseSE3Index), dtype=np.float64) + identity_se3_array[PoseSE3Index.QW] = 1.0 + identity_se3 = PoseSE3.from_array(identity_se3_array) # Check if consistent with absolute-relative-absolute conversion converted_se3_quat = convert_se3_array_between_origins(from_se3, to_se3, random_quat_se3_array) @@ -205,32 +195,32 @@ def test_convert_se3_array_between_origins(self): rel_to_se3_quat = convert_absolute_to_relative_se3_array(to_se3, abs_from_se3_quat) np.testing.assert_allclose( - converted_se3_quat[..., StateSE3Index.XYZ], - rel_to_se3_quat[..., StateSE3Index.XYZ], + converted_se3_quat[..., PoseSE3Index.XYZ], + rel_to_se3_quat[..., PoseSE3Index.XYZ], atol=1e-6, ) np.testing.assert_allclose( - converted_se3_quat[..., StateSE3Index.QUATERNION], - rel_to_se3_quat[..., StateSE3Index.QUATERNION], + converted_se3_quat[..., PoseSE3Index.QUATERNION], + rel_to_se3_quat[..., PoseSE3Index.QUATERNION], atol=1e-6, ) # Check if consistent with absolute conversion to identity origin absolute_se3_quat = convert_se3_array_between_origins(from_se3, identity_se3, random_quat_se3_array) np.testing.assert_allclose( - absolute_se3_quat[..., StateSE3Index.XYZ], - abs_from_se3_quat[..., StateSE3Index.XYZ], + absolute_se3_quat[..., PoseSE3Index.XYZ], + abs_from_se3_quat[..., PoseSE3Index.XYZ], atol=1e-6, ) def test_convert_points_3d_array_between_origins(self): random_points_3d = np.random.rand(10, 3) for _ in range(10): - from_se3 = StateSE3.from_array(self._get_random_quat_se3_array(1)[0]) - to_se3 = StateSE3.from_array(self._get_random_quat_se3_array(1)[0]) - identity_se3_array = np.zeros(len(StateSE3Index), dtype=np.float64) - identity_se3_array[StateSE3Index.QW] = 1.0 - identity_se3 = StateSE3.from_array(identity_se3_array) + from_se3 = PoseSE3.from_array(self._get_random_quat_se3_array(1)[0]) + to_se3 = PoseSE3.from_array(self._get_random_quat_se3_array(1)[0]) + identity_se3_array = np.zeros(len(PoseSE3Index), dtype=np.float64) + identity_se3_array[PoseSE3Index.QW] = 1.0 + identity_se3 = PoseSE3.from_array(identity_se3_array) # Check if consistent with absolute-relative-absolute conversion converted_points_quat = convert_points_3d_array_between_origins(from_se3, to_se3, random_points_3d) @@ -239,12 +229,12 @@ def test_convert_points_3d_array_between_origins(self): np.testing.assert_allclose(converted_points_quat, rel_to_se3_quat, atol=1e-6) # Check if consistent with se3 array conversion - random_se3_poses = np.zeros((random_points_3d.shape[0], len(StateSE3Index)), dtype=np.float64) - random_se3_poses[:, StateSE3Index.XYZ] = random_points_3d - random_se3_poses[:, StateSE3Index.QUATERNION] = np.array([1.0, 0.0, 0.0, 0.0]) # Identity rotation + random_se3_poses = np.zeros((random_points_3d.shape[0], len(PoseSE3Index)), dtype=np.float64) + random_se3_poses[:, PoseSE3Index.XYZ] = random_points_3d + random_se3_poses[:, PoseSE3Index.QUATERNION] = np.array([1.0, 0.0, 0.0, 0.0]) # Identity rotation converted_se3_quat_poses = convert_se3_array_between_origins(from_se3, to_se3, random_se3_poses) np.testing.assert_allclose( - converted_se3_quat_poses[:, StateSE3Index.XYZ], + converted_se3_quat_poses[:, PoseSE3Index.XYZ], converted_points_quat, atol=1e-6, ) @@ -252,8 +242,8 @@ def test_convert_points_3d_array_between_origins(self): # Check if consistent with absolute conversion to identity origin absolute_se3_quat = convert_points_3d_array_between_origins(from_se3, identity_se3, random_points_3d) np.testing.assert_allclose( - absolute_se3_quat[..., StateSE3Index.XYZ], - abs_from_se3_quat[..., StateSE3Index.XYZ], + absolute_se3_quat[..., PoseSE3Index.XYZ], + abs_from_se3_quat[..., PoseSE3Index.XYZ], atol=1e-6, ) @@ -300,7 +290,3 @@ def test_translate_se3_along_body_frame(self): np.testing.assert_allclose(translated_quat.point_3d.array, translated_euler.point_3d.array, atol=1e-6) np.testing.assert_allclose(translated_quat.rotation_matrix, translated_euler.rotation_matrix, atol=1e-6) np.testing.assert_allclose(quat_se3.quaternion.array, translated_quat.quaternion.array, atol=1e-6) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/geometry/utils/test_bounding_box_utils.py b/tests/unit/geometry/utils/test_bounding_box_utils.py index 3b5718ca..dd049b41 100644 --- a/tests/unit/geometry/utils/test_bounding_box_utils.py +++ b/tests/unit/geometry/utils/test_bounding_box_utils.py @@ -1,5 +1,3 @@ -import unittest - import numpy as np import numpy.typing as npt import shapely @@ -8,11 +6,11 @@ BoundingBoxSE3Index, Corners2DIndex, Corners3DIndex, - EulerStateSE3Index, + EulerPoseSE3Index, Point2DIndex, Point3DIndex, ) -from py123d.geometry.se import EulerStateSE3, StateSE3 +from py123d.geometry.pose import EulerPoseSE3, PoseSE3 from py123d.geometry.transform.transform_se3 import translate_se3_along_body_frame from py123d.geometry.utils.bounding_box_utils import ( bbse2_array_to_corners_array, @@ -24,24 +22,23 @@ from py123d.geometry.vector import Vector3D -class TestBoundingBoxUtils(unittest.TestCase): - - def setUp(self): +class TestBoundingBoxUtils: + def setup_method(self): self._num_consistency_checks = 10 self._max_pose_xyz = 100.0 self._max_extent = 200.0 def _get_random_euler_se3_array(self, size: int) -> npt.NDArray[np.float64]: """Generate random SE3 poses""" - random_se3_array = np.zeros((size, len(EulerStateSE3Index)), dtype=np.float64) - random_se3_array[:, EulerStateSE3Index.XYZ] = np.random.uniform( + random_se3_array = np.zeros((size, len(EulerPoseSE3Index)), dtype=np.float64) + random_se3_array[:, EulerPoseSE3Index.XYZ] = np.random.uniform( -self._max_pose_xyz, self._max_pose_xyz, (size, len(Point3DIndex)), ) - random_se3_array[:, EulerStateSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) - random_se3_array[:, EulerStateSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) - random_se3_array[:, EulerStateSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array[:, EulerPoseSE3Index.YAW] = np.random.uniform(-np.pi, np.pi, size) + random_se3_array[:, EulerPoseSE3Index.PITCH] = np.random.uniform(-np.pi / 2, np.pi / 2, size) + random_se3_array[:, EulerPoseSE3Index.ROLL] = np.random.uniform(-np.pi, np.pi, size) return random_se3_array @@ -106,7 +103,7 @@ def test_corners_2d_array_to_polygon_array_one_dim(self): expected_polygon = shapely.geometry.Polygon(corners_array) np.testing.assert_allclose(polygon.area, expected_polygon.area, atol=1e-6) - self.assertTrue(polygon.equals(expected_polygon)) + assert polygon.equals(expected_polygon) def test_corners_2d_array_to_polygon_array_n_dim(self): corners_array = np.array( @@ -131,10 +128,10 @@ def test_corners_2d_array_to_polygon_array_n_dim(self): expected_polygon_2 = shapely.geometry.Polygon(corners_array[1]) np.testing.assert_allclose(polygons[0].area, expected_polygon_1.area, atol=1e-6) - self.assertTrue(polygons[0].equals(expected_polygon_1)) + assert polygons[0].equals(expected_polygon_1) np.testing.assert_allclose(polygons[1].area, expected_polygon_2.area, atol=1e-6) - self.assertTrue(polygons[1].equals(expected_polygon_2)) + assert polygons[1].equals(expected_polygon_2) def test_corners_2d_array_to_polygon_array_zero_dim(self): corners_array = np.zeros((0, 4, 2), dtype=np.float64) @@ -154,7 +151,7 @@ def test_bbse2_array_to_polygon_array_one_dim(self): expected_polygon = shapely.geometry.Polygon(expected_corners) np.testing.assert_allclose(polygon.area, expected_polygon.area, atol=1e-6) - self.assertTrue(polygon.equals(expected_polygon)) + assert polygon.equals(expected_polygon) def test_bbse2_array_to_polygon_array_n_dim(self): bounding_box_se2_array = np.array([1.0, 2.0, 0.0, 4.0, 2.0]) @@ -171,7 +168,7 @@ def test_bbse2_array_to_polygon_array_n_dim(self): for polygon in polygons: np.testing.assert_allclose(polygon.area, expected_polygon.area, atol=1e-6) - self.assertTrue(polygon.equals(expected_polygon)) + assert polygon.equals(expected_polygon) def test_bbse2_array_to_polygon_array_zero_dim(self): bounding_box_se2_array = np.zeros((0, 5), dtype=np.float64) @@ -198,14 +195,14 @@ def test_bbse3_array_to_corners_array_one_dim(self): def test_bbse3_array_to_corners_array_one_dim_rotation(self): for _ in range(self._num_consistency_checks): - se3_state = EulerStateSE3.from_array(self._get_random_euler_se3_array(1)[0]).state_se3 + se3_state = EulerPoseSE3.from_array(self._get_random_euler_se3_array(1)[0]).pose_se3 se3_array = se3_state.array # construct a bounding box bounding_box_se3_array = np.zeros((len(BoundingBoxSE3Index),), dtype=np.float64) length, width, height = np.random.uniform(0.0, self._max_extent, size=3) - bounding_box_se3_array[BoundingBoxSE3Index.STATE_SE3] = se3_array + bounding_box_se3_array[BoundingBoxSE3Index.SE3] = se3_array bounding_box_se3_array[BoundingBoxSE3Index.LENGTH] = length bounding_box_se3_array[BoundingBoxSE3Index.WIDTH] = width bounding_box_se3_array[BoundingBoxSE3Index.HEIGHT] = height @@ -227,13 +224,13 @@ def test_bbse3_array_to_corners_array_n_dim(self): for _ in range(self._num_consistency_checks): N = np.random.randint(1, 20) se3_array = self._get_random_euler_se3_array(N) - se3_state_array = np.array([EulerStateSE3.from_array(arr).state_se3.array for arr in se3_array]) + se3_state_array = np.array([EulerPoseSE3.from_array(arr).pose_se3.array for arr in se3_array]) # construct a bounding box bounding_box_se3_array = np.zeros((N, len(BoundingBoxSE3Index)), dtype=np.float64) lengths, widths, heights = np.random.uniform(0.0, self._max_extent, size=(3, N)) - bounding_box_se3_array[:, BoundingBoxSE3Index.STATE_SE3] = se3_state_array + bounding_box_se3_array[:, BoundingBoxSE3Index.SE3] = se3_state_array bounding_box_se3_array[:, BoundingBoxSE3Index.LENGTH] = lengths bounding_box_se3_array[:, BoundingBoxSE3Index.WIDTH] = widths bounding_box_se3_array[:, BoundingBoxSE3Index.HEIGHT] = heights @@ -249,7 +246,7 @@ def test_bbse3_array_to_corners_array_n_dim(self): np.testing.assert_allclose( corners_array[obj_idx, corner_idx], translate_se3_along_body_frame( - StateSE3.from_array(bounding_box_se3_array[obj_idx, BoundingBoxSE3Index.STATE_SE3]), + PoseSE3.from_array(bounding_box_se3_array[obj_idx, BoundingBoxSE3Index.SE3]), body_translate_vector, ).point_3d.array, atol=1e-6, @@ -260,7 +257,3 @@ def test_bbse3_array_to_corners_array_zero_dim(self): corners_array = bbse3_array_to_corners_array(bounding_box_se3_array) expected_corners = np.zeros((0, 8, 3), dtype=np.float64) np.testing.assert_allclose(corners_array, expected_corners, atol=1e-6) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/unit/geometry/utils/test_rotation_utils.py b/tests/unit/geometry/utils/test_rotation_utils.py index f298af7a..92c4d1b7 100644 --- a/tests/unit/geometry/utils/test_rotation_utils.py +++ b/tests/unit/geometry/utils/test_rotation_utils.py @@ -1,8 +1,8 @@ -import unittest from typing import Tuple import numpy as np import numpy.typing as npt +import pytest from pyquaternion import Quaternion as PyQuaternion from py123d.geometry.geometry_index import EulerAnglesIndex, QuaternionIndex @@ -62,9 +62,8 @@ def _get_rotation_matrix_helper(euler_array: npt.NDArray[np.float64]) -> npt.NDA return R_z @ R_y @ R_x -class TestRotationUtils(unittest.TestCase): - - def setUp(self): +class TestRotationUtils: + def setup_method(self): pass def _get_random_quaternion(self) -> npt.NDArray[np.float64]: @@ -135,11 +134,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) conjugate_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) conjugate_quaternion_array(invalid_quat) @@ -177,11 +176,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) get_euler_array_from_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) get_euler_array_from_quaternion_array(invalid_quat) @@ -221,11 +220,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((0, 3)) # (0, 3) rotation matrix shape (invalid) get_euler_array_from_rotation_matrices(invalid_rot) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3, 3, 8)) # (3, 3, 8) rotation matrix shape (invalid) get_euler_array_from_rotation_matrices(invalid_rot) @@ -246,11 +245,11 @@ def test_get_euler_array_from_rotation_matrix(self): ) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3,)) # (0, 3) rotation matrix shape (invalid) get_euler_array_from_rotation_matrix(invalid_rot) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3, 8)) # (3, 8) rotation matrix shape (invalid) get_euler_array_from_rotation_matrix(invalid_rot) @@ -290,16 +289,15 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) get_q_bar_matrices(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) get_q_bar_matrices(invalid_quat) def test_get_q_matrices(self): - def _test_by_shape(shape: Tuple[int, ...]) -> None: for _ in range(10): N = np.prod(shape) @@ -333,11 +331,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) get_q_matrices(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) get_q_matrices(invalid_quat) @@ -387,16 +385,15 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((0,)) # Zero euler angles (invalid) get_quaternion_array_from_euler_array(invalid_euler) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((3, 8)) # Zero euler angles (invalid) get_quaternion_array_from_euler_array(invalid_euler) def test_get_quaternion_array_from_rotation_matrices(self): - def _test_by_shape(shape: Tuple[int, ...]) -> None: for _ in range(10): N = np.prod(shape) @@ -436,11 +433,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((0, 3)) # (0, 3) rotation matrix shape (invalid) get_quaternion_array_from_rotation_matrices(invalid_rot) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3, 3, 8)) # (3, 3, 8) rotation matrix shape (invalid) get_quaternion_array_from_rotation_matrices(invalid_rot) @@ -465,11 +462,11 @@ def test_get_quaternion_array_from_rotation_matrix(self): ) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3,)) # (0, 3) rotation matrix shape (invalid) get_quaternion_array_from_rotation_matrix(invalid_rot) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_rot = np.zeros((3, 8)) # (3, 8) rotation matrix shape (invalid) get_quaternion_array_from_rotation_matrix(invalid_rot) @@ -504,11 +501,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) normalize_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) normalize_quaternion_array(invalid_quat) @@ -545,11 +542,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((0, 5)) # Zero euler angles (invalid) get_rotation_matrices_from_euler_array(invalid_euler) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((3, 8)) # Zero euler angles (invalid) get_rotation_matrices_from_euler_array(invalid_euler) @@ -589,11 +586,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) get_rotation_matrices_from_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) get_rotation_matrices_from_quaternion_array(invalid_quat) @@ -615,11 +612,11 @@ def test_get_rotation_matrix_from_euler_array(self): ) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((0,)) # Zero euler angles (invalid) get_rotation_matrix_from_euler_array(invalid_euler) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_euler = np.zeros((8,)) # Zero euler angles (invalid) get_rotation_matrix_from_euler_array(invalid_euler) @@ -641,11 +638,11 @@ def test_get_rotation_matrix_from_quaternion_array(self): ) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) get_rotation_matrix_from_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((8,)) # Zero quaternion (invalid) get_rotation_matrix_from_quaternion_array(invalid_quat) @@ -686,11 +683,11 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((0,)) # Zero quaternion (invalid) invert_quaternion_array(invalid_quat) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) invert_quaternion_array(invalid_quat) @@ -735,12 +732,12 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test invalid input - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat1 = np.zeros((0,)) # Zero quaternion (invalid) invalid_quat2 = np.zeros((0,)) # Zero quaternion (invalid) multiply_quaternion_arrays(invalid_quat1, invalid_quat2) - with self.assertRaises(AssertionError): + with pytest.raises(AssertionError): invalid_quat1 = np.zeros((len(QuaternionIndex), 8)) # Zero quaternion (invalid) invalid_quat2 = np.zeros((len(QuaternionIndex), 4)) # Zero quaternion (invalid) multiply_quaternion_arrays(invalid_quat1, invalid_quat2) @@ -762,8 +759,8 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: # Check if each angle is within [-pi, pi] for i in range(N): angle = normalized_angles_flat[i] - self.assertGreaterEqual(angle, -np.pi - 1e-8) - self.assertLessEqual(angle, np.pi + 1e-8) + assert angle >= -np.pi - 1e-8 + assert angle <= np.pi + 1e-8 # Test single-dim shape _test_by_shape((1,)) @@ -776,12 +773,7 @@ def _test_by_shape(shape: Tuple[int, ...]) -> None: _test_by_shape((0,)) # Test float - with self.subTest("Test float input"): - angle = 4 * np.pi - normalized_angle = normalize_angle(angle) - self.assertGreaterEqual(normalized_angle, -np.pi - 1e-8) - self.assertLessEqual(normalized_angle, np.pi + 1e-8) - - -if __name__ == "__main__": - unittest.main() + angle = 4 * np.pi + normalized_angle = normalize_angle(angle) + assert normalized_angle >= -np.pi - 1e-8 + assert normalized_angle <= np.pi + 1e-8 diff --git a/tutorials/01_scene_tutorial.ipynb b/tutorials/01_scene_tutorial.ipynb new file mode 100644 index 00000000..c520bf74 --- /dev/null +++ b/tutorials/01_scene_tutorial.ipynb @@ -0,0 +1,434 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "

\n", + " \n", + " \n", + " \n", + " \"Logo\"\n", + " \n", + "

123D: Scene Tutorial

\n", + "" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from py123d.api import SceneAPI, SceneFilter\n", + "from py123d.api.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", + "from py123d.common.multithreading.worker_parallel import SingleMachineParallelExecutor" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## 1.1 Download Demo Logs\n", + "\n", + "You can download demo logs for 123D as described in the [documentation](https://autonomousvision.github.io/py123d/installation/). After the installation and download, you can start with the tutorial." + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 1.2 Create Scenes by filtering the datasets\n", + "\n", + "\n", + "\n", + "The logs store continuous driving recordings. Scenes in 123D are sequences that are extracted from a log, e.g. given a predefined duration and history.\n", + "\n", + "In the example below, we filter some scenes from all logs with 8 second duration and 8 seconds temporal distance (making the scenes non-overlapping). If `None` is passed to the duration, the scene will contain the complete log.\n", + "\n", + "This `SceneFilter` object is passed to a `SceneBuilder` object to query `SceneAPI`'s from the dataset.\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.sensors import PinholeCamera, PinholeCameraType\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=None,\n", + " log_names=None,\n", + " scene_uuids=None,\n", + " duration_s=7.0,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=7.0,\n", + " pinhole_camera_types=[PinholeCameraType.PCAM_F0],\n", + " shuffle=True,\n", + ")\n", + "worker = SingleMachineParallelExecutor()\n", + "scenes = ArrowSceneBuilder().get_scenes(scene_filter, worker)\n", + "\n", + "dataset_splits = set(scene.log_metadata.split for scene in scenes)\n", + "print(f\"Found {len(scenes)} scenes from {len(dataset_splits)} datasplits:\")\n", + "for split in dataset_splits:\n", + " print(f\" - {split}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 1.2 Inspecting the Scene\n", + "\n", + "Let's inspect a random scenefrom our dataset.\n", + "\n", + "A scene has several different metadata objects attached to it:\n", + "\n", + "`SceneMetadata`: Information how the scene was extracted from the log. Each timestep in the log has universally unique identifier (UUID). The UUID of the initial timestep also serves as identifier for scene filtering." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "scene: SceneAPI = np.random.choice(scenes)\n", + "scene_metadata = scene.scene_metadata\n", + "print(scene_metadata)\n", + "print(\"\\nInitial UUID:\", scene_metadata.initial_uuid)\n", + "print(\"Number of iterations:\", scene_metadata.number_of_iterations)\n", + "print(\"Number of history iterations:\", scene_metadata.number_of_history_iterations)\n", + "print(\"Duration (s):\", scene_metadata.duration_s)\n", + "print(\"Iteration duration (s):\", scene_metadata.iteration_duration_s)" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "`LogMetadata`: Information of the log the scene was extracted from. This object also includes data about the map (if available), or static information of the ego vehicle, e.g. the included sensors and vehicle parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "log_metadata = scene.log_metadata\n", + "log_metadata" + ] + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "`MapMetadata`: If the map is available, this object includes information about the location, wether the map is 3D (`map_has_z`), of the the map is defined per log (`map_is_local`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "map_metadata = scene.map_metadata\n", + "map_metadata" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## 1.3 Retrieving data from the `SceneAPI`\n", + "\n", + "Different datasets might provide different modalities. In general, you can load data using a `scene.get_modality_at_iteration(iteration=...)`\n", + "\n", + "If a modality is not available, the return will be `None`. The `TimePoint` is the only modality that is strictly required to be available in a `Scene\n", + "\n", + "Let's look at some examples:\n" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "### 1.3.1 `TimePoint`\n", + "\n", + "Current time step in microseconds." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.time import TimePoint\n", + "\n", + "iteration = 0\n", + "timepoint: TimePoint = scene.get_timepoint_at_iteration(iteration=iteration)\n", + "print(f\"Time at iteration {iteration}:\", timepoint)" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "### 1.3.2 `EgoStateSE3` \n", + "State of the ego vehicle in 3D with location and orientation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.vehicle_state import EgoStateSE3\n", + "\n", + "if (ego_state := scene.get_ego_state_at_iteration(iteration=iteration)) is not None:\n", + " ego_state: EgoStateSE3\n", + "\n", + " print(\"Vehicle parameters\\t\", ego_state.vehicle_parameters)\n", + "\n", + " # The ego vehicles coordinate system is defined by it's rear-axle / IMU location.\n", + " print(\"Rear axle location:\\t\", ego_state.rear_axle_se3.point_3d)\n", + " print(\"Rear axle orientation:\\t\", ego_state.rear_axle_se3.quaternion)\n", + "\n", + " # You can also use the center pose\n", + " print(\"Center location:\\t\", ego_state.center_se3.point_3d)\n", + " print(\"Center orientation:\\t\", ego_state.center_se3.quaternion)" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "### 1.3.3 `BoxDetectionWrapper`\n", + "\n", + "Object that contains all bounding boxes in the current time step" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.detections import BoxDetectionWrapper\n", + "\n", + "if (box_detections := scene.get_box_detections_at_iteration(iteration=iteration)) is not None:\n", + " box_detections: BoxDetectionWrapper\n", + "\n", + " print(f\"Number of boxes:{len(box_detections)}\")\n", + "\n", + " if len(box_detections) > 0:\n", + " box_detection = box_detections[0]\n", + " print(\"\\nFirst box:\")\n", + " print(\"Dataset Label:\\t\", box_detection.metadata.label)\n", + " print(\"Default Label:\\t\", box_detection.metadata.default_label)\n", + " print(\"Parameters:\\t\", box_detection.bounding_box_se3)" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### 1.3.4 `PinholeCamera`\n", + "Object containing the camera observation with a pinhole model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "available_pinhole_types = scene.available_pinhole_camera_types\n", + "print(\"Available pinhole camera types:\\t\", available_pinhole_types)\n", + "\n", + "if len(available_pinhole_types) > 0:\n", + " camera_type = np.random.choice(available_pinhole_types)\n", + "else:\n", + " camera_type = PinholeCameraType.PCAM_F0 # Front facing camera\n", + "\n", + "# NOTE: If a camera is not available, the return will be None\n", + "if (pinhole_camera := scene.get_pinhole_camera_at_iteration(iteration=iteration, camera_type=camera_type)) is not None:\n", + " pinhole_camera: PinholeCamera\n", + "\n", + " print(f\"\\nCamera type:\\t{camera_type}\")\n", + "\n", + " print(f\"Image shape:\\t{pinhole_camera.image.shape}\")\n", + " print(f\"Intrinsics:\\t{pinhole_camera.metadata.intrinsics}\")\n", + " print(f\"Distortion:\\t{pinhole_camera.metadata.distortion}\")\n", + "\n", + " plt.imshow(pinhole_camera.image)\n", + " plt.title(f\"Camera Type: {camera_type}\")\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "20", + "metadata": {}, + "source": [ + "### 1.3.5 `LiDAR`\n", + "Object containing a point cloud of a single laser scanner." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.sensors import LiDAR, LiDARType\n", + "\n", + "available_lidar_types = scene.available_lidar_types\n", + "print(\"Available LiDAR types:\\t\", available_lidar_types)\n", + "\n", + "if len(available_lidar_types) > 0:\n", + " lidar_type = np.random.choice(available_lidar_types)\n", + "else:\n", + " lidar_type = LiDARType.LIDAR_TOP # Top mounted LiDAR\n", + "\n", + "if (lidar := scene.get_lidar_at_iteration(iteration=iteration, lidar_type=lidar_type)) is not None:\n", + " lidar: LiDAR\n", + "\n", + " print(f\"\\nLiDAR type:\\t{lidar_type}\")\n", + " print(f\"Shape (NxM):\\t{lidar.point_cloud.shape}\")\n", + " print(f\"Features (M):\\t{[(enum.name, enum.value) for enum in lidar.metadata.lidar_index]}\")\n", + "\n", + " xy = lidar.xy\n", + "\n", + " plt.scatter(xy[:, 0], xy[:, 1], s=0.1, alpha=0.25, c=\"black\")\n", + " plt.title(f\"LiDAR Type: {lidar_type}\")\n", + " plt.xlabel(\"X-forward [m]\")\n", + " plt.ylabel(\"Y-left [m]\")\n", + " plt.axis(\"equal\")\n", + "\n", + " range_limit = 100 # meters\n", + " plt.xlim(-range_limit, range_limit)\n", + " plt.ylim(-range_limit, range_limit)\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "22", + "metadata": {}, + "source": [ + "### 1.3.6 `MapAPI`\n", + "\n", + "The `MapAPI` can get retrieved from a scene directly. If the map is available, we plot the map with our default plotting function.\n", + "For further information, you can visit the map or visualization tutorial." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "23", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.api import MapAPI\n", + "from py123d.geometry import Point2D\n", + "from py123d.visualization.matplotlib.observation import add_default_map_on_ax\n", + "from py123d.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", + "\n", + "\n", + "def simple_map_visualization(map_api: MapAPI, center_2d: Point2D, map_radius: float = 100.0):\n", + " \"\"\"Helper to plot the map using matplotlib.\n", + "\n", + " :param map_api: The MapAPI to visualize\n", + " :param center_2d: The center point of the map visualization\n", + " :param map_radius: The radius around the center point to visualize\n", + " \"\"\"\n", + "\n", + " fsize = 8\n", + " _, ax = plt.subplots(figsize=(fsize, fsize))\n", + " add_default_map_on_ax(ax, map_api=map_api, point_2d=center_2d, radius=map_radius)\n", + " add_non_repeating_legend_to_ax(ax)\n", + " ax.set_aspect(\"equal\")\n", + " ax.set_xlim(center_2d.x - map_radius, center_2d.x + map_radius)\n", + " ax.set_ylim(center_2d.y - map_radius, center_2d.y + map_radius)\n", + " plt.show()\n", + "\n", + "\n", + "if (map_api := scene.get_map_api()) is not None:\n", + " map_api: MapAPI\n", + "\n", + " if (ego_state := scene.get_ego_state_at_iteration(iteration=iteration)) is not None:\n", + " center_2d = ego_state.center_se3.point_2d\n", + " else:\n", + " center_2d = Point2D.from_array(np.array([0.0, 0.0]))\n", + "\n", + " print(\"\\nMapAPI is available.\")\n", + " print(\"Map Metadata:\", map_api.map_metadata)\n", + " simple_map_visualization(map_api=map_api, center_2d=center_2d)" + ] + }, + { + "cell_type": "markdown", + "id": "24", + "metadata": {}, + "source": [ + "### 1.3.7 Others:\n", + "\n", + "You can find further modalities in the documentation of [`SceneAPI`](todo)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "py123d_dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tutorials/02_map_tutorial.ipynb b/tutorials/02_map_tutorial.ipynb new file mode 100644 index 00000000..bc0aab2e --- /dev/null +++ b/tutorials/02_map_tutorial.ipynb @@ -0,0 +1,799 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "

\n", + " \n", + " \n", + " \n", + " \"Logo\"\n", + " \n", + "

123D: Map Tutorial

\n", + "" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import List, Optional\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from py123d.api import MapAPI, SceneAPI, SceneFilter\n", + "from py123d.api.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", + "from py123d.common.multithreading.worker_parallel import SingleMachineParallelExecutor\n", + "from py123d.datatypes.map_objects import (\n", + " BaseMapLineObject,\n", + " BaseMapObject,\n", + " BaseMapSurfaceObject,\n", + " Intersection,\n", + " Lane,\n", + " LaneGroup,\n", + " MapLayer,\n", + " RoadEdgeType,\n", + " RoadLineType,\n", + ")\n", + "from py123d.geometry import Point3D, Polyline2D\n", + "from py123d.visualization.color.default import MAP_SURFACE_CONFIG\n", + "from py123d.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", + "\n", + "# Set some default visualization parameters\n", + "DEFAULT_FIGSIZE = (9, 9)\n", + "SURFACE_ALPHA = 0.3" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## 2.1 Download Demo Logs\n", + "\n", + "You can download demo logs for 123D as described in the [documentation](https://autonomousvision.github.io/py123d/installation/). After the installation and download, you can start with the tutorial." + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 2.2 Create Scenes by filtering the datasets\n", + "\n", + "We create some scenes for easy access to some `MapAPI`'s. We use the option `map_api_required=True` to only include scenes/logs with maps." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "scene_filter = SceneFilter(\n", + " split_names=None,\n", + " duration_s=None, # No duration means that the scene will include the complete log.\n", + " shuffle=True,\n", + " map_api_required=True, # Only include scenes/logs with an available map API.\n", + ")\n", + "worker = SingleMachineParallelExecutor()\n", + "scenes = ArrowSceneBuilder().get_scenes(scene_filter, worker)\n", + "\n", + "dataset_splits = set(scene.log_metadata.split for scene in scenes)\n", + "print(f\"Found {len(scenes)} scenes from {len(dataset_splits)} datasplits:\")\n", + "for split in dataset_splits:\n", + " print(f\" - {split}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 2.3 Inspecting the `MapAPI`\n" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "`MapMetadata`: If the map is available, this object includes information about the location, whether the map is 3D (`map_has_z`), of the the map is defined per log (`map_is_local`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "scene: SceneAPI = np.random.choice(scenes)\n", + "map_api: MapAPI = scene.get_map_api()\n", + "\n", + "map_metadata = scene.map_metadata\n", + "map_metadata" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "## 2.4 Querying map objects\n", + "\n", + "There are multiple categories, also called layers, of map objects in 123D. You can get the available layers with:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "map_api.get_available_map_layers()" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "These layers can be divided into:\n", + "- `BaseMapSurfaceObject`: Objects/layers that defined a surface , e.g. a polygon or triangle mesh. Examples are lanes, lane groups, crosswalks, etc.\n", + "- `BaseMapLineObject`: Objects/layers that define a polyline in the map, e.g. a road edge or road line.\n", + "\n", + "All map objects are of type `BaseMapObject` which merely requires a `map_object_id` that is unique in each layer.\n", + "\n", + "The features of a map object depend on the type. You can query map objects with several functions, e.g. given a query point and radius." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "# You can define a query point and radius to search for map objects around that point.\n", + "query_point: Optional[Point3D] = None # e.g. Point(x=0.0, y=0.0, z=0.0)\n", + "query_radius: float = 100.0 # meters\n", + "\n", + "# Otherwise, we use the ego vehicle position, or the origin if no ego state is available.\n", + "if query_point is None:\n", + " if (ego_state := scene.get_ego_state_at_iteration(iteration=0)) is not None:\n", + " query_point = ego_state.center_se3.point_3d\n", + " else:\n", + " query_point = Point3D(0.0, 0.0, 0.0)\n", + "\n", + "\n", + "# For this example, we will query all available map layers.\n", + "map_layers = [\n", + " MapLayer.LANE, # Lanes (surface)\n", + " MapLayer.LANE_GROUP, # Lane groups (surface)\n", + " MapLayer.INTERSECTION, # Intersections (surface)\n", + " MapLayer.CROSSWALK, # Crosswalks (surface)\n", + " MapLayer.WALKWAY, # Walkways (surface)\n", + " MapLayer.CARPARK, # Carparks (surface)\n", + " MapLayer.GENERIC_DRIVABLE, # Generic drivable (surface)\n", + " MapLayer.STOP_ZONE, # Stop zones (surface)\n", + " MapLayer.ROAD_EDGE, # Road edges (lines)\n", + " MapLayer.ROAD_LINE, # Road lines (lines)\n", + "]\n", + "\n", + "# Query map objects: returns a dictionary mapping each map layer to a list of map objects in that layer.\n", + "map_object_dict = map_api.get_map_objects_in_radius(point=query_point, radius=query_radius, layers=map_layers)\n", + "\n", + "print(f\"Map objects found in radius {query_radius} around point {query_point}:\")\n", + "for map_layer, map_objects in map_object_dict.items():\n", + " print(f\"- {map_layer.name}: {len(map_objects)} objects\")" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "Before we inspect the map objects, let's define some helper functions to visualize them." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "def add_map_object_id(ax: plt.Axes, map_object: BaseMapObject) -> None:\n", + " \"\"\"Helper to add the map object ID as text at the centroid of the map object.\"\"\"\n", + " if isinstance(map_object, BaseMapSurfaceObject):\n", + " centroid = map_object.shapely_polygon.centroid\n", + " elif isinstance(map_object, BaseMapLineObject):\n", + " centroid = map_object.polyline_2d.interpolate(0.5, normalized=True).shapely_point\n", + " else:\n", + " raise TypeError(f\"Unsupported map object type of type {type(map_object)}.\")\n", + "\n", + " ax.text(\n", + " centroid.x,\n", + " centroid.y,\n", + " str(map_object.object_id),\n", + " fontsize=8,\n", + " ha=\"center\",\n", + " va=\"center\",\n", + " color=\"black\",\n", + " bbox=dict(boxstyle=\"round,pad=0.3\", facecolor=\"white\", edgecolor=\"black\", alpha=0.7),\n", + " )\n", + "\n", + "\n", + "def add_map_surface_object(\n", + " ax: plt.Axes,\n", + " map_surface_object: BaseMapSurfaceObject,\n", + " add_id: bool = True,\n", + " alpha: float = SURFACE_ALPHA,\n", + " **plot_kwargs,\n", + ") -> None:\n", + " \"\"\"Helper to plot a map surface object.\"\"\"\n", + " x, y = map_surface_object.shapely_polygon.exterior.xy\n", + " ax.fill(x, y, alpha=alpha, **plot_kwargs)\n", + " if add_id:\n", + " add_map_object_id(ax, map_surface_object)\n", + "\n", + "\n", + "def add_map_line_object(ax: plt.Axes, map_line_object: BaseMapLineObject, add_id: bool = True, **plot_kwargs) -> None:\n", + " \"\"\"Helper to plot a map line object.\"\"\"\n", + " polyline_array = map_line_object.polyline_2d.array\n", + " ax.plot(polyline_array[:, 0], polyline_array[:, 1], **plot_kwargs)\n", + " if add_id:\n", + " add_map_object_id(ax, map_line_object)\n", + "\n", + "\n", + "def add_polyline(ax: plt.Axes, polyline: Polyline2D, add_start_end: bool = False, **plot_kwargs) -> None:\n", + " \"\"\"Helper to plot a polyline.\"\"\"\n", + " polyline_array = polyline.array\n", + " ax.plot(polyline_array[:, 0], polyline_array[:, 1], **plot_kwargs)\n", + " if add_start_end:\n", + " ax.plot(polyline.array[0, 0], polyline.array[0, 1], \"o\", label=\"Start\", color=\"black\")\n", + " ax.plot(polyline.array[-1, 0], polyline.array[-1, 1], \"x\", label=\"End\", color=\"black\")\n", + "\n", + "\n", + "def adjust_aspect_custom(ax: plt.Axes) -> None:\n", + " \"\"\"Helper to adjust the aspect ratio of a matplotlib Axes.\"\"\"\n", + " x_limits = ax.get_xlim()\n", + " y_limits = ax.get_ylim()\n", + " x_range = x_limits[1] - x_limits[0]\n", + " y_range = y_limits[1] - y_limits[0]\n", + " max_range = max(x_range, y_range)\n", + "\n", + " x_center = (x_limits[0] + x_limits[1]) / 2\n", + " y_center = (y_limits[0] + y_limits[1]) / 2\n", + "\n", + " ax.set_xlim(x_center - max_range / 2, x_center + max_range / 2)\n", + " ax.set_ylim(y_center - max_range / 2, y_center + max_range / 2)\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "## 2.5 Map Objects in 123D\n", + "\n", + "### 2.5.1 `Lane`\n", + "\n", + "Let's start with the `Lane` object. Each lane can have multiple features, such as polylines from boundaries or the lane center, relational properties that point to neighboring map objects, or other features, such as the speed limit. \n", + "We can sample a lane and have a look:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "# Sample a random lane\n", + "if (lanes := map_object_dict[MapLayer.LANE]) is not None and len(lanes) > 0:\n", + " lane: Lane = np.random.choice(lanes)\n", + "\n", + "if lane is not None:\n", + " fig, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)\n", + "\n", + " centerline = lane.centerline\n", + " right_boundary = lane.right_boundary\n", + " left_boundary = lane.left_boundary\n", + "\n", + " # Plot centerline and boundaries\n", + " add_polyline(ax, centerline, label=\"Centerline\", color=\"blue\", add_start_end=True)\n", + " add_polyline(ax, right_boundary, label=\"Right Boundary\", color=\"red\")\n", + " add_polyline(ax, left_boundary, label=\"Left Boundary\", color=\"green\")\n", + "\n", + " # Plot lane surface\n", + " add_map_surface_object(ax, lane, label=\"Surface\", color=\"grey\")\n", + "\n", + " speed_limit_mps = round(lane.speed_limit_mps, 2) if lane.speed_limit_mps is not None else \"N/A\"\n", + " ax.set_title(f\"Lane {lane.object_id}, speed limit: {speed_limit_mps} m/s\")\n", + " centroid = lane.shapely_polygon.centroid\n", + "\n", + " # Surface of the lane\n", + " adjust_aspect_custom(ax)\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "Lanes can have neighbors, which can either be accessed as IDs or directly. These neighbors include:\n", + "- List of successor / predecessor lanes (`None` if not available)\n", + "- Single lane to the left or right (`None` if not available) " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "# Sample a random lane\n", + "if (lanes := map_object_dict[MapLayer.LANE]) is not None and len(lanes) > 0:\n", + " lane: Lane = np.random.choice(lanes)\n", + "\n", + "if lane is not None:\n", + " size = (8, 8)\n", + "\n", + " fig, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)\n", + "\n", + " # Add the current lane:\n", + " add_map_surface_object(ax, lane, label=f\"Current lane {lane.object_id}\", color=\"grey\")\n", + " add_polyline(ax, lane.centerline, label=\"Centerline\", color=\"grey\", add_start_end=True)\n", + "\n", + " # Add left neighbor lane:\n", + " if lane.left_lane is not None:\n", + " add_map_surface_object(ax, lane.left_lane, label=f\"Left lane {lane.left_lane.object_id}\", color=\"green\")\n", + "\n", + " if lane.right_lane is not None:\n", + " add_map_surface_object(ax, lane.right_lane, label=f\"Right lane {lane.right_lane.object_id}\", color=\"red\")\n", + "\n", + " # Add successor lanes:\n", + " for successor_lane in lane.successors:\n", + " add_map_surface_object(ax, successor_lane, label=f\"Successor lane {successor_lane.object_id}\", color=\"blue\")\n", + "\n", + " for predecessor in lane.predecessors:\n", + " add_map_surface_object(ax, predecessor, label=f\"Predecessor lane {predecessor.object_id}\", color=\"orange\")\n", + "\n", + " adjust_aspect_custom(ax)\n", + " add_non_repeating_legend_to_ax(ax)" + ] + }, + { + "cell_type": "markdown", + "id": "18", + "metadata": {}, + "source": [ + "### 2.5.2 `LaneGroup`" + ] + }, + { + "cell_type": "markdown", + "id": "19", + "metadata": {}, + "source": [ + "A lane can be part of a lane group. Lane groups are sets of lanes that go in the same direction. The lane group can be accessed from the lane directly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "if lane is not None and lane.lane_group is not None:\n", + " lane_group: LaneGroup = lane.lane_group\n", + "\n", + " fig, ax = plt.subplots(figsize=(8, 8))\n", + " ax.set_title(\n", + " f\"Lane Group {lane_group.object_id} includes [{', '.join(str(l.object_id) for l in lane_group.lanes)}]\"\n", + " )\n", + "\n", + " # Original lane\n", + " add_map_surface_object(ax, lane, label=f\"Current Lane {lane.object_id}\", color=\"grey\")\n", + " add_polyline(ax, lane.centerline, label=\"Centerline\", color=\"grey\", add_start_end=True)\n", + "\n", + " # Other lanes in the lane group\n", + " for group_lane in lane_group.lanes:\n", + " if group_lane.object_id != lane.object_id:\n", + " add_map_surface_object(ax, group_lane, label=f\"Other lane {group_lane.object_id}\", color=\"black\")\n", + "\n", + " adjust_aspect_custom(ax)\n", + " add_non_repeating_legend_to_ax(ax)" + ] + }, + { + "cell_type": "markdown", + "id": "21", + "metadata": {}, + "source": [ + "Lane groups are surfaces, with neighboring relationships. Let's sample another lane group and have a look:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "22", + "metadata": {}, + "outputs": [], + "source": [ + "if len(map_object_dict[MapLayer.LANE_GROUP]) > 0:\n", + " lane_group: LaneGroup = np.random.choice(map_object_dict[MapLayer.LANE_GROUP])\n", + "\n", + " fig, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)\n", + " ax.set_title(f\"Lane Group {lane_group.object_id}\")\n", + "\n", + " all_lanes: List[Lane] = lane_group.lanes\n", + "\n", + " # Current lane group\n", + " add_map_surface_object(ax, lane_group, label=f\"Lane Group {lane_group.object_id}\", color=\"grey\")\n", + "\n", + " # Predecessor lane groups\n", + " for predecessor in lane_group.predecessors:\n", + " add_map_surface_object(ax, predecessor, label=f\"Predecessor Lane Group {predecessor.object_id}\", color=\"orange\")\n", + " all_lanes += predecessor.lanes\n", + "\n", + " # Successor lane groups\n", + " for successor in lane_group.successors:\n", + " add_map_surface_object(ax, successor, label=f\"Successor Lane Group {successor.object_id}\", color=\"blue\")\n", + " all_lanes += successor.lanes\n", + "\n", + " # Adding all centerline\n", + " for lane in all_lanes:\n", + " centerline = lane.centerline\n", + " ax.plot(*centerline.array.T[:2], label=\"All centerlines\", color=\"darkgrey\")\n", + "\n", + " adjust_aspect_custom(ax)\n", + " add_non_repeating_legend_to_ax(ax)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "23", + "metadata": {}, + "source": [ + "### 2.5.3 `Intersection`\n", + "\n", + "Intersections are map surfaces that include multiple lane groups. Let's look at a random example:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24", + "metadata": {}, + "outputs": [], + "source": [ + "if len(map_object_dict[MapLayer.INTERSECTION]) > 0:\n", + " intersection: Intersection = np.random.choice(map_object_dict[MapLayer.INTERSECTION])\n", + "\n", + " fig, ax = plt.subplots(ncols=2, figsize=(14, 7))\n", + " ax[0].set_title(f\"Intersection {intersection.object_id}\")\n", + " add_map_surface_object(ax[0], intersection, label=\"Intersection\", color=\"blue\")\n", + "\n", + " lane_groups: List[LaneGroup] = intersection.lane_groups\n", + " ax[1].set_title(f\"Lane Groups of Intersection {intersection.object_id}\")\n", + " for lane_group in lane_groups:\n", + " add_map_surface_object(ax[1], lane_group, label=\"Lane Group\", color=\"grey\")\n", + "\n", + " for ax_ in ax:\n", + " add_non_repeating_legend_to_ax(ax_)\n", + " adjust_aspect_custom(ax_)\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "25", + "metadata": {}, + "source": [ + "### 2.5.4 Other map surfaces\n", + "\n", + "Besides lanes, lane groups, and intersection, the 123D map has a few simpler map surface objects. These include:\n", + "\n", + "- `Crosswalk`\n", + "- `Walkway`\n", + "- `Carpark`\n", + "- `GenericDrivable`\n", + "- `StopZone`\n", + "\n", + "The availability of these objects depend on the dataset. For now, we can just plot all objects we found in our previous query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26", + "metadata": {}, + "outputs": [], + "source": [ + "misc_map_surfaces: List[MapLayer] = [\n", + " MapLayer.CROSSWALK,\n", + " MapLayer.WALKWAY,\n", + " MapLayer.CARPARK,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.STOP_ZONE,\n", + "]\n", + "\n", + "size = 8\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "for map_layer in misc_map_surfaces:\n", + " map_surface_objects: List[BaseMapSurfaceObject] = map_object_dict[map_layer]\n", + "\n", + " # Here, we just use the default 123D colors for each map layer.\n", + " map_layer_color = str(MAP_SURFACE_CONFIG[map_layer].fill_color)\n", + "\n", + " for map_surface_object in map_surface_objects:\n", + " add_map_surface_object(\n", + " ax,\n", + " map_surface_object,\n", + " label=f\"{map_layer.name}\",\n", + " alpha=1.0,\n", + " color=map_layer_color,\n", + " add_id=False,\n", + " )\n", + "\n", + "adjust_aspect_custom(ax)\n", + "add_non_repeating_legend_to_ax(ax)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "27", + "metadata": {}, + "source": [ + "### 2.5.5 `RoadEdge`\n", + "\n", + "In contrast to all previous road objects, we now have a look at map objects that are lines. \n", + "\n", + "Road edges are polylines (either in 2D or 3D) that mark the edge between drivable surfaces and non-drivable areas.\n", + "The concept of road edges is based on the [Waymo Open Datasets (Motion / Perception)](https://waymo.com/open/), where drivable areas are **not** represented by outlines or polygons.\n", + "While we have a method to convert the Waymo map to a more polygon-based representation, this conversion method cannot fully recover the drivable area.\n", + "\n", + "Therefore, we added road edges to the 123D map. The road edges of non-Waymo maps are extracted from the 3D/2D surfaces of the maps. \n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "28", + "metadata": {}, + "source": [ + "\n", + "\n", + "Road edges have a `RoadEdgeType`, which is one of:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "29", + "metadata": {}, + "outputs": [], + "source": [ + "for road_edege_type in RoadEdgeType:\n", + " print(f\"- {road_edege_type}\")" + ] + }, + { + "cell_type": "markdown", + "id": "30", + "metadata": {}, + "source": [ + "We will plot some road edges from our previous query. We will also add other drivable surfaces as polygons." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "31", + "metadata": {}, + "outputs": [], + "source": [ + "drivable_surface_layers = [\n", + " MapLayer.LANE_GROUP,\n", + " MapLayer.INTERSECTION,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + "]\n", + "\n", + "road_edge_color_map = {\n", + " RoadEdgeType.UNKNOWN: \"red\",\n", + " RoadEdgeType.ROAD_EDGE_BOUNDARY: \"orange\",\n", + " RoadEdgeType.ROAD_EDGE_MEDIAN: \"blue\",\n", + "}\n", + "\n", + "fig, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)\n", + "\n", + "for drivable_surface_layer in drivable_surface_layers:\n", + " for drivable_surface_object in map_object_dict[drivable_surface_layer]:\n", + " add_map_surface_object(\n", + " ax,\n", + " drivable_surface_object,\n", + " label=\"All drivable surfaces\",\n", + " color=\"lightgrey\",\n", + " alpha=1.0,\n", + " add_id=False,\n", + " zorder=0,\n", + " )\n", + "\n", + "for road_edge in map_object_dict[MapLayer.ROAD_EDGE]:\n", + " add_map_line_object(\n", + " ax,\n", + " road_edge,\n", + " label=f\"Road Edge: {road_edge.road_edge_type.name}\",\n", + " color=road_edge_color_map.get(road_edge.road_edge_type, \"black\"),\n", + " linewidth=2.0,\n", + " add_id=False,\n", + " zorder=1,\n", + " )\n", + "\n", + "add_non_repeating_legend_to_ax(ax)\n", + "ax.set_xlim(query_point.x - query_radius / 2, query_point.x + query_radius / 2)\n", + "ax.set_ylim(query_point.y - query_radius / 2, query_point.y + query_radius / 2)\n", + "\n", + "# Change background color to see the lines better\n", + "ax.set_facecolor(\"darkseagreen\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "32", + "metadata": {}, + "source": [ + "### 2.5.6 `RoadLine`\n", + "\n", + "Road lines are lane markings that are in the map, either defined by a 3D or 2D polyline. A road line has a `RoadLineType`, which is one of:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33", + "metadata": {}, + "outputs": [], + "source": [ + "for road_line_type in RoadLineType:\n", + " print(f\"- {road_line_type}\")" + ] + }, + { + "cell_type": "markdown", + "id": "34", + "metadata": {}, + "source": [ + "In some datasets, the road line is equivalent to left/right boundaries. A lane can have various road lines along its boundaries, making referencing to lanes difficult. We currently include that as separate map objects." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35", + "metadata": {}, + "outputs": [], + "source": [ + "road_line_color_map = {\n", + " RoadLineType.NONE: \"grey\",\n", + " RoadLineType.UNKNOWN: \"red\",\n", + " RoadLineType.DASH_SOLID_YELLOW: \"yellow\",\n", + " RoadLineType.DASH_SOLID_WHITE: \"white\",\n", + " RoadLineType.DASHED_WHITE: \"white\",\n", + " RoadLineType.DASHED_YELLOW: \"yellow\",\n", + " RoadLineType.DOUBLE_SOLID_YELLOW: \"yellow\",\n", + " RoadLineType.DOUBLE_SOLID_WHITE: \"white\",\n", + " RoadLineType.DOUBLE_DASH_YELLOW: \"yellow\",\n", + " RoadLineType.DOUBLE_DASH_WHITE: \"white\",\n", + " RoadLineType.SOLID_YELLOW: \"yellow\",\n", + " RoadLineType.SOLID_WHITE: \"white\",\n", + " RoadLineType.SOLID_DASH_WHITE: \"white\",\n", + " RoadLineType.SOLID_DASH_YELLOW: \"yellow\",\n", + " RoadLineType.SOLID_BLUE: \"blue\",\n", + "}\n", + "\n", + "road_line_style_map = {\n", + " RoadLineType.DASHED_WHITE: \"--\",\n", + " RoadLineType.DASHED_YELLOW: \"--\",\n", + " RoadLineType.DOUBLE_DASH_YELLOW: \"--\",\n", + " RoadLineType.DOUBLE_DASH_WHITE: \"--\",\n", + "}\n", + "\n", + "road_line_width_map = {\n", + " RoadLineType.DOUBLE_SOLID_YELLOW: 3.0,\n", + " RoadLineType.DOUBLE_SOLID_WHITE: 3.0,\n", + " RoadLineType.DOUBLE_DASH_YELLOW: 3.0,\n", + " RoadLineType.DOUBLE_DASH_WHITE: 3.0,\n", + "}\n", + "\n", + "fig, ax = plt.subplots(figsize=DEFAULT_FIGSIZE)\n", + "\n", + "for drivable_surface_layer in drivable_surface_layers:\n", + " for drivable_surface_object in map_object_dict[drivable_surface_layer]:\n", + " add_map_surface_object(\n", + " ax,\n", + " drivable_surface_object,\n", + " label=\"All drivable surfaces\",\n", + " color=\"lightgrey\",\n", + " alpha=1.0,\n", + " add_id=False,\n", + " zorder=0,\n", + " )\n", + "\n", + "for road_line in map_object_dict[MapLayer.ROAD_LINE]:\n", + " line_color = road_line_color_map.get(road_line.road_line_type, \"black\")\n", + " line_style = road_line_style_map.get(road_line.road_line_type, \"-\")\n", + " line_width = road_line_width_map.get(road_line.road_line_type, 2.0)\n", + "\n", + " add_map_line_object(\n", + " ax,\n", + " road_line,\n", + " label=f\"Road Line: {road_line.road_line_type.name}\",\n", + " color=line_color,\n", + " linestyle=line_style,\n", + " linewidth=line_width,\n", + " add_id=False,\n", + " zorder=1,\n", + " )\n", + "\n", + "add_non_repeating_legend_to_ax(ax)\n", + "ax.set_xlim(query_point.x - query_radius / 2, query_point.x + query_radius / 2)\n", + "ax.set_ylim(query_point.y - query_radius / 2, query_point.y + query_radius / 2)\n", + "\n", + "# Change background color to see the road lines better\n", + "ax.set_facecolor(\"darkseagreen\")\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "36", + "metadata": {}, + "source": [ + "You made it to end. You can repeat the tutorial for different datasets and filtering settings." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "py123d_dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tutorials/03_visualizations_tutorial.ipynb b/tutorials/03_visualizations_tutorial.ipynb new file mode 100644 index 00000000..eee28ea2 --- /dev/null +++ b/tutorials/03_visualizations_tutorial.ipynb @@ -0,0 +1,296 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "

\n", + " \n", + " \n", + " \n", + " \"Logo\"\n", + " \n", + "

123D: Visualization Tutorial

\n", + "" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "from py123d.api import MapAPI, SceneAPI, SceneFilter\n", + "from py123d.api.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", + "from py123d.common.multithreading.worker_parallel import SingleMachineParallelExecutor" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "## 3.1 Download Demo Logs\n", + "\n", + "You can download demo logs for 123D as described in the [documentation](https://autonomousvision.github.io/py123d/installation/). After the installation and download, you can start with the tutorial.\n" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 3.2 Create Scenes by filtering the datasets\n", + "\n", + "As in other tutorials, we first query some scenes fro visualization. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.datatypes.sensors.pinhole_camera import PinholeCameraType\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=None,\n", + " duration_s=7.0, # No duration means that the scene will include the complete log.\n", + " timestamp_threshold_s=0.0,\n", + " shuffle=True,\n", + " map_api_required=False, # Only include scenes/logs with an available map API.\n", + " pinhole_camera_types=[PinholeCameraType.PCAM_F0],\n", + ")\n", + "worker = SingleMachineParallelExecutor()\n", + "scenes = ArrowSceneBuilder().get_scenes(scene_filter, worker)\n", + "\n", + "dataset_splits = set(scene.log_metadata.split for scene in scenes)\n", + "print(f\"Found {len(scenes)} scenes from {len(dataset_splits)} datasplits:\")\n", + "for split in dataset_splits:\n", + " print(f\" - {split}\")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 3.3 Matplotlib\n" + ] + }, + { + "cell_type": "markdown", + "id": "6", + "metadata": {}, + "source": [ + "### 3.3.1 Plots in 2D\n", + "\n", + "Below, we added some standard plot to visualize an iteration of a scene in birds-eye-view:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.visualization.matplotlib.plots import add_scene_on_ax\n", + "\n", + "scene: SceneAPI = np.random.choice(scenes)\n", + "map_api: MapAPI = scene.get_map_api()\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "add_scene_on_ax(ax, scene, iteration=0, radius=80)\n", + "plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "8", + "metadata": {}, + "source": [ + "We can also render an animation of the scene with the `render_scene_animation` function. You can use `mp4` or `gif` formats. The current scene will be saved in `./visualization`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "from py123d.visualization.matplotlib.plots import render_scene_animation\n", + "\n", + "save_path = Path(\"./visualization\")\n", + "save_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + "fps = 1 / scene.log_metadata.timestep_seconds\n", + "\n", + "render_scene_animation(\n", + " scene=scene,\n", + " output_path=save_path,\n", + " start_idx=0,\n", + " end_idx=None,\n", + " step=1,\n", + " fps=fps * 2, # Let's speed it up a bit\n", + " dpi=300,\n", + " format=\"mp4\",\n", + " radius=80,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "10", + "metadata": {}, + "source": [ + "### 3.3.2 Plots of Cameras\n", + "\n", + "Below, we have some plot to visualize camera observations with matplotlib. In the simples form, you can retrieve a random camera and add it to a plot:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.visualization.matplotlib.camera import add_pinhole_camera_ax\n", + "\n", + "iteration = 0\n", + "available_pinhole_cameras = scene.available_pinhole_camera_types\n", + "if len(available_pinhole_cameras) > 1:\n", + " pinhole_camera_type = np.random.choice(available_pinhole_cameras)\n", + " pinhole_camera = scene.get_pinhole_camera_at_iteration(iteration=iteration, camera_type=pinhole_camera_type)\n", + "\n", + " fig, ax = plt.subplots(figsize=(8, 6))\n", + " add_pinhole_camera_ax(ax, pinhole_camera)\n", + " ax.set_title(f\"{pinhole_camera_type} at Iteration {iteration}\")\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "12", + "metadata": {}, + "source": [ + "You can also visualize the bounding boxes in overlaid in the image with:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.visualization.matplotlib.camera import add_box_detections_to_camera_ax\n", + "\n", + "iteration = 0\n", + "available_pinhole_cameras = scene.available_pinhole_camera_types\n", + "if len(available_pinhole_cameras) > 1:\n", + " pinhole_camera_type = np.random.choice(available_pinhole_cameras)\n", + " pinhole_camera = scene.get_pinhole_camera_at_iteration(iteration=iteration, camera_type=pinhole_camera_type)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration=iteration)\n", + " ego_state = scene.get_ego_state_at_iteration(iteration=iteration)\n", + "\n", + " fig, ax = plt.subplots(figsize=(8, 6))\n", + " add_box_detections_to_camera_ax(ax, pinhole_camera, box_detections, ego_state)\n", + " ax.set_title(f\"{pinhole_camera_type} at Iteration {iteration}\")\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "14", + "metadata": {}, + "source": [ + "The same applies for the LiDAR point cloud. Here we add the points to the image of a random camera and LiDAR scanner. The color map reflects the distance to the ego vehicle." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.visualization.matplotlib.camera import add_lidar_to_camera_ax\n", + "\n", + "available_pinhole_cameras = scene.available_pinhole_camera_types\n", + "available_lidars = scene.available_lidar_types\n", + "\n", + "iteration = 0\n", + "if len(available_pinhole_cameras) > 1 and len(available_lidars) > 0:\n", + " pinhole_camera_type = np.random.choice(available_pinhole_cameras)\n", + " pinhole_camera = scene.get_pinhole_camera_at_iteration(iteration=iteration, camera_type=pinhole_camera_type)\n", + " lidar_type = np.random.choice(available_lidars)\n", + " lidar = scene.get_lidar_at_iteration(iteration=iteration, lidar_type=lidar_type)\n", + "\n", + " fig, ax = plt.subplots(figsize=(8, 6))\n", + " add_lidar_to_camera_ax(ax, pinhole_camera, lidar)\n", + " ax.set_title(f\"{pinhole_camera_type} & {lidar_type} at Iteration {iteration}\")\n", + " plt.show()" + ] + }, + { + "cell_type": "markdown", + "id": "16", + "metadata": {}, + "source": [ + "## 3.4. Viser\n", + "\n", + "You can visualize a scene in 3D with our viser viewer. \n", + "You can also run viser with:\n", + "```sh\n", + "py123d-viser scene_filter=...\n", + "```\n", + "By default, you can open viewer via: [`http://localhost:8080`](http://localhost:8080 )\n", + "\n", + "Check out the [viser documentation](https://viser.studio/main/) for further information." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "from py123d.visualization.viser.viser_viewer import ViserViewer\n", + "\n", + "ViserViewer(scenes)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "py123d_dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.12" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}