Merge pull request #62 from UnravelSports/feat/player-ids

no ball label

Author: UnravelSports

unravel/soccer/graphs/graph_converter.py

@@ -676,6 +676,7 @@ def plot(
         anonymous: bool = False,
         plot_type: Literal["pitch_only", "graph_only", "full"] = "full",
         show_label: bool = True,
+        show_ball_label: bool = False,
         show_timestamp: bool = True,
         next_closest_timestamp: bool = False,
     ):
@@ -790,7 +791,8 @@ def plot(
         self._ball_color = ball_color
         self._color_by = color_by
         self._plot_type = plot_type
-        self._show_pitch_label = show_label
+        self._show_label = show_label
+        self._show_ball_label = show_ball_label
         self._show_pitch_timestamp = show_timestamp
         self._next_closest_timestamp = next_closest_timestamp
 
@@ -1110,10 +1112,14 @@ def player_and_ball(frame_data, ax):
                         (
                             self.get_player_by_id(r[Column.OBJECT_ID])["jersey_no"]
                             if r[Column.OBJECT_ID] != Constant.BALL
-                            else Constant.BALL
+                            else Constant.BALL if self._show_ball_label else ""
                         )
                         if not anonymous
-                        else str(i)
+                        else (
+                            str(i)
+                            if r[Column.OBJECT_ID] != Constant.BALL
+                            else Constant.BALL if self._show_ball_label else ""
+                        )
                     ),
                     color=self._ball_color if is_ball else color,
                     fontsize=12,
@@ -1132,7 +1138,7 @@ def player_and_ball(frame_data, ax):
                 )
 
             # Add label and timestamp to pitch if enabled
-            if self._show_pitch_label:
+            if self._show_label:
                 ax.set_xlabel(f"Label: {frame_data['label'][0]}", fontsize=22)
             if self._show_pitch_timestamp:
                 ax.set_title(self._gameclock, fontsize=22)