File size: 6,953 Bytes
6bdfadb |
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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 |
import matplotlib as mpl
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection
import fastf1 as ff1
from PIL import Image
from io import BytesIO
from typing import Union
import json
# Custom types
gp = Union[str, int]
session_type = Union[str, int, None]
def rotate(xy, *, angle):
rot_mat = np.array([[np.cos(angle), np.sin(angle)],
[-np.sin(angle), np.cos(angle)]])
return np.matmul(xy, rot_mat)
def create_track_speed_visualization(session, driver_name: str) -> Image:
weekend = session.event
session.load()
with open("assets/driver_abbreviations.json") as f:
driver_abbreviations = json.load(f)
driver_abbreviation = driver_abbreviations[driver_name]
lap = session.laps.pick_drivers(driver_abbreviation).pick_fastest()
# Get telemetry data
x = lap.telemetry['X'] # values for x-axis
y = lap.telemetry['Y'] # values for y-axis
color = lap.telemetry['Speed'] # value to base color gradient on
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# We create a plot with title and adjust some setting to make it look good.
fig, ax = plt.subplots(sharex=True, sharey=True, figsize=(12, 6.75))
fig.suptitle(f'{weekend["EventName"]} - {driver_name} ', size=24, y=0.97)
# Adjust margins and turn of axis
plt.subplots_adjust(left=0.1, right=0.9, top=0.9, bottom=0.12)
ax.axis('off')
# After this, we plot the data itself.
# Create background track line
ax.plot(lap.telemetry['X'], lap.telemetry['Y'],
color='black', linestyle='-', linewidth=16, zorder=0)
# Create a continuous norm to map from data points to colors
norm = plt.Normalize(color.min(), color.max())
lc = LineCollection(segments, cmap=mpl.colormaps['viridis'], norm=norm,
linestyle='-', linewidth=5)
# Set the values used for colormapping
lc.set_array(color)
# Merge all line segments together
line = ax.add_collection(lc)
# Finally, we create a color bar as a legend.
cbaxes = fig.add_axes([0.25, 0.05, 0.5, 0.05])
normlegend = mpl.colors.Normalize(vmin=color.min(), vmax=color.max())
legend = mpl.colorbar.ColorbarBase(cbaxes, norm=normlegend, cmap=mpl.colormaps['viridis'],
orientation="horizontal")
# Create a PIL image from the plot
fig = plt.gcf()
# Save the figure to a BytesIO buffer and convert to bytes
buf = BytesIO()
fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
buf.seek(0)
# Create PIL image from buffer bytes and close the figure
img_data = buf.getvalue()
plt.close(fig)
buf.close()
# Create new image from the raw bytes
img = Image.open(BytesIO(img_data))
return img
def create_track_corners_visualization(session) -> Image:
lap = session.laps.pick_fastest()
pos = lap.get_pos_data()
circuit_info = session.get_circuit_info()
# Get an array of shape [n, 2] where n is the number of points and the second
# axis is x and y.
track = pos.loc[:, ('X', 'Y')].to_numpy()
# Convert the rotation angle from degrees to radian.
track_angle = circuit_info.rotation / 180 * np.pi
# Rotate and plot the track map.
rotated_track = rotate(track, angle=track_angle)
plt.plot(rotated_track[:, 0], rotated_track[:, 1])
offset_vector = [500, 0] # offset length is chosen arbitrarily to 'look good'
# Iterate over all corners.
for _, corner in circuit_info.corners.iterrows():
# Create a string from corner number and letter
txt = f"{corner['Number']}{corner['Letter']}"
# Convert the angle from degrees to radian.
offset_angle = corner['Angle'] / 180 * np.pi
# Rotate the offset vector so that it points sideways from the track.
offset_x, offset_y = rotate(offset_vector, angle=offset_angle)
# Add the offset to the position of the corner
text_x = corner['X'] + offset_x
text_y = corner['Y'] + offset_y
# Rotate the text position equivalently to the rest of the track map
text_x, text_y = rotate([text_x, text_y], angle=track_angle)
# Rotate the center of the corner equivalently to the rest of the track map
track_x, track_y = rotate([corner['X'], corner['Y']], angle=track_angle)
# Draw a circle next to the track.
plt.scatter(text_x, text_y, color='grey', s=140)
# Draw a line from the track to this circle.
plt.plot([track_x, text_x], [track_y, text_y], color='grey')
# Finally, print the corner number inside the circle.
plt.text(text_x, text_y, txt,
va='center_baseline', ha='center', size='small', color='white')
plt.title(session.event['Location'])
plt.xticks([])
plt.yticks([])
plt.axis('equal')
# Create a PIL image from the plot
fig = plt.gcf()
# Save the figure to a BytesIO buffer and convert to bytes
buf = BytesIO()
fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
buf.seek(0)
# Create PIL image from buffer bytes and close the figure
img_data = buf.getvalue()
plt.close(fig)
buf.close()
# Create new image from the raw bytes
img = Image.open(BytesIO(img_data))
return img
def create_track_gear_visualization(session) -> Image:
lap = session.laps.pick_fastest()
tel = lap.get_telemetry()
x = np.array(tel['X'].values)
y = np.array(tel['Y'].values)
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
gear = tel['nGear'].to_numpy().astype(float)
cmap = plt.cm.get_cmap('viridis', 8)
norm = plt.Normalize(1, 8)
lc_comp = LineCollection(segments, norm=norm, cmap=cmap)
lc_comp.set_array(gear)
lc_comp.set_linewidth(4)
plt.gca().add_collection(lc_comp)
plt.axis('equal')
plt.tick_params(labelleft=False, left=False, labelbottom=False, bottom=False)
plt.suptitle(
f"Fastest Lap Gear Shift Visualization\n"
f"{lap['Driver']} - {session.event['EventName']}"
)
cbar = plt.colorbar(mappable=lc_comp, label="Gear")
cbar.set_ticks(np.arange(1, 9)) # Set ticks at integer positions
cbar.set_ticklabels(np.arange(1, 9)) # Labels from 1 to 8
# Create a PIL image from the plot
fig = plt.gcf()
# Save the figure to a BytesIO buffer and convert to bytes
buf = BytesIO()
fig.savefig(buf, format='png', dpi=150, bbox_inches='tight')
buf.seek(0)
# Create PIL image from buffer bytes and close the figure
img_data = buf.getvalue()
plt.close(fig)
buf.close()
# Create new image from the raw bytes
img = Image.open(BytesIO(img_data))
return img
|