diff --git a/e2e_planner/config/train.yaml b/e2e_planner/config/train.yaml index c069a80d..9b61d257 100644 --- a/e2e_planner/config/train.yaml +++ b/e2e_planner/config/train.yaml @@ -1,5 +1,11 @@ -epochs: 200 -batch_size: 8 -learning_rate: 0.0002 -num_workers: 2 -weight_file: "e2e_model.pt" \ No newline at end of file +epochs: 1000 +batch_size: 16 +learning_rate: 0.0001 +num_workers: 8 +weight_file: "e2e_model.pt" +split_seed: 0 +curve_surprise_weighting: + enabled: true + num_bins: 16 + smoothing: 1.0 + max_weight: 50.0 diff --git a/e2e_planner/e2e_planner/inference_node.py b/e2e_planner/e2e_planner/inference_node.py index 3d1a143a..969f18c7 100644 --- a/e2e_planner/e2e_planner/inference_node.py +++ b/e2e_planner/e2e_planner/inference_node.py @@ -1,4 +1,5 @@ import argparse +import json import rclpy from rclpy.node import Node from rclpy.callback_groups import ReentrantCallbackGroup @@ -22,13 +23,20 @@ from .zed_sdk import ZedSdk from e2e_planner.placenav.place_recognition import PlaceRecognition -from util.yolop_processor import YOLOPv2Processor -from util.preprocessing import MODEL_INPUT_SIZE, center_square_crop, lane_mask_to_tensor_array, overlay_lane_mask +from e2e_planner.util.yolop_processor import YOLOPv2Processor +from e2e_planner.util.preprocessing import MODEL_INPUT_SIZE, center_square_crop, lane_mask_to_tensor_array, overlay_lane_mask -def denormalize_waypoints(normalized: np.ndarray) -> np.ndarray: +NUM_WAYPOINTS = 6 + + +def denormalize_axis(values: np.ndarray, min_value: float, max_value: float) -> np.ndarray: + return (values + 1.0) * 0.5 * (max_value - min_value) + min_value + + +def denormalize_waypoints(normalized: np.ndarray, bounds: dict) -> np.ndarray: denormalized = normalized.copy() - denormalized[0::2] = (normalized[0::2] + 1.0) * 5.0 - denormalized[1::2] = (normalized[1::2] + 1.0) * 3.0 - 3.0 + denormalized[0::2] = denormalize_axis(normalized[0::2], bounds['x_min'], bounds['x_max']) + denormalized[1::2] = denormalize_axis(normalized[1::2], bounds['y_min'], bounds['y_max']) return denormalized class InferenceNode(Node): @@ -41,13 +49,14 @@ def __init__(self, simulator_mode: bool = False) -> None: 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('use_place_recognition', True) self.declare_parameter('yolop_input_size', 256) + self.declare_parameter('yolop_fp16', True) 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_delta', 5.0) self.declare_parameter('placenet_window_lower', -1) - self.declare_parameter('placenet_window_upper', 2) + self.declare_parameter('placenet_window_upper', 10) model_path = self.get_parameter('model_name').value interval_ms = self.get_parameter('interval_ms').value @@ -57,6 +66,7 @@ def __init__(self, simulator_mode: bool = False) -> None: self.command = int(self.get_parameter('default_command').value) self.use_place_recognition = bool(self.get_parameter('use_place_recognition').value) self.yolop_input_size = int(self.get_parameter('yolop_input_size').value) + self.yolop_fp16 = bool(self.get_parameter('yolop_fp16').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) @@ -71,13 +81,14 @@ def __init__(self, simulator_mode: bool = False) -> None: 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) + weight_path = FilePath(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' + self.waypoint_bounds = self._build_waypoint_bounds(weight_path) - if os.path.exists(weight_path): - self.model = torch.jit.load(weight_path, map_location=self.device) + if weight_path.exists(): + self.model = torch.jit.load(str(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: @@ -85,7 +96,7 @@ def __init__(self, simulator_mode: bool = False) -> None: 'Loaded model accepts only image input; navigation command will be ignored.' ) else: - self.get_logger().warn(f'Model file not found: {weight_path}') + self.get_logger().warn(f'Model file not found: {str(weight_path)}') self.model = None self.model_uses_command = False @@ -95,6 +106,7 @@ def __init__(self, simulator_mode: bool = False) -> None: yolop_weight_path, self.device, input_size=self.yolop_input_size, + use_fp16=self.yolop_fp16, ) else: self.get_logger().warn(f'YOLOPv2 model not found: {yolop_weight_path}') @@ -151,10 +163,26 @@ def __init__(self, simulator_mode: bool = False) -> None: callback_group=self.torch_cb_group, ) + def _build_waypoint_bounds(self, weight_path: FilePath) -> dict: + bounds_path = weight_path.with_suffix('.bounds.json') + if not bounds_path.exists(): + raise RuntimeError( + f'Normalization bounds file not found: {bounds_path}\n' + 'Re-train the model to generate it automatically.' + ) + with open(bounds_path, 'r') as f: + bounds = json.load(f) + self.get_logger().info( + f'Waypoint normalization bounds loaded from {bounds_path.name}: ' + f"x=({bounds['x_min']:.3f}, {bounds['x_max']:.3f}), " + f"y=({bounds['y_min']:.3f}, {bounds['y_max']:.3f})" + ) + return bounds + def command_callback(self, msg: UInt8) -> None: self.command = int(msg.data) - def preprocess_command(self, command: int | None = None) -> torch.Tensor: + def preprocess_command(self, command: Optional[int] = 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) @@ -231,11 +259,11 @@ def torch_callback(self) -> None: debug_msg.header = header self.pub_debug_image.publish(debug_msg) - with torch.no_grad(): + with torch.inference_mode(): output = self.run_model(input_tensor, command_tensor) output_normalized = output.cpu().numpy().flatten() - output_denormalized = denormalize_waypoints(output_normalized) + output_denormalized = denormalize_waypoints(output_normalized, self.waypoint_bounds) output_denormalized_tensor = torch.from_numpy(output_denormalized).unsqueeze(0) path_raw_msg = self.create_path_from_output(output_denormalized_tensor, header) diff --git a/e2e_planner/e2e_planner/placenav/place_recognition.py b/e2e_planner/e2e_planner/placenav/place_recognition.py index bc157114..c4710ec9 100644 --- a/e2e_planner/e2e_planner/placenav/place_recognition.py +++ b/e2e_planner/e2e_planner/placenav/place_recognition.py @@ -45,8 +45,7 @@ def _compute_distances(self, query_feature): def _initialize_belief(self, query_feature): dists = self._compute_distances(query_feature) descriptor_quantiles = np.quantile(dists, [0.025, 0.975]) - denom = descriptor_quantiles[1] - descriptor_quantiles[0] - self.lambda1 = np.log(self.delta) / denom if denom > 1e-6 else 1.0 + self.lambda1 = np.log(self.delta) / (descriptor_quantiles[1] - descriptor_quantiles[0]) self.belief = np.exp(-self.lambda1 * dists) self.belief /= self.belief.sum() @@ -70,11 +69,7 @@ def _update_belief(self, query_feature): self.belief[:self.window_lower] = 0.0 self.belief *= self._observation_likelihood(query_feature) - belief_sum = self.belief.sum() - if belief_sum <= 0.0: - self._initialize_belief(query_feature) - else: - self.belief /= belief_sum + self.belief /= self.belief.sum() def get_recognition(self, image_tensor): image_tensor = image_tensor.to(self.device, dtype=torch.float32) diff --git a/e2e_planner/e2e_planner/util/__init__.py b/e2e_planner/e2e_planner/util/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/e2e_planner/e2e_planner/util/__init__.py @@ -0,0 +1 @@ + diff --git a/e2e_planner/e2e_planner/util/preprocessing.py b/e2e_planner/e2e_planner/util/preprocessing.py new file mode 100644 index 00000000..9ec4d2a3 --- /dev/null +++ b/e2e_planner/e2e_planner/util/preprocessing.py @@ -0,0 +1,56 @@ +from typing import Tuple + +import cv2 +import numpy as np + + +PLACENET_CROP_SIZE = 288 +MODEL_INPUT_SIZE = (85, 85) + + +def center_square_crop(image: np.ndarray, crop_size: int = PLACENET_CROP_SIZE) -> np.ndarray: + height, width = image.shape[:2] + if height < crop_size or width < crop_size: + raise ValueError(f'Image is smaller than {crop_size}x{crop_size}: {width}x{height}') + + top = (height - crop_size) // 2 + left = (width - crop_size) // 2 + return image[top:top + crop_size, left:left + crop_size] + + +def color_mask_to_binary(mask_image: np.ndarray) -> np.ndarray: + if mask_image.ndim == 2: + return (mask_image > 0).astype(np.uint8) + + red_mask = ( + (mask_image[:, :, 2] > 200) + & (mask_image[:, :, 0] < 50) + & (mask_image[:, :, 1] < 50) + ) + bright_mask = mask_image.max(axis=2) > 127 + return (red_mask | bright_mask).astype(np.uint8) + + +def preprocess_lane_mask(mask: np.ndarray) -> np.ndarray: + binary_mask = color_mask_to_binary(mask) + height, width = binary_mask.shape[:2] + if height < PLACENET_CROP_SIZE or width < PLACENET_CROP_SIZE: + return cv2.resize(binary_mask, MODEL_INPUT_SIZE, interpolation=cv2.INTER_NEAREST) + + cropped_mask = center_square_crop(binary_mask) + return cv2.resize(cropped_mask, MODEL_INPUT_SIZE, interpolation=cv2.INTER_NEAREST) + + +def lane_mask_to_tensor_array(mask: np.ndarray) -> np.ndarray: + return preprocess_lane_mask(mask).astype(np.float32) + + +def overlay_lane_mask(image_bgr: np.ndarray, processed_mask: np.ndarray) -> np.ndarray: + height, width = image_bgr.shape[:2] + if height < PLACENET_CROP_SIZE or width < PLACENET_CROP_SIZE: + cropped_image = image_bgr + else: + cropped_image = center_square_crop(image_bgr) + debug_image = cv2.resize(cropped_image, MODEL_INPUT_SIZE, interpolation=cv2.INTER_AREA) + debug_image[processed_mask == 1] = [0, 0, 255] + return debug_image diff --git a/e2e_planner/e2e_planner/util/slit_aug.py b/e2e_planner/e2e_planner/util/slit_aug.py new file mode 100644 index 00000000..eebdfdfd --- /dev/null +++ b/e2e_planner/e2e_planner/util/slit_aug.py @@ -0,0 +1,38 @@ +from typing import List, Tuple + +import numpy as np + + +def crop_images(image: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + center_crop = image[:, 40:440] + right_crop = image[:, 80:480] + left_crop = image[:, 0:400] + return center_crop, right_crop, left_crop + + +def rotate_waypoints(waypoints: List[List[float]], angle: float) -> List[List[float]]: + cos_theta = np.cos(angle) + sin_theta = np.sin(angle) + rotation_matrix = np.array([[cos_theta, -sin_theta], + [sin_theta, cos_theta]]) + + rotated_waypoints = [] + for waypoint in waypoints: + rotated = rotation_matrix @ np.array(waypoint) + rotated_waypoints.append(rotated.tolist()) + + return rotated_waypoints + + +def augment(image: np.ndarray, waypoints: List[List[float]]) -> List[Tuple[np.ndarray, List[List[float]]]]: + center_crop, right_crop, left_crop = crop_images(image) + + center_waypoints = waypoints + right_waypoints = rotate_waypoints(waypoints, 0.1745) + left_waypoints = rotate_waypoints(waypoints, -0.1745) + + return [ + (center_crop, center_waypoints), + (right_crop, right_waypoints), + (left_crop, left_waypoints) + ] diff --git a/e2e_planner/e2e_planner/util/yolop_processor.py b/e2e_planner/e2e_planner/util/yolop_processor.py new file mode 100644 index 00000000..c8e395b5 --- /dev/null +++ b/e2e_planner/e2e_planner/util/yolop_processor.py @@ -0,0 +1,97 @@ +from pathlib import Path +from typing import Optional, Tuple + +import cv2 +import numpy as np +import torch + + +class YOLOPv2Processor: + def __init__(self, model_path: Path, device: torch.device, input_size: int = 640, use_fp16: bool = False): + self.device = device + self.input_shape = (input_size, input_size) + self.use_fp16 = use_fp16 and device.type == 'cuda' + + if model_path.exists(): + self.model = torch.jit.load(str(model_path), map_location=device) + self.model.to(device) + if self.use_fp16: + self.model.half() + self.model.eval() + else: + raise FileNotFoundError(f'YOLOPv2 model not found: {model_path}') + + def letterbox(self, img: np.ndarray, new_shape: Tuple[int, int], color: Tuple[int, int, int] = (114, 114, 114), stride: int = 32) -> Tuple[np.ndarray, float, Tuple[float, float]]: + shape = img.shape[:2] + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] + dw, dh = np.mod(dw, stride) / 2, np.mod(dh, stride) / 2 + + if shape[::-1] != new_unpad: + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + + img = cv2.copyMakeBorder( + img, top, bottom, left, right, + cv2.BORDER_CONSTANT, value=color + ) + + return img, r, (dw, dh) + + def lane_line_mask(self, ll: torch.Tensor) -> np.ndarray: + if ll.shape[1] == 1: + ll_seg_mask = (ll[:, 0] > 0.5).int() + else: + ll_seg_mask = torch.argmax(ll, dim=1).int() + return ll_seg_mask.squeeze().cpu().numpy() + + def _restore_original_size( + self, + mask: np.ndarray, + original_shape: Tuple[int, int], + ratio: float, + pad: Tuple[float, float], + ) -> np.ndarray: + original_h, original_w = original_shape + pad_left, pad_top = pad + top = max(int(round(pad_top - 0.1)), 0) + left = max(int(round(pad_left - 0.1)), 0) + unpad_h = int(round(original_h * ratio)) + unpad_w = int(round(original_w * ratio)) + + unpadded = mask[top:top + unpad_h, left:left + unpad_w] + return cv2.resize(unpadded, (original_w, original_h), interpolation=cv2.INTER_NEAREST) + + def process_image(self, image: np.ndarray, target_size: Optional[Tuple[int, int]] = None) -> np.ndarray: + original_shape = image.shape[:2] + img_resized, ratio, (pad_left, pad_top) = self.letterbox(image, self.input_shape) + + img = img_resized.astype(np.float32) / 255.0 + img = torch.from_numpy(np.transpose(img, (2, 0, 1))).unsqueeze(0).to(self.device) + if self.use_fp16: + img = img.half() + + with torch.no_grad(): + outputs = self.model(img) + [pred, anchor_grid], seg, ll = outputs + + ll_seg_mask = self.lane_line_mask(ll) + original_mask = self._restore_original_size( + ll_seg_mask, + original_shape, + ratio, + (pad_left, pad_top), + ) + + if target_size is None: + return original_mask + + return cv2.resize( + original_mask, + target_size, + interpolation=cv2.INTER_NEAREST + ) diff --git a/e2e_planner/launch/e2e_planner.launch.py b/e2e_planner/launch/e2e_planner.launch.py index 05f8a6d2..45da6cee 100644 --- a/e2e_planner/launch/e2e_planner.launch.py +++ b/e2e_planner/launch/e2e_planner.launch.py @@ -22,13 +22,14 @@ def launch_setup(context, *args, **kwargs): 'image_topic': '/image_raw', 'debug_mode': True, 'default_command': 1, - 'use_place_recognition': False, + 'use_place_recognition': True, 'yolop_input_size': 256, + 'yolop_fp16': True, 'placenet_model_name': 'placenet.pt', 'topomap_dir_name': 'topomap', - 'placenet_delta': 10.0, + 'placenet_delta': 5.0, 'placenet_window_lower': -1, - 'placenet_window_upper': 2, + 'placenet_window_upper': 10, }] ) diff --git a/e2e_planner/scripts/binarize_dataset.py b/e2e_planner/scripts/binarize_dataset.py index 787f3d17..0fd600c9 100644 --- a/e2e_planner/scripts/binarize_dataset.py +++ b/e2e_planner/scripts/binarize_dataset.py @@ -21,7 +21,7 @@ from tqdm import tqdm sys.path.insert(0, str(Path(__file__).parent)) -from util.yolop_processor import YOLOPv2Processor +from e2e_planner.util.yolop_processor import YOLOPv2Processor def resolve_weights_path(weights_arg): @@ -29,8 +29,8 @@ def resolve_weights_path(weights_arg): return Path(weights_arg) candidates = [ + Path(__file__).resolve().parent.parent / 'weights/yolopv2.pt', Path.home() / 'ros2_ws/install/e2e_planner/share/e2e_planner/weights/yolopv2.pt', - Path(__file__).parent.parent / 'weights/yolopv2.pt', ] return next((p for p in candidates if p.exists()), None) diff --git a/e2e_planner/scripts/create_data.py b/e2e_planner/scripts/create_data.py index 10686438..d69f468a 100755 --- a/e2e_planner/scripts/create_data.py +++ b/e2e_planner/scripts/create_data.py @@ -1,9 +1,10 @@ #!/usr/bin/env python3 import argparse +from dataclasses import dataclass import rclpy from rclpy.node import Node -from sensor_msgs.msg import Image, Joy, PointCloud2 +from sensor_msgs.msg import Image, Imu, Joy, NavSatFix, NavSatStatus, PointCloud2 from geometry_msgs.msg import PoseWithCovarianceStamped from nav_msgs.msg import Odometry from cv_bridge import CvBridge @@ -13,22 +14,33 @@ import csv import copy import sys +import re +import subprocess +import threading from pathlib import Path from collections import deque -from typing import Optional, List, Tuple, Deque, Union -from tf_transformations import euler_from_quaternion +from typing import Any, Optional, List, Tuple, Deque, Union from rclpy.qos import qos_profile_sensor_data -import sensor_msgs_py.point_cloud2 as pc2 -import pymap3d as pm -try: - import pyzed.sl as sl - ZED_SDK_AVAILABLE = True -except ImportError: - ZED_SDK_AVAILABLE = False + +sl = None + + +def _load_zed_sdk(): + global sl + if sl is not None: + return sl + + try: + import pyzed.sl as zed_sl + except ImportError as exc: + raise RuntimeError('ZED SDK not available. Install pyzed package.') from exc + + sl = zed_sl + return sl SAMPLE_INTERVAL = 0.2 WAYPOINT_INTERVAL = 0.5 -NUM_WAYPOINTS = 10 +NUM_WAYPOINTS = 6 DEFAULT_COMMAND = 1 COMMAND_LABELS = { 0: 'roadside', @@ -37,30 +49,68 @@ 3: 'right', } -PoseSample = Union[PoseWithCovarianceStamped, Tuple[float, float, float]] +@dataclass +class GnssPose: + latitude: float + longitude: float + altitude: float + yaw: float + +PoseSample = Union[PoseWithCovarianceStamped, GnssPose, Tuple[float, float, float]] + +def yaw_from_quaternion(x: float, y: float, z: float, w: float) -> float: + siny_cosp = 2.0 * (w * z + x * y) + cosy_cosp = 1.0 - 2.0 * (y * y + z * z) + return float(np.arctan2(siny_cosp, cosy_cosp)) + +def default_save_dir() -> Path: + current_file = Path(__file__).resolve() + for parent in current_file.parents: + source_data_dir = parent / 'src' / 'aiformula' / 'e2e_planner' / 'data' + if source_data_dir.is_dir(): + return source_data_dir + package_data_dir = parent / 'data' + if parent.name == 'e2e_planner' and package_data_dir.is_dir(): + return package_data_dir + return current_file.parent.parent / 'data' class Sample: - def __init__(self, image: np.ndarray, timestamp: float, reference_pose: PoseSample, command: int, point_cloud: Optional[np.ndarray] = None): + def __init__( + self, + image: np.ndarray, + timestamp: float, + reference_pose: PoseSample, + reference_pose_source: str, + command: int, + point_cloud: Optional[np.ndarray] = None, + ): self.image: np.ndarray = image self.timestamp: float = timestamp self.reference_pose: PoseSample = reference_pose + self.reference_pose_source: str = reference_pose_source self.command: int = command self.point_cloud: Optional[np.ndarray] = point_cloud self.waypoints: List[Tuple[float, float]] = [] + self.debug_waypoints: List[dict] = [] self.target_times: List[float] = [timestamp + WAYPOINT_INTERVAL * (i + 1) for i in range(NUM_WAYPOINTS)] class DataCollectionNode(Node): def __init__(self, simulator_mode: bool = False) -> None: super().__init__('data_collection_node') - package_root = Path(__file__).parent.parent self.declare_parameter('simulator_mode', simulator_mode) self.declare_parameter('sdk_flag', True) - self.declare_parameter('save_dir', str(package_root / 'data')) + self.declare_parameter('save_dir', str(default_save_dir())) self.declare_parameter('image_topic', '/image_raw') self.declare_parameter('pointcloud_topic', '/zed/zed_node/pointcloud') + self.declare_parameter('real_pose_source', 'gnss') self.declare_parameter('pose_topic', '/vectornav/pose') + self.declare_parameter('gnss_topic', '/vectornav/gnss') + self.declare_parameter('imu_topic', '/vectornav/imu') self.declare_parameter('odom_topic', '/odom') + self.declare_parameter('ground_truth_pose_topic', '/world/car_world/pose/info') + self.declare_parameter('ground_truth_frame', 'ai_car1') + self.declare_parameter('use_ground_truth_pose', True) self.declare_parameter('joy_topic', '/joy') self.declare_parameter('save_interval_sec', SAMPLE_INTERVAL) self.declare_parameter('toggle_button_index', 2) @@ -72,8 +122,14 @@ def __init__(self, simulator_mode: bool = False) -> None: self.save_base_dir = Path(self.get_parameter('save_dir').value) self.image_topic = self.get_parameter('image_topic').value self.pointcloud_topic = self.get_parameter('pointcloud_topic').value + self.real_pose_source = str(self.get_parameter('real_pose_source').value).lower() self.pose_topic = self.get_parameter('pose_topic').value + self.gnss_topic = self.get_parameter('gnss_topic').value + self.imu_topic = self.get_parameter('imu_topic').value self.odom_topic = self.get_parameter('odom_topic').value + self.ground_truth_pose_topic = self.get_parameter('ground_truth_pose_topic').value + self.ground_truth_frame = self.get_parameter('ground_truth_frame').value + self.use_ground_truth_pose = bool(self.get_parameter('use_ground_truth_pose').value) self.joy_topic = self.get_parameter('joy_topic').value self.save_interval = float(self.get_parameter('save_interval_sec').value) self.toggle_button_index = int(self.get_parameter('toggle_button_index').value) @@ -83,27 +139,31 @@ def __init__(self, simulator_mode: bool = False) -> None: self.bridge: CvBridge = CvBridge() self.latest_image: Optional[Image] = None self.latest_pose: Optional[PoseSample] = None + self.latest_pose_source: str = 'none' + self.latest_gnss: Optional[NavSatFix] = None + self.latest_imu_yaw: Optional[float] = None self.latest_pointcloud: Optional[PointCloud2] = None self.latest_command: int = DEFAULT_COMMAND + self.ground_truth_process: Optional[subprocess.Popen] = None + self.ground_truth_thread: Optional[threading.Thread] = None + self._stop_ground_truth_reader = threading.Event() self.samples: List[Sample] = [] - self.pose_history: Deque[Tuple[float, PoseSample]] = deque() - self.collected_data: List[Tuple[np.ndarray, List[Tuple[float, float]], int, Optional[np.ndarray]]] = [] + self.pose_history: Deque[Tuple[float, PoseSample, str]] = deque() + self.collected_data: List[Tuple[np.ndarray, List[Tuple[float, float]], int, Optional[np.ndarray], List[dict]]] = [] self.last_sample_time: Optional[float] = None self.is_paused: bool = True self.prev_toggle_button_state: int = 0 - self.zed_camera: Optional[sl.Camera] = None - self.zed_image: Optional[sl.Mat] = None - self.zed_point_cloud: Optional[sl.Mat] = None - self.zed_pose: Optional[sl.Pose] = None - self.zed_runtime_params: Optional[sl.RuntimeParameters] = None + self.zed_camera: Optional[Any] = None + self.zed_image: Optional[Any] = None + self.zed_point_cloud: Optional[Any] = None + self.zed_pose: Optional[Any] = None + self.zed_runtime_params: Optional[Any] = None if self.sdk_flag_: - if not ZED_SDK_AVAILABLE: - self.get_logger().error('ZED SDK not available. Install pyzed package.') - raise RuntimeError('ZED SDK not available') + _load_zed_sdk() self._initialize_zed_camera() else: image_topic = self.image_topic if self.simulator_mode else '/zed/zed_node/rgb/image_rect_color' @@ -112,9 +172,21 @@ def __init__(self, simulator_mode: bool = False) -> None: self.create_subscription(PointCloud2, self.pointcloud_topic, self.pointcloud_callback, qos_profile_sensor_data) if self.simulator_mode: + if self.use_ground_truth_pose: + self._start_ground_truth_pose_reader() self.create_subscription(Odometry, self.odom_topic, self.odom_callback, qos_profile_sensor_data) else: - self.create_subscription(PoseWithCovarianceStamped, self.pose_topic, self.pose_callback, qos_profile_sensor_data) + if self.real_pose_source == 'gnss': + self.create_subscription(NavSatFix, self.gnss_topic, self.gnss_callback, qos_profile_sensor_data) + self.create_subscription(Imu, self.imu_topic, self.imu_callback, qos_profile_sensor_data) + self.get_logger().info( + f'Using GNSS pose from {self.gnss_topic} with yaw from {self.imu_topic}' + ) + elif self.real_pose_source == 'pose': + self.create_subscription(PoseWithCovarianceStamped, self.pose_topic, self.pose_callback, qos_profile_sensor_data) + self.get_logger().info(f'Using pose from {self.pose_topic}') + else: + raise ValueError('real_pose_source must be "gnss" or "pose"') self.create_subscription(Joy, self.joy_topic, self.joy_callback, 10) self.create_timer(0.1, self.timer_callback) @@ -166,16 +238,145 @@ def image_callback(self, msg: Image) -> None: def pose_callback(self, msg: PoseWithCovarianceStamped) -> None: self.latest_pose = msg + self.latest_pose_source = 'pose' + + def gnss_callback(self, msg: NavSatFix) -> None: + if msg.status.status < NavSatStatus.STATUS_FIX: + return + if not all(np.isfinite([msg.latitude, msg.longitude, msg.altitude])): + return + + self.latest_gnss = msg + self._update_latest_gnss_pose() + + def imu_callback(self, msg: Imu) -> None: + q = msg.orientation + self.latest_imu_yaw = yaw_from_quaternion(q.x, q.y, q.z, q.w) + self._update_latest_gnss_pose() + + def _update_latest_gnss_pose(self) -> None: + if self.latest_gnss is None or self.latest_imu_yaw is None: + return + self.latest_pose = GnssPose( + latitude=float(self.latest_gnss.latitude), + longitude=float(self.latest_gnss.longitude), + altitude=float(self.latest_gnss.altitude), + yaw=float(self.latest_imu_yaw), + ) + self.latest_pose_source = 'gnss' def odom_callback(self, msg: Odometry) -> None: + if self.simulator_mode and self.use_ground_truth_pose and self.latest_pose is not None: + return q = msg.pose.pose.orientation - _, _, yaw = euler_from_quaternion([q.x, q.y, q.z, q.w]) + yaw = yaw_from_quaternion(q.x, q.y, q.z, q.w) self.latest_pose = (msg.pose.pose.position.x, msg.pose.pose.position.y, yaw) + self.latest_pose_source = 'odom' + + def _start_ground_truth_pose_reader(self) -> None: + command = ['ign', 'topic', '-e', '-t', self.ground_truth_pose_topic] + try: + self.ground_truth_process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + ) + except OSError as exc: + self.get_logger().error(f'Failed to start Gazebo ground truth reader: {exc}') + return + + self.ground_truth_thread = threading.Thread( + target=self._read_ground_truth_pose_loop, + daemon=True, + ) + self.ground_truth_thread.start() + self.get_logger().info( + f'Using Gazebo ground truth pose from {self.ground_truth_pose_topic}, frame "{self.ground_truth_frame}"' + ) + + def _read_ground_truth_pose_loop(self) -> None: + if self.ground_truth_process is None or self.ground_truth_process.stdout is None: + return + + in_pose = False + brace_depth = 0 + pose_lines: List[str] = [] + + for line in self.ground_truth_process.stdout: + if self._stop_ground_truth_reader.is_set(): + break + + stripped = line.strip() + if not in_pose and stripped == 'pose {': + in_pose = True + brace_depth = 1 + pose_lines = [line] + continue + + if not in_pose: + continue + + pose_lines.append(line) + brace_depth += line.count('{') - line.count('}') + if brace_depth == 0: + self._update_ground_truth_pose_from_block(''.join(pose_lines)) + in_pose = False + pose_lines = [] + + def _update_ground_truth_pose_from_block(self, block: str) -> None: + name_match = re.search(r'name:\s*"([^"]+)"', block) + if name_match is None or not self._is_ground_truth_frame(name_match.group(1)): + return + + position_block = self._extract_named_block(block, 'position') + orientation_block = self._extract_named_block(block, 'orientation') + x = self._extract_float_field(position_block, 'x', 0.0) + y = self._extract_float_field(position_block, 'y', 0.0) + qx = self._extract_float_field(orientation_block, 'x', 0.0) + qy = self._extract_float_field(orientation_block, 'y', 0.0) + qz = self._extract_float_field(orientation_block, 'z', 0.0) + qw = self._extract_float_field(orientation_block, 'w', 1.0) + yaw = yaw_from_quaternion(qx, qy, qz, qw) + frame_name = name_match.group(1) + self.latest_pose = (x, y, yaw) + self.latest_pose_source = f'ground_truth:{frame_name}' + + def _extract_named_block(self, text: str, name: str) -> str: + match = re.search(rf'{name}\s*\{{(.*?)\n\s*\}}', text, re.DOTALL) + return match.group(1) if match else '' + + def _extract_float_field(self, text: str, name: str, default: float) -> float: + match = re.search(rf'\b{name}:\s*([-+0-9.eE]+)', text) + return float(match.group(1)) if match else default + + def _is_ground_truth_frame(self, child_frame: str) -> bool: + frame = self.ground_truth_frame.strip('/') + candidates = { + frame, + f'{frame}/chassis', + f'{frame}::chassis', + } + return ( + child_frame.strip('/') in candidates + or child_frame.endswith(f'/{frame}') + or child_frame.endswith(f'::{frame}') + or child_frame.endswith(f'/{frame}/chassis') + or child_frame.endswith(f'::{frame}::chassis') + ) + + def destroy_node(self) -> bool: + self._stop_ground_truth_reader.set() + if self.ground_truth_process is not None and self.ground_truth_process.poll() is None: + self.ground_truth_process.terminate() + return super().destroy_node() def pointcloud_callback(self, msg: PointCloud2) -> None: self.latest_pointcloud = msg def _convert_pointcloud2_to_array(self, pointcloud_msg: PointCloud2) -> np.ndarray: + import sensor_msgs_py.point_cloud2 as pc2 points_list = [[point[0], point[1], point[2], point[3]] for point in pc2.read_points(pointcloud_msg, skip_nans=True, field_names=("x", "y", "z", "rgb"))] return np.array(points_list, dtype=np.float32) @@ -218,7 +419,7 @@ def timer_callback(self) -> None: current_time = time.time() if self.latest_pose is not None: - self.pose_history.append((current_time, copy.deepcopy(self.latest_pose))) + self.pose_history.append((current_time, copy.deepcopy(self.latest_pose), self.latest_pose_source)) if self.sdk_flag_: image, point_cloud = self._capture_data_from_zed() @@ -226,7 +427,14 @@ def timer_callback(self) -> None: return if self.last_sample_time is None or current_time - self.last_sample_time >= self.save_interval: - sample = Sample(image, current_time, copy.deepcopy(self.latest_pose), self.latest_command, point_cloud) + sample = Sample( + image, + current_time, + copy.deepcopy(self.latest_pose), + self.latest_pose_source, + self.latest_command, + point_cloud, + ) self.samples.append(sample) self.last_sample_time = current_time else: @@ -237,7 +445,14 @@ def timer_callback(self) -> None: point_cloud = None if not self.simulator_mode and self.latest_pointcloud is not None: point_cloud = self._convert_pointcloud2_to_array(self.latest_pointcloud) - sample = Sample(cv_image, current_time, copy.deepcopy(self.latest_pose), self.latest_command, point_cloud) + sample = Sample( + cv_image, + current_time, + copy.deepcopy(self.latest_pose), + self.latest_pose_source, + self.latest_command, + point_cloud, + ) self.samples.append(sample) self.last_sample_time = current_time @@ -246,7 +461,13 @@ def timer_callback(self) -> None: completed_samples = [sample for sample in self.samples if len(sample.waypoints) == NUM_WAYPOINTS] for sample in completed_samples: - self.collected_data.append((sample.image, sample.waypoints, sample.command, sample.point_cloud)) + self.collected_data.append(( + sample.image, + sample.waypoints, + sample.command, + sample.point_cloud, + sample.debug_waypoints, + )) self.get_logger().info(f'🟡Collected data #{len(self.collected_data)}') self.samples = [sample for sample in self.samples if len(sample.waypoints) < NUM_WAYPOINTS] @@ -256,17 +477,26 @@ def timer_callback(self) -> None: def collect_waypoints_for_sample(self, sample: Sample) -> None: for i in range(len(sample.waypoints), NUM_WAYPOINTS): target_time = sample.target_times[i] - pose = self.find_closest_pose(target_time) - if pose is not None: - x, y = self.transform_to_robot_frame(sample.reference_pose, pose) + pose_record = self.find_closest_pose(target_time) + if pose_record is not None: + pose, pose_source = pose_record + x, y, debug_row = self.transform_to_robot_frame(sample.reference_pose, pose) + debug_row.update({ + 'waypoint_index': i, + 'sample_time': sample.timestamp, + 'target_time': target_time, + 'reference_pose_source': sample.reference_pose_source, + 'target_pose_source': pose_source, + }) sample.waypoints.append((x, y)) + sample.debug_waypoints.append(debug_row) else: break - def find_closest_pose(self, target_time: float) -> Optional[PoseSample]: - for t, pose in self.pose_history: + def find_closest_pose(self, target_time: float) -> Optional[Tuple[PoseSample, str]]: + for t, pose, pose_source in self.pose_history: if t >= target_time: - return pose + return pose, pose_source return None def cleanup_pose_history(self) -> None: @@ -279,7 +509,7 @@ def cleanup_pose_history(self) -> None: while self.pose_history and self.pose_history[0][0] < min_target_time: self.pose_history.popleft() - def transform_to_robot_frame(self, reference_pose: PoseSample, current_pose: PoseSample) -> Tuple[float, float]: + def transform_to_robot_frame(self, reference_pose: PoseSample, current_pose: PoseSample) -> Tuple[float, float, dict]: if self.simulator_mode: x0, y0, yaw0 = reference_pose x, y, _ = current_pose @@ -287,15 +517,57 @@ def transform_to_robot_frame(self, reference_pose: PoseSample, current_pose: Pos dy = y - y0 x_robot = np.cos(yaw0) * dx + np.sin(yaw0) * dy y_robot = -np.sin(yaw0) * dx + np.cos(yaw0) * dy - return x_robot, y_robot + debug_row = { + 'ref_x': x0, + 'ref_y': y0, + 'ref_yaw': yaw0, + 'cur_x': x, + 'cur_y': y, + 'cur_yaw': current_pose[2], + 'dx_world': dx, + 'dy_world': dy, + 'x_robot': x_robot, + 'y_robot': y_robot, + } + return x_robot, y_robot, debug_row + + if isinstance(reference_pose, GnssPose) and isinstance(current_pose, GnssPose): + import pymap3d as pm + + e, n, _ = pm.geodetic2enu( + current_pose.latitude, + current_pose.longitude, + current_pose.altitude, + reference_pose.latitude, + reference_pose.longitude, + reference_pose.altitude, + ) + yaw0 = reference_pose.yaw + x_robot = -e * np.sin(yaw0) + n * np.cos(yaw0) + y_robot = -e * np.cos(yaw0) - n * np.sin(yaw0) + + debug_row = { + 'ref_x': reference_pose.latitude, + 'ref_y': reference_pose.longitude, + 'ref_yaw': yaw0, + 'cur_x': current_pose.latitude, + 'cur_y': current_pose.longitude, + 'cur_yaw': current_pose.yaw, + 'dx_world': e, + 'dy_world': n, + 'x_robot': x_robot, + 'y_robot': y_robot, + } + return x_robot, y_robot, debug_row x0_ecef = reference_pose.pose.pose.position.x y0_ecef = reference_pose.pose.pose.position.y z0_ecef = reference_pose.pose.pose.position.z + import pymap3d as pm lat0, lon0, alt0 = pm.ecef2geodetic(x0_ecef, y0_ecef, z0_ecef) q0 = reference_pose.pose.pose.orientation - _, _, yaw0 = euler_from_quaternion([q0.x, q0.y, q0.z, q0.w]) + yaw0 = yaw_from_quaternion(q0.x, q0.y, q0.z, q0.w) xi_ecef = current_pose.pose.pose.position.x yi_ecef = current_pose.pose.pose.position.y @@ -305,7 +577,24 @@ def transform_to_robot_frame(self, reference_pose: PoseSample, current_pose: Pos x_robot = -e * np.sin(yaw0) + n * np.cos(yaw0) y_robot = -e * np.cos(yaw0) - n * np.sin(yaw0) - return x_robot, y_robot + debug_row = { + 'ref_x': x0_ecef, + 'ref_y': y0_ecef, + 'ref_yaw': yaw0, + 'cur_x': xi_ecef, + 'cur_y': yi_ecef, + 'cur_yaw': yaw_from_quaternion( + current_pose.pose.pose.orientation.x, + current_pose.pose.pose.orientation.y, + current_pose.pose.pose.orientation.z, + current_pose.pose.pose.orientation.w, + ), + 'dx_world': e, + 'dy_world': n, + 'x_robot': x_robot, + 'y_robot': y_robot, + } + return x_robot, y_robot, debug_row def save_data(self) -> None: if len(self.collected_data) == 0: @@ -317,17 +606,20 @@ def save_data(self) -> None: images_dir = dataset_dir / 'images' path_dir = dataset_dir / 'path' commands_dir = dataset_dir / 'commands' + debug_dir = dataset_dir / 'debug_waypoints' pointclouds_dir = dataset_dir / 'pointclouds' images_dir.mkdir(parents=True, exist_ok=True) path_dir.mkdir(parents=True, exist_ok=True) commands_dir.mkdir(parents=True, exist_ok=True) + debug_dir.mkdir(parents=True, exist_ok=True) pointclouds_dir.mkdir(parents=True, exist_ok=True) - for idx, (image, waypoints, command, point_cloud) in enumerate(self.collected_data, start=1): + for idx, (image, waypoints, command, point_cloud, debug_waypoints) in enumerate(self.collected_data, start=1): image_path = images_dir / f'{idx:05d}.png' waypoints_path = path_dir / f'{idx:05d}.csv' command_path = commands_dir / f'{idx:05d}.csv' + debug_path = debug_dir / f'{idx:05d}.csv' pointcloud_path = pointclouds_dir / f'{idx:05d}.npy' cv2.imwrite(str(image_path), image) @@ -342,6 +634,13 @@ def save_data(self) -> None: csv_writer = csv.writer(csvfile) csv_writer.writerow([command]) + if debug_waypoints: + with open(str(debug_path), 'w', newline='') as csvfile: + fieldnames = list(debug_waypoints[0].keys()) + csv_writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + csv_writer.writeheader() + csv_writer.writerows(debug_waypoints) + if point_cloud is not None: np.save(str(pointcloud_path), point_cloud) diff --git a/e2e_planner/scripts/create_topomap.py b/e2e_planner/scripts/create_topomap.py index d1884fd8..857f3408 100644 --- a/e2e_planner/scripts/create_topomap.py +++ b/e2e_planner/scripts/create_topomap.py @@ -9,7 +9,7 @@ import yaml from torchvision import transforms -from util.preprocessing import MODEL_INPUT_SIZE, center_square_crop +from e2e_planner.util.preprocessing import MODEL_INPUT_SIZE, center_square_crop class TopomapGenerator: @@ -19,14 +19,14 @@ class TopomapGenerator: 2: 'left', 3: 'right', } - SAVED_STEP = 10 + SAVED_STEP = 5 def __init__(self, dataset_path): self.dataset_root = Path(dataset_path) self.image_dir = self.dataset_root / 'images' self.command_dir = self.dataset_root / 'commands' - script_dir = Path(__file__).parent + script_dir = Path(__file__).resolve().parent self.package_root = script_dir.parent self.topomap_dir = self.package_root / 'config' / 'topomap' self.topomap_images_dir = self.topomap_dir / 'images' @@ -57,13 +57,11 @@ def _load_command(self, image_path): with command_path.open('r', newline='') as f: return int(float(next(csv.reader(f))[0])) - def _preprocess_image(self, image_path): + def _center_crop(self, image_path): image = cv2.imread(str(image_path), cv2.IMREAD_COLOR) if image is None: raise ValueError(f'Failed to read image: {image_path}') - - cropped_image = center_square_crop(image) - return cv2.resize(cropped_image, MODEL_INPUT_SIZE, interpolation=cv2.INTER_AREA) + return center_square_crop(image) def extract_feature(self, image): image_tensor = self.placenet_transform(cv2.cvtColor(image, cv2.COLOR_BGR2RGB)).unsqueeze(0) @@ -83,15 +81,20 @@ def build_nodes(self): if command not in self.COMMAND_TO_ACTION: raise ValueError(f'Unsupported command value: {command}') - processed_image = self._preprocess_image(image_path) + # センタークロップ（288x288）を特徴抽出と保存用画像の両方のベースにする + cropped_image = self._center_crop(image_path) + + # 保存用は cv2.resize で 85x85 に縮小 + save_image = cv2.resize(cropped_image, MODEL_INPUT_SIZE, interpolation=cv2.INTER_AREA) output_image_name = f'img{idx + 1:05d}.png' output_image_path = self.topomap_images_dir / output_image_name - cv2.imwrite(str(output_image_path), processed_image) + cv2.imwrite(str(output_image_path), save_image) + # 特徴抽出は 288x288 のまま渡し、transforms.Resize で縮小する（推論時と同じ経路） nodes.append({ 'id': idx, 'image': output_image_name, - 'feature': self.extract_feature(processed_image), + 'feature': self.extract_feature(cropped_image), 'action': self.COMMAND_TO_ACTION[command], }) diff --git a/e2e_planner/scripts/network.py b/e2e_planner/scripts/network.py index f16a5482..d4893289 100644 --- a/e2e_planner/scripts/network.py +++ b/e2e_planner/scripts/network.py @@ -3,7 +3,7 @@ class Network(nn.Module): - def __init__(self, num_waypoints: int = 10, num_branches: int = 4): + def __init__(self, num_waypoints: int = 6, num_branches: int = 4): super(Network, self).__init__() self.conv1 = nn.Conv2d(1, 32, kernel_size=8, stride=4) diff --git a/e2e_planner/scripts/train.py b/e2e_planner/scripts/train.py index 0c1843a9..6e931a7c 100755 --- a/e2e_planner/scripts/train.py +++ b/e2e_planner/scripts/train.py @@ -1,53 +1,170 @@ #!/usr/bin/env python3 +import json import sys import yaml +from datetime import datetime +from pathlib import Path + import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader, random_split from torch.utils.tensorboard import SummaryWriter import cv2 import csv -from pathlib import Path import numpy as np -from typing import Tuple +from typing import Dict, List, Tuple from tqdm import tqdm -from network import Network -from util.preprocessing import lane_mask_to_tensor_array -NUM_WAYPOINTS = 10 +try: + from e2e_collector.network import Network +except ImportError: + from network import Network + +from e2e_planner.util.preprocessing import lane_mask_to_tensor_array + +NUM_WAYPOINTS = 6 NUM_BRANCHES = 4 DEFAULT_COMMAND = 1 +METADATA_CURVE_SCORE_INDEX = 7 +METADATA_CURVE_BIN_INDEX = 8 +METADATA_SAMPLE_WEIGHT_INDEX = 9 + + +def normalize_axis(values: torch.Tensor, min_value: float, max_value: float) -> torch.Tensor: + return (values - min_value) / (max_value - min_value) * 2.0 - 1.0 + + +def denormalize_axis(values: torch.Tensor, min_value: float, max_value: float) -> torch.Tensor: + return (values + 1.0) * 0.5 * (max_value - min_value) + min_value + + +def denormalize_waypoints(waypoints: torch.Tensor, bounds: Dict[str, float]) -> torch.Tensor: + denormalized = waypoints.clone() + denormalized[..., 0::2] = denormalize_axis(waypoints[..., 0::2], bounds['x_min'], bounds['x_max']) + denormalized[..., 1::2] = denormalize_axis(waypoints[..., 1::2], bounds['y_min'], bounds['y_max']) + return denormalized class E2EDataset(Dataset): - def __init__(self, dataset_path: Path): + def __init__( + self, + dataset_path: Path, + surprise_enabled: bool = False, + surprise_num_bins: int = 16, + surprise_smoothing: float = 1.0, + surprise_max_weight: float = 50.0, + ): self.dataset_path = dataset_path self.mask_images_dir = dataset_path / 'mask_images' self.path_dir = dataset_path / 'path' self.command_dir = dataset_path / 'commands' self.mask_files = sorted(list(self.mask_images_dir.glob('*.png'))) + self.surprise_enabled = surprise_enabled + self.surprise_num_bins = max(2, surprise_num_bins) + self.surprise_smoothing = max(0.0, surprise_smoothing) + self.surprise_max_weight = max(1.0, surprise_max_weight) + self.normalization_bounds = self._build_normalization_bounds() + ( + self.curve_scores, + self.curve_bins, + self.sample_weights, + self.bin_entropy_nats, + ) = self._build_curve_surprise_weights() def __len__(self) -> int: return len(self.mask_files) - def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + def _read_waypoints(self, csv_file: Path) -> List[List[float]]: + with open(csv_file, 'r') as f: + reader = csv.DictReader(f) + waypoints = [[float(row['x']), float(row['y'])] for row in reader] + if len(waypoints) < NUM_WAYPOINTS: + raise ValueError(f'Expected at least {NUM_WAYPOINTS} waypoints in {csv_file}, got {len(waypoints)}') + return waypoints[:NUM_WAYPOINTS] + + def _curve_score(self, waypoints: List[List[float]]) -> float: + raw_waypoints = np.array(waypoints, dtype=np.float32) + max_abs_y_index = int(np.argmax(np.abs(raw_waypoints[:, 1]))) + return float(raw_waypoints[max_abs_y_index, 1]) + + def _build_normalization_bounds(self) -> Dict[str, float]: + if len(self.mask_files) == 0: + raise RuntimeError( + f'No training masks found in {self.mask_images_dir}. ' + 'Cannot compute waypoint normalization bounds.' + ) + + waypoints = [] + for mask_file in self.mask_files: + waypoints.extend(self._read_waypoints(self.path_dir / f'{mask_file.stem}.csv')) + + waypoint_array = np.asarray(waypoints, dtype=np.float32) + return { + 'x_min': float(np.min(waypoint_array[:, 0])), + 'x_max': float(np.max(waypoint_array[:, 0])), + 'y_min': float(np.min(waypoint_array[:, 1])), + 'y_max': float(np.max(waypoint_array[:, 1])), + } + + def _build_curve_surprise_weights( + self, + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]: + curve_scores = np.zeros(len(self.mask_files), dtype=np.float32) + bin_indices = np.zeros(len(self.mask_files), dtype=np.float32) + weights = np.ones(len(self.mask_files), dtype=np.float32) + + if not self.surprise_enabled or len(self.mask_files) == 0: + return curve_scores, bin_indices, weights, 0.0 + + for idx, mask_file in enumerate(self.mask_files): + waypoints = self._read_waypoints(self.path_dir / f'{mask_file.stem}.csv') + curve_scores[idx] = self._curve_score(waypoints) + + # Signed curve_score の経験分布から自己情報量 -log p をサンプル重みにする + hist, bin_edges = np.histogram(curve_scores, bins=self.surprise_num_bins) + denom = float(hist.sum()) + self.surprise_smoothing * self.surprise_num_bins + probs = (hist.astype(np.float64) + self.surprise_smoothing) / denom + surprise_per_bin = (-np.log(probs)).astype(np.float32) + + indices = np.clip( + np.digitize(curve_scores, bin_edges[1:-1]), + 0, + self.surprise_num_bins - 1, + ) + bin_indices = indices.astype(np.float32) + weights = surprise_per_bin[indices] + weights = np.minimum(weights, np.float32(self.surprise_max_weight)) + mean_weight = float(weights.mean()) + if mean_weight > 0.0: + weights = weights / mean_weight # 平均=1 で全体の loss スケールを保つ + entropy_nats = float(-np.sum(probs * np.log(probs))) + return curve_scores, bin_indices, weights.astype(np.float32), entropy_nats + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, str, torch.Tensor]: mask_file = self.mask_files[idx] csv_file = self.path_dir / f'{mask_file.stem}.csv' command_file = self.command_dir / f'{mask_file.stem}.csv' mask_bgr = cv2.imread(str(mask_file), cv2.IMREAD_COLOR) + mask_nonzero = int(np.count_nonzero(mask_bgr)) - with open(csv_file, 'r') as f: - reader = csv.DictReader(f) - waypoints = [[float(row['x']), float(row['y'])] for row in reader] + waypoints = self._read_waypoints(csv_file) mask_normalized = lane_mask_to_tensor_array(mask_bgr) mask_tensor = torch.from_numpy(mask_normalized).unsqueeze(0) waypoints_tensor = torch.tensor(waypoints, dtype=torch.float32).flatten() - waypoints_tensor[0::2] = waypoints_tensor[0::2] / 5.0 - 1.0 - waypoints_tensor[1::2] = (waypoints_tensor[1::2] + 3.0) / 3.0 - 1.0 + waypoints_tensor[0::2] = normalize_axis( + waypoints_tensor[0::2], + self.normalization_bounds['x_min'], + self.normalization_bounds['x_max'], + ) + waypoints_tensor[1::2] = normalize_axis( + waypoints_tensor[1::2], + self.normalization_bounds['y_min'], + self.normalization_bounds['y_max'], + ) command = DEFAULT_COMMAND if command_file.exists(): @@ -58,7 +175,24 @@ def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tenso command_tensor = torch.zeros(NUM_BRANCHES, dtype=torch.float32) command_tensor[command] = 1.0 - return mask_tensor, waypoints_tensor, command_tensor + raw_waypoints = np.array(waypoints, dtype=np.float32) + metadata = torch.tensor( + [ + float(command), + float(mask_nonzero), + float(raw_waypoints[-1, 0]), + float(raw_waypoints[-1, 1]), + float(np.max(np.abs(raw_waypoints[:, 1]))), + float(np.min(raw_waypoints[:, 0])), + float(np.max(raw_waypoints[:, 0])), + float(self.curve_scores[idx]), + float(self.curve_bins[idx]), + float(self.sample_weights[idx]), + ], + dtype=torch.float32, + ) + + return mask_tensor, waypoints_tensor, command_tensor, mask_file.stem, metadata class Config: def __init__(self, config_path: Path, package_root: Path): @@ -72,11 +206,18 @@ def __init__(self, config_path: Path, package_root: Path): self.learning_rate = config_dict['learning_rate'] self.num_workers = config_dict['num_workers'] self.weight_file = config_dict['weight_file'] + self.split_seed = int(config_dict.get('split_seed', 0)) + surprise_cfg = config_dict.get('curve_surprise_weighting', {}) + self.surprise_enabled = bool(surprise_cfg.get('enabled', False)) + self.surprise_num_bins = max(2, int(surprise_cfg.get('num_bins', 16))) + self.surprise_smoothing = float(surprise_cfg.get('smoothing', 1.0)) + self.surprise_max_weight = float(surprise_cfg.get('max_weight', 50.0)) self.weights_dir = package_root / 'weights' self.weights_dir.mkdir(exist_ok=True) self.logs_dir = package_root / 'runs' + self.logs_dir.mkdir(exist_ok=True) self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @@ -84,11 +225,25 @@ class Trainer: def __init__(self, dataset_path: Path, config: Config): self.config = config self.device = config.device + self.dataset_path = dataset_path - dataset = E2EDataset(dataset_path) + dataset = E2EDataset( + dataset_path, + surprise_enabled=config.surprise_enabled, + surprise_num_bins=config.surprise_num_bins, + surprise_smoothing=config.surprise_smoothing, + surprise_max_weight=config.surprise_max_weight, + ) + self.normalization_bounds = dataset.normalization_bounds + if len(dataset) == 0: + raise RuntimeError( + f'No training masks found in {dataset_path / "mask_images"}. ' + 'Run binarize_dataset.py before train.py.' + ) train_size = int(0.8 * len(dataset)) val_size = len(dataset) - train_size - train_dataset, val_dataset = random_split(dataset, [train_size, val_size]) + split_generator = torch.Generator().manual_seed(config.split_seed) + train_dataset, val_dataset = random_split(dataset, [train_size, val_size], generator=split_generator) self.train_loader = DataLoader( train_dataset, @@ -106,23 +261,50 @@ def __init__(self, dataset_path: Path, config: Config): self.model = Network(num_waypoints=NUM_WAYPOINTS, num_branches=NUM_BRANCHES).to(self.device) self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=config.learning_rate) self.mseloss = nn.MSELoss() - self.writer = SummaryWriter(log_dir=str(config.logs_dir)) + run_name = f'{datetime.now().strftime("%Y%m%d_%H%M%S")}_{dataset_path.name}' + self.run_dir = config.logs_dir / run_name + self.writer = SummaryWriter(log_dir=str(self.run_dir)) self.best_val_loss = float('inf') + self.last_val_details = [] + self.last_val_rmse_m = 0.0 print(f'Using device: {self.device}') print(f'Train size: {len(train_dataset)}, Val size: {len(val_dataset)}') + print( + 'Waypoint normalization bounds: ' + f"x=({self.normalization_bounds['x_min']:.3f}, {self.normalization_bounds['x_max']:.3f}), " + f"y=({self.normalization_bounds['y_min']:.3f}, {self.normalization_bounds['y_max']:.3f})" + ) + if config.surprise_enabled: + print( + 'Curve surprise weighting: ' + f'num_bins={config.surprise_num_bins}, ' + f'smoothing={config.surprise_smoothing:.3f}, ' + f'max_weight={config.surprise_max_weight:.3f}, ' + f'entropy={dataset.bin_entropy_nats:.3f} nats, ' + f'weight_range={float(np.min(dataset.sample_weights)):.3f}-{float(np.max(dataset.sample_weights)):.3f}' + ) + print(f'TensorBoard log dir: {self.run_dir}') + print(f'View with: tensorboard --logdir {config.logs_dir}') + + def weighted_mse_loss(self, outputs: torch.Tensor, targets: torch.Tensor, sample_weights: torch.Tensor) -> torch.Tensor: + sample_losses = torch.mean((outputs - targets) ** 2, dim=1) + return torch.mean(sample_losses * sample_weights) def validate(self) -> float: self.model.eval() total_loss = 0.0 + total_rmse_m = 0.0 + self.last_val_details = [] with torch.no_grad(): pbar = tqdm(self.val_loader, desc='Validation') - for images, waypoints, commands in pbar: + for images, waypoints, commands, sample_names, metadata in pbar: images = images.to(self.device) waypoints = waypoints.to(self.device) commands = commands.to(self.device) + metadata = metadata.to(self.device) outputs = self.model(images, commands) @@ -130,44 +312,136 @@ def validate(self) -> float: total_loss += loss.item() pbar.set_postfix({'loss': f'{loss.item():.6f}'}) + sample_losses = torch.mean((outputs - waypoints) ** 2, dim=1) + output_meters = denormalize_waypoints(outputs, self.normalization_bounds) + target_meters = denormalize_waypoints(waypoints, self.normalization_bounds) + sample_rmse_meters = torch.sqrt(torch.mean((output_meters - target_meters) ** 2, dim=1)) + total_rmse_m += sample_rmse_meters.mean().item() + + for i, sample_name in enumerate(sample_names): + self.last_val_details.append({ + 'sample': sample_name, + 'loss': float(sample_losses[i].item()), + 'rmse_m': float(sample_rmse_meters[i].item()), + 'command': int(metadata[i, 0].item()), + 'mask_nonzero': int(metadata[i, 1].item()), + 'last_x_3s': float(metadata[i, 2].item()), + 'last_y_3s': float(metadata[i, 3].item()), + 'max_abs_y_3s': float(metadata[i, 4].item()), + 'min_x_3s': float(metadata[i, 5].item()), + 'max_x_3s': float(metadata[i, 6].item()), + 'curve_y_m': float(metadata[i, METADATA_CURVE_SCORE_INDEX].item()), + 'curve_bin': int(metadata[i, METADATA_CURVE_BIN_INDEX].item()), + 'sample_weight': float(metadata[i, METADATA_SAMPLE_WEIGHT_INDEX].item()), + }) + + self.last_val_rmse_m = total_rmse_m / len(self.val_loader) return total_loss / len(self.val_loader) + def write_val_details(self, epoch: int) -> None: + if not self.last_val_details: + return + + sorted_details = sorted(self.last_val_details, key=lambda row: row['loss'], reverse=True) + csv_path = self.run_dir / f'val_loss_epoch_{epoch:04d}.csv' + fieldnames = [ + 'rank', + 'sample', + 'loss', + 'rmse_m', + 'command', + 'mask_nonzero', + 'last_x_3s', + 'last_y_3s', + 'max_abs_y_3s', + 'min_x_3s', + 'max_x_3s', + 'curve_y_m', + 'curve_bin', + 'sample_weight', + ] + with open(csv_path, 'w', newline='') as f: + writer = csv.DictWriter(f, fieldnames=fieldnames) + writer.writeheader() + for rank, row in enumerate(sorted_details, start=1): + writer.writerow({'rank': rank, **row}) + + worst = sorted_details[:5] + worst_summary = ', '.join(f"{row['sample']}:{row['loss']:.4f}" for row in worst) + print(f'Validation details saved: {csv_path}') + print(f'Worst val samples: {worst_summary}') + def save_checkpoint(self, val_loss: float) -> None: if val_loss < self.best_val_loss: self.best_val_loss = val_loss weight_path = self.config.weights_dir / self.config.weight_file scripted_model = torch.jit.script(self.model) scripted_model.save(str(weight_path)) + bounds_path = weight_path.with_suffix('.bounds.json') + with open(bounds_path, 'w') as f: + json.dump(self.normalization_bounds, f, indent=2) print(f'Best model saved: {weight_path} (val_loss: {val_loss:.6f})') def train(self, epochs: int) -> None: + self.writer.add_text('run/dataset_path', str(self.dataset_path), 0) + self.writer.add_text('run/device', str(self.device), 0) + self.writer.add_scalar('Dataset/train_size', len(self.train_loader.dataset), 0) + self.writer.add_scalar('Dataset/val_size', len(self.val_loader.dataset), 0) + for epoch in range(1, epochs + 1): self.model.train() total_train_loss = 0.0 + total_train_rmse_m = 0.0 + total_train_weight = 0.0 pbar = tqdm(self.train_loader, desc=f'Epoch {epoch} [Train]') - for images, waypoints, commands in pbar: + for images, waypoints, commands, _, metadata in pbar: images = images.to(self.device) waypoints = waypoints.to(self.device) commands = commands.to(self.device) + metadata = metadata.to(self.device) + sample_weights = metadata[:, METADATA_SAMPLE_WEIGHT_INDEX] self.optimizer.zero_grad() outputs = self.model(images, commands) - loss = self.mseloss(outputs, waypoints) + loss = self.weighted_mse_loss(outputs, waypoints, sample_weights) loss.backward() self.optimizer.step() total_train_loss += loss.item() - pbar.set_postfix({'loss': f'{loss.item():.6f}'}) + with torch.no_grad(): + total_train_weight += sample_weights.mean().item() + output_meters = denormalize_waypoints(outputs, self.normalization_bounds) + target_meters = denormalize_waypoints(waypoints, self.normalization_bounds) + rmse_meters = torch.sqrt(torch.mean((output_meters - target_meters) ** 2, dim=1)) + total_train_rmse_m += rmse_meters.mean().item() + pbar.set_postfix({ + 'loss': f'{loss.item():.6f}', + 'w': f'{sample_weights.mean().item():.3f}', + }) train_loss = total_train_loss / len(self.train_loader) + train_rmse_m = total_train_rmse_m / len(self.train_loader) + train_weight_mean = total_train_weight / len(self.train_loader) val_loss = self.validate() self.writer.add_scalar('Loss/train', train_loss, epoch) self.writer.add_scalar('Loss/val', val_loss, epoch) - - print(f'Epoch [{epoch}/{epochs}], Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}') + self.writer.add_scalar('RMSE_m/train', train_rmse_m, epoch) + self.writer.add_scalar('RMSE_m/val', self.last_val_rmse_m, epoch) + if self.config.surprise_enabled: + self.writer.add_scalar('CurveSurprise/train_weight_mean', train_weight_mean, epoch) + self.writer.add_scalar('LearningRate', self.optimizer.param_groups[0]['lr'], epoch) + self.writer.add_scalar('Best/val_loss', min(self.best_val_loss, val_loss), epoch) + self.writer.flush() + + print( + f'Epoch [{epoch}/{epochs}], ' + f'Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}, ' + f'Train RMSE: {train_rmse_m:.3f} m, Val RMSE: {self.last_val_rmse_m:.3f} m' + ) + self.write_val_details(epoch) self.save_checkpoint(val_loss) @@ -183,7 +457,7 @@ def main() -> None: print(f'Dataset path does not exist: {dataset_path}') sys.exit(1) - script_dir = Path(__file__).parent + script_dir = Path(__file__).resolve().parent package_root = script_dir.parent config_path = package_root / 'config' / 'train.yaml' diff --git a/e2e_planner/scripts/util/yolop_processor.py b/e2e_planner/scripts/util/yolop_processor.py index 48d124f9..4e879312 100644 --- a/e2e_planner/scripts/util/yolop_processor.py +++ b/e2e_planner/scripts/util/yolop_processor.py @@ -8,13 +8,16 @@ class YOLOPv2Processor: - def __init__(self, model_path: Path, device: torch.device, input_size: int = 640): + def __init__(self, model_path: Path, device: torch.device, input_size: int = 640, use_fp16: bool = False): self.device = device self.input_shape = (input_size, input_size) + self.use_fp16 = use_fp16 and device.type == 'cuda' if model_path.exists(): self.model = torch.jit.load(str(model_path), map_location=device) self.model.to(device) + if self.use_fp16: + self.model.half() self.model.eval() else: raise FileNotFoundError(f'YOLOPv2 model not found: {model_path}') @@ -70,8 +73,10 @@ def process_image(self, image: np.ndarray, target_size: Optional[Tuple[int, int] img = img_resized.astype(np.float32) / 255.0 img = torch.from_numpy(np.transpose(img, (2, 0, 1))).unsqueeze(0).to(self.device) + if self.use_fp16: + img = img.half() - with torch.no_grad(): + with torch.inference_mode(): outputs = self.model(img) [pred, anchor_grid], seg, ll = outputs diff --git a/e2e_planner/setup.py b/e2e_planner/setup.py index 131a0b17..f70828f1 100644 --- a/e2e_planner/setup.py +++ b/e2e_planner/setup.py @@ -1,28 +1,34 @@ from setuptools import setup import os -from glob import glob +from pathlib import Path package_name = 'e2e_planner' +package_root = Path(__file__).resolve().parent + + +def source_glob(pattern): + cwd = Path.cwd() + return [os.path.relpath(path, cwd) for path in package_root.glob(pattern)] # scripts/ を e2e_collector パッケージとして公開し、 # data_collector を ros2 run から実行可能にする setup( name=package_name, version='0.0.1', - packages=[package_name, package_name + '.placenav', 'util', 'e2e_collector'], + packages=[package_name, package_name + '.placenav', package_name + '.util', 'e2e_collector'], package_dir={ - 'util': 'scripts/util', 'e2e_collector': 'scripts', }, data_files=[ ('share/ament_index/resource_index/packages', ['resource/' + package_name]), ('share/' + package_name, ['package.xml']), - (os.path.join('share', package_name, 'launch'), glob('launch/*.launch.py')), - (os.path.join('share', package_name, 'config'), glob('config/*.yaml')), - (os.path.join('share', package_name, 'config', 'topomap'), glob('config/topomap/*.yaml')), - (os.path.join('share', package_name, 'config', 'topomap', 'images'), glob('config/topomap/images/*.png')), - (os.path.join('share', package_name, 'weights'), glob('weights/*.pt')), + (os.path.join('share', package_name, 'launch'), source_glob('launch/*.launch.py')), + (os.path.join('share', package_name, 'config'), source_glob('config/*.yaml')), + (os.path.join('share', package_name, 'config', 'topomap'), source_glob('config/topomap/*.yaml')), + (os.path.join('share', package_name, 'config', 'topomap', 'images'), source_glob('config/topomap/images/*.png')), + (os.path.join('share', package_name, 'weights'), source_glob('weights/*.pt')), + (os.path.join('share', package_name, 'weights'), source_glob('weights/*.json')), ], install_requires=['setuptools'], zip_safe=True, @@ -36,6 +42,10 @@ 'data_collector = e2e_collector.create_data:main', # 実機用ウェイポイント推論 'inference_node = e2e_planner.inference_node:main', + # 学習・データセット前処理 + 'train = e2e_collector.train:main', + 'binarize_dataset = e2e_collector.binarize_dataset:main', + 'create_topomap = e2e_collector.create_topomap:main', ], }, ) diff --git a/simulator/CMakeLists.txt b/simulator/CMakeLists.txt index c11cfea1..bf97a648 100644 --- a/simulator/CMakeLists.txt +++ b/simulator/CMakeLists.txt @@ -20,10 +20,14 @@ find_package(ament_cmake REQUIRED) # uncomment the following section in order to fill in # further dependencies manually. # find_package( REQUIRED) -install( - DIRECTORY launch world models config - DESTINATION share/${PROJECT_NAME} -) +foreach(dir launch world models config) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${dir}") + install( + DIRECTORY ${dir} + DESTINATION share/${PROJECT_NAME} + ) + endif() +endforeach() install( PROGRAMS