Skip to content

Commit

Permalink
Fix for single individual
Browse files Browse the repository at this point in the history
  • Loading branch information
sfmig committed Sep 12, 2024
1 parent f4f46df commit f58a968
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions movement/io/load_bboxes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit f58a968

Please sign in to comment.