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