Adding a transform function to overlay teams

kloppy/domain/services/transformers/dataset.py

@@ -1,8 +1,9 @@
 import warnings
 from dataclasses import fields, replace
 
+from kloppy.domain.models.common import Player
 from kloppy.domain.models.tracking import PlayerData
-from typing import Union, Optional
+from typing import Dict, Union, Optional
 
 from kloppy.domain import (
     AttackingDirection,
@@ -37,6 +38,9 @@ def __init__(
         to_coordinate_system: Optional[CoordinateSystem] = None,
         to_pitch_dimensions: Optional[PitchDimensions] = None,
         to_orientation: Optional[Orientation] = None,
+        overlay_teams: bool = False,
+        horizontal_mirror: bool = False,
+        vertical_mirror: bool = False,
     ):
         if (
             from_pitch_dimensions
@@ -73,9 +77,15 @@ 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
+        self._horizontal_mirror = horizontal_mirror
+        self._vertical_mirror = vertical_mirror
         if (
             from_orientation
             and not to_orientation
@@ -128,19 +138,31 @@ def flip_point(
         if not point:
             return None
 
-        x_base = self._to_pitch_dimensions.x_dim.to_base(point.x)
-        y_base = self._to_pitch_dimensions.y_dim.to_base(point.y)
+        x = self.mirror_horizontally(point.x)
+        y = self.mirror_vertically(point.y)
+
+        if isinstance(point, Point3D):
+            return Point3D(x=x, y=y, z=point.z)
+        else:
+            return Point(x=x, y=y)
+
+    def mirror_horizontally(self, x_point: float) -> float:
+        x_base = self._to_pitch_dimensions.x_dim.to_base(x_point)
 
         x_base = 1 - x_base
-        y_base = 1 - y_base
 
         x = self._to_pitch_dimensions.x_dim.from_base(x_base)
+
+        return x
+
+    def mirror_vertically(self, y_point: float) -> float:
+        y_base = self._to_pitch_dimensions.y_dim.to_base(y_point)
+
+        y_base = 1 - y_base
+
         y = self._to_pitch_dimensions.y_dim.from_base(y_base)
 
-        if isinstance(point, Point3D):
-            return Point3D(x=x, y=y, z=point.z)
-        else:
-            return Point(x=x, y=y)
+        return y
 
     def __needs_flip(
         self,
@@ -185,6 +207,16 @@ def transform_frame(self, frame: Frame) -> Frame:
         elif self._needs_pitch_dimensions_change:
             frame = self.__change_frame_dimensions(frame)
 
+        elif self._overlay_teams:
+            frame = self.transform_frame_overlay_teams(frame)
+
+        elif self._vertical_mirror or self._horizontal_mirror:
+            frame = self.transform_mirror(
+                frame=frame,
+                vertical_mirror=self._vertical_mirror,
+                horizontal_mirror=self._horizontal_mirror,
+            )
+
         # Flip frame based on orientation
         if self._needs_orientation_change:
             if self.__needs_flip(
@@ -308,6 +340,79 @@ def __flip_frame(self, frame: Frame):
             statistics=frame.statistics,
         )
 
+    def _needs_mirror(self, player: Player, team: Team | None):
+        return (team and team.team_id == player.team.team_id) or not team
+
+    def transform_mirror(
+        self,
+        frame: Frame,
+        team: Team | None = None,
+        vertical_mirror: bool = False,
+        horizontal_mirror: bool = False,
+    ):
+        players_data = {}
+
+        players_data = {
+            player: PlayerData(
+                coordinates=Point(
+                    x=(
+                        self.mirror_horizontally(player_data.coordinates.x)
+                        if horizontal_mirror
+                        else player_data.coordinates.x
+                    ),
+                    y=(
+                        self.mirror_vertically(player_data.coordinates.y)
+                        if vertical_mirror
+                        else player_data.coordinates.y
+                    ),
+                ),
+                distance=player_data.distance,
+                speed=player_data.speed,
+                other_data=player_data.other_data,
+            )
+            for player, player_data in frame.players_data.items()
+            if self._needs_mirror(player, team)
+        }
+
+        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=frame.ball_coordinates,
+            players_data={**frame.players_data, **players_data},
+        )
+
+    def transform_frame_overlay_teams(self, frame: Frame):
+
+        if frame.attacking_direction == AttackingDirection.RTL:
+            left_team = [
+                team
+                for team in frame.dataset.metadata.teams
+                if team.team_id != frame.ball_owning_team.team_id
+            ][0]
+            frame = self.transform_mirror(
+                frame, left_team, vertical_mirror=True, horizontal_mirror=True
+            )
+        else:
+            frame = self.transform_mirror(
+                frame,
+                frame.ball_owning_team,
+                vertical_mirror=True,
+                horizontal_mirror=True,
+            )
+            frame.ball_coordinates = Point(
+                x=self.mirror_horizontally(frame.ball_coordinates.x),
+                y=self.mirror_vertically(frame.ball_coordinates.y),
+            )
+
+        return frame
+
     def transform_event(self, event: Event) -> Event:
         # Change coordinate system
         if self._needs_coordinate_system_change:
@@ -375,11 +480,17 @@ def transform_dataset(
         to_pitch_dimensions: Optional[PitchDimensions] = None,
         to_orientation: Optional[Orientation] = None,
         to_coordinate_system: Optional[CoordinateSystem] = None,
+        overlay_teams: bool = False,
+        horizontal_mirror: bool = False,
+        vertical_mirror: 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
+            and horizontal_mirror is False
+            and vertical_mirror is False
         ):
             return dataset
 
@@ -391,8 +502,34 @@ 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 horizontal_mirror or vertical_mirror:
+            transformer = cls(
+                from_pitch_dimensions=dataset.metadata.pitch_dimensions,
+                from_orientation=dataset.metadata.orientation,
+                to_orientation=to_orientation,
+                to_pitch_dimensions=to_pitch_dimensions,
+                horizontal_mirror=horizontal_mirror,
+                vertical_mirror=vertical_mirror,
+            )
+            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 +555,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,9 @@ def transform(
     to_coordinate_system: Optional[
         Union[CoordinateSystem, Provider, str]
     ] = None,
+    overlay_teams: bool = False,
+    horizontal_mirror: bool = False,
+    vertical_mirror: bool = False,
 ) -> Dataset:
     # convert raw orientation to object
     if to_orientation is not None and isinstance(to_orientation, str):
@@ -43,4 +46,7 @@ def transform(
         to_orientation=to_orientation,
         to_coordinate_system=to_coordinate_system,
         to_pitch_dimensions=to_pitch_dimensions,
+        overlay_teams=overlay_teams,
+        horizontal_mirror=horizontal_mirror,
+        vertical_mirror=vertical_mirror,
     )