animate.py

from csv import reader
from dataclasses import dataclass
from math import radians

from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from scipath import Profile, create_cubic_path_2d

from kbm import KinematicBicycleModel
from libs import CarDescription, StanleyController


class Simulation:
    def __init__(self):
        fps = 50.0

        self.dt = 1 / fps
        self.map_size_x = 70
        self.map_size_y = 40
        self.frames = 2500
        self.loop = False


class Path:
    def __init__(self):
        # Get path to waypoints.csv
        with open("data/waypoints.csv", newline="") as f:
            rows = list(reader(f, delimiter=","))

        points = [(float(row[0]), float(row[1])) for row in rows[1:]]
        path, self.pyaw, _ = create_cubic_path_2d(points, profile=Profile.NO_CURVATURE)
        self.px, self.py = path.T


class Car:
    def __init__(self, init_x, init_y, init_yaw, px, py, pyaw, delta_time):
        # Model parameters
        self.x = init_x
        self.y = init_y
        self.yaw = init_yaw
        self.delta_time = delta_time
        self.time = 0.0
        self.velocity = 0.0
        self.wheel_angle = 0.0
        self.angular_velocity = 0.0
        max_steer = radians(33)
        wheelbase = 2.96

        # Acceleration parameters
        target_velocity = 10.0
        self.time_to_reach_target_velocity = 5.0
        self.required_acceleration = target_velocity / self.time_to_reach_target_velocity

        # Tracker parameters
        self.px = px
        self.py = py
        self.pyaw = pyaw
        self.k = 8.0
        self.ksoft = 1.0
        self.kyaw = 0.01
        self.ksteer = 0.0
        self.crosstrack_error = None
        self.target_id = None

        # Description parameters
        self.colour = "black"
        overall_length = 4.97
        overall_width = 1.964
        tyre_diameter = 0.4826
        tyre_width = 0.265
        axle_track = 1.7
        rear_overhang = 0.5 * (overall_length - wheelbase)

        self.bicycle_model = KinematicBicycleModel(
            wheelbase=wheelbase,
            max_steer=max_steer,
            delta_time=self.delta_time,
        )

        self.tracker = StanleyController(
            self.k,
            self.ksoft,
            self.kyaw,
            self.ksteer,
            max_steer,
            wheelbase,
            self.px,
            self.py,
            self.pyaw,
        )

        self.description = CarDescription(
            overall_length,
            overall_width,
            rear_overhang,
            tyre_diameter,
            tyre_width,
            axle_track,
            wheelbase,
        )

    def get_required_acceleration(self) -> float:
        self.time += self.delta_time
        return self.required_acceleration

    def plot_car(self):
        return self.description.plot_car(self.x, self.y, self.yaw, self.wheel_angle)

    def drive(self):
        acceleration = 0 if self.time > self.time_to_reach_target_velocity else self.get_required_acceleration()

        self.wheel_angle, self.target_id, self.crosstrack_error = self.tracker.stanley_control(
            self.x,
            self.y,
            self.yaw,
            self.velocity,
            self.wheel_angle,
        )

        vehicle_state = self.bicycle_model.compute_state(
            x=self.x,
            y=self.y,
            yaw=self.yaw,
            steer=self.wheel_angle,
            velocity=self.velocity,
            acceleration=acceleration,
        )

        self.x = vehicle_state["x"]
        self.y = vehicle_state["y"]
        self.yaw = vehicle_state["yaw"]
        self.velocity = vehicle_state["velocity"]

        print(f"Cross-track term: {self.crosstrack_error}{' '*10}", end="\r")


@dataclass
class Fargs:
    ax: plt.Axes
    sim: Simulation
    path: Path
    car: Car
    car_outline: plt.Line2D
    front_right_wheel: plt.Line2D
    front_left_wheel: plt.Line2D
    rear_right_wheel: plt.Line2D
    rear_left_wheel: plt.Line2D
    rear_axle: plt.Line2D
    annotation: plt.Annotation
    target: plt.Line2D


