235 lines
8.1 KiB
Python
235 lines
8.1 KiB
Python
import numpy as np
|
||
import torch
|
||
import random
|
||
|
||
from metadrive.type import MetaDriveType
|
||
|
||
|
||
def _static_obbs_overlap(car_a, car_b):
|
||
"""
|
||
Check if two static vehicle OBBs overlap (2D SAT).
|
||
car_a, car_b: dicts with "begin" (x, y), "heading" (rad), "length", "width".
|
||
begin is center; half-extents are length/2, width/2.
|
||
"""
|
||
def _get_corners(car):
|
||
cx, cy = car["begin"][0], car["begin"][1]
|
||
h = float(car["heading"])
|
||
L2 = float(car["length"]) / 2.0
|
||
W2 = float(car["width"]) / 2.0
|
||
ux, uy = np.cos(h), np.sin(h)
|
||
vx, vy = -np.sin(h), np.cos(h)
|
||
return np.array([
|
||
[cx + L2 * ux + W2 * vx, cy + L2 * uy + W2 * vy],
|
||
[cx + L2 * ux - W2 * vx, cy + L2 * uy - W2 * vy],
|
||
[cx - L2 * ux - W2 * vx, cy - L2 * uy - W2 * vy],
|
||
[cx - L2 * ux + W2 * vx, cy - L2 * uy + W2 * vy],
|
||
])
|
||
|
||
def _get_axes(car):
|
||
h = float(car["heading"])
|
||
return [
|
||
np.array([np.cos(h), np.sin(h)]),
|
||
np.array([-np.sin(h), np.cos(h)]),
|
||
]
|
||
|
||
corners_a = _get_corners(car_a)
|
||
corners_b = _get_corners(car_b)
|
||
axes = _get_axes(car_a) + _get_axes(car_b)
|
||
|
||
for axis in axes:
|
||
proj_a = corners_a @ axis
|
||
proj_b = corners_b @ axis
|
||
min_a, max_a = proj_a.min(), proj_a.max()
|
||
min_b, max_b = proj_b.min(), proj_b.max()
|
||
if max_a < min_b or max_b < min_a:
|
||
return False
|
||
return True
|
||
|
||
|
||
def _deduplicate_background_by_collision(background_vehicles):
|
||
"""
|
||
Merge static tracks that collide (same physical vehicle). Build collision graph,
|
||
find connected components, keep one representative per component (min show_time, then scenario_id).
|
||
"""
|
||
if not background_vehicles:
|
||
return background_vehicles
|
||
items = list(background_vehicles.items())
|
||
n = len(items)
|
||
# Build adjacency by index
|
||
parent = list(range(n))
|
||
|
||
def find(i):
|
||
if parent[i] != i:
|
||
parent[i] = find(parent[i])
|
||
return parent[i]
|
||
|
||
def union(i, j):
|
||
pi, pj = find(i), find(j)
|
||
if pi != pj:
|
||
parent[pi] = pj
|
||
|
||
for i in range(n):
|
||
for j in range(i + 1, n):
|
||
if _static_obbs_overlap(items[i][1], items[j][1]):
|
||
union(i, j)
|
||
|
||
# Representative per component: index with min (show_time, scenario_id)
|
||
comp_rep = {}
|
||
for i in range(n):
|
||
r = find(i)
|
||
sid, car = items[i][0], items[i][1]
|
||
key = (car.get("show_time", 0), sid)
|
||
if r not in comp_rep or key < comp_rep[r][0]:
|
||
comp_rep[r] = (key, sid, car)
|
||
|
||
return {sid: car for (_, sid, car) in comp_rep.values()}
|
||
|
||
|
||
def is_on_lane(pos, map_manager, threshold=2.0):
|
||
"""Check if a position is on a valid lane (within lateral tolerance)."""
|
||
if map_manager is None or map_manager.current_map is None:
|
||
return True
|
||
try:
|
||
lane, _ = map_manager.current_map.road_network.get_closest_lane_index(pos, return_lane=True)
|
||
if lane is None:
|
||
return False
|
||
long, lat = lane.local_coordinates(pos)
|
||
width = lane.width
|
||
if abs(lat) <= (width / 2 + threshold):
|
||
return True
|
||
return False
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
def filter_traffic_tracks_to_birth_lists(
|
||
current_traffic_data,
|
||
sdc_scenario_id,
|
||
map_manager,
|
||
*,
|
||
lane_threshold=5.0,
|
||
static_displacement_threshold=5.0,
|
||
static_speed_threshold=1.0,
|
||
return_stats=False,
|
||
deduplicate_static_by_collision=True,
|
||
):
|
||
"""
|
||
Filter traffic tracks into controlled (car_birth_info_list) and background lists.
|
||
|
||
- controlled (car_birth_info_list): 非 SDC、类型 VEHICLE、至少一帧 valid、在车道内、且非静态
|
||
（位移/速度超过阈值）。用于策略控制或专家回放，spawn 时机为 show_time == round。
|
||
- background (background_vehicles): 同上但在车道内且判定为静态（位移 < 5m、速度 < 1 m/s）。
|
||
仅作场景占位与观测邻居，spawn 时机为 show_time == round，按 valid 在 step 中移除。
|
||
|
||
Returns (car_birth_info_list, background_vehicles, obj_to_clean) or, if return_stats=True,
|
||
(car_birth_info_list, background_vehicles, obj_to_clean, stats_dict).
|
||
stats_dict: n_total, n_no_valid, n_off_lane, n_static, n_controlled.
|
||
"""
|
||
car_birth_info_list = []
|
||
background_vehicles = {}
|
||
obj_to_clean = []
|
||
n_total = 0
|
||
n_no_valid = 0
|
||
n_off_lane = 0
|
||
n_static = 0
|
||
|
||
for scenario_id, track in current_traffic_data.items():
|
||
if scenario_id == sdc_scenario_id:
|
||
continue
|
||
if track["type"] != MetaDriveType.VEHICLE:
|
||
continue
|
||
|
||
n_total += 1
|
||
obj_to_clean.append(scenario_id)
|
||
valid = track["state"]["valid"]
|
||
if not valid.any():
|
||
n_no_valid += 1
|
||
continue
|
||
|
||
first_show = int(np.argmax(valid))
|
||
last_show = len(valid) - 1 - int(np.argmax(valid[::-1]))
|
||
mid_show = (first_show + last_show) // 2
|
||
|
||
start_pos = track["state"]["position"][first_show]
|
||
is_valid_track = True
|
||
if not is_on_lane(start_pos, map_manager, threshold=lane_threshold):
|
||
mid_pos = track["state"]["position"][mid_show]
|
||
if not is_on_lane(mid_pos, map_manager, threshold=lane_threshold):
|
||
is_valid_track = False
|
||
|
||
if not is_valid_track:
|
||
n_off_lane += 1
|
||
continue
|
||
|
||
positions = track["state"]["position"][valid.astype(bool)]
|
||
velocities = track["state"]["velocity"][valid.astype(bool)]
|
||
total_displacement = 0.0
|
||
max_speed = 0.0
|
||
if len(positions) > 1:
|
||
total_displacement = float(np.linalg.norm(positions[-1] - positions[0]))
|
||
max_speed = float(np.max(np.linalg.norm(velocities, axis=1)))
|
||
is_static = total_displacement < static_displacement_threshold and max_speed < static_speed_threshold
|
||
|
||
if is_static:
|
||
n_static += 1
|
||
background_vehicles[scenario_id] = {
|
||
"id": track["metadata"]["object_id"],
|
||
"show_time": first_show,
|
||
"begin": (
|
||
float(track["state"]["position"][first_show, 0]),
|
||
float(track["state"]["position"][first_show, 1]),
|
||
),
|
||
"heading": float(track["state"]["heading"][first_show]),
|
||
"end": (
|
||
float(track["state"]["position"][last_show, 0]),
|
||
float(track["state"]["position"][last_show, 1]),
|
||
),
|
||
"scenario_id": scenario_id,
|
||
"length": track["state"]["length"][first_show],
|
||
"width": track["state"]["width"][first_show],
|
||
"valid": valid,
|
||
}
|
||
continue
|
||
|
||
car_birth_info_list.append({
|
||
"id": track["metadata"]["object_id"],
|
||
"show_time": first_show,
|
||
"begin": (
|
||
float(track["state"]["position"][first_show, 0]),
|
||
float(track["state"]["position"][first_show, 1]),
|
||
),
|
||
"heading": float(track["state"]["heading"][first_show]),
|
||
"end": (
|
||
float(track["state"]["position"][last_show, 0]),
|
||
float(track["state"]["position"][last_show, 1]),
|
||
),
|
||
"scenario_id": scenario_id,
|
||
"length": track["state"]["length"][first_show],
|
||
"width": track["state"]["width"][first_show],
|
||
})
|
||
|
||
if deduplicate_static_by_collision and background_vehicles:
|
||
background_vehicles = _deduplicate_background_by_collision(background_vehicles)
|
||
|
||
if return_stats:
|
||
stats = {
|
||
"n_total": n_total,
|
||
"n_no_valid": n_no_valid,
|
||
"n_off_lane": n_off_lane,
|
||
"n_static": n_static,
|
||
"n_controlled": len(car_birth_info_list),
|
||
}
|
||
return car_birth_info_list, background_vehicles, obj_to_clean, stats
|
||
return car_birth_info_list, background_vehicles, obj_to_clean
|
||
|
||
|
||
def set_seed(seed):
|
||
if seed == -1:
|
||
seed = np.random.randint(0, 10000)
|
||
print('Random seed: {}'.format(seed))
|
||
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
torch.cuda.manual_seed(seed)
|
||
torch.cuda.manual_seed_all(seed) |