From 4d5a4de406c767f7daaaaf4f22b82acac3e2eff9 Mon Sep 17 00:00:00 2001 From: Victor Joos Date: Mon, 26 Feb 2024 17:39:54 +0100 Subject: [PATCH] Create json a bit faster and less memory intensive --- .../wrappers/datasets/soccernet/soccernet_game_state.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py b/tracklab/wrappers/datasets/soccernet/soccernet_game_state.py index 17eb1ae8..d4db6b8d 100644 --- 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