Spaces:
Running
Running
File size: 2,784 Bytes
5769ee4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
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()
|