import logging
import os
import shutil
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib import animation
from matplotlib.lines import Line2D
from matplotlib.patches import Circle, Ellipse, FancyArrowPatch, FancyBboxPatch, Polygon, Rectangle
from IPython.display import HTML, Markdown, display
from gridiron import Play, Player, Field, offense, defense, routes as r, style, draw_marker, draw_frame
from gridiron.field import FIELD_WIDTH, HASH_NEAR, HASH_FAR, MIDDLE
logging.getLogger("matplotlib.font_manager").setLevel(logging.ERROR)
style.apply()
PDF = os.environ.get("RTG_FORMAT", "").lower() in ("pdf", "typst", "print")
GREEN = "#0f7a55" # darker step of the aqua highlight: "legal" on the light field
RED = style.SERIES[7] # "foul"
PURPLE = style.SERIES[6] # the tackle box
ZEBRA = style.INK # officials: white disc, black ring
BOX_W = 2 * 1.6 + 0.8 # outside edge of a tackle's normal position, yards from the ball
def O(key, pos, d, w, label=""):
"""An offensive (blue) player at (d, w): d = yards downfield of the line, w = yards to the offense's right."""
return Player(key, pos, "off", d, w, label=label)
def D(key, pos, d, w, label=""):
"""A defensive (orange) player."""
return Player(key, pos, "def", d, w, label=label)
def XY(fld, p, d, w):
"""(d, w) in the offense's frame -> plot coordinates on a drawn field."""
x, y = fld.xy(*p.to_abs(d, w))
return float(x), float(y)
def say(fld, p, d, w, text, color=style.INK_2, fs=6.8, ha="center", weight="normal", box=True, rot=0):
"""A short on-field label at (d, w), on a light pad so lines underneath don't cut it."""
x, y = XY(fld, p, d, w)
kw = dict(boxstyle="round,pad=0.18", fc=style.SURFACE, ec="none", alpha=0.92) if box else None
fld.ax.text(x, y, text, fontsize=fs, color=color, ha=ha, va="center", zorder=12, fontweight=weight,
linespacing=1.12, bbox=kw, rotation=rot)
def path_arrow(fld, p, pts, color=style.BALL, lw=1.4, ls=(0, (2.5, 2)), head=9, z=7, rad=0.0):
"""A dashed arrow through (d, w) points: a ball's flight or a loose ball's bounce."""
xy = [XY(fld, p, *q) for q in pts]
if len(xy) > 2:
xs, ys = zip(*xy[:-1])
fld.ax.plot(xs, ys, color=color, lw=lw, ls=ls, zorder=z)
fld.ax.add_patch(FancyArrowPatch(xy[-2], xy[-1], arrowstyle="-|>", mutation_scale=head, color=color,
lw=lw, ls=ls, zorder=z, connectionstyle=f"arc3,rad={rad}"))
def xmark(fld, p, d, w, color=style.DEFENSE_DARK, ms=8):
x, y = XY(fld, p, d, w)
fld.ax.plot(x, y, marker="x", ms=ms, mew=2.2, color=color, zorder=11)
def ball_at(fld, p, d, w):
x, y = XY(fld, p, d, w)
draw_marker(fld.ax, x, y, "ball", "", fld.pts_per_yard())
def no_numbers(ax):
"""Drop the painted yard numbers: on zoomed diagrams they compete with the labels."""
for t in list(ax.texts):
if t.get_text().isdigit():
t.remove()
def official(ax, x, y, label, ppy, r=1.1, alpha=1.0):
"""An official: white disc with a black ring (offense = blue disc, defense = orange square)."""
ax.add_patch(Circle((x, y), r, fc="white", ec=ZEBRA, lw=1.6, zorder=9, alpha=alpha))
fs = max(4.6, min(8.5, r * ppy * (1.15 if len(label) == 1 else 0.92)))
ax.text(x, y, label, ha="center", va="center", fontsize=fs, color=ZEBRA, fontweight="bold", zorder=10,
alpha=alpha)
def marker_legend(ax, y=-0.01, extra=()):
h = [Line2D([], [], marker="o", ls="", markerfacecolor=style.OFFENSE, markeredgecolor="white", ms=8,
label="Offense"),
Line2D([], [], marker="D", ls="", markerfacecolor=style.DEFENSE_TINT, markeredgecolor=style.DEFENSE,
ms=7, label="Defense"),
Line2D([], [], marker="o", ls="", markerfacecolor="white", markeredgecolor=ZEBRA, mew=1.6, ms=9,
label="Official")] + list(extra)
ax.legend(handles=h, loc="upper center", bbox_to_anchor=(0.5, y), ncol=len(h), fontsize=7.2,
frameon=False, handletextpad=0.4, columnspacing=1.2)
def tackle_box(ax, fld, p, back, alpha=0.10):
"""The shaded tackle box (the rulebook's pocket area), from the line back to depth `back`."""
xy = [XY(fld, p, *q) for q in [(-0.7, -BOX_W), (-0.7, BOX_W), (back, BOX_W), (back, -BOX_W)]]
ax.add_patch(Polygon(xy, closed=True, fc=PURPLE, alpha=alpha, ec=PURPLE, lw=0.9, ls=(0, (3, 2)),
zorder=1.4))
# --- flowchart pieces ---------------------------------------------------------------------------
def fbox(ax, x, y, w, h, title, body="", fc=style.SURFACE, ec=style.BASELINE, tfs=7.8, bfs=6.7, lw=1.0):
ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle="round,pad=0.02,rounding_size=1.1",
fc=fc, ec=ec, lw=lw, zorder=3))
if body:
ax.text(x, y + h / 2 - 1.6, title, ha="center", va="top", fontsize=tfs, fontweight="bold",
color=style.INK, zorder=4, linespacing=1.1)
nl = title.count("\n") + 1
ax.text(x, y + h / 2 - 1.9 - 2.9 * nl, body, ha="center", va="top", fontsize=bfs, color=style.INK_2,
zorder=4, linespacing=1.18)
else:
ax.text(x, y, title, ha="center", va="center", fontsize=tfs, fontweight="bold", color=style.INK,
zorder=4, linespacing=1.15)
def farrow(ax, p0, p1, text=None, color=style.INK_2, toff=(0.0, 0.0), ha="center", fs=6.8):
ax.add_patch(FancyArrowPatch(p0, p1, arrowstyle="-|>", mutation_scale=10, color=color, lw=1.0,
shrinkA=1, shrinkB=1, zorder=2))
if text:
ax.text((p0[0] + p1[0]) / 2 + toff[0], (p0[1] + p1[1]) / 2 + toff[1], text, ha=ha, va="center",
fontsize=fs, color=color, fontweight="bold", zorder=5,
bbox=dict(boxstyle="round,pad=0.12", fc=style.SURFACE, ec="none"))
def canvas(w_in, h_in, xmax=100, ymax=None):
fig, ax = plt.subplots(figsize=(w_in, h_in))
ymax = ymax or xmax * h_in / w_in
ax.set_xlim(0, xmax); ax.set_ylim(0, ymax); ax.set_axis_off(); ax.grid(False)
return fig, ax
# --- scenes: simulated tracking plus extra actors (officials), as a video or a frame strip ------
def kf_interp(kf, t):
kf = np.asarray(kf, dtype=float)
return np.array([np.interp(t, kf[:, 0], kf[:, 1]), np.interp(t, kf[:, 0], kf[:, 2])])
def draw_scene(ax, play, trk, t, window, lateral, extra=None, decor=None, trails=True, clock=True):
"""One moment of a play on `ax`: players (and ball) from tracking `trk`, extra actors from keyframes."""
ylo, yhi = sorted((play.ball_y - lateral[1], play.ball_y - lateral[0]))
fld = Field(ax=ax, orient="vertical", los=play.los, window=window, to_go=play.to_go, lateral=(ylo, yhi))
no_numbers(ax)
ppy = fld.pts_per_yard()
if decor:
decor(fld, t)
times = trk.groupby("frameId")["time"].first()
fid = int((times - t).abs().idxmin())
cur = trk[trk.frameId == fid]
hist = trk[trk.frameId <= fid] if trails else None
draw_frame(fld, cur, hist, ppy)
for key, (lab, kf) in (extra or {}).items():
ts = np.linspace(min(np.asarray(kf)[:, 0].min(), 0), t, 30)
pts = np.array([kf_interp(kf, s) for s in ts])
xs, ys = zip(*[XY(fld, play, d, w) for d, w in pts])
if trails:
ax.plot(xs, ys, color=ZEBRA, lw=0.9, alpha=0.35, zorder=3)
official(ax, xs[-1], ys[-1], lab, ppy, r=1.15)
for art in list(ax.patches) + list(ax.texts) + list(ax.lines):
art.set_clip_on(True) # players who ran out of the window don't spill outside it
if clock:
ax.set_title("snap" if abs(t) < 1e-6 else f"{t:+.1f} s", fontsize=8.6, color=style.INK_2)
return fld
def show_scene(play, trk, times, window, lateral, extra=None, decor=None, notes=None, ncols=2, width_in=6.5):
"""HTML: a video of the scene. PDF: a frame strip at `times`, each panel with its note."""
trk = trk.copy()
trk["_id"] = trk["displayName"].astype(str)
if PDF:
n = len(times)
nrows = int(np.ceil(n / ncols))
span_w, span_d = lateral[1] - lateral[0], window[1] - window[0]
pw = width_in / ncols
fig, axes = plt.subplots(nrows, ncols, figsize=(width_in, nrows * (pw * span_d / span_w + 0.35)),
squeeze=False)
for ax in axes.flat[n:]:
ax.axis("off")
for i, (ax, t) in enumerate(zip(axes.flat, times)):
fld = draw_scene(ax, play, trk, t, window, lateral, extra, decor)
if notes:
notes(fld, i, t)
fig.tight_layout(h_pad=1.0, w_pad=0.6)
plt.show()
return
tmax = float(trk["time"].max())
t0 = float(trk["time"].min())
ts = np.round(np.arange(max(t0, -0.5), tmax + 1e-9, 0.1), 2)
span_w, span_d = lateral[1] - lateral[0], window[1] - window[0]
fig, ax = plt.subplots(figsize=(5.6, 5.6 * span_d / span_w))
hold = {}
def update(t):
ax.clear()
fld = draw_scene(ax, play, trk, float(t), window, lateral, extra, decor)
if notes:
# show the note of the last panel time already passed
k = max([i for i, tt in enumerate(times) if tt <= t + 1e-6] or [0])
notes(fld, k, float(t))
return []
anim = animation.FuncAnimation(fig, update, frames=ts, interval=100, blit=False)
if shutil.which("ffmpeg"):
html = anim.to_html5_video(embed_limit=40).replace("<video ", '<video style="max-width:100%;height:auto" ')
else:
html = anim.to_jshtml(default_mode="once")
plt.close(fig)
display(HTML(html))
def md_table(rows, header, align):
lines = ["| " + " | ".join(header) + " |", "|" + "|".join(align) + "|"]
lines += ["| " + " | ".join(map(str, row)) + " |" for row in rows]
return Markdown("\n".join(lines))
