Upload logic.py with huggingface_hub

logic.py

@@ -0,0 +1,495 @@
+# 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