From 0bcb77c2ade4e5cf41317f5cd6914153398495ad Mon Sep 17 00:00:00 2001 From: Tigran Grigoryan Date: Thu, 26 Sep 2024 20:53:30 +0400 Subject: [PATCH] 4161 - Fix TRex dataframe cleanup function Fix 'clean_df' function in 'trex/df_preprocessing.py' module which misbehaves not selecting any existing columns to delete in order to cleanup dataframe for visualization. Resolves: #4161 Signed-off-by: Tigran Grigoryan --- .../experimental/trt-engine-explorer/trex/df_preprocessing.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/experimental/trt-engine-explorer/trex/df_preprocessing.py b/tools/experimental/trt-engine-explorer/trex/df_preprocessing.py index e6126706..e98e987a 100644 --- a/tools/experimental/trt-engine-explorer/trex/df_preprocessing.py +++ b/tools/experimental/trt-engine-explorer/trex/df_preprocessing.py @@ -254,7 +254,8 @@ def drop_columns(df: pd.DataFrame, columns: list): def clean_df(df: pd.DataFrame, inplace=True) -> pd.DataFrame: clean_io(df) - columns = set([col for col_list in layer_attributes.keys() for col in col_list]) + lists = [list(layer.keys()) for layer in layer_attributes.values()] + columns = {key for list_ in lists for key in list_} drop_columns(df, columns) df.fillna(0, inplace=inplace) return df