Feat/yolop shortcut

.gitignore

@@ -20,3 +20,7 @@ simulator_old/
 
 
 #例外ファイルを指定
+zed-ros2-wrapper/
+
+#トポマップデータ (環境依存のため除外)
+e2e_planner/config/topomap/

e2e_planner/e2e_planner/inference_node.py

@@ -1,24 +1,29 @@
+import argparse
 import rclpy
 from rclpy.node import Node
 from rclpy.callback_groups import ReentrantCallbackGroup
 from rclpy.executors import MultiThreadedExecutor
 from sensor_msgs.msg import Image, PointCloud2
-from std_msgs.msg import Header
+from std_msgs.msg import Header, UInt8
 from nav_msgs.msg import Path
 from geometry_msgs.msg import PoseStamped, Pose, Point
 from cv_bridge import CvBridge
 import cv2
 import torch
+from torchvision import transforms
 import numpy as np
 import os
+import sys
 from pathlib import Path as FilePath
 from rclpy.qos import qos_profile_system_default, qos_profile_sensor_data
 from ament_index_python.packages import get_package_share_directory
 from scipy.interpolate import splprep, splev
-from typing import Tuple
+from typing import Optional, Tuple
 from .zed_sdk import ZedSdk
 
+from e2e_planner.placenav.place_recognition import PlaceRecognition
 from util.yolop_processor import YOLOPv2Processor
+from util.preprocessing import center_square_crop, lane_mask_to_tensor_array, overlay_lane_mask
 
 def denormalize_waypoints(normalized: np.ndarray) -> np.ndarray:
     denormalized = normalized.copy()
@@ -27,86 +32,199 @@ def denormalize_waypoints(normalized: np.ndarray) -> np.ndarray:
     return denormalized
 
 class InferenceNode(Node):
-    def __init__(self) -> None:
+    def __init__(self, simulator_mode: bool = False) -> None:
         super().__init__('inference_node')
 
         self.declare_parameter('model_name', 'model.pt')
         self.declare_parameter('interval_ms', 100)
+        self.declare_parameter('sim_flag', simulator_mode)
+        self.declare_parameter('image_topic', '/image_raw')
+        self.declare_parameter('debug_mode', True)
+        self.declare_parameter('default_command', 1)
+        self.declare_parameter('use_place_recognition', False)
+        self.declare_parameter('placenet_model_name', 'placenet.pt')
+        self.declare_parameter('topomap_dir_name', 'topomap')
+        self.declare_parameter('placenet_delta', 10.0)
+        self.declare_parameter('placenet_window_lower', -1)
+        self.declare_parameter('placenet_window_upper', 2)
 
         model_path = self.get_parameter('model_name').value
         interval_ms = self.get_parameter('interval_ms').value
+        self.sim_flag = bool(self.get_parameter('sim_flag').value)
+        image_topic = self.get_parameter('image_topic').value
+        self.debug_mode = bool(self.get_parameter('debug_mode').value)
+        self.command = int(self.get_parameter('default_command').value)
+        self.use_place_recognition = bool(self.get_parameter('use_place_recognition').value)
+        placenet_model_name = self.get_parameter('placenet_model_name').value
+        topomap_dir_name = self.get_parameter('topomap_dir_name').value
+        self.placenet_delta = float(self.get_parameter('placenet_delta').value)
+        self.placenet_window_lower = int(self.get_parameter('placenet_window_lower').value)
+        self.placenet_window_upper = int(self.get_parameter('placenet_window_upper').value)
 
         self.bridge = CvBridge()
-        self.device = torch.device('cuda')
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.zed = None
 
         self.cv_image = None
-
-        self.sim_flag = False
+        self.cv_header: Optional[Header] = None
 
         package_share_directory = get_package_share_directory('e2e_planner')
         weight_path = os.path.join(package_share_directory, 'weights', model_path)
         yolop_weight_path = FilePath(package_share_directory) / 'weights' / 'yolopv2.pt'
+        placenet_weight_path = FilePath(package_share_directory) / 'weights' / placenet_model_name
+        topomap_path = FilePath(package_share_directory) / 'config' / topomap_dir_name / 'topomap.yaml'
 
         if os.path.exists(weight_path):
             self.model = torch.jit.load(weight_path, map_location=self.device)
             self.model.eval()
+            self.model_uses_command = self._model_uses_command(self.model)
+            if not self.model_uses_command:
+                self.get_logger().warn(
+                    'Loaded model accepts only image input; navigation command will be ignored.'
+                )
         else:
             self.get_logger().warn(f'Model file not found: {weight_path}')
             self.model = None
+            self.model_uses_command = False
 
