Add projection utility and improve standard plots (#21)

* Add bool to apply the proj information right after loding

* Add functionality to plot function

* Use better marker (rectangle) for traffic lights

* Formatting

* Check if the recording has a projection defined

* Change to OP version 0.2.2

* Set apply_proj parameter to false in specific examples

* Set apply_proj to False

* Add execution count to metadata to pass pipeline

* Remove old if statement

---------

Co-authored-by: kallenberg <oliver.kallenberg@ika.rwth-aachen.de>

Author: oliverkallenberg

docs/notebooks/tutorial.ipynb

@@ -27,14 +27,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import omega_prime\n",
     "\n",
     "r = omega_prime.Recording.from_file(\n",
-    "    \"../../example_files/pedestrian.osi\", map_path=\"../../example_files/fabriksgatan.xodr\"\n",
+    "    \"../../example_files/pedestrian.osi\", map_path=\"../../example_files/fabriksgatan.xodr\", apply_proj=False\n",
     ")\n",
     "r.to_mcap(\"example.mcap\")"
    ]
@@ -194,7 +194,7 @@
     }
    ],
    "source": [
-    "r = omega_prime.Recording.from_file(\"example.mcap\", parse_map=True)\n",
+    "r = omega_prime.Recording.from_file(\"example.mcap\", parse_map=True, apply_proj=False)\n",
     "ax = r.plot()\n",
     "ax.set_xlim(-20, 50)\n",
     "ax.set_ylim(-75, 20)"

docs/notebooks/tutorial_locator.ipynb

@@ -19,12 +19,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "r = omega_prime.Recording.from_file(\n",
-    "    \"../../example_files/pedestrian.osi\", map_path=\"../../example_files/fabriksgatan.xodr\"\n",
+    "    \"../../example_files/pedestrian.osi\", map_path=\"../../example_files/fabriksgatan.xodr\", apply_proj=False\n",
     ")"
    ]
   },

omega_prime/recording.py

@@ -584,6 +584,7 @@ def from_file(
         validate: bool = False,
         parse_map: bool = False,
         step_size: float = 0.01,
+        apply_proj: bool = True,
         **kwargs,
     ) -> "Recording":
         """Load a Recording from a file. Supports `.parquet`, `.osi` and `.mcap` files.
@@ -594,6 +595,7 @@ def from_file(
             validate (bool): Whether to validate the data against the schema.
             parse_map (bool): Whether to create python objects from the map data or just load it.
             step_size (float): Step size for map parsing, if applicable (Used for ASAM OpenDRIVE).
+            apply_proj (bool): Whether to apply projection transformations to the recording's moving object data.
 
         Returns:
             Recording (Recording): The loaded Recording object.
@@ -624,6 +626,12 @@ def from_file(
             r.map = map
         elif r.map is None:
             warn(f"No map could be found: {map_parsing}")
+
+        if r.projections and apply_proj:
+            try:
+                r.apply_projections()
+            except Exception:
+                warn("Failed to apply projections.")
         return r
 
     def to_file(self, filepath):
@@ -824,28 +832,56 @@ def interpolate(self, new_nanos: list[int] | None = None, hz: float | None = Non
         new_df = pl.concat(new_dfs)
         return self.__init__(df=new_df, map=self.map, host_vehicle_idx=self.host_vehicle_idx)
 
-    def plot(self, ax=None, legend=False) -> plt.Axes:
+    def _create_legend(self, ax):
+        handles, labels = ax.get_legend_handles_labels()
+        host_label = f"{self.host_vehicle_idx} - HV"
+
+        def sort_key(item):
+            label = item[1]
+            if label == host_label:
+                return (-1, -1)
+            try:
+                return (0, int(label))
+            except ValueError:
+                return (0, float("inf"))  # non-numeric labels go last
+
+        items = sorted(zip(handles, labels), key=sort_key)
+        handles, labels = zip(*items)
+        ax.legend(handles, labels, loc="center left", bbox_to_anchor=(1, 0.5))
+        return ax
+
+    def plot(self, ax=None, legend=False, mvs_plt_type: str = "scatter") -> plt.Axes:
         "Generate a static plot of the recording using Matplotlib. Plots the map (if available), moving objects, and traffic light states."
         if ax is None:
             fig, ax = plt.subplots(1, 1)
             ax.set_aspect(1)
         if self.map:
             self.map.plot(ax)
-        self.plot_mvs(ax=ax)
+        self.plot_mvs(ax=ax, mvs_plt_type=mvs_plt_type)
         self.plot_tl(ax=ax)
         if legend:
-            ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+            ax = self._create_legend(ax)
         return ax
 
-    def plot_mvs(self, ax=None, legend=False):
+    def plot_mvs(self, ax=None, legend=False, mvs_plt_type: str = "scatter"):
         "Generate a static plot of the moving objects in the recording using Matplotlib."
         if ax is None:
             fig, ax = plt.subplots(1, 1)
             ax.set_aspect(1)
-        for [idx], mv in self._df["idx", "x", "y"].group_by("idx"):
-            ax.plot(*mv["x", "y"], c="red", alpha=0.5, label=str(idx))
+        plot_fn = {"scatter": ax.scatter, "plot": ax.plot}.get(mvs_plt_type)
+        if plot_fn is None:
+            raise ValueError("`mvs_plt_type` must be one of: 'scatter', 'plot'.")
+
+        plot_df = self._df["idx", "x", "y"]
+        base_kwargs = {"alpha": 0.5}
+        for [idx], mv in plot_df.group_by("idx"):
+            if idx == self.host_vehicle_idx:
+                ax.plot(*mv["x", "y"], c="red", label=f"{idx} - HV")
+                continue
+            plot_fn(*mv["x", "y"], label=str(idx), **base_kwargs)
+
         if legend:
-            ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+            ax = self._create_legend(ax)
         return ax
 
     def plot_tl(self, ax=None):
@@ -866,7 +902,7 @@ def plot_tl(self, ax=None):
                 ax.plot(
                     x,
                     y,
-                    marker="o",
+                    marker="s",
                     label=f"Traffic Light {tl_dict[tl].id.value}",
                     c="blue",
                     alpha=0.7,

pyproject.toml

@@ -48,7 +48,7 @@ dependencies = [
     'tqdm_joblib',
     'filelock>=3.18.0'
 ]
-version = "0.2.1"
+version = "0.2.2"
 
 [project.urls]
 Homepage = "https://github.com/ika-rwth-aachen/omega-prime"