Add vehicle of interest coloring.

CHANGELOG.md

@@ -12,6 +12,7 @@ Copy and pasting the git commit messages is __NOT__ enough.
 ### Added
 - Added support for the [Argoverse 2 Motion Forecasting Dataset](https://www.argoverse.org/av2.html#forecasting-link) (see `scenarios/argoverse`)
 - Added `waymo_open_dataset` as a module at the SMARTS repo level, to be able to load waymo scenarios without any external packages
+- Added vehicle of interest coloring through scenario studio. This lets the scenario color vehicles that match a certain pattern of vehicle id.
 ### Changed
 ### Deprecated
 ### Fixed

scenarios/sumo/intersections/4lane_t/scenario.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 
+from smarts.core.colors import Colors
 from smarts.sstudio import gen_scenario
 from smarts.sstudio import types as t
 
@@ -99,6 +100,7 @@
         bubbles=bubbles,
         ego_missions=ego_missions,
         social_agent_missions=social_agent_missions,
+        scenario_metadata=t.ScenarioMetadata(r".*-1.*", Colors.Yellow),
     ),
     output_dir=Path(__file__).parent,
 )

smarts/core/renderer.py

@@ -25,10 +25,12 @@
 import importlib.resources as pkg_resources
 import logging
 import os
+import re
 from enum import IntEnum
 from pathlib import Path
+from re import Pattern
 from threading import Lock
-from typing import NamedTuple
+from typing import NamedTuple, Optional, Union
 
 import gltf
 from direct.showbase.ShowBase import ShowBase
@@ -54,7 +56,7 @@
 )
 
 from . import glsl, models
-from .colors import SceneColors
+from .colors import Colors, SceneColors
 from .coordinates import Pose
 from .masks import RenderMasks
 from .scenario import Scenario
@@ -196,6 +198,8 @@ def __init__(self, simid: str, debug_mode: DEBUG_MODE = DEBUG_MODE.ERROR):
         # Note: Each instance of the SMARTS simulation will have its own Renderer,
         # but all Renderer objects share the same ShowBaseInstance.
         self._showbase_instance: _ShowBaseInstance = _ShowBaseInstance()
+        self._interest_filter: Optional[Pattern] = None
+        self._interest_color: Optional[Union[Colors, SceneColors]] = None
 
     @property
     def id(self):
@@ -304,7 +308,14 @@ def setup(self, scenario: Scenario):
                     "Resolution", self._showbase_instance.getSize()
                 )
             self._dashed_lines_np = dashed_lines_np
-
+        if scenario_metadata := scenario.metadata:
+            if interest_pattern := scenario_metadata.get(
+                "actor_of_interest_re_filter", None
+            ):
+                self._interest_filter = re.compile(interest_pattern)
+                self._interest_color = scenario_metadata.get(
+                    "actor_of_interest_color", SceneColors.SocialAgent
+                )
         self._is_setup = True
 
     def render(self):
@@ -335,12 +346,32 @@ def destroy(self):
     def __del__(self):
         self.destroy()
 
-    def create_vehicle_node(self, glb_model: str, vid: str, color, pose: Pose):
+    def set_interest(self, interest_filter: Pattern, interest_color: Colors):
+        """Sets the color of all vehicles that have ids that match the given pattern.
+
+        Args:
+            interest_filter (Pattern): The regular expression pattern to match.
+            interest_color (Colors): The color that the vehicle should show as.
+        """
+        assert isinstance(interest_filter, Pattern)
+        self._interest_filter = interest_filter
+        self._interest_color = interest_color
+
+    def create_vehicle_node(
+        self, glb_model: str, vid: str, color: Union[Colors, SceneColors], pose: Pose
+    ):
         """Create a vehicle node."""
         with pkg_resources.path(models, glb_model) as path:
             node_path = self._showbase_instance.loader.loadModel(str(path.absolute()))
         node_path.setName("vehicle-%s" % vid)
-        node_path.setColor(color)
+        if (
+            self._interest_filter is not None
+            and self._interest_color is not None
+            and self._interest_filter.match(vid)
+        ):
+            node_path.setColor(self._interest_color.value)
+        else:
+            node_path.setColor(color.value)
         pos, heading = pose.as_panda3d()
         node_path.setPosHpr(*pos, heading, 0, 0)
         node_path.hide(RenderMasks.DRIVABLE_AREA_HIDE)
@@ -350,7 +381,7 @@ def begin_rendering_vehicle(self, vid: str, is_agent: bool):
         """Add the vehicle node to the scene graph"""
         vehicle_path = self._vehicle_nodes.get(vid, None)
         if not vehicle_path:
