Feat/shannon surprise

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"
+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

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,21 +81,22 @@ 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:
                 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.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)

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)

e2e_planner/e2e_planner/util/__init__.py

@@ -0,0 +1 @@
+

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

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)
+    ]

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
+        )

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,
         }]
     )