Spaces:
Running
Running
| import os | |
| import matplotlib.pyplot as plt | |
| from mmcv import Config | |
| import numpy as np | |
| from pytorch_lightning.utilities.seed import seed_everything | |
| from risk_biased.scene_dataset.scene import RandomScene, RandomSceneParams | |
| from risk_biased.scene_dataset.scene_plotter import ScenePlotter | |
| from risk_biased.utils.cost import ( | |
| DistanceCostNumpy, | |
| DistanceCostParams, | |
| TTCCostNumpy, | |
| TTCCostParams, | |
| ) | |
| if __name__ == "__main__": | |
| working_dir = os.path.dirname(os.path.realpath(os.path.join(__file__, ".."))) | |
| config_path = os.path.join( | |
| working_dir, "..", "risk_biased", "config", "learning_config.py" | |
| ) | |
| config = Config.fromfile(config_path) | |
| if config.seed is not None: | |
| seed_everything(config.seed) | |
| ped_speed = 2 | |
| is_torch = False | |
| fig, ax = plt.subplots( | |
| 3, 4, sharex=True, sharey=True, tight_layout=True, subplot_kw={"aspect": 1} | |
| ) | |
| scene_params = RandomSceneParams.from_config(config) | |
| scene_params.batch_size = 1 | |
| for ii in range(9): | |
| test_scene = RandomScene( | |
| scene_params, | |
| is_torch=is_torch, | |
| ) | |
| dist_cost = DistanceCostNumpy(DistanceCostParams.from_config(config)) | |
| ttc_cost = TTCCostNumpy(TTCCostParams.from_config(config)) | |
| nx = 1000 | |
| ny = 100 | |
| x, y = np.meshgrid( | |
| np.linspace(-test_scene.ego_length, test_scene.road_length, nx), | |
| np.linspace(test_scene.bottom, test_scene.top, ny), | |
| ) | |
| i = 2 - (int(ii >= 6) + int(ii >= 3)) | |
| j = ii % 3 | |
| vx = float(ii % 3 - 1) | |
| vy = float((ii >= 6)) - float(ii <= 2) | |
| print(f"horizontal velocity {vx}") | |
| print(f"vertical velocity {vy}") | |
| norm = np.maximum(np.sqrt(vx * vx + vy * vy), 1) | |
| vx = vx / norm * np.ones([nx * ny, 1]) | |
| vy = vy / norm * np.ones([nx * ny, 1]) | |
| v_ped = ped_speed * np.stack((vx, vy), -1) | |
| v_ego = np.array([[[test_scene.ego_ref_speed, 0]]]) | |
| p_init = np.stack((x, y), -1).reshape((nx * ny, 2)) | |
| p_final = p_init + v_ped[:, 0, :] * test_scene.time_scene | |
| len_traj = 30 | |
| ped_trajs = np.linspace(p_init, p_final, len_traj, axis=1) | |
| ego_traj = np.linspace( | |
| [[0, 0]], | |
| [test_scene.ego_ref_speed * test_scene.time_scene, 0], | |
| len_traj, | |
| axis=1, | |
| ) | |
| cost, _ = ttc_cost(ego_traj, ped_trajs, v_ego, v_ped) | |
| cost = cost.reshape(ny, nx) | |
| colorbar = ax[i][j].pcolormesh(x, y, cost, cmap="RdBu") | |
| plotter = ScenePlotter(test_scene, ax=ax[i][j]) | |
| plotter.plot_road() | |
| fig.subplots_adjust(wspace=0.1, hspace=0.1) | |
| fig.tight_layout() | |
| fig.colorbar(colorbar, ax=ax.ravel().tolist()) | |
| for a in ax[:, -1]: | |
| a.remove() | |
| plt.show() | |