From f58a9681bf108b3578b015118f8fb21f4df3cc4c Mon Sep 17 00:00:00 2001 From: sfmig <33267254+sfmig@users.noreply.github.com> Date: Thu, 12 Sep 2024 14:18:45 +0100 Subject: [PATCH] Fix for single individual --- movement/io/load_bboxes.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/movement/io/load_bboxes.py b/movement/io/load_bboxes.py index 6971de0f..70d95135 100644 --- a/movement/io/load_bboxes.py +++ b/movement/io/load_bboxes.py @@ -327,7 +327,7 @@ def from_via_tracks_file( shape_array=bboxes_arrays["shape_array"], confidence_array=bboxes_arrays["confidence_array"], individual_names=[ - f"id_{id}" for id in bboxes_arrays["ID_array"].squeeze() + f"id_{id.item()}" for id in bboxes_arrays["ID_array"] ], frame_array=( bboxes_arrays["frame_array"] @@ -399,8 +399,11 @@ def _numpy_arrays_from_via_tracks_file(file_path: Path) -> dict: df[map_key_to_columns[key]].to_numpy(), indices_id_switch, # indices along axis=0 ) + array_dict[key] = np.stack(list_arrays, axis=1) - array_dict[key] = np.stack(list_arrays, axis=1).squeeze() + # squeeze only last dimension if it is 1 + if array_dict[key].shape[-1] == 1: + array_dict[key] = array_dict[key].squeeze(axis=-1) # Add remaining arrays to dict array_dict["ID_array"] = df["ID"].unique().reshape(-1, 1)