Add speed parsing for SecondSpectrum

PySport · Jun 30, 2023 · 637fad9 · 637fad9
1 parent c5e5f5f
commit 637fad9
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 1 deletion.
diff --git a/kloppy/domain/models/tracking.py b/kloppy/domain/models/tracking.py
@@ -24,6 +24,7 @@ class Frame(DataRecord):
     players_data: Dict[Player, PlayerData]
     other_data: Dict[str, Any]
     ball_coordinates: Point3D
+    ball_speed: int
 
     @property
     def record_id(self) -> int:

diff --git a/kloppy/domain/services/transformers/dataset.py b/kloppy/domain/services/transformers/dataset.py
@@ -190,6 +190,7 @@ def __change_frame_coordinate_system(self, frame: Frame):
             ball_coordinates=self.__change_point_coordinate_system(
                 frame.ball_coordinates
             ),
+            ball_speed=frame.ball_speed,
             players_data={
                 key: PlayerData(
                     coordinates=self.__change_point_coordinate_system(

diff --git a/kloppy/infra/serializers/tracking/secondspectrum.py b/kloppy/infra/serializers/tracking/secondspectrum.py
@@ -64,8 +64,10 @@ def _frame_from_framedata(cls, teams, period, frame_data):
             ball_coordinates = Point3D(
                 float(ball_x), float(ball_y), float(ball_z)
             )
+            ball_speed = frame_data["ball"]["speed"]
         else:
             ball_coordinates = None
+            ball_speed = None
 
         ball_state = BallState.ALIVE if frame_data["live"] else BallState.DEAD
         ball_owning_team = (
@@ -78,6 +80,7 @@ def _frame_from_framedata(cls, teams, period, frame_data):
 
                 jersey_no = player_data["number"]
                 x, y, _ = player_data["xyz"]
+                speed = player_data["speed"]
                 player = team.get_player_by_jersey_number(jersey_no)
 
                 if not player:
@@ -89,13 +92,14 @@ def _frame_from_framedata(cls, teams, period, frame_data):
                     team.players.append(player)
 
                 players_data[player] = PlayerData(
-                    coordinates=Point(float(x), float(y))
+                    coordinates=Point(float(x), float(y)), speed=speed
                 )
 
         return Frame(
             frame_id=frame_id,
             timestamp=frame_timestamp,
             ball_coordinates=ball_coordinates,
+            ball_speed=ball_speed,
             ball_state=ball_state,
             ball_owning_team=ball_owning_team,
             players_data=players_data,

diff --git a/kloppy/tests/test_secondspectrum.py b/kloppy/tests/test_secondspectrum.py
@@ -106,6 +106,10 @@ def test_correct_normalized_deserialization(
         assert dataset.records[0].players_coordinates[home_player] == Point(
             x=0.4146981051733674, y=0.9144718866065965
         )
+        assert (
+            dataset.records[0].players_data[home_player].speed
+            == 6.578958220040129
+        )
 
         # Check normalised pitch dimensions
         pitch_dimensions = dataset.metadata.pitch_dimensions