Create a new transform function to overlay teams

Author: douglaspampu

kloppy/domain/services/transformers/dataset.py

@@ -37,6 +37,7 @@ def __init__(
         to_coordinate_system: Optional[CoordinateSystem] = None,
         to_pitch_dimensions: Optional[PitchDimensions] = None,
         to_orientation: Optional[Orientation] = None,
+        overlay_teams: bool = False
     ):
         if (
             from_pitch_dimensions
@@ -73,9 +74,13 @@ def __init__(
                     "You must specify the source CoordinateSystem when specifying the target CoordinateSystem"
                 )
             self._to_pitch_dimensions = to_coordinate_system.pitch_dimensions
+        else:
+            self._to_pitch_dimensions = self._from_pitch_dimensions
 
         self._from_orientation = from_orientation
         self._to_orientation = to_orientation
+
+        self._overlay_teams = overlay_teams
         if (
             from_orientation
             and not to_orientation
@@ -184,6 +189,9 @@ def transform_frame(self, frame: Frame) -> Frame:
         # Change dimensions
         elif self._needs_pitch_dimensions_change:
             frame = self.__change_frame_dimensions(frame)
+        
+        elif self._overlay_teams:
+            frame = self.transform_frame_overlay_teams(frame)
 
         # Flip frame based on orientation
         if self._needs_orientation_change:
@@ -308,6 +316,37 @@ def __flip_frame(self, frame: Frame):
             statistics=frame.statistics,
         )
 
+    def _get_overlay_players_coordinates(self, player_data:PlayerData, player_team:Team, ball_owning_team:Team, attacking_direction:AttackingDirection):
+        if attacking_direction == AttackingDirection.RTL:
+            if player_team != ball_owning_team:
+                player_data.coordinates = self.flip_point(player_data.coordinates)
+        else:
+            if player_team == ball_owning_team:
+                player_data.coordinates = self.flip_point(player_data.coordinates)
+
+        return player_data
+
+    def transform_frame_overlay_teams(self, frame: Frame):
+        players_data = {player:self._get_overlay_players_coordinates(player_data, player.team, frame.ball_owning_team, frame.attacking_direction) for player, player_data in frame.players_data.items()}
+
+        ball_coordinates = frame.ball_coordinates
+        if frame.attacking_direction != AttackingDirection.RTL:
+            ball_coordinates = self.flip_point(ball_coordinates)
+
+        return Frame(
+            # doesn't change
+            timestamp=frame.timestamp,
+            frame_id=frame.frame_id,
+            ball_owning_team=frame.ball_owning_team,
+            ball_state=frame.ball_state,
+            period=frame.period,
+            other_data=frame.other_data,
+            statistics=frame.statistics,
+            # changes
+            ball_coordinates=ball_coordinates,
+            players_data=players_data,
+        )
+    
     def transform_event(self, event: Event) -> Event:
         # Change coordinate system
         if self._needs_coordinate_system_change:
@@ -375,11 +414,13 @@ def transform_dataset(
         to_pitch_dimensions: Optional[PitchDimensions] = None,
         to_orientation: Optional[Orientation] = None,
         to_coordinate_system: Optional[CoordinateSystem] = None,
+        overlay_teams: bool = False
     ) -> Dataset:
         if (
             to_pitch_dimensions is None
             and to_orientation is None
             and to_coordinate_system is None
+            and overlay_teams is False
         ):
             return dataset
 
@@ -391,8 +432,20 @@ def transform_dataset(
                     "Cannot transform to BALL_OWNING_TEAM orientation when "
                     "dataset doesn't contain ball owning team data"
                 )
-
-        if to_pitch_dimensions is not None:
+        if overlay_teams:
+            transformer = cls(
+                from_pitch_dimensions=dataset.metadata.pitch_dimensions,
+                from_orientation=dataset.metadata.orientation,
+                to_orientation=to_orientation,
+                to_pitch_dimensions=to_pitch_dimensions,
+                overlay_teams=overlay_teams
+            )
+            metadata = replace(
+                dataset.metadata,
+                pitch_dimensions=to_pitch_dimensions,
+                orientation=to_orientation,
+            )
+        elif to_pitch_dimensions is not None:
             # Transform the pitch dimensions and optionally the orientation
             transformer = cls(
                 from_pitch_dimensions=dataset.metadata.pitch_dimensions,
@@ -418,7 +471,6 @@ def transform_dataset(
                 dataset.metadata,
                 coordinate_system=to_coordinate_system,
                 pitch_dimensions=to_coordinate_system.pitch_dimensions,
-                orientation=to_orientation,
             )
 
         else:

kloppy/helpers.py

@@ -18,6 +18,7 @@ def transform(
     to_coordinate_system: Optional[
         Union[CoordinateSystem, Provider, str]
     ] = None,
+    overlay_teams:bool = False
 ) -> Dataset:
     # convert raw orientation to object
     if to_orientation is not None and isinstance(to_orientation, str):
@@ -43,4 +44,5 @@ def transform(
         to_orientation=to_orientation,
         to_coordinate_system=to_coordinate_system,
         to_pitch_dimensions=to_pitch_dimensions,
+        overlay_teams=overlay_teams
     )

kloppy/tests/test_helpers.py

@@ -29,6 +29,88 @@
 
 
 class TestHelpers:
+
+    def _get_tracking_dataset_multiple_players(self):
+        home_team = Team(team_id="home", name="home", ground=Ground.HOME)
+        away_team = Team(team_id="away", name="away", ground=Ground.AWAY)
+        teams = [home_team, away_team]
+
+        periods = [
+            Period(
+                id=1,
+                start_timestamp=0.0,
+                end_timestamp=10.0,
+            ),
+            Period(
+                id=2,
+                start_timestamp=15.0,
+                end_timestamp=25.0,
+            ),
+        ]
+        metadata = Metadata(
+            flags=(DatasetFlag.BALL_OWNING_TEAM),
+            pitch_dimensions=NormalizedPitchDimensions(
+                x_dim=Dimension(0, 100),
+                y_dim=Dimension(-50, 50),
+                pitch_length=105,
+                pitch_width=68,
+            ),
+            orientation=Orientation.HOME_AWAY,
+            frame_rate=25,
+            periods=periods,
+            teams=teams,
+            score=None,
+            provider=None,
+            coordinate_system=None,
+            date="2024-05-19T13:30:00",
+            game_week="35",
+            game_id="2374516",
+        )
+
+        tracking_data = TrackingDataset(
+            metadata=metadata,
+            records=[
+                create_frame(
+                    frame_id=1,
+                    timestamp=0.1,
+                    ball_owning_team=teams[0],
+                    ball_state=None,
+                    period=periods[0],
+                    players_data={},
+                    other_data=None,
+                    ball_coordinates=Point3D(x=100, y=-50, z=0),
+                ),
+                create_frame(
+                    frame_id=2,
+                    timestamp=0.2,
+                    ball_owning_team=teams[1],
+                    ball_state=None,
+                    period=periods[1],
+                    players_data={
+                        Player(
+                            team=home_team, player_id="home_1", jersey_no=1
+                        ): PlayerData(
+                            coordinates=Point(x=15, y=35),
+                            distance=0.03,
+                            speed=10.5,
+                            other_data={"extra_data": 1},
+                        ),
+                        Player(
+                            team=away_team, player_id="away_1", jersey_no=1
+                        ): PlayerData(
+                            coordinates=Point(x=15, y=35),
+                            distance=0.03,
+                            speed=10.5,
+                            other_data={"extra_data": 1},
+                        )
+                    },
+                    other_data={"extra_data": 1},
+                    ball_coordinates=Point3D(x=0, y=50, z=1),
+                ),
+            ],
+        )
+        return tracking_data
+
     def _get_tracking_dataset(self):
         home_team = Team(team_id="home", name="home", ground=Ground.HOME)
         away_team = Team(team_id="away", name="away", ground=Ground.AWAY)
@@ -371,6 +453,45 @@ def test_transform_event_data_freeze_frame(self, base_dir):
         assert coordinates.x == 1 - coordinates_transformed.x
         assert coordinates.y == 1 - coordinates_transformed.y
 
+    def test_transform_overlay_teams_frames(self, base_dir):
+        dataset = tracab.load(
+            meta_data=base_dir / "files/tracab_meta.xml",
+            raw_data=base_dir / "files/tracab_raw.dat",
+            only_alive=False,
+            coordinates="tracab",
+        )
+
+        transformed_frame = dataset.transform(overlay_teams=True)
+
+        ## Assert frame with attacking direction right-to-left
+
+        home_players_coordinates = [coordinates for player, coordinates in dataset.frames[0].players_coordinates.items() if dataset.metadata.teams[0].get_player_by_id(player.player_id)]
+        converted_home_players_coordinates = [coordinates for player, coordinates in transformed_frame.frames[0].players_coordinates.items() if dataset.metadata.teams[0].get_player_by_id(player.player_id)]
+
+        away_players_coordinates = [coordinates for player, coordinates in dataset.frames[0].players_coordinates.items() if dataset.metadata.teams[1].get_player_by_id(player.player_id)]
+        converted_away_players_coordinates = [coordinates for player, coordinates in transformed_frame.frames[0].players_coordinates.items() if dataset.metadata.teams[1].get_player_by_id(player.player_id)]
+
+        assert all([a==b] for a, b in zip(home_players_coordinates, converted_home_players_coordinates))
+
+        assert all([a!=b] for a, b in zip(away_players_coordinates, converted_away_players_coordinates))
+
+        assert dataset.frames[0].ball_coordinates == transformed_frame.frames[0].ball_coordinates
+
+        ## Assert frame with attacking direction with left-to-right
+
+        home_players_coordinates = [coordinates for player, coordinates in dataset.frames[4].players_coordinates.items() if dataset.metadata.teams[0].get_player_by_id(player.player_id)]
+        converted_home_players_coordinates = [coordinates for player, coordinates in transformed_frame.frames[4].players_coordinates.items() if dataset.metadata.teams[0].get_player_by_id(player.player_id)]
+
+        away_players_coordinates = [coordinates for player, coordinates in dataset.frames[4].players_coordinates.items() if dataset.metadata.teams[1].get_player_by_id(player.player_id)]
+        converted_away_players_coordinates = [coordinates for player, coordinates in transformed_frame.frames[4].players_coordinates.items() if dataset.metadata.teams[1].get_player_by_id(player.player_id)]
+
+        assert all([a!=b] for a, b in zip(home_players_coordinates, converted_home_players_coordinates))
+
+        assert all([a==b] for a, b in zip(away_players_coordinates, converted_away_players_coordinates))
+
+        assert dataset.frames[4].ball_coordinates != transformed_frame.frames[4].ball_coordinates
+
+
     def test_to_pandas(self):
         tracking_data = self._get_tracking_dataset()