Skip to content

Commit

Permalink
Set dim order in Validators
Browse files Browse the repository at this point in the history
  • Loading branch information
@lochhh
lochhh committed Nov 26, 2024
1 parent de2414a commit 7dd22c8
Show file tree
Hide file tree
Showing 7 changed files with 113 additions and 87 deletions.
15 changes: 7 additions & 8 deletions movement/io/load_bboxes.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -398,8 +398,7 @@ 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).squeeze()
array_dict[key] = np.stack(list_arrays, axis=-1).squeeze()

# Transform position_array to represent centroid of bbox,
# rather than top-left corner
Expand Down Expand Up @@ -629,21 +628,21 @@ def _ds_from_valid_data(data: ValidBboxesDataset) -> xr.Dataset:
time_unit = "seconds"

# Convert data to an xarray.Dataset
# with dimensions ('time', 'individuals', 'space')
# with dimensions ('time', 'space', 'individuals')
DIM_NAMES = ValidBboxesDataset.DIM_NAMES
n_space = data.position_array.shape[-1]
n_space = data.position_array.shape[1]
return xr.Dataset(
data_vars={
"position": xr.DataArray(data.position_array, dims=DIM_NAMES),
"shape": xr.DataArray(data.shape_array, dims=DIM_NAMES),
"confidence": xr.DataArray(
data.confidence_array, dims=DIM_NAMES[:-1]
data.confidence_array, dims=DIM_NAMES[:1] + DIM_NAMES[2:]
),
},
coords={
DIM_NAMES[0]: time_coords,
DIM_NAMES[1]: data.individual_names,
DIM_NAMES[2]: ["x", "y", "z"][:n_space],
DIM_NAMES[1]: ["x", "y", "z"][:n_space],
DIM_NAMES[2]: data.individual_names,
},
attrs={
"fps": data.fps,
Expand All @@ -652,4 +651,4 @@ def _ds_from_valid_data(data: ValidBboxesDataset) -> xr.Dataset:
"source_file": None,
"ds_type": "bboxes",
},
).transpose("time", "space", "individuals")
)
49 changes: 26 additions & 23 deletions movement/io/load_poses.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,15 +194,17 @@ def from_dlc_style_df(
df.columns.get_level_values("bodyparts").unique().to_list()
)

