Skip to content

Commit d579957

Browse files
MichaelSchuldesMichaelSchuldes
committed
switch from pandas to polars
1 parent e223fc9 commit d579957

File tree

10 files changed

+238
-131
lines changed

10 files changed

+238
-131
lines changed

‎omega_prime/asam_odr/opendriveconverter/elements/lane.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -341,7 +341,7 @@ def set_lanes(
341341
)
342342
# set lane class (none, intersection or roundabout)
343343
if lane_class["intersection"] or lane_class["roundabout"]:
344-
my_lane.classification = betterosi.LaneClassificationType.TYPE_INTERSECTION
344+
my_lane.type = betterosi.LaneClassificationType.TYPE_INTERSECTION
345345

346346
# add lane to road
347347
my_road.lanes.update({len(my_road.lanes): my_lane})

‎omega_prime/cli.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from pathlib import Path
22
from typing import Annotated
33

4-
import pandas as pd
4+
import polars as pl
55
import typer
66

77
import omega_prime
@@ -40,7 +40,7 @@ def from_csv(
4040
validate: bool = True,
4141
skip_odr_parse: bool = False,
4242
):
43-
df = pd.read_csv(input)
43+
df = pl.read_csv(input)
4444
r = omega_prime.Recording(df, validate=validate)
4545
if odr is not None:
4646
r.map = omega_prime.asam_odr.MapOdr.from_file(odr, skip_parse=skip_odr_parse)
@@ -51,7 +51,7 @@ def from_csv(
5151
def validate(
5252
input: Annotated[Path, typer.Argument(help="Path to omega file to validate", exists=True, dir_okay=False)],
5353
):
54-
omega_prime.Recording.from_file(input)
54+
omega_prime.Recording.from_file(input, validate=True)
5555
print(f"File {input} is valid.")
5656

5757

‎omega_prime/converters/lxd.py‎

Lines changed: 51 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import multiprocessing as mp
1010
from ..asam_odr import MapOdr
1111
from ..recording import Recording
12-
12+
import polars as pl
1313

1414
__all__ = ["convert_lxd"]
1515
logger.configure(handlers=[{"sink": sys.stdout, "level": "WARNING"}])
@@ -29,10 +29,6 @@
2929
pedestrians = {"pedestrian": betterosi.MovingObjectType.TYPE_PEDESTRIAN}
3030

3131

32-
def wrap_angle(angle):
33-
return (angle + np.pi) % (2 * np.pi) - np.pi
34-
35-
3632
class DatasetConverter:
3733
def __init__(self, dataset_dir: Path) -> None:
3834
self._dataset = Dataset(dataset_dir)
@@ -45,62 +41,70 @@ def get_recording_opendrive_path(self, recording_id: int) -> Path:
4541

4642
def rec2df(self, recording_id):
4743
rec = self._dataset.get_recording(recording_id)
48-
meta = rec._tracks_meta_data
4944
dt = 1 / rec.get_meta_data("frameRate")
50-
meta["type"] = (
51-
meta["class"]
52-
.apply(
53-
lambda x: betterosi.MovingObjectType.TYPE_VEHICLE
54-
if x in vehicles
55-
else betterosi.MovingObjectType.TYPE_PEDESTRIAN
45+
46+
meta = rec._tracks_meta_data
47+
meta = meta.with_columns(
48+
pl.col("class")
49+
.map_elements(
50+
(
51+
lambda x: betterosi.MovingObjectType.TYPE_VEHICLE
52+
if x in vehicles
53+
else betterosi.MovingObjectType.TYPE_PEDESTRIAN
54+
),
55+
return_dtype=int,
5656
)
57-
.values
58-
)
59-
meta.loc[:, "role"] = meta["class"].apply(
60-
lambda x: betterosi.MovingObjectVehicleClassificationRole.ROLE_CIVIL if x in vehicles else -1
57+
.alias("type"),
58+
pl.col("class")
59+
.map_elements(
60+
(lambda x: betterosi.MovingObjectVehicleClassificationRole.ROLE_CIVIL if x in vehicles else -1),
61+
return_dtype=int,
62+
)
63+
.alias("role"),
64+
pl.col("class")
65+
.map_elements((lambda x: vehicles[x] if x in vehicles else -1), return_dtype=int)
66+
.alias("subtype"),
6167
)
62-
meta["subtype"] = meta["class"].apply(lambda x: vehicles[x] if x in vehicles else -1)
63-
meta = meta.rename(columns={"trackId": "idx"})
68+
meta = meta.rename({"trackId": "idx"})
6469

6570
tracks = rec._get_tracks_data()
6671
tracks = tracks.rename(
67-
columns={
72+
{
6873
"xCenter": "x",
6974
"yCenter": "y",
7075
"xVelocity": "vel_x",
7176
"yVelocity": "vel_y",
7277
"xAcceleration": "acc_x",
7378
"yAcceleration": "acc_y",
7479
"trackId": "idx",
75-
"width": "width",
76-
"length": "length",
80+
# "width": "width",
81+
# "length": "length",
7782
}
7883
)
79-
for k in ["acc_z", "z", "vel_z", "roll", "pitch"]:
80-
tracks[k] = 0.0
81-
tracks["height"] = 2.0
82-
tracks["yaw"] = wrap_angle(np.deg2rad(tracks["heading"]))
83-
tracks["total_nanos"] = np.array(tracks["frame"] * dt * NANOS_PER_SEC, dtype=int)
84-
tracks = tracks.merge(meta[["role", "idx", "type", "subtype"]], on="idx")
85-
86-
tracks.loc[
87-
(tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
88-
& (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
89-
"width",
90-
] = 0.8
91-
tracks.loc[
92-
(tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
93-
& (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
94-
"length",
95-
] = 2
96-
tracks.loc[
97-
(tracks["type"] == betterosi.MovingObjectType.TYPE_VEHICLE)
98-
& (tracks["subtype"] == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE),
99-
"height",
100-
] = 1.9
101-
tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "width"] = 0.5
102-
tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "length"] = 0.5
103-
tracks.loc[(tracks["type"] == betterosi.MovingObjectType.TYPE_PEDESTRIAN), "height"] = 1.8
84+
tracks = tracks.join(meta.select(["idx", "role", "type", "subtype"]), on="idx", how="left")
85+
is_vehicle = pl.col("type") == betterosi.MovingObjectType.TYPE_VEHICLE
86+
is_bicycle = pl.col("subtype") == betterosi.MovingObjectVehicleClassificationType.TYPE_BICYCLE
87+
is_pedestrian = pl.col("type") == betterosi.MovingObjectType.TYPE_PEDESTRIAN
88+
tracks = tracks.with_columns(
89+
[pl.lit(0.0).alias(k) for k in ["acc_z", "z", "vel_z", "roll", "pitch"]]
90+
+ [
91+
(((pl.col("heading") + np.pi) % (2 * np.pi)) - np.pi).alias("yaw"),
92+
(pl.col("frame") * dt * NANOS_PER_SEC).cast(pl.Int64).alias("total_nanos"),
93+
pl.when(is_vehicle & is_bicycle)
94+
.then(0.8)
95+
.when(is_pedestrian)
96+
.then(0.5)
97+
.otherwise(pl.col("width"))
98+
.alias("width"),
99+
pl.when(is_vehicle & is_bicycle)
100+
.then(2.0)
101+
.when(is_pedestrian)
102+
.then(0.5)
103+
.otherwise(pl.col("length"))
104+
.alias("length"),
105+
pl.when(is_vehicle & is_bicycle).then(1.9).when(is_pedestrian).then(1.8).otherwise(2.0).alias("height"),
106+
]
107+
)
104108

105109
return tracks
106110

@@ -109,7 +113,7 @@ def convert_recording(args):
109113
converter, recording_id, out_filename = args
110114
tracks = converter.rec2df(recording_id)
111115
xodr_path = converter.get_recording_opendrive_path(recording_id)
112-
rec = Recording(df=tracks, map=MapOdr.from_file(xodr_path))
116+
rec = Recording(df=tracks, map=MapOdr.from_file(xodr_path), validate=False)
113117
rec.to_mcap(out_filename)
114118

115119

‎omega_prime/map.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,6 @@ def create(cls, lane: betterosi.Lane):
9595
subtype=betterosi.LaneClassificationSubtype(lane.classification.subtype),
9696
successor_ids=[],
9797
predecessor_ids=[],
98-
right_boundary_id=None,
99-
left_boundary_id=None,
10098
)
10199

102100
def plot(self, ax: plt.Axes):

0 commit comments

Comments
 (0)