def animate(frame, fargs):
    ax = fargs.ax
    sim = fargs.sim
    path = fargs.path
    car = fargs.car
    car_outline = fargs.car_outline
    front_right_wheel = fargs.front_right_wheel
    front_left_wheel = fargs.front_left_wheel
    rear_right_wheel = fargs.rear_right_wheel
    rear_left_wheel = fargs.rear_left_wheel
    rear_axle = fargs.rear_axle
    annotation = fargs.annotation
    target = fargs.target

    # Camera tracks car
    ax.set_xlim(car.x - sim.map_size_x, car.x + sim.map_size_x)
    ax.set_ylim(car.y - sim.map_size_y, car.y + sim.map_size_y)

    # Drive and draw car
    car.drive()
    outline_plot, fr_plot, rr_plot, fl_plot, rl_plot = car.plot_car()
    car_outline.set_data(*outline_plot)
    front_right_wheel.set_data(*fr_plot)
    rear_right_wheel.set_data(*rr_plot)
    front_left_wheel.set_data(*fl_plot)
    rear_left_wheel.set_data(*rl_plot)
    rear_axle.set_data([car.x], [car.y])

    # Show car's target
    target.set_data([path.px[car.target_id]], [path.py[car.target_id]])

    # Annotate car's coordinate above car
    annotation.set_text(f"{car.x:.1f}, {car.y:.1f}")
    annotation.set_position((car.x, car.y + 5))

    plt.title(f"{sim.dt*frame:.2f}s", loc="right")
    plt.xlabel(f"Speed: {car.velocity:.2f} m/s", loc="left")
    # plt.savefig(f'image/visualisation_{frame:03}.png', dpi=300)

    return (
        car_outline,
        front_right_wheel,
        rear_right_wheel,
        front_left_wheel,
        rear_left_wheel,
        rear_axle,
        target,
    )


def main():
    sim = Simulation()
    path = Path()
    car = Car(path.px[0], path.py[0], path.pyaw[0], path.px, path.py, path.pyaw, sim.dt)

    interval = sim.dt * 10**3

    fig = plt.figure()
    ax = plt.axes()
    ax.set_aspect("equal")

    road = plt.Circle((0, 0), 50, color="gray", fill=False, linewidth=30)
    ax.add_patch(road)
    ax.plot(path.px, path.py, "--", color="gold")

    empty = ([], [])
    (target,) = ax.plot(*empty, "+r")
    (car_outline,) = ax.plot(*empty, color=car.colour)
    (front_right_wheel,) = ax.plot(*empty, color=car.colour)
    (rear_right_wheel,) = ax.plot(*empty, color=car.colour)
    (front_left_wheel,) = ax.plot(*empty, color=car.colour)
    (rear_left_wheel,) = ax.plot(*empty, color=car.colour)
    (rear_axle,) = ax.plot(car.x, car.y, "+", color=car.colour, markersize=2)
    annotation = ax.annotate(
        f"{car.x:.1f}, {car.y:.1f}",
        xy=(car.x, car.y + 5),
        color="black",
        annotation_clip=False,
    )

    fargs = [
        Fargs(
            ax=ax,
            sim=sim,
            path=path,
            car=car,
            car_outline=car_outline,
            front_right_wheel=front_right_wheel,
            front_left_wheel=front_left_wheel,
            rear_right_wheel=rear_right_wheel,
            rear_left_wheel=rear_left_wheel,
            rear_axle=rear_axle,
            annotation=annotation,
            target=target,
        )
    ]

    _ = FuncAnimation(
        fig,
        animate,
        frames=sim.frames,
        init_func=lambda: None,
        fargs=fargs,
        interval=interval,
        repeat=sim.loop,
    )
    # anim.save('animation.gif', writer='imagemagick', fps=50)

    plt.grid()
    plt.show()


if __name__ == "__main__":
    main()