# reshape the data into (n_frames, n_individuals, n_keypoints, 3)
# where the last axis contains "x", "y", "likelihood"
tracks_with_scores = df.to_numpy().reshape(
(-1, len(individual_names), len(keypoint_names), 3)
# reshape the data into (n_frames, 3, n_keypoints, n_individuals)
# where the second axis contains "x", "y", "likelihood"
tracks_with_scores = (
df.to_numpy()
.reshape((-1, len(individual_names), len(keypoint_names), 3))
.transpose(0, 3, 2, 1)
)

return from_numpy(
position_array=tracks_with_scores[:, :, :, :-1],
confidence_array=tracks_with_scores[:, :, :, -1],
position_array=tracks_with_scores[:, :-1, :, :],
confidence_array=tracks_with_scores[:, -1, :, :],
individual_names=individual_names,
keypoint_names=keypoint_names,
fps=fps,
Expand Down Expand Up @@ -460,10 +462,10 @@ def _ds_from_sleap_analysis_file(
file = ValidHDF5(file_path, expected_datasets=["tracks"])

with h5py.File(file.path, "r") as f:
# transpose to shape: (n_frames, n_tracks, n_keypoints, n_space)
tracks = f["tracks"][:].transpose((3, 0, 2, 1))
# transpose to shape: (n_frames, n_space, n_keypoints, n_tracks)
tracks = f["tracks"][:].transpose((3, 1, 2, 0))
# Create an array of NaNs for the confidence scores
scores = np.full(tracks.shape[:-1], np.nan)
scores = np.full(tracks.shape[:1] + tracks.shape[2:], np.nan)
individual_names = [n.decode() for n in f["track_names"][:]] or None
if individual_names is None:
log_warning(
Expand All @@ -472,9 +474,9 @@ def _ds_from_sleap_analysis_file(
"default individual name."
)
# If present, read the point-wise scores,
# and transpose to shape: (n_frames, n_tracks, n_keypoints)
# and transpose to shape: (n_frames, n_keypoints, n_tracks)
if "point_scores" in f:
scores = f["point_scores"][:].transpose((2, 0, 1))
scores = f["point_scores"][:].T
return from_numpy(
position_array=tracks.astype(np.float32),
confidence_array=scores.astype(np.float32),
Expand Down Expand Up @@ -516,8 +518,8 @@ def _ds_from_sleap_labels_file(
"default individual name."
)
return from_numpy(
position_array=tracks_with_scores[:, :, :, :-1],
confidence_array=tracks_with_scores[:, :, :, -1],
position_array=tracks_with_scores[:, :-1, :, :],
confidence_array=tracks_with_scores[:, -1, :, :],
individual_names=individual_names,
keypoint_names=[kp.name for kp in labels.skeletons[0].nodes],
fps=fps,
Expand All @@ -538,7 +540,8 @@ def _sleap_labels_to_numpy(labels: Labels) -> np.ndarray:
Returns
-------
numpy.ndarray
A NumPy array containing pose tracks and confidence scores.
A NumPy array containing pose tracks and confidence scores,
with shape ``(n_frames, 3, n_nodes, n_tracks)``.
Notes
-----
Expand Down Expand Up @@ -568,7 +571,7 @@ def _sleap_labels_to_numpy(labels: Labels) -> np.ndarray:
skeleton = labels.skeletons[-1] # Assume project only uses last skeleton
n_nodes = len(skeleton.nodes)
n_frames = int(last_frame - first_frame + 1)
tracks = np.full((n_frames, n_tracks, n_nodes, 3), np.nan, dtype="float32")
tracks = np.full((n_frames, 3, n_nodes, n_tracks), np.nan, dtype="float32")

for lf in lfs:
i = int(lf.frame_idx - first_frame)
Expand All @@ -584,12 +587,12 @@ def _sleap_labels_to_numpy(labels: Labels) -> np.ndarray:
# Use user-labelled instance if available
if user_track_instances:
inst = user_track_instances[-1]
tracks[i, j] = np.hstack(
tracks[i, ..., j] = np.hstack(
(inst.numpy(), np.full((n_nodes, 1), np.nan))
)
).T
elif predicted_track_instances:
inst = predicted_track_instances[-1]
tracks[i, j] = inst.numpy(scores=True)
tracks[i, ..., j] = inst.numpy(scores=True).T
return tracks


Expand Down Expand Up @@ -679,7 +682,7 @@ def _ds_from_valid_data(data: ValidPosesDataset) -> xr.Dataset:
"""
n_frames = data.position_array.shape[0]
n_space = data.position_array.shape[-1]
n_space = data.position_array.shape[1]

# Create the time coordinate, depending on the value of fps
time_coords = np.arange(n_frames, dtype=int)
Expand All @@ -694,14 +697,14 @@ def _ds_from_valid_data(data: ValidPosesDataset) -> xr.Dataset:
data_vars={
"position": xr.DataArray(data.position_array, dims=DIM_NAMES),
"confidence": xr.DataArray(
data.confidence_array, dims=DIM_NAMES[:-1]
data.confidence_array, dims=DIM_NAMES[:1] + DIM_NAMES[2:]
),
},
coords={
DIM_NAMES[0]: time_coords,
DIM_NAMES[3]: ["x", "y", "z"][:n_space],
DIM_NAMES[1]: ["x", "y", "z"][:n_space],
DIM_NAMES[2]: data.keypoint_names,
DIM_NAMES[1]: data.individual_names,
DIM_NAMES[3]: data.individual_names,
},
attrs={
"fps": data.fps,
Expand All @@ -710,4 +713,4 @@ def _ds_from_valid_data(data: ValidPosesDataset) -> xr.Dataset:
"source_file": None,
"ds_type": "poses",
},
).transpose("time", "space", "keypoints", "individuals")
)
58 changes: 38 additions & 20 deletions movement/validators/datasets.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,32 +143,39 @@ class ValidPosesDataset:
)

# Class variables
DIM_NAMES: ClassVar[tuple] = ("time", "individuals", "keypoints", "space")
DIM_NAMES: ClassVar[tuple] = ("time", "space", "keypoints", "individuals")
VAR_NAMES: ClassVar[tuple] = ("position", "confidence")

# Add validators
@position_array.validator
def _validate_position_array(self, attribute, value):
_validate_type_ndarray(value)
if value.ndim != 4:
n_dims = value.ndim
if n_dims != 4:
raise log_error(
ValueError,
f"Expected '{attribute.name}' to have 4 dimensions, "
f"but got {value.ndim}.",
f"but got {n_dims}.",
)
if value.shape[-1] not in [2, 3]:
space_dim_shape = value.shape[1]
if space_dim_shape not in [2, 3]:
raise log_error(
ValueError,
f"Expected '{attribute.name}' to have 2 or 3 spatial "
f"dimensions, but got {value.shape[-1]}.",
f"dimensions, but got {space_dim_shape}.",
)

@confidence_array.validator
def _validate_confidence_array(self, attribute, value):
if value is not None:
_validate_type_ndarray(value)
# Expected shape is the same as position_array,
# but without the `space` dim
expected_shape = (
self.position_array.shape[:1] + self.position_array.shape[2:]
)
_validate_array_shape(
attribute, value, expected_shape=self.position_array.shape[:-1]
attribute, value, expected_shape=expected_shape
)

@individual_names.validator
Expand All @@ -178,7 +185,7 @@ def _validate_individual_names(self, attribute, value):
_validate_list_length(attribute, value, 1)
else:
_validate_list_length(
attribute, value, self.position_array.shape[1]
attribute, value, self.position_array.shape[-1]
)

@keypoint_names.validator
Expand All @@ -187,25 +194,28 @@ def _validate_keypoint_names(self, attribute, value):

def __attrs_post_init__(self):
"""Assign default values to optional attributes (if None)."""
position_array_shape = self.position_array.shape
if self.confidence_array is None:
self.confidence_array = np.full(
(self.position_array.shape[:-1]), np.nan, dtype="float32"
position_array_shape[:1] + position_array_shape[2:],
np.nan,
dtype="float32",
)
log_warning(
"Confidence array was not provided."
"Setting to an array of NaNs."
)
if self.individual_names is None:
self.individual_names = [
f"individual_{i}" for i in range(self.position_array.shape[1])
f"individual_{i}" for i in range(position_array_shape[-1])
]
log_warning(
"Individual names were not provided. "
f"Setting to {self.individual_names}."
)
if self.keypoint_names is None:
self.keypoint_names = [
f"keypoint_{i}" for i in range(self.position_array.shape[2])
f"keypoint_{i}" for i in range(position_array_shape[2])
]
log_warning(
"Keypoint names were not provided. "
Expand Down Expand Up @@ -297,28 +307,30 @@ class ValidBboxesDataset:
validator=validators.optional(validators.instance_of(str)),
)

DIM_NAMES: ClassVar[tuple] = ("time", "individuals", "space")
DIM_NAMES: ClassVar[tuple] = ("time", "space", "individuals")
VAR_NAMES: ClassVar[tuple] = ("position", "shape", "confidence")

# Validators
@position_array.validator
@shape_array.validator
def _validate_position_and_shape_arrays(self, attribute, value):
_validate_type_ndarray(value)
# check last dimension (spatial) has 2 coordinates
# check `space` dim (at idx 1) has 2 coordinates
n_expected_spatial_coordinates = 2
if value.shape[-1] != n_expected_spatial_coordinates:
n_spatial_coordinates = value.shape[1]
if n_spatial_coordinates != n_expected_spatial_coordinates:
raise log_error(
ValueError,
f"Expected '{attribute.name}' to have 2 spatial coordinates, "
f"but got {value.shape[-1]}.",
f"Expected '{attribute.name}' to have "
f"{n_expected_spatial_coordinates} spatial coordinates, "
f"but got {n_spatial_coordinates}.",
)

@individual_names.validator
def _validate_individual_names(self, attribute, value):
if value is not None:
_validate_list_length(
attribute, value, self.position_array.shape[1]
attribute, value, self.position_array.shape[-1]
)
# check n_individual_names are unique
# NOTE: combined with the requirement above, we are enforcing
Expand All @@ -335,8 +347,13 @@ def _validate_individual_names(self, attribute, value):
def _validate_confidence_array(self, attribute, value):
if value is not None:
_validate_type_ndarray(value)
# Expected shape is the same as position_array,
# but without the `space` dim
expected_shape = (
self.position_array.shape[:1] + self.position_array.shape[2:]
)
_validate_array_shape(
attribute, value, expected_shape=self.position_array.shape[:-1]
attribute, value, expected_shape=expected_shape
)

@frame_array.validator
Expand Down Expand Up @@ -364,10 +381,11 @@ def __attrs_post_init__(self):
If no individual names are provided, assign them unique IDs per frame,
starting with 0 ("id_0").
"""
position_array_shape = self.position_array.shape
# assign default confidence_array
if self.confidence_array is None:
self.confidence_array = np.full(
(self.position_array.shape[:-1]),
position_array_shape[:1] + position_array_shape[2:],
np.nan,
dtype="float32",
)
Expand All @@ -378,7 +396,7 @@ def __attrs_post_init__(self):
# assign default individual_names
if self.individual_names is None:
self.individual_names = [
f"id_{i}" for i in range(self.position_array.shape[1])
f"id_{i}" for i in range(position_array_shape[-1])
]
log_warning(
"Individual names for the bounding boxes "
Expand All @@ -388,7 +406,7 @@ def __attrs_post_init__(self):
)
# assign default frame_array
if self.frame_array is None:
n_frames = self.position_array.shape[0]
n_frames = position_array_shape[0]
self.frame_array = np.arange(n_frames).reshape(-1, 1)
log_warning(
"Frame numbers were not provided. "
Expand Down
Loading

0 comments on commit 7dd22c8

Please sign in to comment.