feat(tracab): decoupled Tracab dat / json from meta data file types (#364)

---------

Co-authored-by: UnravelSports [JB] <jors@unravelsports.com>
Co-authored-by: Pieter Robberechts <pieter.robberechts@kuleuven.be>

Author: 3 people

kloppy/_providers/tracab.py

@@ -1,15 +1,12 @@
-from typing import Optional, Union, Type
-
+import warnings
+from typing import Optional
 
 from kloppy.domain import TrackingDataset
-from kloppy.infra.serializers.tracking.tracab.tracab_dat import (
-    TRACABDatDeserializer,
-)
-from kloppy.infra.serializers.tracking.tracab.tracab_json import (
-    TRACABJSONDeserializer,
+from kloppy.infra.serializers.tracking.tracab.deserializer import (
+    TRACABDeserializer,
     TRACABInputs,
 )
-from kloppy.io import FileLike, open_as_file, get_file_extension
+from kloppy.io import FileLike, open_as_file
 
 
 def load(
@@ -18,17 +15,67 @@ def load(
     sample_rate: Optional[float] = None,
     limit: Optional[int] = None,
     coordinates: Optional[str] = None,
-    only_alive: Optional[bool] = True,
+    only_alive: bool = True,
     file_format: Optional[str] = None,
 ) -> TrackingDataset:
-    if file_format == "dat":
-        deserializer_class = TRACABDatDeserializer
-    elif file_format == "json":
-        deserializer_class = TRACABJSONDeserializer
-    else:
-        deserializer_class = identify_deserializer(meta_data, raw_data)
-
-    deserializer = deserializer_class(
+    """
+    Load TRACAB tracking data.
+
+    Args:
+        meta_data: A JSON or XML feed containing the meta data.
+        raw_data: A JSON or dat feed containing the raw tracking data.
+        sample_rate: Sample the data at a specific rate.
+        limit: Limit the number of frames to load to the first `limit` frames.
+        coordinates: The coordinate system to use.
+        only_alive: Only include frames in which the game is not paused.
+        file_format:
+            Deprecated. The format will be inferred based on the file extensions.
+
+    Returns:
+        The parsed tracking data.
+
+    Notes:
+        Tracab distributes its metadata in various formats. Kloppy tries to
+        infer automatically which format applies. Currently, kloppy supports
+        the following formats:
+
+        - **Flat XML structure**:
+
+            <root>
+                <GameID>13331</GameID>
+                <CompetitionID>55</CompetitionID>
+                ...
+            </root>
+
+        - **Hierarchical XML structure**:
+
+            <match iId="1" ...>
+                <period iId="1" iStartFrame="1848508" iEndFrame="1916408"/>
+                ...
+            </match>
+
+        - **JSON structure**:
+
+            {
+                "GameID": 1,
+                "CompetitionID": 1,
+                "SeasonID": 2023,
+                ...
+            }
+
+        If parsing fails for a supported format or you encounter an unsupported
+        structure, please create an issue on the kloppy GitHub repository
+        with a sample of the problematic data.
+    """
+    # file format is deprecated
+    if file_format is not None:
+        warnings.warn(
+            "file_format is deprecated. This is now automatically infered.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    deserializer = TRACABDeserializer(
         sample_rate=sample_rate,
         limit=limit,
         coordinate_system=coordinates,
@@ -40,22 +87,3 @@ def load(
         return deserializer.deserialize(
             inputs=TRACABInputs(meta_data=meta_data_fp, raw_data=raw_data_fp)
         )
-
-
-def identify_deserializer(
-    meta_data: FileLike,
-    raw_data: FileLike,
-) -> Union[Type[TRACABDatDeserializer], Type[TRACABJSONDeserializer]]:
-    meta_data_extension = get_file_extension(meta_data)
-    raw_data_extension = get_file_extension(raw_data)
-
-    if meta_data_extension == ".xml" and raw_data_extension == ".dat":
-        deserializer = TRACABDatDeserializer
-    elif meta_data_extension == ".json" and raw_data_extension == ".json":
-        deserializer = TRACABJSONDeserializer
-    else:
-        raise ValueError(
-            "Tracab file format could not be recognized, please specify"
-        )
-
-    return deserializer

kloppy/infra/serializers/tracking/metrica_epts/reader.py

@@ -1,12 +1,11 @@
 import re
-from typing import List, Iterator, IO
 from datetime import timedelta
-
+from typing import IO, Iterator, List
 
 from .models import (
-    PlayerChannel,
     DataFormatSpecification,
     EPTSMetadata,
+    PlayerChannel,
     Sensor,
 )
 

kloppy/infra/serializers/tracking/tracab/__init__.py

@@ -1,9 +1,6 @@
-from .common import TRACABInputs
-from .tracab_dat import TRACABDatDeserializer
-from .tracab_json import TRACABJSONDeserializer
+from .deserializer import TRACABDeserializer, TRACABInputs
 
 __all__ = [
-    "TRACABDatDeserializer",
-    "TRACABJSONDeserializer",
+    "TRACABDeserializer",
     "TRACABInputs",
 ]

kloppy/infra/serializers/tracking/tracab/common.py

@@ -0,0 +1,110 @@
+import logging
+import warnings
+from typing import IO, NamedTuple, Optional, Union
+
+from kloppy.domain import (
+    AttackingDirection,
+    DatasetFlag,
+    Metadata,
+    Orientation,
+    Provider,
+    TrackingDataset,
+    attacking_direction_from_frame,
+)
+from kloppy.utils import performance_logging
+
+from ..deserializer import TrackingDataDeserializer
+from .parsers import get_metadata_parser, get_raw_data_parser
+
+logger = logging.getLogger(__name__)
+
+
+class TRACABInputs(NamedTuple):
+    meta_data: IO[bytes]
+    raw_data: IO[bytes]
+
+
+class TRACABDeserializer(TrackingDataDeserializer[TRACABInputs]):
+    def __init__(
+        self,
+        limit: Optional[int] = None,
+        sample_rate: Optional[float] = None,
+        coordinate_system: Optional[Union[str, Provider]] = None,
+        only_alive: bool = True,
+    ):
+        super().__init__(limit, sample_rate, coordinate_system)
+        self.only_alive = only_alive
+
+    @property
+    def provider(self) -> Provider:
+        return Provider.TRACAB
+
+    def deserialize(self, inputs: TRACABInputs) -> TrackingDataset:
+        with performance_logging("Loading metadata", logger=logger):
+            metadata_parser = get_metadata_parser(inputs.meta_data)
+            (
+                pitch_length,
+                pitch_width,
+            ) = metadata_parser.extract_pitch_dimensions()
+            teams = metadata_parser.extract_lineups()
+            periods = metadata_parser.extract_periods()
+            frame_rate = metadata_parser.extract_frame_rate()
+            date = metadata_parser.extract_date()
+            game_id = metadata_parser.extract_game_id()
+            orientation = metadata_parser.extract_orientation()
+
+        transformer = self.get_transformer(
+            pitch_length=pitch_length, pitch_width=pitch_width
+        )
+
+        with performance_logging("Loading data", logger=logger):
+            raw_data_parser = get_raw_data_parser(
+                inputs.raw_data, periods, teams, frame_rate
+            )
+            frames = []
+            for n, frame in enumerate(
+                raw_data_parser.extract_frames(
+                    self.sample_rate, self.only_alive
+                )
+            ):
+                frame = transformer.transform_frame(frame)
+                frames.append(frame)
+
+                if self.limit and n + 1 >= (self.limit / self.sample_rate):
+                    break
+
+        if orientation is None:
+            try:
+                first_frame = next(
+                    frame for frame in frames if frame.period.id == 1
+                )
+                orientation = (
+                    Orientation.HOME_AWAY
+                    if attacking_direction_from_frame(first_frame)
+                    == AttackingDirection.LTR
+                    else Orientation.AWAY_HOME
+                )
+            except StopIteration:
+                warnings.warn(
+                    "Could not determine orientation of dataset, defaulting to NOT_SET"
+                )
+                orientation = Orientation.NOT_SET
+
+        metadata = Metadata(
+            teams=list(teams),
+            periods=periods,
+            pitch_dimensions=transformer.get_to_coordinate_system().pitch_dimensions,
+            score=None,
+            frame_rate=frame_rate,
+            orientation=orientation,
+            provider=Provider.TRACAB,
+            flags=DatasetFlag.BALL_OWNING_TEAM | DatasetFlag.BALL_STATE,
+            coordinate_system=transformer.get_to_coordinate_system(),
+            date=date,
+            game_id=game_id,
+        )
+
+        return TrackingDataset(
+            records=frames,
+            metadata=metadata,
+        )

kloppy/infra/serializers/tracking/tracab/deserializer.py

@@ -0,0 +1,62 @@
+from typing import IO, List, Optional, Tuple
+
+from lxml import objectify
+
+from kloppy.domain import Period, Team
+
+from .metadata.base import TracabMetadataParser
+from .metadata.flat_xml import TracabFlatXMLMetadataParser
+from .metadata.hierarchical_xml import TracabHierarchicalXMLMetadataParser
+from .metadata.json import TracabJSONMetadataParser
+from .raw_data.base import TracabDataParser
+from .raw_data.dat import TracabDatParser
+from .raw_data.json import TracabJSONParser
+
+
+def get_metadata_parser(
+    feed: IO[bytes], feed_format: Optional[str] = None
+) -> TracabMetadataParser:
+    # infer the data format if not provided
+    if feed_format is None:
+        if feed.read(1).decode("utf-8")[0] == "<":
+            feed.seek(0)
+            meta_data = objectify.fromstring(feed.read())
+            if hasattr(meta_data, "match"):
+                feed_format = "HIERARCHICAL_XML"
+            else:
+                feed_format = "FLAT_XML"
+        else:
+            feed_format = "JSON"
+        feed.seek(0)
+
+    if feed_format.upper() == "JSON":
+        return TracabJSONMetadataParser(feed)
+    elif feed_format.upper() == "FLAT_XML":
+        return TracabFlatXMLMetadataParser(feed)
+    elif feed_format.upper() == "HIERARCHICAL_XML":
+        return TracabHierarchicalXMLMetadataParser(feed)
+    else:
+        raise ValueError(f"Unknown metadata feed format {feed_format}")
+
+
+def get_raw_data_parser(
+    feed: IO[bytes],
+    periods: List[Period],
+    teams: Tuple[Team, Team],
+    frame_rate: int,
+    feed_format: Optional[str] = None,
+) -> TracabDataParser:
+    # infer the data format if not provided
+    if feed_format is None:
+        if feed.read(1).decode("utf-8")[0] == "{":
+            feed_format = "JSON"
+        else:
+            feed_format = "DAT"
+        feed.seek(0)
+
+    if feed_format.upper() == "DAT":
+        return TracabDatParser(feed, periods, teams, frame_rate)
+    elif feed_format.upper() == "JSON":
+        return TracabJSONParser(feed, periods, teams, frame_rate)
+    else:
+        raise ValueError(f"Unknown raw data feed format {feed_format}")

kloppy/infra/serializers/tracking/tracab/parsers/__init__.py

@@ -0,0 +1,58 @@
+"""Base class for all parsers that can handle Tracab metadata files.
+
+A parser reads a single metadata file and should extend the 'TracabMetadataParser'
+class to extract the data about periods, lineups, pitch dimensions, etc.
+"""
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import IO, List, Optional, Tuple
+
+from kloppy.domain import Orientation, Period, Score, Team
+
+
+class TracabMetadataParser(ABC):
+    """Extract data from a tracab metadata file."""
+
+    def __init__(self, feed: IO[bytes]) -> None:
+        """Initialize the parser with the data stream.
+
+        Args:
+            feed : The metadata of a game to parse.
+        """
+
+    @abstractmethod
+    def extract_periods(self) -> List[Period]:
+        """Extract the periods of the game."""
+
+    def extract_score(self) -> Optional[Score]:
+        """Extract the game's score."""
+        return None
+
+    def extract_date(self) -> Optional[datetime]:
+        """Extract the game's date."""
+        return None
+
+    def extract_game_week(self) -> Optional[str]:
+        """Extract the game week."""
+        return None
+
+    def extract_game_id(self) -> Optional[str]:
+        """Extract the game's id."""
+        return None
+
+    @abstractmethod
+    def extract_lineups(self) -> Tuple[Team, Team]:
+        """Extract the home and away team."""
+
+    @abstractmethod
+    def extract_pitch_dimensions(self) -> Tuple[float, float]:
+        """Extract the pitch size as (length, width)."""
+
+    @abstractmethod
+    def extract_frame_rate(self) -> int:
+        """Extract the tracking data's frame rate."""
+
+    def extract_orientation(self) -> Optional[Orientation]:
+        """Extract the orientation of the data."""
+        return None