Create json a bit faster and less memory intensive

TrackingLaboratory · Feb 26, 2024 · 4d5a4de · 4d5a4de
1 parent 2daf968
commit 4d5a4de
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py b/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import time
 import zipfile
 import numpy as np
 import pandas as pd
@@ -84,7 +85,6 @@ def save_for_eval(self,
     @staticmethod
     def soccernet_encoding(dataframe: pd.DataFrame, supercategory):
         dataframe["supercategory"] = supercategory
-        dataframe = dataframe.map(lambda x: x.tolist() if isinstance(x, np.ndarray) else x)
         dataframe = dataframe.replace({np.nan: None})
         if supercategory == "object":
             # Remove detections that don't have mandatory columns
@@ -100,10 +100,7 @@ def soccernet_encoding(dataframe: pd.DataFrame, supercategory):
             )
             dataframe = dataframe.rename(columns={"bbox_ltwh": "bbox_image", "jersey_number": "jersey"})
             dataframe["track_id"] = dataframe["track_id"]
-            dataframe["attributes"] = dataframe.apply(
-                lambda x: x[x.index.intersection(["role", "jersey", "team"])].to_dict(),
-                axis=1
-            )
+            dataframe["attributes"] = [{"role": x.get("role"), "jersey": x.get("jersey"), "team": x.get("team")} for n, x in dataframe.iterrows()]
             dataframe["id"] = dataframe.index
             dataframe = dataframe[dataframe.columns.intersection(
                 ["id", "image_id", "video_id", "track_id", "supercategory",
@@ -127,6 +124,8 @@ def soccernet_encoding(dataframe: pd.DataFrame, supercategory):
         dataframe["video_id"] = dataframe["video_id"].apply(str)
         dataframe["image_id"] = dataframe["image_id"].apply(str)
         dataframe["id"] = dataframe["id"].apply(str)
+        dataframe = dataframe.map(
+            lambda x: x.tolist() if isinstance(x, np.ndarray) else x)
         return dataframe