From 340e5f6e1337c5046e1b0b2139791e1875701c4d Mon Sep 17 00:00:00 2001
From: Pieter Robberechts <pieter.robberechts@kuleuven.be>
Date: Tue, 17 Dec 2024 11:26:37 +0100
Subject: [PATCH] fix(SkillCorner): load data error for game ids 2269 and 3442

Fixes #192
---
 .../infra/serializers/tracking/skillcorner.py | 24 +++++++++++--------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/kloppy/infra/serializers/tracking/skillcorner.py b/kloppy/infra/serializers/tracking/skillcorner.py
index b5cc0306..d67b6d96 100644
--- a/kloppy/infra/serializers/tracking/skillcorner.py
+++ b/kloppy/infra/serializers/tracking/skillcorner.py
@@ -1,15 +1,15 @@
+import json
 import logging
-from datetime import timedelta, timezone
-from dateutil.parser import parse
 import warnings
-from typing import NamedTuple, IO, Optional, Union, Dict
 from collections import Counter
-import numpy as np
-import json
+from datetime import timedelta, timezone
 from pathlib import Path
+from typing import IO, Dict, NamedTuple, Optional, Union
+
+import numpy as np
+from dateutil.parser import parse
 
 from kloppy.domain import (
-    attacking_direction_from_frame,
     AttackingDirection,
     DatasetFlag,
     Frame,
@@ -18,6 +18,7 @@
     Orientation,
     Period,
     Player,
+    PlayerData,
     Point,
     Point3D,
     PositionType,
@@ -25,7 +26,7 @@
     Score,
     Team,
     TrackingDataset,
-    PlayerData,
+    attacking_direction_from_frame,
 )
 from kloppy.infra.serializers.tracking.deserializer import (
     TrackingDataDeserializer,
@@ -133,15 +134,18 @@ def _get_frame_data(
             track_id = frame_record.get("track_id", None)
             group_name = frame_record.get("group_name", None)
 
-            if trackable_object == ball_id:
-                group_name = "ball"
+            if trackable_object == ball_id or group_name == "balls":
+                group_name = "balls"
                 z = frame_record.get("z")
                 if z is not None:
                     z = float(z)
                 ball_coordinates = Point3D(x=float(x), y=float(y), z=z)
                 continue
 
-            elif trackable_object in referee_dict.keys():
+            elif (
+                trackable_object in referee_dict.keys()
+                or group_name == "referee"
+            ):
                 group_name = "referee"
                 continue  # Skip Referee Coords