Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

update coordinatesdataframe.py #134

Open
wants to merge 2 commits into
base: develop
Choose a base branch
from
Open

update coordinatesdataframe.py #134

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
bb1bdc1
update coordinatesdataframe.py
shunsuke-iwashita Nov 15, 2023
287d167
update coordinatesdataframe.py
shunsuke-iwashita Nov 15, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 69 additions & 18 deletions sportslabkit/dataframe/coordinatesdataframe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,12 @@
from typing import Any

import cv2
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib.animation import FuncAnimation
from matplotlib.lines import Line2D
from mplsoccer import Pitch
from numpy.typing import ArrayLike, NDArray

Expand Down Expand Up @@ -244,6 +247,7 @@ def visualize_frame(
home_kwargs: dict[str, Any] | None = None,
away_kwargs: dict[str, Any] | None = None,
save_kwargs: dict[str, Any] | None = None,
sports: str = "soccer"
):
"""Visualize a single frame.

Expand All @@ -262,6 +266,7 @@ def visualize_frame(
home_kwargs: Keyword arguments specifically for the home team markers.
away_kwargs: Keyword arguments specifically for the away team markers.
save_kwargs: Keyword arguments for the save function.
sports: Name of sports. Defaults to "soccer".

Note:
`marker_kwargs` will be used for all markers but will be overwritten by `ball_kwargs`, `home_kwargs` and `away_kwargs`. All keyword arguments are passed to `plt.plot`. `save_kwargs` are passed to `plt.savefig`.
Expand Down Expand Up @@ -309,16 +314,38 @@ def visualize_frame(
df_ball = _df[ball_key]
df_home = _df[home_key]
df_away = _df[away_key]
pitch = Pitch(
pitch_color="black",
line_color=(0.3, 0.3, 0.3),
pitch_type="custom",
pitch_length=105,
pitch_width=68,
label=False,
)
if sports == "soccer":
pitch = Pitch(
pitch_color="black",
line_color=(0.3, 0.3, 0.3),
pitch_type="custom",
pitch_length=105,
pitch_width=68,
label=False,
)

fig, ax = pitch.draw(figsize=(8, 5.2))

elif sports == "ultimate_3on3":
plt.style.use('dark_background')
fig, ax = plt.subplots()

pitch_length = 54
pitch_width = 22
end_line_left = 10
end_line_right = 44

rect = patches.Rectangle((0, 0), pitch_length, pitch_width, edgecolor='#606060', lw=1.5, fill=False)
ax.add_patch(rect)

line1 = Line2D([end_line_left, end_line_left], [0, pitch_width], color='#606060', lw=1.5)
ax.add_line(line1)

line2 = Line2D([end_line_right, end_line_right], [0, pitch_width], color='#606060', lw=1.5)
ax.add_line(line2)

fig, ax = pitch.draw(figsize=(8, 5.2))
ax.set_aspect('equal')
ax.axis('off')

ax.plot(
df_ball.loc[:, (slice(None), "x")],
Expand Down Expand Up @@ -350,6 +377,7 @@ def visualize_frames(
home_kwargs: dict[str, Any] | None = None,
away_kwargs: dict[str, Any] | None = None,
save_kwargs: dict[str, Any] | None = None,
sports: str = "soccer",
):
"""Visualize multiple frames using matplotlib.animation.FuncAnimation.

Expand All @@ -370,6 +398,7 @@ def visualize_frames(
home_kwargs: Keyword arguments specifically for the home team markers.
away_kwargs: Keyword arguments specifically for the away team markers.
save_kwargs: Keyword arguments for the save function.
sports: Name of sports. Defaults to "soccer".

Note:
`marker_kwargs` will be used for all markers but will be overwritten by `ball_kwargs`, `home_kwargs` and `away_kwargs`. All keyword arguments are passed to `plt.plot`. `save_kwargs` are passed to `FuncAnimation.save`.
Expand Down Expand Up @@ -428,16 +457,38 @@ def visualize_frames(
df_ball = _df[ball_key]
df_home = _df[home_key]
df_away = _df[away_key]
pitch = Pitch(
pitch_color="black",
line_color=(0.3, 0.3, 0.3),
pitch_type="custom",
pitch_length=105,
pitch_width=68,
label=False,
)
if sports == "soccer":
pitch = Pitch(
pitch_color="black",
line_color=(0.3, 0.3, 0.3),
pitch_type="custom",
pitch_length=105,
pitch_width=68,
label=False,
)

fig, ax = pitch.draw(figsize=(8, 5.2))

elif sports == "ultimate_3on3":
plt.style.use('dark_background')
fig, ax = plt.subplots()

pitch_length = 54
pitch_width = 22
end_line_left = 10
end_line_right = 44

rect = patches.Rectangle((0, 0), pitch_length, pitch_width, edgecolor='#606060', lw=1.5, fill=False)
ax.add_patch(rect)

line1 = Line2D([end_line_left, end_line_left], [0, pitch_width], color='#606060', lw=1.5)
ax.add_line(line1)

line2 = Line2D([end_line_right, end_line_right], [0, pitch_width], color='#606060', lw=1.5)
ax.add_line(line2)

fig, ax = pitch.draw(figsize=(8, 5.2))
ax.set_aspect('equal')
ax.axis('off')

ball, *_ = ax.plot([], [], **_ball_kwargs)
away, *_ = ax.plot([], [], **_away_kwargs)
Expand Down