-            self._log.warning(f"Renderer ignoring invalid vehicle id: {vid}")
+            self._log.warning("Renderer ignoring invalid vehicle id: %s", vid)
             return
         # TAI: consider reparenting hijacked vehicles too?
         vehicle_path.reparentTo(self._vehicles_np if is_agent else self._root_np)
@@ -359,7 +390,7 @@ def update_vehicle_node(self, vid: str, pose: Pose):
         """Move the specified vehicle node."""
         vehicle_path = self._vehicle_nodes.get(vid, None)
         if not vehicle_path:
-            self._log.warning(f"Renderer ignoring invalid vehicle id: {vid}")
+            self._log.warning("Renderer ignoring invalid vehicle id: %s", vid)
             return
         pos, heading = pose.as_panda3d()
         vehicle_path.setPosHpr(*pos, heading, 0, 0)

smarts/core/scenario.py

@@ -125,6 +125,7 @@ def __init__(
             self._traffic_specs = [os.path.join(traffic_path, route)]
         self._missions = missions or {}
         self._bubbles = Scenario._discover_bubbles(scenario_root)
+        self._metadata = Scenario._discover_metadata(scenario_root)
         self._social_agents = social_agents or {}
         self._surface_patches = surface_patches
         self._log_dir = self._resolve_log_dir(log_dir)
@@ -547,6 +548,16 @@ def _discover_bubbles(scenario_root):
             bubbles = pickle.load(f)
             return bubbles
 
+    @staticmethod
+    def _discover_metadata(scenario_root):
+        path = os.path.join(scenario_root, "build", "scenario_metadata.pkl")
+        if not os.path.exists(path):
+            return dict()
+
+        with open(path, "rb") as f:
+            metadata = pickle.load(f)
+            return metadata
+
     def set_ego_missions(self, ego_missions: Dict[str, Mission]):
         """Replaces the ego missions within the scenario.
         Args:
@@ -1053,3 +1064,12 @@ def traffic_history(self) -> Optional[TrafficHistory]:
     def scenario_hash(self) -> str:
         """A hash of the scenario."""
         return self._scenario_hash
+
+    @property
+    def metadata(self) -> Dict:
+        """Scenario metadata values.
+
+        Returns:
+            Dict: The values.
+        """
+        return self._metadata or {}

smarts/core/smarts.py

@@ -1074,7 +1074,6 @@ def _pybullet_provider_sync(self, provider_state: ProviderState):
                         vehicle_state=vehicle,
                         actor_id=vehicle_id,
                         vehicle_id=vehicle_id,
-                        vehicle_config_type=vehicle.vehicle_config_type,
                     )
 
                 if not vehicle.updated:

smarts/core/tests/test_renderers.py

@@ -124,7 +124,7 @@ def test_renderer(self):
             heading_=Heading(math.pi * 0.91),
         )
         self._rdr.create_vehicle_node(
-            "simple_car.glb", self._vid, SceneColors.SocialVehicle.value, pose
+            "simple_car.glb", self._vid, SceneColors.SocialVehicle, pose
         )
         self._rdr.begin_rendering_vehicle(self._vid, is_agent=False)
         for s in range(self._num_steps):

smarts/core/vehicle.py

@@ -106,43 +106,43 @@ class VehicleConfig:
 VEHICLE_CONFIGS = {
     "passenger": VehicleConfig(
         vehicle_type="car",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=3.68, width=1.47, height=1.4),
         glb_model="simple_car.glb",
     ),
     "bus": VehicleConfig(
         vehicle_type="bus",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=7, width=2.25, height=3),
         glb_model="bus.glb",
     ),
     "coach": VehicleConfig(
         vehicle_type="coach",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=8, width=2.4, height=3.5),
         glb_model="coach.glb",
     ),
     "truck": VehicleConfig(
         vehicle_type="truck",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=5, width=1.91, height=1.89),
         glb_model="truck.glb",
     ),
     "trailer": VehicleConfig(
         vehicle_type="trailer",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=10, width=2.5, height=4),
         glb_model="trailer.glb",
     ),
     "pedestrian": VehicleConfig(
         vehicle_type="pedestrian",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=0.5, width=0.5, height=1.6),
         glb_model="pedestrian.glb",
     ),
     "motorcycle": VehicleConfig(
         vehicle_type="motorcycle",
-        color=SceneColors.SocialVehicle.value,
+        color=SceneColors.SocialVehicle,
         dimensions=Dimensions(length=2.5, width=1, height=1.4),
         glb_model="motorcycle.glb",
     ),
@@ -252,7 +252,7 @@ def renderer(self):  # type: ignore
     #     self._chassis.speed = speed
 
     @property
-    def vehicle_color(self) -> Union[SceneColors, Tuple, None]:
+    def vehicle_color(self) -> Union[SceneColors, None]:
         """The color of this vehicle (generally used for rendering purposes.)"""
         self._assert_initialized()
         return self._color
@@ -390,9 +390,7 @@ def build_agent_vehicle(
         else:
             start_pose = Pose.from_center(start.position, start.heading)
 
-        vehicle_color = (
-            SceneColors.Agent.value if trainable else SceneColors.SocialAgent.value
-        )
+        vehicle_color = SceneColors.Agent if trainable else SceneColors.SocialAgent
 
         if agent_interface.vehicle_type == "sedan":
             urdf_name = "vehicle"
@@ -442,12 +440,11 @@ def build_agent_vehicle(
         return vehicle
 
     @staticmethod
-    def build_social_vehicle(
-        sim, vehicle_id, vehicle_state, vehicle_config_type
-    ) -> "Vehicle":
+    def build_social_vehicle(sim, vehicle_id, vehicle_state: VehicleState) -> "Vehicle":
         """Create a new unassociated vehicle."""
         dims = Dimensions.copy_with_defaults(
-            vehicle_state.dimensions, VEHICLE_CONFIGS[vehicle_config_type].dimensions
+            vehicle_state.dimensions,
+            VEHICLE_CONFIGS[vehicle_state.vehicle_config_type].dimensions,
         )
         chassis = BoxChassis(
             pose=vehicle_state.pose,
@@ -456,7 +453,9 @@ def build_social_vehicle(
             bullet_client=sim.bc,
         )
         return Vehicle(
-            id=vehicle_id, chassis=chassis, vehicle_config_type=vehicle_config_type
+            id=vehicle_id,
+            chassis=chassis,
+            vehicle_config_type=vehicle_state.vehicle_config_type,
         )
 
     @staticmethod

smarts/core/vehicle_index.py

@@ -734,14 +734,16 @@ def _enfranchise_actor(
 
     @clear_cache
     def build_social_vehicle(
-        self, sim, vehicle_state, actor_id, vehicle_config_type, vehicle_id=None
+        self, sim, vehicle_state, actor_id, vehicle_id=None
     ) -> Vehicle:
         """Build an entirely new vehicle for a social agent."""
         if vehicle_id is None:
             vehicle_id = gen_id()
 
         vehicle = Vehicle.build_social_vehicle(
-            sim, vehicle_id, vehicle_state, vehicle_config_type
+            sim,
+            vehicle_id,
+            vehicle_state,
         )
 
         vehicle_id, actor_id = _2id(vehicle_id), _2id(actor_id)

smarts/sstudio/genscenario.py

@@ -28,7 +28,7 @@
 import os
 import pickle
 import sqlite3
-from dataclasses import replace
+from dataclasses import asdict, replace
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
@@ -92,6 +92,9 @@ def _build_graph(scenario: types.Scenario, base_dir: str) -> Dict[str, Any]:
             artifact_path = os.path.join(base_dir, f"{dataset.name}.shf")
             graph["traffic_histories"].append(artifact_path)
 
+    if scenario.scenario_metadata is not None:
+        graph["scenario_metadata"] = [os.path.join(base_dir, "scenario_metadata.pkl")]
+
     return graph
 
 
@@ -328,6 +331,18 @@ def gen_scenario(
             )
             _update_artifacts(db_conn, artifact_paths, obj_hash)
 
+    # Scenario metadata
+    artifact_paths = build_graph["scenario_metadata"]
+    obj_hash = pickle_hash(scenario.scenario_metadata, True)
+    if _needs_build(
+        db_conn, scenario.scenario_metadata, artifact_paths, obj_hash, map_needs_rebuild
+    ):
+        with timeit("scenario_metadata", logger.info):
+            gen_metadata(
+                scenario=output_dir,
+                scenario_metadata=scenario.scenario_metadata,
+            )
+
 
 def gen_map(scenario: str, map_spec: types.MapSpec, output_dir: Optional[str] = None):
     """Saves a map spec to file."""
@@ -658,3 +673,16 @@ def gen_traffic_histories(
                 )
         output_dir = os.path.join(scenario, "build")
         genhistories.import_dataset(hdsr, output_dir, map_bbox)
+
+
+def gen_metadata(scenario: str, scenario_metadata: types.ScenarioMetadata):
+    """Generate the metadata for the scenario
+
+    Args:
+        scenario (str):The scenario directory
+        scenario_metadata (types.ScenarioMetadata): _description_
+    """
+    _check_if_called_externally()
+    output_path = os.path.join(scenario, "build", "scenario_metadata.pkl")
+    with open(output_path, "wb") as f:
+        pickle.dump(asdict(scenario_metadata), f)