def transform(self, batch):
# (2,) agent coordinates in image frame of reference
agent_position_image = np.array([
self.cfg['raster_params']['ego_center'][0] * self.W,
self.cfg['raster_params']['ego_center'][1] * self.W
])
# (1,) meters per pixel
r = self.cfg['raster_params']['pixel_size'][0]
# initially positions are given in agent frame of reference but with wrong rotation
# transform into image frame of reference using rotation and affine transformation
# after that tranform into agent frame of reference back
batch["target_positions"] = (transform_points(
batch["target_positions"] + batch["centroid"],
batch["world_to_image"]) - agent_position_image) * r
batch['history_positions'] = (transform_points(
batch['history_positions'] + batch["centroid"],
batch["world_to_image"]) - agent_position_image) * r
return (batch["image"].astype(np.float32),
batch["target_positions"].astype(np.float32),
batch["target_availabilities"].astype(np.float32),
batch['history_positions'].astype(np.float32),
batch['history_yaws'].astype(np.float32),
batch["centroid"].astype(np.float32),
batch["world_to_image"].astype(np.float32))
def check_performance(dataset, name="", num_samples=64 * 20, random_order=False):
with utils.timeit_context(f"iterate {name} dataset"):
sample = dataset[63]
# print("image shape", sample["image"]["image_sem"].shape, sample["image"]["image_sem"].dtype)
print("Keys:", sample.keys())
target_positions = sample["target_positions"]
target_positions_world = transform_points(target_positions, sample["world_from_agent"])
# output_mask = sample["output_mask"]
img = dataset.rasterizer.to_rgb(sample["image"].transpose(1, 2, 0))
plt.imshow(img)
agents_history = sample["agents_history"]
cur_frame_positions = agents_history[-1, :, :2] * 100.0
cur_frame_velocity = agents_history[-1, :, 2:4] * 10.0
cur_frame_positions_img = transform_points(cur_frame_positions, sample["raster_from_agent"])
plt.scatter(cur_frame_positions_img[:, 0], cur_frame_positions_img[:, 1])
plt.scatter(cur_frame_positions_img[:, 0] + cur_frame_velocity[:, 0] * 1.0,
cur_frame_positions_img[:, 1] + cur_frame_velocity[:, 1] * 1.0,
c='red')
plt.show()
nb_samples = len(dataset)
for i in tqdm(range(num_samples)):
if random_order:
sample = dataset[np.random.randint(0, nb_samples)]
else:
sample = dataset[i]
target_positions = sample["target_positions"]
def get_box_nodes(self, history_frames: np.ndarray,
history_agents: List[np.ndarray],
agent: Optional[np.ndarray],
raster_from_world: np.ndarray) -> Dict[str, List]:
"""
remove node not in the latest
"""
all_agents_info = []
ego_info = []
# loop over historical frames
for i, (frame, agents) in enumerate(zip(history_frames,
history_agents)):
agents = filter_agents_by_labels(agents,
self.filter_agents_threshold)
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
agents = np.concatenate([agents, av_agent])
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
agents = self.__class__.remove_agents_by_track_id(
agents, track_id=agent["track_id"])
agent_ego["centroid"] = transform_points(agent_ego["centroid"],
raster_from_world)
agents["centroid"] = transform_points(agents["centroid"],
raster_from_world)
all_agents_info.append(agents)
ego_info.append(agent_ego)
return {
self.__class__.AGENTS: all_agents_info,
self.__class__.EGO: ego_info
}
def test_transform_points_transpose_equivalence() -> None:
input_points = np.random.rand(10, 3)
input_points_t = input_points.transpose(1, 0)
# Generate some random transformation matrix
tf = np.identity(4)
tf[:3, :3] = transforms3d.euler.euler2mat(np.random.rand(),
np.random.rand(),
np.random.rand())
tf[3, :3] = np.random.rand(3)
output_points = transform_points(input_points, tf)
output_points_t = transform_points_transposed(input_points_t, tf)
np.testing.assert_array_equal(output_points.transpose(1, 0),
output_points_t)
tf_inv = np.linalg.inv(tf)
input_points_recovered = transform_points(output_points, tf_inv)
input_points_t_recovered = transform_points_transposed(
output_points_t, tf_inv)
np.testing.assert_almost_equal(input_points_recovered,
input_points,
decimal=10)
np.testing.assert_almost_equal(input_points_t_recovered,
input_points_t,
decimal=10)
def test_wrong_input_shape(self) -> None:
tf = np.eye(4)
with self.assertRaises(ValueError):
points = np.zeros((3, 10))
transform_points(points, tf)
with self.assertRaises(ValueError):
points = np.zeros((10, 3))
transform_points_transposed(points, tf)
def traj_from_item(item):
history = transform_points(
item['history_positions'][item['history_availabilities'].astype(bool)],
item['world_from_agent'])
future = transform_points(
item['target_positions'][item['target_availabilities'].astype(bool)],
item['world_from_agent'])
return Traj(xx0=history[:, 0],
yy0=history[:, 1],
xx1=future[:, 0],
yy1=future[:, 1])
def test_transform_points_broadcast() -> None:
# transform a batch of points with the same matrix, should be the same a transforming each
tf = np.random.randn(3, 3)
batch_points = np.random.randn(16, 50, 2)
output_points = transform_points(batch_points, tf)
expected_points = []
for points in batch_points:
expected_points.append(transform_points(points, tf))
expected_points = np.stack(expected_points)
assert np.allclose(output_points, expected_points, atol=1e-5)
def rasterize_box(
self,
history_frames: np.ndarray,
history_agents: List[np.ndarray],
history_tl_faces: List[np.ndarray],
agent: Optional[np.ndarray] = None,
svg=False, svg_args=None,
) -> th.Union[dict]:
# all frames are drawn relative to this one"
frame = history_frames[0]
if agent is None:
ego_translation_m = history_frames[0]["ego_translation"]
ego_yaw_rad = rotation33_as_yaw(frame["ego_rotation"])
else:
ego_translation_m = np.append(agent["centroid"], history_frames[0]["ego_translation"][-1])
ego_yaw_rad = agent["yaw"]
svg_args = svg_args or dict()
raster_from_world = self.render_context.raster_from_world(ego_translation_m, ego_yaw_rad)
raster_size = self.render_context.raster_size_px
# this ensures we always end up with fixed size arrays, +1 is because current time is also in the history
res = dict(ego=list(), agents=defaultdict(list))
for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):
# print('history index', i)
agents = filter_agents_by_labels(agents, self.filter_agents_threshold)
# note the cast is for legacy support of dataset before April 2020
av_agent = get_ego_as_agent(frame).astype(agents.dtype)
if agent is None:
add_agents(res['agents'], av_agent)
res['ego'].append(av_agent[0]["centroid"][:2])
else:
agent_ego = filter_agents_by_track_id(agents, agent["track_id"])
if len(agent_ego) == 0: # agent not in this history frame
add_agents(res['agents'], np.append(agents, av_agent))
else: # add av to agents and remove the agent from agents
agents = agents[agents != agent_ego[0]]
add_agents(res['agents'], np.append(agents, av_agent))
res['ego'].append(agent_ego[0]["centroid"][:2])
tolerance = svg_args.get('tolerance', 20.)
_ego = normalize_line(
cv2_subpixel(transform_points(np.array(res['ego']).reshape((-1, 2)), raster_from_world)))
res['ego'] = crop_tensor(_ego, raster_size)
ego_grad = calc_max_grad(res['ego'])
res['agents'] = [normalize_line(cv2_subpixel(transform_points(np.array(path).reshape((-1, 2)), raster_from_world))
) for idx, path in res['agents'].items()]
res['agents'] = [
crop_tensor(path, raster_size) for path in res['agents'] if not is_noisy(path, ego_grad, tolerance)]
res['agents'] = [path for path in res['agents'] if len(path)]
if svg:
res['path'] = torch.cat([linear_path_to_tensor(path, svg_args.get('pad_val', -1)) for path in res['agents']
] + [linear_path_to_tensor(res['ego'], svg_args.get('pad_val', -1))], 0)
res['path_type'] = [path_type_to_number('agent')] * len(res['agents']) + [path_type_to_number('ego')]
return res
def draw_velocity(
raster_size: Tuple[int, int],
world_to_image_space: np.ndarray,
agents: np.ndarray,
color: Union[int, Tuple[int, int, int]],
velocity_scale: float = 1.0,
thickness: int = 1,
tipLength: float = 0.1,
) -> np.ndarray:
"""
Draw multiple boxes in one sweep over the image.
Boxes corners are extracted from agents, and the coordinates are projected in the image plane.
Finally, cv2 draws the boxes.
Args:
raster_size (Tuple[int, int]): Desired output image size
world_to_image_space (np.ndarray): 3x3 matrix to convert from world to image coordinated
agents (np.ndarray): array of agents to be drawn
color (Union[int, Tuple[int, int, int]]): single int or RGB color
velocity_scale (float): how long velocity vector is drawn. bigger value will draw longer vector.
thickness (int): thickness of arrow
tipLength (float): arrow's tip length
Returns:
np.ndarray: the image with agents rendered. RGB if color RGB, otherwise GRAY
"""
if isinstance(color, int):
im = np.zeros(raster_size, dtype=np.uint8)
else:
im = np.zeros(raster_size + (3,), dtype=np.uint8)
# (n_agents, 2)
world_coords = agents["centroid"][:, :2]
world_velocity = agents["velocity"]
# print("world_coords", world_coords.shape, world_coords)
# print("world_velocity", world_velocity.shape, world_velocity)
image_coords = transform_points(world_coords, world_to_image_space)
image_coords_next = transform_points(world_coords + world_velocity * velocity_scale, world_to_image_space)
# image_velocity = transform_points(world_velocity, world_to_image_space) # This does NOT work!
# print("image_coords", image_coords.shape, image_coords)
# print("image_velocity", image_velocity.shape, image_velocity)
# fillPoly wants polys in a sequence with points inside as (x,y)
pt1 = image_coords.astype(np.int64)
pt2 = image_coords_next.astype(np.int64)
# print("pt1", pt1.shape, pt1)
# print("pt2", pt2.shape, pt2)
for i in range(len(pt1)):
cv2.arrowedLine(im, tuple(pt1[i]), tuple(pt2[i]), color, thickness=thickness, tipLength=tipLength)
return im
def test_transform_batch_points() -> None:
# transform batch and singular elements one by one, results should match
# note: we use random here as the validity of transform is checked below
tfs = np.random.randn(16, 3, 3)
batch_points = np.random.randn(16, 50, 2)
output_points = transform_points(batch_points, tfs)
expected_points = []
for points, tf in zip(batch_points, tfs):
expected_points.append(transform_points(points, tf))
expected_points = np.stack(expected_points)
assert np.allclose(output_points, expected_points, atol=1e-5)
def _get_in_out_as_trajectories(
in_out: UnrollInputOutput) -> Tuple[np.ndarray, np.ndarray]:
"""Convert the input (log-replayed) and output (simulated) trajectories into world space.
Apply availability on the log-replayed one
:param in_out: an UnrollInputOutput object
:return: the replayed and simulated trajectory as numpy arrays
"""
replay_traj = transform_points(in_out.inputs["target_positions"],
in_out.inputs["world_from_agent"])
replay_traj = replay_traj[in_out.inputs["target_availabilities"] > 0]
sim_traj = transform_points(in_out.outputs["positions"],
in_out.inputs["world_from_agent"])
return replay_traj, sim_traj
def render_semantic_map(
self, center_in_world: np.ndarray, raster_from_world: np.ndarray, tl_faces: np.ndarray = None,
tl_face_color=True
) -> th.Union[torch.Tensor, dict]:
"""Renders the semantic map at given x,y coordinates.
Args:
center_in_world (np.ndarray): XY of the image center in world ref system
raster_from_world (np.ndarray):
Returns:
th.Union[torch.Tensor, dict]
"""
# filter using half a radius from the center
raster_radius = float(np.linalg.norm(self.raster_size * self.pixel_size)) / 2
# get active traffic light faces
if tl_face_color:
active_tl_ids = set(filter_tl_faces_by_status(tl_faces, "ACTIVE")["face_id"].tolist())
# setup canvas
raster_size = self.render_context.raster_size_px
res = dict(path=list(), path_type=list())
for idx in elements_within_bounds(center_in_world, self.bounds_info["lanes"]["bounds"], raster_radius):
lane = self.proto_API[self.bounds_info["lanes"]["ids"][idx]].element.lane
# get image coords
lane_coords = self.proto_API.get_lane_coords(self.bounds_info["lanes"]["ids"][idx])
xy_left = cv2_subpixel(transform_points(lane_coords["xyz_left"][:, :2], raster_from_world))
xy_right = cv2_subpixel(transform_points(lane_coords["xyz_right"][:, :2], raster_from_world))
xy_left = normalize_line(xy_left)
xy_right = normalize_line(xy_right)
lane_type = "black" # no traffic light face is controlling this lane
if tl_face_color:
lane_tl_ids = set([MapAPI.id_as_str(la_tc) for la_tc in lane.traffic_controls])
for tl_id in lane_tl_ids.intersection(active_tl_ids):
if self.proto_API.is_traffic_face_colour(tl_id, "red"):
lane_type = "red"
elif self.proto_API.is_traffic_face_colour(tl_id, "green"):
lane_type = "green"
elif self.proto_API.is_traffic_face_colour(tl_id, "yellow"):
lane_type = "yellow"
for vector in [xy_left, xy_right]:
vector = crop_tensor(vector, raster_size)
if len(vector):
res['path'].append(vector)
res['path_type'].append(path_type_to_number(lane_type))
return res
def run_prediction(predictor, data_loader):
predictor.eval()
pred_coords_list = []
confidences_list = []
timestamps_list = []
track_id_list = []
with torch.no_grad():
dataiter = tqdm(data_loader)
for data in dataiter:
image = data["image"].to(device)
pred, confidences = predictor(image)
world_from_agents = data["world_from_agent"].cpu().numpy()
centroid = data["centroid"].cpu().numpy()
pred = pred.cpu().numpy()
for i in range(pred.shape[0]):
for j in range(pred.shape[1]):
pred[i, j, :] = transform_points(
pred[i, j, ], world_from_agents[i]) - centroid[i, :2]
pred_coords_list.append(pred)
confidences_list.append(confidences.cpu().numpy().copy())
timestamps_list.append(data["timestamp"].numpy().copy())
track_id_list.append(data["track_id"].numpy().copy())
timestamps = np.concatenate(timestamps_list)
track_ids = np.concatenate(track_id_list)
coords = np.concatenate(pred_coords_list)
confs = np.concatenate(confidences_list)
return timestamps, track_ids, coords, confs
def generate_mask_coords(agent_centroid, ego_centroid, world_from_agent, mask):
"""agent_centroid: (1,2) arr
ego_centroid: (,2 arr)"""
assert mask.shape[1] % 2 == 1 #masks must be odd number shaped
assert mask.shape[2] % 2 == 1 #masks must be odd number shaped
ego_centroid = torch.from_numpy(np.expand_dims(ego_centroid, axis=0))
ego_centroid = transform_points(ego_centroid, world_from_agent)
rel_centroid = (agent_centroid.squeeze() - ego_centroid.numpy().squeeze())
#for now, let's generate a 13x13 mask. Could choose long mask (e.x. 13x3) if we
#input ego heading to determine whether x or y becomes the long end.
#1 square == ~15 feet
map_x_uncentered = rel_centroid[0] // 15
map_y_uncentered = rel_centroid[1] // 15
#Adjust coords to put ego vehicle in center, and flip x and y for python array indexing
map_x_center = mask.shape[1] // 2 + 1
map_y_center = mask.shape[2] // 2 + 1
map_x_centered = map_y_center - map_y_uncentered
map_y_centered = map_x_center + map_x_uncentered
if map_x_centered < 0 or map_x_centered >= mask.shape[1]:
return None
if map_y_centered < 0 or map_y_centered >= mask.shape[2]:
return None
else:
return [int(map_x_centered), int(map_y_centered)]
def get_frame(self, scene_index: int, state_index: int, track_id: Optional[int] = None) -> dict:
"""
A utility function to get the rasterisation and trajectory target for a given agent in a given frame
Args:
scene_index (int): the index of the scene in the zarr
state_index (int): a relative frame index in the scene
track_id (Optional[int]): the agent to rasterize or None for the AV
Returns:
dict: the rasterised image, the target trajectory (position and yaw) along with their availability,
the 2D matrix to center that agent, the agent track (-1 if ego) and the timestamp
"""
frames = self.dataset.frames[get_frames_slice_from_scenes(self.dataset.scenes[scene_index])]
data = self.sample_function(state_index, frames, self.dataset.agents, self.dataset.tl_faces, track_id)
# 0,1,C -> C,0,1
image = data["image"].transpose(2, 0, 1)
history_positions = np.array(data["history_positions"], dtype=np.float32)
history_yaws = np.array(data["history_yaws"], dtype=np.float32)
timestamp = frames[state_index]["timestamp"]
track_id = np.int64(-1 if track_id is None else track_id) # always a number to avoid crashing torch
target_positions = np.array(data["target_positions"], dtype=np.float32)
target_yaws = np.array(data["target_yaws"], dtype=np.float32)
target_availabilities = data["target_availabilities"]
# Get target_positions from ground truth if self.gt is available
if self.gt is not None:
assert str(track_id) + str(
timestamp) in self.gt, 'self.gt (ground truth) does not contain requested track_id/timestamp combination. We have got a problem somewhere!'
target_positions = np.array(self.gt[str(track_id) + str(timestamp)][0], dtype=np.float32)
target_positions = transform_points(target_positions + data['centroid'][:2], data['agent_from_world'])
target_availabilities = np.array(self.gt[str(track_id) + str(timestamp)][1], dtype=np.float32)
ego_center = data['ego_center'] if 'ego_center' in data else np.array(self.cfg['raster_params']['ego_center'])
return {
"image": image,
"target_positions": target_positions,
"target_yaws": target_yaws,
"target_availabilities": target_availabilities,
"history_positions": history_positions,
"history_yaws": history_yaws,
"history_availabilities": data["history_availabilities"],
"world_to_image": data["world_to_image"],
"raster_from_world": data["raster_from_world"],
"raster_from_agent": data["raster_from_agent"],
"world_from_agent": data["world_from_agent"],
"agent_from_world": data["agent_from_world"],
"track_id": track_id,
"timestamp": timestamp,
"centroid": data["centroid"],
"ego_center": ego_center,
"yaw": data["yaw"],
"extent": data["extent"],
}
def update_agents(dataset: SimulationDataset, frame_idx: int,
input_dict: Dict[str, np.ndarray],
output_dict: Dict[str, np.ndarray]) -> None:
"""Update the agents in frame_idx (across scenes) using agents_output_dict
:param dataset: the simulation dataset
:param frame_idx: index of the frame to modify
:param input_dict: the input to the agent model
:param output_dict: the output of the agent model
:return:
"""
agents_update_dict: Dict[Tuple[int, int], np.ndarray] = {}
world_from_agent = input_dict["world_from_agent"]
yaw = input_dict["yaw"]
pred_trs = transform_points(output_dict["positions"][:, :1],
world_from_agent)[:, 0]
pred_yaws = yaw + output_dict["yaws"][:, 0, 0]
next_agents = np.zeros(len(yaw), dtype=AGENT_DTYPE)
next_agents["centroid"] = pred_trs
next_agents["yaw"] = pred_yaws
next_agents["track_id"] = input_dict["track_id"]
next_agents["extent"] = input_dict["extent"]
next_agents["label_probabilities"][:, PERCEPTION_LABEL_TO_INDEX[
"PERCEPTION_LABEL_CAR"]] = 1
for scene_idx, next_agent in zip(input_dict["scene_index"],
next_agents):
agents_update_dict[(scene_idx,
next_agent["track_id"])] = np.expand_dims(
next_agent, 0)
dataset.set_agents(frame_idx, agents_update_dict)
def unpack_deltas_cm(
self, dx: Sequence[int], dy: Sequence[int], dz: Sequence[int], frame: GeoFrame
) -> np.ndarray:
"""
Get coords in world reference system (local ENU->ECEF->world).
See the protobuf annotations for additional information about how coordinates are stored
Args:
dx (Sequence[int]): X displacement in centimeters in local ENU
dy (Sequence[int]): Y displacement in centimeters in local ENU
dz (Sequence[int]): Z displacement in centimeters in local ENU
frame (GeoFrame): geo-location information for the local ENU. It contains lat and long origin of the frame
Returns:
np.ndarray: array of shape (Nx3) with XYZ coordinates in world ref system
"""
x = np.cumsum(np.asarray(dx) / 100)
y = np.cumsum(np.asarray(dy) / 100)
z = np.cumsum(np.asarray(dz) / 100)
frame_lat, frame_lng = self._undo_e7(frame.origin.lat_e7), self._undo_e7(
frame.origin.lng_e7
)
xyz = np.stack(pm.enu2ecef(x, y, z, frame_lat, frame_lng, 0), axis=-1)
xyz = transform_points(xyz, self.ecef_to_world)
return xyz
def test_transform_points() -> None:
tf = np.asarray([[1.0, 0, 100], [0, 0.5, 50], [0, 0, 1]])
points = np.array([[0, 10], [10, 0], [10, 10]])
expected_point = np.array([[100, 55], [110, 50], [110, 55]])
output_points = transform_points(points, tf)
np.testing.assert_array_equal(output_points, expected_point)
def test_transform_points() -> None:
# 2D points (e.g. "world_from_agent")
tf = np.asarray(
[
[-1.08912448e00, 2.25029062e00, 6.03325482e03],
[-2.25029062e00, -1.08912448e00, -1.28582624e03],
[0.0, 0.0, 1.0],
]
)
points = np.array([[0, 10], [10, 0], [10, 10]])
expected_points = np.array([[6055.757726, -1296.717485], [6022.363575, -1308.329146], [6044.866481, -1319.220391]])
output_points = transform_points(points, tf)
np.testing.assert_allclose(output_points, expected_points)
# 3D points (e.g. "world_to_ecef")
tf = np.asarray(
[
[0.846617444, 0.323463078, -0.422623402, -2698767.44],
[-0.532201938, 0.514559352, -0.672301845, -4293151.58],
[-3.05311332e-16, 0.794103464, 0.6077826, 3855164.76],
[0.0, 0.0, 0.0, 1.0],
]
)
points = np.array([[0, 10, 0], [10, 0, 0], [0, 0, 10], [10, 10, 0], [10, 10, 10], [0, 10, 10]])
expected_points = np.array(
[
[-2698764.20536922, -4293146.43440648, 3855172.70103464],
[-2698758.97382556, -4293156.90201938, 3855164.76],
[-2698771.66623402, -4293158.30301845, 3855170.837826],
[-2698755.73919478, -4293151.75642586, 3855172.70103464],
[-2698759.9654288, -4293158.47944431, 3855178.77886064],
[-2698768.43160324, -4293153.15742493, 3855178.77886064],
]
)
output_points = transform_points(points, tf)
np.testing.assert_allclose(output_points, expected_points)
# original test
tf = np.asarray([[1.0, 0, 100], [0, 0.5, 50], [0, 0, 1]])
points = np.array([[0, 10], [10, 0], [10, 10]])
expected_points = np.array([[100, 55], [110, 50], [110, 55]])
output_points = transform_points(points, tf)
np.testing.assert_array_equal(output_points, expected_points)
def visualize_trajectory(dataset, index, title="target_positions movement with draw_trajectory"):
data = dataset[index]
im = data["image"].transpose(1, 2, 0)
im = dataset.rasterizer.to_rgb(im)
target_positions_pixels = transform_points(data["target_positions"] + data["centroid"][:2], data["world_to_image"])
draw_trajectory(im, target_positions_pixels, TARGET_POINTS_COLOR, radius=1, yaws=data["target_yaws"])
plt.title(title)
plt.imshow(im[::-1])
plt.show()
def validation(model, device):
model.eval()
torch.set_grad_enabled(False)
# store information for evaluation
future_coords_offsets_pd = []
timestamps = []
agent_ids = []
confidences_list = []
val_dataloader = load_val_data()
# num_iter = iter(val_dataloader)
# progress_bar = tqdm(range(cfg["train_params"]["val_num_steps"]))
progress_bar = tqdm(val_dataloader)
for data in progress_bar:
# data = next(num_iter)
preds, confs = forward(data, model, device)
# convert agent coordinates into world offsets
preds = preds.cpu().numpy()
confs = confs.cpu().numpy()
world_from_agents = data["world_from_agent"].numpy()
centroids = data["centroid"].numpy()
coords_offset = []
for pred, world_from_agent, centroid in zip(preds, world_from_agents,
centroids):
for mode in range(3):
pred[mode] = transform_points(pred[mode],
world_from_agent) - centroid[:2]
coords_offset.append(pred)
confidences_list.append(confs)
future_coords_offsets_pd.append(np.stack(coords_offset))
timestamps.append(data["timestamp"].numpy().copy())
agent_ids.append(data["track_id"].numpy().copy())
pred_path = f"{gettempdir()}/pred.csv"
write_pred_csv(pred_path,
timestamps=np.concatenate(timestamps),
track_ids=np.concatenate(agent_ids),
coords=np.concatenate(future_coords_offsets_pd),
confs=np.concatenate(confidences_list))
eval_base_path = '/home/axot/lyft/data/scenes/validate_chopped_31'
eval_gt_path = str(Path(eval_base_path) / "gt.csv")
metrics = compute_metrics_csv(eval_gt_path, pred_path,
[neg_multi_log_likelihood, time_displace])
for metric_name, metric_mean in metrics.items():
print(metric_name, metric_mean)
return metrics['neg_multi_log_likelihood']
def visualizeAV(self):
data = self.dataset[50]
im = data["image"].transpose(1, 2, 0)
im = self.dataset.rasterizer.to_rgb(im)
target_positions_pixels = transform_points(data["target_positions"] + data["centroid"][:2],
data["world_to_image"])
draw_trajectory(im, target_positions_pixels, data["target_yaws"], TARGET_POINTS_COLOR)
plt.imshow(im[::-1])
plt.show()
def plot_image(data_point, rasterizer):
im = data_point["image"].transpose(1, 2, 0)
im = rasterizer.to_rgb(im)
target_positions_pixels = transform_points(data_point["target_positions"],
data_point["raster_from_agent"])
draw_trajectory(im,
target_positions_pixels,
TARGET_POINTS_COLOR,
yaws=data_point["target_yaws"])
fig, ax = plt.subplots(figsize=(5, 5))
ax.imshow(im[::-1])
def __getitem__(self, index):
sample = super().__getitem__(index)
mask = np.zeros(sample["image"].shape[1:], dtype=np.uint8)
points = transform_points(sample["target_positions"],
sample["raster_from_agent"])
points = points[sample["target_availabilities"].astype(bool)]
draw_trajectory(mask, points, (255))
mask = torch.from_numpy((mask / 255.0).astype(np.float32)).unsqueeze(0)
sample["mask"] = mask
sample["square_category"] = torch.tensor(
to_flatten_square_idx(sample)).long()
return sample
def visualize_trajectory(dataset, index):
data = dataset[index]
im = data['image'].transpose(1, 2, 0)
im = dataset.rasterizer.to_rgb(im)
target_position_pixels = transform_points(
data['target_positions'] + data['centroid'][:2],
data['world_to_image'])
draw_trajectory(im, target_position_pixels, data['target_yaws'],
TARGET_POINTS_COLOR)
plt.imshow(im[::-1])
plt.show()
def test_transform_to_image_space_2d(self) -> None:
image_shape = (200, 200)
pixel_size = np.asarray((1.0, 0.5))
offset = np.asarray((0, -2))
input_points = np.array([[0, 0], [10, 10], [-10, -10]])
expected_output_points = np.array([[100, 104], [110, 124], [90, 84]])
tf = world_to_image_pixels_matrix(image_shape, pixel_size, offset)
output_points = transform_points(input_points, tf)
np.testing.assert_array_equal(output_points, expected_output_points)
def get_simple_element_info(self, center_in_world: np.ndarray,
raster_from_world: np.ndarray,
raster_radius: float, element_key: str):
element_indexes = elements_within_bounds(
center_in_world, self.bounds_info[element_key]["bounds"],
raster_radius)
for idx in element_indexes:
xyz = self.get_polyline_coords(
self.bounds_info[element_key]["ids"][idx])
xy_pos = transform_points(xyz[:, :2], raster_from_world)
yield {"centroid": xy_pos}
def plt_show_agent_map(self, idx):
zarr_dataset = self.chunked_dataset("scenes/train.zarr")
agent_dataset = AgentDataset(self.cfg, zarr_dataset, self.rast)
data = agent_dataset[idx]
im = data["image"].transpose(1, 2, 0)
im = self.rast.to_rgb(im)
target_positions_pixels = transform_points(
data["target_positions"] + data["centroid"][:2],
data["world_to_image"])
draw_trajectory(im, target_positions_pixels, TARGET_POINTS_COLOR, 1,
data["target_yaws"])
plt.imshow(im[::-1])
plt.savefig("filename.png")
def trajectory_stat(
history_positions: np.array,
target_positions: np.array,
centroid: np.array,
world_to_image: np.array,
) -> Any:
history_pixels = transform_points(history_positions + centroid,
world_to_image)
history_pixels -= history_pixels[0]
history_y_change = history_pixels[np.argmax(np.abs(history_pixels[:, 1])),
1]
history_x_change = history_pixels[np.argmax(np.abs(history_pixels[:, 0])),
0]
target_pixels = transform_points(target_positions + centroid,
world_to_image)
target_pixels -= target_pixels[0]
target_y_change = target_pixels[np.argmax(np.abs(target_pixels[:, 1])), 1]
target_x_change = target_pixels[np.argmax(np.abs(target_pixels[:, 0])), 0]
hist_diff = np.linalg.norm(np.diff(history_positions, axis=0), axis=1)
history_speed = hist_diff.sum() / history_positions.shape[0]
history_acceleration = (hist_diff[-1] - hist_diff[0]) / hist_diff.shape[0]
target_diff = np.linalg.norm(np.diff(target_positions, axis=0), axis=1)
target_speed = target_diff.sum() / target_positions.shape[0]
target_acceleration = (target_diff[-1] -
target_diff[0]) / target_diff.shape[0]
total_acceleration = (target_diff[-1] - hist_diff[0]) / (
target_diff.shape[0] + hist_diff.shape[0])
return ('history_y_change', history_y_change), ('history_x_change', history_x_change), \
('target_y_change', target_y_change), ('target_x_change', target_x_change), \
('history_speed', history_speed), ('history_acceleration', history_acceleration), \
('target_speed', target_speed), ('target_acceleration', target_acceleration), \
('total_acceleration', total_acceleration)
def draw_single_image(
rasterizer,
image: np.array,
centroid: np.array,
world_to_image: np.array,
target_positions: np.array,
target_yaws: np.array,
predicted_positions: Optional[np.array] = None,
predicted_yaws: Optional[np.array] = None,
target_color: Optional[tuple] = TARGET_POINTS_COLOR,
predicted_color: Optional[tuple] = PREDICTED_POINTS_COLOR,
) -> torch.Tensor:
"""
Produce a single RGB representation of the rasterized input image and its corresponding position prediction
:param rasterizer:
:param image:
:param centroid:
:param world_to_image:
:param target_positions:
:param target_yaws:
:param predicted_positions:
:param predicted_yaws:
:param target_color:
:param predicted_color:
:return:
"""
predicted_yaws = predicted_yaws if predicted_yaws is not None else target_yaws
im = _set_image_type(rasterizer.to_rgb(image.cpu().data.numpy().transpose(1, 2, 0))) # Todo enhance
draw_trajectory(im, transform_points(
target_positions.cpu().data.numpy() + centroid[:2].cpu().data.numpy(), world_to_image.cpu().data.numpy()),
target_yaws.cpu().data.numpy(), target_color)
if predicted_positions is not None:
draw_trajectory(im, transform_points(
predicted_positions.cpu().data.numpy() + centroid[:2].cpu().data.numpy(),
world_to_image.cpu().data.numpy()), predicted_yaws.cpu().data.numpy(), predicted_color)
return np.uint8(im).transpose(2, 0, 1)
