track visualization for speed, corner and gears

app.py

@@ -13,53 +13,26 @@ from utils.constants import (
 # Variables
 CURRENT_YEAR = datetime.datetime.now().year
 
-def driver_championship_score(driver_name: str) -> str:
-    """
-    Get the championship score for the given driver.
-
-    Args:
-        driver_name (str): The driver's name
-
-    Returns:
-        int: The driver's championship score
-    """
-    return f"Driver {driver_name} has {random.randint(0, 100)} championship points"
-
-
-def driver_position(driver_name: str) -> str:
-    """
-    Get the current position of the given driver.
-
-    Args:
-        driver_name (str): The driver's name
-
-    Returns:
-        str: The driver's current position
-    """
-    return f"Driver {driver_name} is in position {random.randint(1, 20)}"
-
 
 # Load in driver names
 with open("assets/driver_names.json") as f:
     driver_names = json.load(f)["drivers"]
+with open("assets/constructor_details.json") as f:
+    constructor_names = [constructor["constructor"] for constructor in json.load(f)["constructors"]]
 
-iface1 = gr.Interface(
-    fn=driver_championship_score,
-    inputs=gr.Dropdown(driver_names),
-    outputs="text",
-    title="[WIP] Driver Championship Score",
-    description="Get the championship score for the given driver"
-)
 
-iface2 = gr.Interface(
-    fn=driver_position,
-    inputs=gr.Dropdown(driver_names),
+iface_driver_championship_standings = gr.Interface(
+    fn=tools.driver_championship_score,
+    inputs=[
+        gr.Number(label="Calendar year", value=CURRENT_YEAR, minimum=1950, maximum=CURRENT_YEAR),
+        gr.Dropdown(driver_names)
+    ],
     outputs="text",
-    title="[WIP] Driver Position",
-    description="Get the current position of the given driver"
+    title="Driver Championship Standings",
+    description="Get the championship standings for the given driver"
 )
 
-iface3 = gr.Interface(
+iface_event_info = gr.Interface(
     fn=tools.get_event_info,
     inputs=[
         gr.Number(label="Calendar year", value=CURRENT_YEAR, minimum=1950, maximum=CURRENT_YEAR),
@@ -71,7 +44,7 @@ iface3 = gr.Interface(
     description="Get the Grand Prix event info for a specific race week"
 )
 
-iface4 = gr.Interface(
+iface_season_calendar = gr.Interface(
     fn=tools.get_season_calendar,
     inputs=[
         gr.Number(label="Calendar year", value=CURRENT_YEAR, minimum=1950, maximum=CURRENT_YEAR),
@@ -81,9 +54,30 @@ iface4 = gr.Interface(
     description="Get the season calendar information for the given year"
 )
 
+iface_track_visualization = gr.Interface(
+    fn=tools.track_visualization,
+    inputs=[
+        gr.Number(label="Calendar year", value=CURRENT_YEAR, minimum=1950, maximum=CURRENT_YEAR),
+        gr.Textbox(label="Grand Prix", placeholder="Ex: Monaco", info="The name of the GP/country/location (Fuzzy matching supported)"),
+        gr.Radio(["speed", "corners", "gear"], label="Visualization type", value="speed"),
+        gr.Dropdown(driver_names)
+    ],
+    outputs="image",
+    title="[WIP] Track Visualization",
+    description="Get the track visualization for the given Grand Prix"
+)
+
+
+# TODO swap to this format
+named_interfaces = {
+    "driver_championship_standings": iface_driver_championship_standings,
+    "event_info": iface_event_info,
+    "season_calendar": iface_season_calendar,
+    "track_visualization": iface_track_visualization
+}
 
-interface_list = [iface1, iface2, iface3, iface4]
-tab_names = ["[WIP] Driver Championship Score", "[WIP] Driver Position", "Event Info", "Season Calendar"]
+interface_list = [iface_driver_championship_standings, iface_event_info, iface_season_calendar, iface_track_visualization]
+tab_names = ["Driver Championship Standings", "Event Info", "Season Calendar", "Track Visualization"]
 
 
 # Combine all the interfaces into a single TabbedInterface

tools.py

@@ -18,7 +18,8 @@ from utils.constants import AVAILABLE_SESSION_TYPES
 from fastf1.core import Session
 from typing import Union
 
-from utils import parser_utils
+from utils import parser_utils, track_utils
+from PIL import Image
 
 # Custom types
 gp = Union[str, int]
@@ -56,10 +57,37 @@ def get_event_info(year: int, round: gp, format: str) -> str:
 def get_constructor_standings(year: int) -> str:
     pass
 
-def get_driver_standings(year): 
+def get_end_of_season_standings(year): 
     pass
 
 
+def driver_position(driver_name: str) -> str:
+    """
+    Get the current position of the given driver.
+
+    Args:
+        driver_name (str): The driver's name
+
+    Returns:
+        str: The driver's current position
+    """
+    return f"Driver {driver_name} is in position {random.randint(1, 20)}"
+
+def driver_championship_score(year: int, driver_name: str) -> str:
+    pass
+
+def track_visualization(year: int, round: gp, visualization_type: str, driver_name: str) -> Image:
+
+    session = get_session(year, round, "race")
+    session.load()
+
+    if visualization_type == "speed":
+        return track_utils.create_track_speed_visualization(session, driver_name)
+    elif visualization_type == "corners":
+        return track_utils.create_track_corners_visualization(session)
+    elif visualization_type == "gear":
+        return track_utils.create_track_gear_visualization(session)
+
 
 if __name__ == "__main__":
     session = get_session(2024, 1, "fp1")

utils/track_utils.py

@@ -0,0 +1,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
+