p42_perception_node.py

#!/usr/bin/env python3
"""
ROS2 Perception Node - Worker process for perception compute.

Subscribes to localization and ground-truth agents from the in-process world,
runs the perception strategy, and publishes tracked objects.
"""
import argparse
import json
import logging
import os
import sys

import rclpy
from rclpy.node import Node
from std_msgs.msg import String

from avlite import ExecutionSettings, PerceptionModel, PerceptionStrategy, SensorFrame
from avlite.c10_perception.c11_perception_model import EgoState
from avlite.c50_common.c56_fps_tracker import FpsTracker
from .p46_autoware_converters import (
    AUTOWARE_AVAILABLE,
    ego_state_from_kinematic_state,
    agents_from_bounding_boxes,
)
from .spawn_commands import agents_from_perception_json
from .lidar_commands import parse_lidar_cloud
from .pm_snapshot_commands import encode_pm_snapshot
from .settings import PluginSettings

log = logging.getLogger(__name__)

if AUTOWARE_AVAILABLE:
    from autoware_auto_msgs.msg import VehicleKinematicState
    from autoware_auto_msgs.msg import BoundingBoxArray, BoundingBox


class PerceptionNode(Node):
    """ROS2 worker that runs perception and publishes tracked objects."""

    def __init__(
        self,
        perception_model: PerceptionModel,
        perception: PerceptionStrategy | None = None,
        perception_dt: float | None = None,
    ):
        super().__init__("avlite_perception")

        self.settings = PluginSettings
        self.pm = perception_model
        self.perception = perception
        self.use_autoware = AUTOWARE_AVAILABLE and self.settings.use_autoware_msgs

        perception_dt = perception_dt if perception_dt is not None else ExecutionSettings.c40_perception_dt
        self.declare_parameter("perception_dt", perception_dt)
        perception_dt = self.get_parameter("perception_dt").get_parameter_value().double_value
        pace_perception = bool(getattr(self.settings, "pace_perception", True))

        self._fps_tracker = FpsTracker()
        self._shutdown = False
        self._latest_lidar = None

        self._setup_subscriptions()
        if self.use_autoware:
            self.objects_pub = self.create_publisher(
                BoundingBoxArray,
                self.settings.perception_topic,
                10,
            )
        else:
            self.objects_pub = self.create_publisher(String, self.settings.perception_topic, 10)
        self.viz_pub = self.create_publisher(String, self.settings.perception_viz_topic, 10)

        timer_period = perception_dt if pace_perception else 0.001
        self.timer = self.create_timer(timer_period, self._perception_tick)
        self.get_logger().info(
            f"PerceptionNode worker started (autoware={self.use_autoware}, "
            f"{1.0 / timer_period:.1f} Hz, pace_perception={pace_perception})"
        )

    def _setup_subscriptions(self) -> None:
        if self.use_autoware:
            self.create_subscription(
                VehicleKinematicState,
                self.settings.localization_topic,
                self._on_localization,
                10,
            )
            self.create_subscription(
                BoundingBoxArray,
                self.settings.world_gt_topic,
                self._on_world_gt,
                10,
            )
        else:
            self.create_subscription(
                String,
                self.settings.localization_topic,
                self._on_localization_json,
                10,
            )
            self.create_subscription(
                String,
                self.settings.world_gt_topic,
                self._on_world_gt_json,
                10,
            )
        self.create_subscription(
            String,
            self.settings.lidar_topic,
            self._on_lidar_json,
            10,
        )

    def _on_lidar_json(self, msg: String) -> None:
        try:
            self._latest_lidar = parse_lidar_cloud(msg.data)
        except json.JSONDecodeError:
            pass

    def _on_localization(self, msg: "VehicleKinematicState") -> None:
        if self.pm is None:
            return
        ego_state_from_kinematic_state(msg, self.pm.ego_vehicle)

    def _on_localization_json(self, msg: String) -> None:
        if self.pm is None:
            return
        try:
            data = json.loads(msg.data)
            ego = self.pm.ego_vehicle
            ego.x = data.get("x", ego.x)
            ego.y = data.get("y", ego.y)
            ego.theta = data.get("theta", ego.theta)
            ego.velocity = data.get("velocity", ego.velocity)
        except json.JSONDecodeError:
            pass

    def _on_world_gt(self, msg: "BoundingBoxArray") -> None:
        if self.pm is not None:
            self.pm.agent_vehicles = agents_from_bounding_boxes(msg)

    def _on_world_gt_json(self, msg: String) -> None:
        if self.pm is None:
            return
        try:
            data = json.loads(msg.data)
            self.pm.agent_vehicles = agents_from_perception_json(data)
        except json.JSONDecodeError:
            pass

    def _run_perception(self) -> None:
        if self.perception is None or self.pm is None:
            return
        try:
            sensors = SensorFrame(lidar=self._latest_lidar)
            self.perception.perceive(perception_model=self.pm, sensors=sensors)
        except Exception as e:
            log.debug("Perception worker step failed: %s", e)

    def _perception_tick(self) -> None:
        if self._shutdown or not rclpy.ok():
            return
        self._run_perception()
        self._publish_tracked_objects()
        self._publish_pm_snapshot()
        self._fps_tracker.tick()

    def _publish_pm_snapshot(self) -> None:
        if self.pm is None:
            return
        try:
            msg = String()
            msg.data = encode_pm_snapshot(self.pm)
            self.viz_pub.publish(msg)
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Failed to publish PM snapshot: {e}")

    def _publish_tracked_objects(self) -> None:
        if self.pm is None:
            return
        agents = getattr(self.pm, "agent_vehicles", []) or []
        try:
            if self.use_autoware:
                msg = BoundingBoxArray()
                msg.header.stamp = self.get_clock().now().to_msg()
                msg.header.frame_id = self.settings.map_frame
                for agent in agents:
                    box = BoundingBox()
                    box.centroid.x = float(agent.x)
                    box.centroid.y = float(agent.y)
                    box.centroid.z = 0.0
                    box.size.x = float(getattr(agent, "length", 4.5))
                    box.size.y = float(getattr(agent, "width", 2.0))
                    box.size.z = 1.5
                    box.heading = float(agent.theta)
                    box.velocity = float(agent.velocity)
                    box.vehicle_label = 1
                    msg.boxes.append(box)
            else:
                msg = String()
                objects_list = []
                for agent in agents:
                    objects_list.append({
                        "id": getattr(agent, "agent_id", 0),
                        "x": float(agent.x),
                        "y": float(agent.y),
                        "theta": float(agent.theta),
                        "velocity": float(agent.velocity),
                    })
                msg.data = json.dumps({"objects": objects_list})
            self.objects_pub.publish(msg)
        except (rclpy.exceptions.InvalidHandle, RuntimeError):
            pass
        except Exception as e:
            if not self._shutdown:
                self.get_logger().error(f"Failed to publish perception: {e}")

    def destroy_node(self):
        self._shutdown = True
        if self.timer:
            self.timer.cancel()
        super().destroy_node()


def _parse_args(argv=None):
    parser = argparse.ArgumentParser(description="AVLite perception ROS worker")
    parser.add_argument(
        "--profile",
        default=os.environ.get("AVLITE_PROFILE", "default"),
    )
    return parser.parse_args(argv)


def main(args=None):
    parsed = _parse_args(args)
    os.environ["AVLITE_PROFILE"] = parsed.profile

    from .p48_node_bootstrap import attach_worker_logging, bootstrap_role, spin_node

    logging.basicConfig(level=logging.WARNING)
    boot = bootstrap_role("perception", profile=parsed.profile)
    if boot.pm is None:
        log.error("Perception bootstrap failed")
        sys.exit(1)

    rclpy.init(args=None)
    node = PerceptionNode(
        perception_model=boot.pm,
        perception=boot.perception,
    )
    attach_worker_logging(node, "avlite.c10_perception")
    try:
        spin_node(node)
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    main()