-        self.yolop_processor = YOLOPv2Processor(yolop_weight_path, self.device)
-        self.zed = ZedSdk(self, self.sim_flag)
+        self.yolop_processor = None
+        if yolop_weight_path.exists():
+            self.yolop_processor = YOLOPv2Processor(yolop_weight_path, self.device)
+        else:
+            self.get_logger().warn(f'YOLOPv2 model not found: {yolop_weight_path}')
+        self.place_recognition = None
+        self.placenet_transform = None
+        if self.use_place_recognition:
+            if not placenet_weight_path.exists():
+                self.get_logger().warn(f'Place recognition model not found: {placenet_weight_path}')
+                self.use_place_recognition = False
+            elif not topomap_path.exists():
+                self.get_logger().warn(f'Topomap not found: {topomap_path}')
+                self.use_place_recognition = False
+            else:
+                self.place_recognition = PlaceRecognition(
+                    placenet_weight_path,
+                    topomap_path,
+                    device=self.device,
+                    delta=self.placenet_delta,
+                    window_lower=self.placenet_window_lower,
+                    window_upper=self.placenet_window_upper,
+                )
+                self.placenet_transform = transforms.Compose([
+                    transforms.ToTensor(),
+                    transforms.Resize((85, 85), antialias=True),
+                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                ])
+        self.create_subscription(UInt8, '/command', self.command_callback, qos_profile_system_default)
 
         self.pub_raw = self.create_publisher(Path, 'e2e_planner/path_raw', qos_profile_system_default)
         self.pub = self.create_publisher(Path, 'e2e_planner/path', qos_profile_system_default)
-        self.pub_pointcloud = self.create_publisher(PointCloud2, '/zed/zed_node/pointcloud', qos_profile_sensor_data)
-        self.pub_debug_image = self.create_publisher(Image, 'e2e_planner/debug_image', qos_profile_system_default)
-        self.get_logger().info('Debug mode enabled: publishing preprocessed images to e2e_planner/debug_image')
+        self.pub_pointcloud = None
+        if self.debug_mode:
+            self.pub_debug_image = self.create_publisher(Image, 'e2e_planner/debug_image', qos_profile_system_default)
+            self.get_logger().info('Debug mode enabled: publishing preprocessed images to e2e_planner/debug_image')
 
         self.torch_cb_group = ReentrantCallbackGroup()
         self.zed_cb_group = ReentrantCallbackGroup()
 
+        if self.sim_flag:
+            self.create_subscription(Image, image_topic, self.image_callback, qos_profile_sensor_data)
+            self.get_logger().info(f'Using simulator image topic: {image_topic}')
+        else:
+            self.zed = ZedSdk(self, self.sim_flag)
+            self.pub_pointcloud = self.create_publisher(PointCloud2, '/zed/zed_node/pointcloud', qos_profile_sensor_data)
+            self.zed_timer = self.create_timer(
+                10.0 / 1000.0,
+                self.zed_sensor_callback,
+                callback_group=self.zed_cb_group,
+            )
+
         self.torch_timer = self.create_timer(
             interval_ms / 1000.0,
             self.torch_callback,
             callback_group=self.torch_cb_group,
         )
-        self.zed_timer = self.create_timer(
-            10.0 / 1000.0,
-            self.zed_sensor_callback,
-            callback_group=self.zed_cb_group,
-        )
+
+    def command_callback(self, msg: UInt8) -> None:
+        self.command = int(msg.data)
+
+    def preprocess_command(self, command: int | None = None) -> torch.Tensor:
+        command_tensor = torch.zeros((1, 4), device=self.device, dtype=torch.float32)
+        command_idx = self.command if command is None else int(command)
+        command_idx = min(max(command_idx, 0), 3)
+        command_tensor[0, command_idx] = 1.0
+        return command_tensor
+
+    def _model_uses_command(self, model) -> bool:
+        try:
+            schema = str(model.forward.schema)
+        except RuntimeError:
+            return True
+        return 'Tensor cmd' in schema or 'Tensor command' in schema
+
+    def run_model(self, input_tensor: torch.Tensor, command_tensor: torch.Tensor) -> torch.Tensor:
+        if self.model_uses_command:
+            return self.model(input_tensor, command_tensor)
+        return self.model(input_tensor)
+
+    def preprocess_placenet_image(self, image: np.ndarray) -> torch.Tensor:
+        bgr_image = self._to_bgr(image)
+        cropped_image = center_square_crop(bgr_image)
+        rgb_image = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2RGB)
+        placenet_image_tensor = self.placenet_transform(rgb_image).unsqueeze(0)
+        return placenet_image_tensor.to(self.device, dtype=torch.float32)
+
+    def recognize_command(self, image: np.ndarray) -> int:
+        if self.place_recognition is None:
+            return self.command
+
+        placenet_image_tensor = self.preprocess_placenet_image(image)
+        command, idx = self.place_recognition.get_recognition(placenet_image_tensor)
+        self.command = int(command)
+        self.get_logger().info(f'place recognition: command={command}, idx={idx}')
+        return int(command)
+
+    def _to_bgr(self, image: np.ndarray) -> np.ndarray:
+        if image.shape[2] == 4:
+            return cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
+        return image
 
     def preprocess_image(self, image: np.ndarray) -> Tuple[torch.Tensor, np.ndarray]:
