284 lines
9.0 KiB
Python
284 lines
9.0 KiB
Python
import fastf1
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
import os
|
|
import json
|
|
import re
|
|
from typing import List, Dict, Any, Optional, Tuple
|
|
|
|
# Import for type hints
|
|
import fastf1.core
|
|
|
|
def rotate(xy, *, angle):
|
|
"""Rotate coordinates by the given angle."""
|
|
rot_mat = np.array([[np.cos(angle), np.sin(angle)],
|
|
[-np.sin(angle), np.cos(angle)]])
|
|
return np.matmul(xy, rot_mat)
|
|
|
|
def get_track_coordinates(session: fastf1.core.Session):
|
|
"""Extract track coordinates from session data."""
|
|
try:
|
|
# Get the fastest lap and position data
|
|
lap = session.laps.pick_fastest()
|
|
pos = lap.get_pos_data()
|
|
|
|
# Get circuit information
|
|
circuit_info = session.get_circuit_info()
|
|
|
|
# Get track coordinates
|
|
track = pos.loc[:, ('X', 'Y')].to_numpy()
|
|
|
|
# Convert the rotation angle from degrees to radian
|
|
track_angle = getattr(circuit_info, 'rotation', 0) / 180 * np.pi
|
|
|
|
# Rotate the track coordinates
|
|
rotated_track = rotate(track, angle=track_angle)
|
|
|
|
# Create circuit_info dictionary with safe attribute access
|
|
circuit_info_dict = {
|
|
'name': session.event['EventName'],
|
|
'country': session.event['Country'],
|
|
'city': session.event['Location'],
|
|
'length': getattr(circuit_info, 'length', None),
|
|
'rotation': getattr(circuit_info, 'rotation', None),
|
|
}
|
|
|
|
# Create event_info dictionary
|
|
event_info_dict = {
|
|
'name': session.event['EventName'],
|
|
'year': session.event['EventDate'].year,
|
|
'official_name': session.event.get('OfficialEventName', session.event['EventName']),
|
|
'location': session.event['Location']
|
|
}
|
|
|
|
return {
|
|
'coordinates': rotated_track.tolist(),
|
|
'circuit_info': circuit_info_dict,
|
|
'event_info': event_info_dict
|
|
}
|
|
except Exception as e:
|
|
print(f"Error extracting track coordinates: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
return None
|
|
|
|
def create_geojson(track_data):
|
|
"""Create a GeoJSON representation of the track."""
|
|
# Extract data from track_data
|
|
coordinates = track_data['coordinates']
|
|
circuit_info = track_data['circuit_info']
|
|
event_info = track_data['event_info']
|
|
|
|
# Create a simplified ID based on country and year
|
|
country_code = get_country_code(circuit_info['country'])
|
|
year = event_info['year']
|
|
circuit_id = f"{country_code.lower()}-{year}"
|
|
|
|
# Ensure length is an integer if available, or None if not
|
|
circuit_length = None
|
|
if circuit_info.get('length'):
|
|
try:
|
|
circuit_length = int(circuit_info['length'])
|
|
except (ValueError, TypeError):
|
|
pass
|
|
|
|
# Create GeoJSON structure
|
|
geojson = {
|
|
"type": "FeatureCollection",
|
|
"name": circuit_id,
|
|
"features": [
|
|
{
|
|
"type": "Feature",
|
|
"properties": {
|
|
"id": circuit_id,
|
|
"Location": event_info['location'],
|
|
"Name": circuit_info['name'],
|
|
"opened": None, # Not available from API
|
|
"firstgp": None, # Not available from API
|
|
"length": circuit_length,
|
|
"altitude": None # Not available from API
|
|
},
|
|
"geometry": {
|
|
"type": "LineString",
|
|
"coordinates": coordinates
|
|
}
|
|
}
|
|
]
|
|
}
|
|
|
|
# Calculate bounding box if coordinates are available
|
|
if coordinates and len(coordinates) > 0:
|
|
x_coords = [coord[0] for coord in coordinates]
|
|
y_coords = [coord[1] for coord in coordinates]
|
|
bbox = [min(x_coords), min(y_coords), max(x_coords), max(y_coords)]
|
|
|
|
# Add bounding box to GeoJSON
|
|
geojson["bbox"] = bbox
|
|
geojson["features"][0]["bbox"] = bbox
|
|
|
|
return geojson
|
|
|
|
def get_country_code(country_name):
|
|
"""Get a two-letter country code from a country name."""
|
|
country_codes = {
|
|
"Australia": "AU",
|
|
"Austria": "AT",
|
|
"Azerbaijan": "AZ",
|
|
"Bahrain": "BH",
|
|
"Belgium": "BE",
|
|
"Brazil": "BR",
|
|
"Canada": "CA",
|
|
"China": "CN",
|
|
"Denmark": "DK",
|
|
"France": "FR",
|
|
"Germany": "DE",
|
|
"Hungary": "HU",
|
|
"Italy": "IT",
|
|
"Japan": "JP",
|
|
"Mexico": "MX",
|
|
"Monaco": "MC",
|
|
"Netherlands": "NL",
|
|
"Portugal": "PT",
|
|
"Qatar": "QA",
|
|
"Russia": "RU",
|
|
"Saudi Arabia": "SA",
|
|
"Singapore": "SG",
|
|
"Spain": "ES",
|
|
"Sweden": "SE",
|
|
"Switzerland": "CH",
|
|
"United Arab Emirates": "AE",
|
|
"United Kingdom": "GB",
|
|
"United States": "US",
|
|
"USA": "US",
|
|
"Abu Dhabi": "AE", # Special case for events
|
|
"Great Britain": "GB",
|
|
"UK": "GB",
|
|
"Emilia Romagna": "IT", # Special case for events
|
|
"Styria": "AT", # Special case for events
|
|
"Miami": "US",
|
|
"Las Vegas": "US"
|
|
}
|
|
|
|
return country_codes.get(country_name, "XX")
|
|
|
|
def sanitize_filename(name):
|
|
"""Convert a string to a valid filename."""
|
|
# Replace special characters with spaces
|
|
name = re.sub(r'[\\/*?:"<>|]', ' ', name)
|
|
# Replace multiple spaces with single space
|
|
name = re.sub(r'\s+', ' ', name).strip()
|
|
return name
|
|
|
|
def save_svg(coordinates, output_path):
|
|
"""Generate and save an SVG representation of the circuit."""
|
|
if not coordinates:
|
|
print("No coordinates available to create SVG")
|
|
return False
|
|
|
|
# Extract x and y coordinates
|
|
x_coords = [coord[0] for coord in coordinates]
|
|
y_coords = [coord[1] for coord in coordinates]
|
|
|
|
# Normalize coordinates to fit in SVG
|
|
min_x, max_x = min(x_coords), max(x_coords)
|
|
min_y, max_y = min(y_coords), max(y_coords)
|
|
|
|
# Add some padding
|
|
x_padding = (max_x - min_x) * 0.05
|
|
y_padding = (max_y - min_y) * 0.05
|
|
|
|
min_x -= x_padding
|
|
max_x += x_padding
|
|
min_y -= y_padding
|
|
max_y += y_padding
|
|
|
|
# SVG dimensions
|
|
svg_width = 800
|
|
svg_height = int(svg_width * (max_y - min_y) / (max_x - min_x))
|
|
|
|
# Create SVG header
|
|
svg = [
|
|
f'<?xml version="1.0" encoding="UTF-8" standalone="no"?>',
|
|
f'<svg width="{svg_width}" height="{svg_height}" viewBox="0 0 {svg_width} {svg_height}" xmlns="http://www.w3.org/2000/svg">',
|
|
]
|
|
|
|
# Create path data
|
|
path_data = []
|
|
for i, (x, y) in enumerate(coordinates):
|
|
# Convert coordinates to SVG coordinates
|
|
svg_x = svg_width * (x - min_x) / (max_x - min_x)
|
|
# Flip y-axis because SVG has origin at top-left
|
|
svg_y = svg_height * (1 - (y - min_y) / (max_y - min_y))
|
|
|
|
if i == 0:
|
|
path_data.append(f"M{svg_x:.2f},{svg_y:.2f}")
|
|
else:
|
|
path_data.append(f"L{svg_x:.2f},{svg_y:.2f}")
|
|
|
|
# Add path to SVG
|
|
svg.append(f' <path d="{" ".join(path_data)}" fill="none" stroke="black" stroke-width="2"/>')
|
|
svg.append('</svg>')
|
|
|
|
# Write SVG to file
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
with open(output_path, 'w', encoding='utf-8') as f:
|
|
f.write('\n'.join(svg))
|
|
|
|
return True
|
|
|
|
def save_png(coordinates, output_path):
|
|
"""Generate and save a PNG representation of the circuit."""
|
|
if not coordinates:
|
|
print("No coordinates available to create PNG")
|
|
return False
|
|
|
|
# Extract x and y coordinates
|
|
x_coords = [coord[0] for coord in coordinates]
|
|
y_coords = [coord[1] for coord in coordinates]
|
|
|
|
# Create figure with transparent background
|
|
fig = plt.figure(figsize=(10, 8), facecolor="none")
|
|
ax = fig.add_subplot(111)
|
|
|
|
# Plot the circuit with black line
|
|
ax.plot(x_coords, y_coords, 'k-', linewidth=2)
|
|
|
|
# Remove grid, axes, and border
|
|
ax.grid(False)
|
|
ax.set_xticks([])
|
|
ax.set_yticks([])
|
|
ax.spines['top'].set_visible(False)
|
|
ax.spines['right'].set_visible(False)
|
|
ax.spines['bottom'].set_visible(False)
|
|
ax.spines['left'].set_visible(False)
|
|
|
|
# Keep aspect ratio equal
|
|
ax.set_aspect('equal')
|
|
|
|
# Add padding around the circuit
|
|
ax.margins(0.1)
|
|
|
|
# Save with transparent background
|
|
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
|
fig.savefig(output_path, transparent=True, bbox_inches='tight', pad_inches=0.1)
|
|
plt.close(fig)
|
|
|
|
return True
|
|
|
|
# Additional utility functions for interactive mode
|
|
def get_available_seasons():
|
|
"""Get a list of available F1 seasons."""
|
|
from datetime import datetime
|
|
current_year = datetime.now().year
|
|
# F1 API has data from 1950 to current year
|
|
return list(range(1950, current_year + 1))
|
|
|
|
def get_events_for_season(year):
|
|
"""Get all events for a specific season."""
|
|
try:
|
|
schedule = fastf1.get_event_schedule(year)
|
|
return schedule
|
|
except Exception as e:
|
|
print(f"Error fetching events for {year}: {str(e)}")
|
|
return None
|