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)