297 lines
9.2 KiB
Python
297 lines
9.2 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.animation as ani
|
|
import matplotlib.font_manager as fon
|
|
import sys
|
|
import math
|
|
|
|
# default setting of figures
|
|
plt.rcParams["mathtext.fontset"] = 'stix' # math fonts
|
|
plt.rcParams['xtick.direction'] = 'in' # x axis in
|
|
plt.rcParams['ytick.direction'] = 'in' # y axis in
|
|
plt.rcParams["font.size"] = 10
|
|
plt.rcParams['axes.linewidth'] = 1.0 # axis line width
|
|
plt.rcParams['axes.grid'] = True # make grid
|
|
|
|
def coordinate_transformation_in_angle(positions, base_angle):
|
|
'''
|
|
Transformation the coordinate in the angle
|
|
|
|
Parameters
|
|
-------
|
|
positions : numpy.ndarray
|
|
this parameter is composed of xs, ys
|
|
should have (2, N) shape
|
|
base_angle : float [rad]
|
|
|
|
Returns
|
|
-------
|
|
traslated_positions : numpy.ndarray
|
|
the shape is (2, N)
|
|
|
|
'''
|
|
if positions.shape[0] != 2:
|
|
raise ValueError('the input data should have (2, N)')
|
|
|
|
positions = np.array(positions)
|
|
positions = positions.reshape(2, -1)
|
|
|
|
rot_matrix = [[np.cos(base_angle), np.sin(base_angle)],
|
|
[-1*np.sin(base_angle), np.cos(base_angle)]]
|
|
|
|
rot_matrix = np.array(rot_matrix)
|
|
|
|
translated_positions = np.dot(rot_matrix, positions)
|
|
|
|
return translated_positions
|
|
|
|
def square_make_with_angles(center_x, center_y, size, angle):
|
|
'''
|
|
Create square matrix with angle line matrix(2D)
|
|
|
|
Parameters
|
|
-------
|
|
center_x : float in meters
|
|
the center x position of the square
|
|
center_y : float in meters
|
|
the center y position of the square
|
|
size : float in meters
|
|
the square's half-size
|
|
angle : float in radians
|
|
|
|
Returns
|
|
-------
|
|
square xs : numpy.ndarray
|
|
lenght is 5 (counterclockwise from right-up)
|
|
square ys : numpy.ndarray
|
|
length is 5 (counterclockwise from right-up)
|
|
angle line xs : numpy.ndarray
|
|
angle line ys : numpy.ndarray
|
|
'''
|
|
|
|
# start with the up right points
|
|
# create point in counterclockwise
|
|
square_xys = np.array([[size, 0.5 * size], [-size, 0.5 * size], [-size, -0.5 * size], [size, -0.5 * size], [size, 0.5 * size]])
|
|
trans_points = coordinate_transformation_in_angle(square_xys.T, -angle) # this is inverse type
|
|
trans_points += np.array([[center_x], [center_y]])
|
|
|
|
square_xs = trans_points[0, :]
|
|
square_ys = trans_points[1, :]
|
|
|
|
angle_line_xs = [center_x, center_x + math.cos(angle) * size]
|
|
angle_line_ys = [center_y, center_y + math.sin(angle) * size]
|
|
|
|
return square_xs, square_ys, np.array(angle_line_xs), np.array(angle_line_ys)
|
|
|
|
|
|
def circle_make_with_angles(center_x, center_y, radius, angle):
|
|
'''
|
|
Create circle matrix with angle line matrix
|
|
|
|
Parameters
|
|
-------
|
|
center_x : float
|
|
the center x position of the circle
|
|
center_y : float
|
|
the center y position of the circle
|
|
radius : float
|
|
angle : float [rad]
|
|
|
|
Returns
|
|
-------
|
|
circle xs : numpy.ndarray
|
|
circle ys : numpy.ndarray
|
|
angle line xs : numpy.ndarray
|
|
angle line ys : numpy.ndarray
|
|
'''
|
|
|
|
point_num = 100 # 分解能
|
|
|
|
circle_xs = []
|
|
circle_ys = []
|
|
|
|
for i in range(point_num + 1):
|
|
circle_xs.append(center_x + radius * math.cos(i*2*math.pi/point_num))
|
|
circle_ys.append(center_y + radius * math.sin(i*2*math.pi/point_num))
|
|
|
|
angle_line_xs = [center_x, center_x + math.cos(angle) * radius]
|
|
angle_line_ys = [center_y, center_y + math.sin(angle) * radius]
|
|
|
|
return np.array(circle_xs), np.array(circle_ys), np.array(angle_line_xs), np.array(angle_line_ys)
|
|
|
|
|
|
class AnimDrawer():
|
|
"""create animation of path and robot
|
|
|
|
Attributes
|
|
------------
|
|
cars :
|
|
anim_fig : figure of matplotlib
|
|
axis : axis of matplotlib
|
|
|
|
"""
|
|
def __init__(self, objects):
|
|
"""
|
|
Parameters
|
|
------------
|
|
objects : list of objects
|
|
|
|
Notes
|
|
---------
|
|
lead_history_states, lead_history_predict_states, traj_ref, history_traj_ref, history_angle_ref
|
|
"""
|
|
self.car_history_state = objects[0]
|
|
self.target = objects[1]
|
|
|
|
self.history_target_x = [self.target[:, 0]]
|
|
self.history_target_y = [self.target[:, 1]]
|
|
|
|
self.history_xs = [self.car_history_state[:, 0]]
|
|
self.history_ys = [self.car_history_state[:, 1]]
|
|
self.history_ths = [self.car_history_state[:, 2]]
|
|
|
|
# setting up figure
|
|
self.anim_fig = plt.figure(dpi=150)
|
|
self.axis = self.anim_fig.add_subplot(111)
|
|
|
|
# imgs
|
|
self.car_imgs = []
|
|
self.traj_imgs = []
|
|
|
|
def draw_anim(self, interval=10):
|
|
"""draw the animation and save
|
|
|
|
Parameteres
|
|
-------------
|
|
interval : int, optional
|
|
animation's interval time, you should link the sampling time of systems
|
|
default is 50 [ms]
|
|
"""
|
|
self._set_axis()
|
|
self._set_img()
|
|
|
|
self.skip_num = 1
|
|
frame_num = int((len(self.history_xs[0])-1) / self.skip_num)
|
|
|
|
animation = ani.FuncAnimation(self.anim_fig, self._update_anim, interval=interval, frames=frame_num)
|
|
|
|
# self.axis.legend()
|
|
print('save_animation?')
|
|
shuold_save_animation = int(input())
|
|
|
|
if shuold_save_animation:
|
|
print('animation_number?')
|
|
num = int(input())
|
|
animation.save('animation_{0}.mp4'.format(num), writer='ffmpeg')
|
|
# animation.save("Sample.gif", writer = 'imagemagick') # gif保存
|
|
|
|
plt.show()
|
|
|
|
def _set_axis(self):
|
|
""" initialize the animation axies
|
|
"""
|
|
# (1) set the axis name
|
|
self.axis.set_xlabel(r'$\it{x}$ [m]')
|
|
self.axis.set_ylabel(r'$\it{y}$ [m]')
|
|
self.axis.set_aspect('equal', adjustable='box')
|
|
|
|
LOW_MARGIN = 2
|
|
HIGH_MARGIN = 2
|
|
|
|
self.axis.set_xlim(np.min(self.history_xs) - LOW_MARGIN, np.max(self.history_xs) + HIGH_MARGIN)
|
|
self.axis.set_ylim(np.min(self.history_ys) - LOW_MARGIN, np.max(self.history_ys) + HIGH_MARGIN)
|
|
|
|
def _set_img(self):
|
|
""" initialize the imgs of animation
|
|
this private function execute the make initial imgs for animation
|
|
"""
|
|
# object imgs
|
|
obj_color_list = ["k", "k", "m", "m"]
|
|
obj_styles = ["solid", "solid", "solid", "solid"]
|
|
|
|
for i in range(len(obj_color_list)):
|
|
temp_img, = self.axis.plot([], [], color=obj_color_list[i], linestyle=obj_styles[i])
|
|
self.car_imgs.append(temp_img)
|
|
|
|
traj_color_list = ["k", "b"]
|
|
|
|
for i in range(len(traj_color_list)):
|
|
temp_img, = self.axis.plot([],[], color=traj_color_list[i], linestyle="dashed")
|
|
self.traj_imgs.append(temp_img)
|
|
|
|
temp_img, = self.axis.plot([],[], "*", color="b")
|
|
self.traj_imgs.append(temp_img)
|
|
|
|
def _update_anim(self, i):
|
|
"""the update animation
|
|
this function should be used in the animation functions
|
|
|
|
Parameters
|
|
------------
|
|
i : int
|
|
time step of the animation
|
|
the sampling time should be related to the sampling time of system
|
|
|
|
Returns
|
|
-----------
|
|
object_imgs : list of img
|
|
traj_imgs : list of img
|
|
"""
|
|
i = int(i * self.skip_num)
|
|
|
|
# self._draw_set_axis(i)
|
|
self._draw_car(i)
|
|
self._draw_traj(i)
|
|
# self._draw_prediction(i)
|
|
|
|
return self.car_imgs, self.traj_imgs,
|
|
|
|
def _draw_set_axis(self, i):
|
|
"""
|
|
"""
|
|
# (2) set the xlim and ylim
|
|
LOW_MARGIN = 20
|
|
HIGH_MARGIN = 20
|
|
OVER_LOOK = 50
|
|
self.axis.set_xlim(np.min(self.history_xs[0][i : i + OVER_LOOK]) - LOW_MARGIN, np.max(self.history_xs[0][i : i + OVER_LOOK]) + HIGH_MARGIN)
|
|
self.axis.set_ylim(np.min(self.history_ys[0][i : i + OVER_LOOK]) - LOW_MARGIN, np.max(self.history_ys[0][i : i + OVER_LOOK]) + HIGH_MARGIN)
|
|
|
|
def _draw_car(self, i):
|
|
"""
|
|
This private function is just divided thing of
|
|
the _update_anim to see the code more clear
|
|
|
|
Parameters
|
|
------------
|
|
i : int
|
|
time step of the animation
|
|
the sampling time should be related to the sampling time of system
|
|
"""
|
|
# cars
|
|
object_x, object_y, angle_x, angle_y = square_make_with_angles(self.history_xs[0][i],
|
|
self.history_ys[0][i],
|
|
1.0,
|
|
self.history_ths[0][i])
|
|
|
|
self.car_imgs[0].set_data([object_x, object_y])
|
|
self.car_imgs[1].set_data([angle_x, angle_y])
|
|
|
|
def _draw_traj(self, i):
|
|
"""
|
|
This private function is just divided thing of
|
|
the _update_anim to see the code more clear
|
|
|
|
Parameters
|
|
------------
|
|
i : int
|
|
time step of the animation
|
|
the sampling time should be related to the sampling time of system
|
|
"""
|
|
# car
|
|
self.traj_imgs[0].set_data(self.history_xs[0][:i], self.history_ys[0][:i])
|
|
|
|
# goal
|
|
self.traj_imgs[-1].set_data(self.history_target_x[0][i-1], self.history_target_y[0][i-1])
|
|
|
|
# traj_ref
|
|
self.traj_imgs[1].set_data(self.history_target_x[0][:i], self.history_target_y[0][:i]) |