Long Example Usage

Long Example

import os.path
from pathlib import Path

from commonroad.common.file_reader import CommonRoadFileReader
from commonroad.planning.planning_problem import PlanningProblem
from commonroad_route_planner.reference_path import ReferencePath
from commonroad_route_planner.fast_api.fast_api import generate_reference_path_from_lanelet_network_and_planning_problem

from commonroad_idm_planner.configuration.planner_config import IDMConfig, IDMConfigFactory

# typing
from typing import Union

from commonroad_idm_planner.idm_planner import IDMPlanner
from commonroad_idm_planner.idm_trajectory import IDMTrajectory
from commonroad_idm_planner.util.collision.collision_status import CollisionStatus
from commonroad_idm_planner.util.pp_util import PlanningProblemFactory
from commonroad_idm_planner.util.visualization.visualize_scenario import visualize_idm_trajectory, make_gif
from commonroad_idm_planner.util.visualization.visualize_physics import visualize_physics
from commonroad_idm_planner.util.visualization.visualize_inputs import visualize_inputs
from commonroad_idm_planner.idm_path import IDMPath, IDMPathFactory


def main(
        scenario_path: Union[Path, str],
        show_plots: bool = True,
        save_img: bool = False,
        save_dir: Union[str, Path] = None,
        save_gif: bool = False
) -> None:
    pp_factory: PlanningProblemFactory = PlanningProblemFactory()

    scenario, planning_problem_set = CommonRoadFileReader(scenario_path).open()
    planning_problem: PlanningProblem = list(planning_problem_set.planning_problem_dict.values())[0]

    config: IDMConfig = IDMConfigFactory().generate_default_config()

    reference_path: ReferencePath = generate_reference_path_from_lanelet_network_and_planning_problem(
        lanelet_network=scenario.lanelet_network,
        planning_problem=planning_problem
    )

    idm_path: IDMPath = IDMPathFactory().generate_idm_path_from_cr_route_planner_reference_path(
        cr_reference_path=reference_path
    )

    idm_planner = IDMPlanner(
        scenario=scenario,
        planning_problem=planning_problem,
        config=config,
        idm_path=idm_path
    )

    trajectory_one_cycle: IDMTrajectory = idm_planner.plan()

    collision_status: CollisionStatus = idm_planner.check_collision()
    if collision_status.collision_detected:
        idm_planner.logger.info(f"{collision_status}")


    if not idm_planner.is_goal_reached:
        for replan_cycle_idx in range(6):
            new_pp: PlanningProblem = pp_factory.planning_problem_from_idm_state(
                idm_state=trajectory_one_cycle.final_state,
                goal_region=planning_problem.goal,
                planning_problem_id=planning_problem.planning_problem_id + 20000 + replan_cycle_idx
            )

            new_rp: ReferencePath = generate_reference_path_from_lanelet_network_and_planning_problem(
                lanelet_network=scenario.lanelet_network,
                planning_problem=new_pp
            )

            new_idm_path: IDMPath = IDMPathFactory().generate_idm_path_from_cr_route_planner_reference_path(
                cr_reference_path=new_rp
            )

            trajectory_one_cycle: IDMTrajectory = idm_planner.re_plan(
                planning_problem=new_pp,
                idm_path=new_idm_path
            )

            collision_status: CollisionStatus = idm_planner.check_collision()

            if collision_status.collision_detected:
                idm_planner.logger.info(f"{collision_status}")

            if idm_planner.is_goal_reached:
                idm_planner.logger.info(
                    f"Goal is reached:  \n"
                    f"  planner goal state: {idm_planner.goal_state} \n"
                )
                # prune scenario trajectory so last state is goal state
                idm_planner.prune_scenario_trajectory()
                break

    # calculate reconstructed input trajectory from scenario trajectory
    idm_planner.calc_steering_angle_for_scenario_traj(
        dt=scenario.dt,
        wheelbase_rear_to_cog=2.5
    )
    idm_planner.calc_input_trajectory_from_scenario_trajectory()

    if show_plots:
        print(f"generating images")

        if save_img:
            save_path = os.path.join(save_dir, str(scenario.scenario_id))
            save_path_physics = os.path.join(save_dir, str(scenario.scenario_id), "physics")
            save_path_input = os.path.join(save_dir, str(scenario.scenario_id), "rec_input")
        else:
            save_path = None
            save_path_physics = None
            save_path_input = None

        visualize_idm_trajectory(
            scenario=scenario,
            planning_problem=planning_problem,
            idm_path=idm_path,
            idm_trajectory=idm_planner.trajectory_over_scenario,
            idm_config=config,
            save_img=save_img,
            save_path=save_path
        )


        if save_img and save_gif:
            make_gif(
                path_to_img_dir=save_path,
                scenario_name=str(scenario.scenario_id),
                num_imgs=idm_planner.trajectory_over_scenario.num_time_steps
            )

        visualize_physics(
            idm_trajectory=idm_planner.trajectory_over_scenario,
            idm_config=config,
            save_img=save_img,
            delta_t=scenario.dt,
            save_path=save_path_physics
        )

        if idm_planner.input_trajectory_over_scenario is not None:
            visualize_inputs(
                input_list=idm_planner.input_trajectory_over_scenario,
                save_img=save_img,
                save_path=save_path_input
            )

if __name__ == "__main__":

    dir_path = "PATH/TO/SCENARIO/FOLDER"
    filenames = sorted(os.listdir(dir_path))

    for file_name in filenames:
        print(f'Planning scenario {file_name}')

        scenario_path = os.path.join(dir_path, file_name)
        save_dir = "PATH/TO/SAVE/IMGS
        save_gif: bool = True
        main(
            scenario_path=scenario_path,
            save_img=True,
            save_dir=save_dir,
            save_gif=save_gif
        )