# logic.py import os import json import cv2 import math import numpy as np from collections import deque from pathlib import Path import torch from torch.utils.data import Dataset from torchvision import transforms # --- Constants & Configs --- WAYPOINT_SCALE_FACTOR = 5.0 T1_FUTURE_TIME = 1.0 T2_FUTURE_TIME = 2.0 TRACKER_FREQUENCY = 10 MERGE_PERCENT = 0.4 PIXELS_PER_METER = 8 MAX_DISTANCE = 32 IMG_SIZE = MAX_DISTANCE * PIXELS_PER_METER * 2 EGO_CAR_X = IMG_SIZE // 2 EGO_CAR_Y = IMG_SIZE - (4.0 * PIXELS_PER_METER) reweight_array = np.ones((20, 20, 7)) last_valid_waypoints = None last_valid_theta = 0.0 class ControllerConfig: turn_KP, turn_KI, turn_KD, turn_n = 1.0, 0.1, 0.1, 20 speed_KP, speed_KI, speed_KD, speed_n = 0.5, 0.05, 0.1, 20 max_speed, max_throttle, clip_delta = 6.0, 0.75, 0.25 collision_buffer, detect_threshold = [0.0, 0.0], 0.04 brake_speed, brake_ratio = 0.4, 1.1 # --- Data Handling --- transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) lidar_transform = transforms.Compose([ transforms.Resize((112, 112)), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]), ]) class LMDriveDataset(Dataset): def __init__(self, data_dir, transform=None, lidar_transform=None): self.data_dir = Path(data_dir) self.transform = transform self.lidar_transform = lidar_transform self.samples = [] measurement_dir = self.data_dir / "measurements" image_dir = self.data_dir / "rgb_full" measurement_files = sorted([f for f in os.listdir(measurement_dir) if f.endswith(".json")]) image_files = sorted([f for f in os.listdir(image_dir) if f.endswith(".jpg")]) num_samples = min(len(measurement_files), len(image_files)) for i in range(num_samples): frame_id = i measurement_path = str(measurement_dir / f"{frame_id:04d}.json") image_name = f"{frame_id:04d}.jpg" image_path = str(image_dir / image_name) if not os.path.exists(measurement_path) or not os.path.exists(image_path): continue with open(measurement_path, "r") as f: measurements_data = json.load(f) self.samples.append({ "image_path": image_path, "measurement_path": measurement_path, "frame_id": frame_id, "measurements": measurements_data }) def __len__(self): return len(self.samples) def __getitem__(self, idx): sample = self.samples[idx] full_image = cv2.imread(sample["image_path"]) if full_image is None: raise ValueError(f"Failed to load image: {sample['image_path']}") full_image = cv2.cvtColor(full_image, cv2.COLOR_BGR2RGB) front_image = full_image[:600, :800] left_image = full_image[600:1200, :800] right_image = full_image[1200:1800, :800] center_image = full_image[1800:2400, :800] front_image_tensor = self.transform(front_image) left_image_tensor = self.transform(left_image) right_image_tensor = self.transform(right_image) center_image_tensor = self.transform(center_image) lidar_path = str(self.data_dir / "lidar" / f"{sample['frame_id']:04d}.png") lidar = cv2.imread(lidar_path) if lidar is None: lidar = np.zeros((112, 112, 3), dtype=np.uint8) else: if len(lidar.shape) == 2: lidar = cv2.cvtColor(lidar, cv2.COLOR_GRAY2BGR) lidar = cv2.cvtColor(lidar, cv2.COLOR_BGR2RGB) lidar_tensor = self.lidar_transform(lidar) measurements_data = sample["measurements"] x = measurements_data.get("x", 0.0) y = measurements_data.get("y", 0.0) theta = measurements_data.get("theta", 0.0) speed = measurements_data.get("speed", 0.0) steer = measurements_data.get("steer", 0.0) throttle = measurements_data.get("throttle", 0.0) brake = int(measurements_data.get("brake", False)) command = measurements_data.get("command", 0) is_junction = int(measurements_data.get("is_junction", False)) should_brake = int(measurements_data.get("should_brake", 0)) x_command = measurements_data.get("x_command", 0.0) y_command = measurements_data.get("y_command", 0.0) target_point = torch.tensor([x_command, y_command], dtype=torch.float32) measurements = torch.tensor( [x, y, theta, speed, steer, throttle, brake, command, is_junction, should_brake], dtype=torch.float32 ) return { "rgb": front_image_tensor, "rgb_left": left_image_tensor, "rgb_right": right_image_tensor, "rgb_center": center_image_tensor, "lidar": lidar_tensor, "measurements": measurements, "target_point": target_point } # --- Object Tracking --- class TrackedObject: def __init__(self): self.last_step = 0 self.last_pos = [0, 0] self.historical_pos = deque(maxlen=10) self.historical_steps = deque(maxlen=10) self.historical_features = deque(maxlen=10) def update(self, step, object_info): self.last_step = step self.last_pos = object_info[:2] self.historical_pos.append(self.last_pos) self.historical_steps.append(step) if len(object_info) > 2: self.historical_features.append(object_info[2]) class Tracker: def __init__(self, frequency=10): self.tracks = [] self.alive_ids = [] self.frequency = frequency def update_and_predict(self, det_data, pos, theta, frame_num): det_data_weighted = det_data * reweight_array detected_objects = find_peak_box(det_data_weighted) objects_info = [] R = np.array([[np.cos(-theta), -np.sin(-theta)], [np.sin(-theta), np.cos(-theta)]]) for obj in detected_objects: i, j = obj['coords'] obj_data = obj['raw_data'] center_y, center_x = convert_grid_to_xy(i, j) center_x += obj_data[1] center_y += obj_data[2] loc = R.T.dot(np.array([center_x, center_y])) objects_info.append([loc[0] + pos[0], loc[1] + pos[1], obj_data[1:]]) updates_ids = self._update(objects_info, frame_num) speed_results, heading_results = self._predict(updates_ids) for k, poi in enumerate(updates_ids): i, j = poi if heading_results[k] is not None: factor = MERGE_PERCENT * 0.1 det_data[i, j, 3] = heading_results[k] * factor + det_data[i, j, 3] * (1 - factor) if speed_results[k] is not None: factor = MERGE_PERCENT * 0.1 det_data[i, j, 6] = speed_results[k] * factor + det_data[i, j, 6] * (1 - factor) return det_data def _update(self, objects_info, step): latest_ids = [] if len(self.tracks) == 0: for object_info in objects_info: to = TrackedObject() to.update(step, object_info) self.tracks.append(to) latest_ids.append(len(self.tracks) - 1) else: matched_ids = set() for idx, object_info in enumerate(objects_info): min_id, min_error = -1, float('inf') pos_x, pos_y = object_info[:2] for _id in self.alive_ids: if _id in matched_ids: continue track_pos = self.tracks[_id].last_pos distance = np.sqrt((track_pos[0] - pos_x)**2 + (track_pos[1] - pos_y)**2) if distance < 2.0 and distance < min_error: min_error = distance min_id = _id if min_id != -1: self.tracks[min_id].update(step, objects_info[idx]) latest_ids.append(min_id) matched_ids.add(min_id) else: to = TrackedObject() to.update(step, object_info) self.tracks.append(to) latest_ids.append(len(self.tracks) - 1) self.alive_ids = [i for i, track in enumerate(self.tracks) if track.last_step > step - 6] return latest_ids def _predict(self, updates_ids): speed_results, heading_results = [], [] for each_id in updates_ids: to = self.tracks[each_id] avg_speed, avg_heading = [], [] for feature in to.historical_features: avg_speed.append(feature[2]) avg_heading.append(feature[:2]) if len(avg_speed) < 2: speed_results.append(None) heading_results.append(None) continue avg_speed = np.mean(avg_speed) avg_heading = np.mean(np.stack(avg_heading), axis=0) yaw_angle = get_yaw_angle(avg_heading) heading_results.append((4 - yaw_angle / np.pi) % 2) speed_results.append(avg_speed) return speed_results, heading_results # --- Control System --- class PIDController: def __init__(self, K_P=1.0, K_I=0.0, K_D=0.0, n=20): self._K_P = K_P self._K_I = K_I self._K_D = K_D self._window = deque([0 for _ in range(n)], maxlen=n) self._max = 0.0 self._min = 0.0 def step(self, error): self._window.append(error) self._max = max(self._max, abs(error)) self._min = -abs(self._max) if len(self._window) >= 2: integral = np.mean(self._window) derivative = self._window[-1] - self._window[-2] else: integral = 0.0 derivative = 0.0 return self._K_P * error + self._K_I * integral + self._K_D * derivative class InterfuserController: def __init__(self, config): self.turn_controller = PIDController(config.turn_KP, config.turn_KI, config.turn_KD, config.turn_n) self.speed_controller = PIDController(config.speed_KP, config.speed_KI, config.speed_KD, config.speed_n) self.config = config self.collision_buffer = np.array(config.collision_buffer) self.detect_threshold = config.detect_threshold self.stop_steps = 0 self.forced_forward_steps = 0 self.red_light_steps = 0 self.block_red_light = 0 self.in_stop_sign_effect = False self.block_stop_sign_distance = 0 self.stop_sign_timer = 0 self.stop_sign_trigger_times = 0 def run_step(self, speed, waypoints, junction, traffic_light_state, stop_sign, meta_data): if speed < 0.2: self.stop_steps += 1 else: self.stop_steps = max(0, self.stop_steps - 10) if speed < 0.06 and self.in_stop_sign_effect: self.in_stop_sign_effect = False if junction < 0.3: self.stop_sign_trigger_times = 0 if traffic_light_state > 0.7: self.red_light_steps += 1 else: self.red_light_steps = 0 if self.red_light_steps > 1000: self.block_red_light = 80 self.red_light_steps = 0 if self.block_red_light > 0: self.block_red_light -= 1 traffic_light_state = 0.01 if stop_sign < 0.6 and self.block_stop_sign_distance < 0.1: self.in_stop_sign_effect = True self.block_stop_sign_distance = 2.0 self.stop_sign_trigger_times = 3 self.block_stop_sign_distance = max(0, self.block_stop_sign_distance - 0.05 * speed) if self.block_stop_sign_distance < 0.1: if self.stop_sign_trigger_times > 0: self.block_stop_sign_distance = 2.0 self.stop_sign_trigger_times -= 1 self.in_stop_sign_effect = True aim = (waypoints[1] + waypoints[0]) / 2.0 angle = np.degrees(np.pi / 2 - np.arctan2(aim[1], aim[0])) / 90 if speed < 0.01: angle = 0 steer = self.turn_controller.step(angle) steer = np.clip(steer, -1.0, 1.0) desired_speed = self.config.max_speed downsampled_waypoints = downsample_waypoints(waypoints) safe_dis = get_max_safe_distance(meta_data, downsampled_waypoints, t=0, collision_buffer=self.collision_buffer, threshold=self.detect_threshold) if traffic_light_state > 0.5 or (stop_sign > 0.6 and self.stop_sign_timer < 20): desired_speed = 0.0 if stop_sign > 0.6: self.stop_sign_timer += 1 else: self.stop_sign_timer = 0 brake = speed > desired_speed * self.config.brake_ratio delta = np.clip(desired_speed - speed, 0.0, self.config.clip_delta) throttle = self.speed_controller.step(delta) throttle = np.clip(throttle, 0.0, self.config.max_throttle) meta_info_1 = f"speed: {speed:.2f}, target_speed: {desired_speed:.2f}" meta_info_2 = f"on_road_prob: {junction:.2f}, red_light_prob: {traffic_light_state:.2f}, stop_sign_prob: {1 - stop_sign:.2f}" meta_info_3 = f"stop_steps: {self.stop_steps}, block_stop_sign_distance: {self.block_stop_sign_distance:.1f}" if self.stop_steps > 1200: self.forced_forward_steps = 12 self.stop_steps = 0 if self.forced_forward_steps > 0: throttle = 0.8 brake = False self.forced_forward_steps -= 1 if self.in_stop_sign_effect: throttle = 0 brake = True return steer, throttle, brake, (meta_info_1, meta_info_2, meta_info_3, safe_dis) # --- Visualization & Helper Functions --- class DisplayInterface: def __init__(self, width=1200, height=600): self._width = width self._height = height def run_interface(self, data): dashboard = np.zeros((self._height, self._width, 3), dtype=np.uint8) font = cv2.FONT_HERSHEY_SIMPLEX dashboard[:, :800] = cv2.resize(data.get('camera_view'), (800, 600)) dashboard[:400, 800:1200] = cv2.resize(data['map_t0'], (400, 400)) dashboard[400:600, 800:1000] = cv2.resize(data['map_t1'], (200, 200)) dashboard[400:600, 1000:1200] = cv2.resize(data['map_t2'], (200, 200)) cv2.line(dashboard, (800, 0), (800, 600), (255, 255, 255), 2) cv2.line(dashboard, (800, 400), (1200, 400), (255, 255, 255), 2) cv2.line(dashboard, (1000, 400), (1000, 600), (255, 255, 255), 2) y_pos = 40 for key, text in data['text_info'].items(): cv2.putText(dashboard, text, (820, y_pos), font, 0.6, (255, 255, 255), 1) y_pos += 30 y_pos += 10 for t, counts in data['object_counts'].items(): count_str = f"{t}: C={counts['car']} B={counts['bike']} P={counts['pedestrian']}" cv2.putText(dashboard, count_str, (820, y_pos), font, 0.5, (255, 255, 255), 1) y_pos += 20 cv2.putText(dashboard, "t0", (1160, 30), font, 0.8, (0, 255, 255), 2) cv2.putText(dashboard, "t1", (960, 430), font, 0.8, (0, 255, 255), 2) cv2.putText(dashboard, "t2", (1160, 430), font, 0.8, (0, 255, 255), 2) return dashboard def get_yaw_angle(forward_vector): forward_vector = forward_vector / np.linalg.norm(forward_vector) return math.atan2(forward_vector[1], forward_vector[0]) def ensure_rgb(image): if len(image.shape) == 2 or image.shape[2] == 1: return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) return image def convert_grid_to_xy(i, j): return (j - 9.5) * 1.6, (19.5 - i) * 1.6 def find_peak_box(data): det_data = np.zeros((22, 22, 7)) det_data[1:21, 1:21] = data detected_objects = [] for i in range(1, 21): for j in range(1, 21): if det_data[i, j, 0] > 0.6 and (det_data[i, j, 0] > det_data[i, j - 1, 0] and det_data[i, j, 0] > det_data[i, j + 1, 0] and det_data[i, j, 0] > det_data[i - 1, j, 0] and det_data[i, j, 0] > det_data[i + 1, j, 0]): length, width, confidence = det_data[i, j, 4], det_data[i, j, 5], det_data[i, j, 0] obj_type = 'unknown' if length > 4.0: obj_type = 'car' elif length / width > 1.5: obj_type = 'bike' else: obj_type = 'pedestrian' detected_objects.append({'coords': (i - 1, j - 1), 'type': obj_type, 'confidence': confidence, 'raw_data': det_data[i, j]}) return detected_objects def add_rect(img, loc, ori, box, value, color): center_x = int(loc[0] * PIXELS_PER_METER + MAX_DISTANCE * PIXELS_PER_METER) center_y = int(loc[1] * PIXELS_PER_METER + MAX_DISTANCE * PIXELS_PER_METER) size_px = (int(box[0] * PIXELS_PER_METER), int(box[1] * PIXELS_PER_METER)) angle_deg = -np.degrees(math.atan2(ori[1], ori[0])) box_points = cv2.boxPoints(((center_x, center_y), size_px, angle_deg)) cv2.fillConvexPoly(img, np.int32(box_points), [int(x * value) for x in color]) return img def render(det_data, t=0): CLASS_COLORS = {'car': (0, 0, 255), 'bike': (0, 255, 0), 'pedestrian': (255, 0, 0), 'unknown': (128, 128, 128)} det_weighted = det_data * reweight_array detected_objects = find_peak_box(det_weighted) counts = {cls: 0 for cls in CLASS_COLORS.keys()} [counts.update({obj['type']: counts.get(obj['type'], 0) + 1}) for obj in detected_objects] img = np.zeros((IMG_SIZE, IMG_SIZE, 3), np.uint8) for obj in detected_objects: i, j = obj['coords'] obj_data = obj['raw_data'] speed, theta = obj_data[6], obj_data[3] * np.pi ori = np.array([math.cos(theta), math.sin(theta)]) loc_x = obj_data[1] + t * speed * ori[0] + convert_grid_to_xy(i, j)[0] loc_y = obj_data[2] - t * speed * ori[1] + convert_grid_to_xy(i, j)[1] box = np.array([obj_data[4], obj_data[5]]) * (1.5 if obj['type'] == 'pedestrian' else 1.0) add_rect(img, np.array([loc_x, loc_y]), ori, box, obj['confidence'], CLASS_COLORS[obj['type']]) return img, counts def render_self_car(loc, ori, box, pixels_per_meter=PIXELS_PER_METER): img = np.zeros((IMG_SIZE, IMG_SIZE, 3), np.uint8) center_x = int(loc[0] * pixels_per_meter + MAX_DISTANCE * pixels_per_meter) center_y = int(loc[1] * pixels_per_meter + MAX_DISTANCE * pixels_per_meter) size_px = (int(box[0] * pixels_per_meter), int(box[1] * pixels_per_meter)) angle_deg = -np.degrees(math.atan2(ori[1], ori[0])) box_points = np.int32(cv2.boxPoints(((center_x, center_y), size_px, angle_deg))) cv2.fillConvexPoly(img, box_points, (0, 255, 255)) return img def render_waypoints(waypoints, pixels_per_meter=PIXELS_PER_METER): global last_valid_waypoints img = np.zeros((IMG_SIZE, IMG_SIZE, 3), np.uint8) current_waypoints = waypoints if waypoints is not None and len(waypoints) > 2 else last_valid_waypoints if current_waypoints is not None: last_valid_waypoints = current_waypoints for i, point in enumerate(current_waypoints): px = int(EGO_CAR_X + point[1] * pixels_per_meter) py = int(EGO_CAR_Y - point[0] * pixels_per_meter) color = (0, 0, 255) if i == len(current_waypoints) - 1 else (0, 255, 0) cv2.circle(img, (px, py), 4, color, -1) return img def collision_detections(map1, map2, threshold=0.04): if len(map2.shape) == 3 and map2.shape[2] == 3: map2 = cv2.cvtColor(map2, cv2.COLOR_BGR2GRAY) assert map1.shape == map2.shape overlap_map = (map1 > 0.01) & (map2 > 0.01) return float(np.sum(overlap_map)) / np.sum(map2 > 0) < threshold def get_max_safe_distance(meta_data, downsampled_waypoints, t, collision_buffer, threshold): surround_map = meta_data.reshape(20, 20, 7)[..., :3][..., 0] if np.sum(surround_map) < 1: return np.linalg.norm(downsampled_waypoints[-3]) hero_bounding_box = np.array([2.45, 1.0]) + collision_buffer safe_distance = 0.0 for i in range(len(downsampled_waypoints) - 2): aim = (downsampled_waypoints[i + 1] + downsampled_waypoints[i + 2]) / 2.0 loc, ori = downsampled_waypoints[i], aim - downsampled_waypoints[i] self_car_map = render_self_car(loc=loc, ori=ori, box=hero_bounding_box, pixels_per_meter=PIXELS_PER_METER) self_car_map_gray = cv2.cvtColor(cv2.resize(self_car_map, (20, 20)), cv2.COLOR_BGR2GRAY) if not collision_detections(surround_map, self_car_map_gray, threshold): break safe_distance = max(safe_distance, np.linalg.norm(loc)) return safe_distance def downsample_waypoints(waypoints, precision=0.2): downsampled_waypoints = [] last_waypoint = np.array([0.0, 0.0]) for i in range(len(waypoints)): now_waypoint = waypoints[i] dis = np.linalg.norm(now_waypoint - last_waypoint) if dis > precision: interval = int(dis / precision) move_vector = (now_waypoint - last_waypoint) / (interval + 1) for j in range(interval): downsampled_waypoints.append(last_waypoint + move_vector * (j + 1)) downsampled_waypoints.append(now_waypoint) last_waypoint = now_waypoint return downsampled_waypoints