-        if self.sim_flag:
-            mask = ((image[:, :, 2] > 200) & (image[:, :, 0] < 50) & (image[:, :, 1] < 50)).astype(np.uint8)
-            mask = cv2.resize(mask, (64, 48))
-        else:
-            bgr_image = cv2.cvtColor(image, cv2.COLOR_BGRA2BGR)
-            mask = self.yolop_processor.process_image(bgr_image, (64, 48))
+        bgr_image = self._to_bgr(image)
+        if self.yolop_processor is None:
+            raise RuntimeError('YOLOPv2 processor is not available.')
+        mask = self.yolop_processor.process_image(bgr_image)
 
-        mask_normalized = mask.astype(np.float32)
+        mask_normalized = lane_mask_to_tensor_array(mask)
         tensor = torch.from_numpy(mask_normalized).unsqueeze(0).unsqueeze(0)
-        return tensor.to(self.device), mask
+        return tensor.to(self.device), mask_normalized
+
+    def image_callback(self, msg: Image) -> None:
+        self.cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
+        self.cv_header = msg.header
 
     def torch_callback(self) -> None:
-        if self.cv_image is None:
+        if self.cv_image is None or self.model is None or self.yolop_processor is None:
             return
-        header = Header()
-        header.stamp = self.get_clock().now().to_msg()
-        header.frame_id = 'base_link'
+        if self.cv_header is None:
+            header = Header()
+            header.stamp = self.get_clock().now().to_msg()
+            header.frame_id = 'base_link'
+        else:
+            header = self.cv_header
 
         input_tensor, mask = self.preprocess_image(self.cv_image)
+        command = self.recognize_command(self.cv_image)
+        command_tensor = self.preprocess_command(command)
 
-        resized_input = cv2.resize(cv2.cvtColor(self.cv_image, cv2.COLOR_BGRA2BGR), (64, 48))
-        resized_input[mask == 1] = [0, 0, 255]
-        debug_msg = self.bridge.cv2_to_imgmsg(resized_input, encoding='bgr8')
-        debug_msg.header = header
-        self.pub_debug_image.publish(debug_msg)
+        if self.debug_mode:
+            debug_image = overlay_lane_mask(self._to_bgr(self.cv_image), mask)
+            debug_msg = self.bridge.cv2_to_imgmsg(debug_image, encoding='bgr8')
+            debug_msg.header = header
+            self.pub_debug_image.publish(debug_msg)
 
         with torch.no_grad():
-            output = self.model(input_tensor)
+            output = self.run_model(input_tensor, command_tensor)
 
         output_normalized = output.cpu().numpy().flatten()
         output_denormalized = denormalize_waypoints(output_normalized)
@@ -123,6 +241,7 @@ def zed_sensor_callback(self) -> None:
             return
 
         self.cv_image = self.zed.get_image()
+        self.cv_header = None
 
         header = Header()
         header.stamp = self.get_clock().now().to_msg()
@@ -132,7 +251,8 @@ def zed_sensor_callback(self) -> None:
         if pointcloud_msg is None:
             pointcloud_msg = PointCloud2()
             pointcloud_msg.header = header
-        self.pub_pointcloud.publish(pointcloud_msg)
+        if self.pub_pointcloud is not None:
+            self.pub_pointcloud.publish(pointcloud_msg)
 
     def apply_bspline_smoothing(self, output: torch.Tensor, header) -> Path:
         waypoints = output.cpu().numpy().reshape(-1, 2)
@@ -166,9 +286,17 @@ def create_path_from_output(self, output: torch.Tensor, header) -> Path:
         return path_msg
 
 
+def _parse_args(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument('--sim', action='store_true', help='Use simulator image topic instead of ZED SDK')
+    source_args = sys.argv[1:] if args is None else list(args)
+    return parser.parse_known_args(source_args)
+
+
 def main(args=None) -> None:
-    rclpy.init(args=args)
-    node = InferenceNode()
+    cli_args, ros_args = _parse_args(args)
+    rclpy.init(args=[sys.argv[0], *ros_args])
+    node = InferenceNode(simulator_mode=cli_args.sim)
     executor = MultiThreadedExecutor(num_threads=2)
     executor.add_node(node)
     executor.spin()