From 7dd22c8efb691d2fd13af5d63953134960fccd9b Mon Sep 17 00:00:00 2001 From: lochhh Date: Tue, 26 Nov 2024 19:05:20 +0000 Subject: [PATCH] Set dim order in Validators --- movement/io/load_bboxes.py | 15 +++-- movement/io/load_poses.py | 49 ++++++++-------- movement/validators/datasets.py | 58 ++++++++++++------- tests/conftest.py | 41 +++++++------ tests/test_unit/test_load_bboxes.py | 18 +++--- tests/test_unit/test_load_poses.py | 11 ++-- .../test_datasets_validators.py | 8 +-- 7 files changed, 113 insertions(+), 87 deletions(-) diff --git a/movement/io/load_bboxes.py b/movement/io/load_bboxes.py index 5a75306d8..5e65e6fe3 100644 --- a/movement/io/load_bboxes.py +++ b/movement/io/load_bboxes.py @@ -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 @@ -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, @@ -652,4 +651,4 @@ def _ds_from_valid_data(data: ValidBboxesDataset) -> xr.Dataset: "source_file": None, "ds_type": "bboxes", }, - ).transpose("time", "space", "individuals") + ) diff --git a/movement/io/load_poses.py b/movement/io/load_poses.py index 34fc6aa8b..b46ba81b7 100644 --- a/movement/io/load_poses.py +++ b/movement/io/load_poses.py @@ -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, @@ -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( @@ -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), @@ -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, @@ -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 ----- @@ -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) @@ -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 @@ -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) @@ -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, @@ -710,4 +713,4 @@ def _ds_from_valid_data(data: ValidPosesDataset) -> xr.Dataset: "source_file": None, "ds_type": "poses", }, - ).transpose("time", "space", "keypoints", "individuals") + ) diff --git a/movement/validators/datasets.py b/movement/validators/datasets.py index 99a68c102..43dace1db 100644 --- a/movement/validators/datasets.py +++ b/movement/validators/datasets.py @@ -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 @@ -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 @@ -187,9 +194,12 @@ 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." @@ -197,7 +207,7 @@ def __attrs_post_init__(self): ) 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. " @@ -205,7 +215,7 @@ def __attrs_post_init__(self): ) 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. " @@ -297,7 +307,7 @@ 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 @@ -305,20 +315,22 @@ class ValidBboxesDataset: @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 @@ -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 @@ -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", ) @@ -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 " @@ -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. " diff --git a/tests/conftest.py b/tests/conftest.py index baedc31b3..b234f7ffe 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -223,10 +223,10 @@ def valid_bboxes_arrays_all_zeros(): ValidBboxesDataset. """ # define the shape of the arrays - n_frames, n_individuals, n_space = (10, 2, 2) + n_frames, n_space, n_individuals = (10, 2, 2) # build a valid array for position or shape with all zeros - valid_bbox_array_all_zeros = np.zeros((n_frames, n_individuals, n_space)) + valid_bbox_array_all_zeros = np.zeros((n_frames, n_space, n_individuals)) # return as a dict return { @@ -252,21 +252,23 @@ def valid_bboxes_arrays(): - Individual 1 at frames 2, 3 """ # define the shape of the arrays - n_frames, n_individuals, n_space = (10, 2, 2) + n_frames, n_space, n_individuals = (10, 2, 2) # build a valid array for position # make bbox with id_i move along x=((-1)**(i))*y line from the origin # if i is even: along x = y line # if i is odd: along x = -y line # moving one unit along each axis in each frame - position = np.empty((n_frames, n_individuals, n_space)) + position = np.empty((n_frames, n_space, n_individuals)) for i in range(n_individuals): - position[:, i, 0] = np.arange(n_frames) - position[:, i, 1] = (-1) ** i * np.arange(n_frames) + position[:, 0, i] = np.arange(n_frames) + position[:, 1, i] = (-1) ** i * np.arange(n_frames) # build a valid array for constant bbox shape (60, 40) constant_shape = (60, 40) # width, height in pixels - shape = np.tile(constant_shape, (n_frames, n_individuals, 1)) + shape = np.tile(constant_shape, (n_frames, n_individuals, 1)).transpose( + 0, 2, 1 + ) # build an array of confidence values, all 0.9 confidence = np.full((n_frames, n_individuals), 0.9) @@ -304,12 +306,14 @@ def valid_bboxes_dataset( data_vars={ "position": xr.DataArray(position_array, dims=dim_names), "shape": xr.DataArray(shape_array, dims=dim_names), - "confidence": xr.DataArray(confidence_array, dims=dim_names[:-1]), + "confidence": xr.DataArray( + confidence_array, dims=dim_names[:1] + dim_names[2:] + ), }, coords={ dim_names[0]: np.arange(n_frames), - dim_names[1]: [f"id_{id}" for id in range(n_individuals)], - dim_names[2]: ["x", "y"], + dim_names[1]: ["x", "y"], + dim_names[2]: [f"id_{id}" for id in range(n_individuals)], }, attrs={ "fps": None, @@ -367,8 +371,8 @@ def _valid_position_array(array_type): n_individuals = 1 x_points = np.repeat(base * base, n_individuals * n_keypoints) y_points = np.repeat(base * 4, n_individuals * n_keypoints) - position_array = np.ravel(np.column_stack((x_points, y_points))) - return position_array.reshape(n_frames, n_individuals, n_keypoints, 2) + position_array = np.vstack((x_points, y_points)) + return position_array.reshape(n_frames, 2, n_keypoints, n_individuals) return _valid_position_array @@ -383,7 +387,9 @@ def valid_poses_dataset(valid_position_array, request): except AttributeError: array_format = "multi_individual_array" position_array = valid_position_array(array_format) - n_frames, n_individuals, n_keypoints = position_array.shape[:3] + n_frames, n_keypoints, n_individuals = ( + position_array.shape[:1] + position_array.shape[2:] + ) return xr.Dataset( data_vars={ "position": xr.DataArray(position_array, dims=dim_names), @@ -391,8 +397,8 @@ def valid_poses_dataset(valid_position_array, request): np.repeat( np.linspace(0.1, 1.0, n_frames), n_individuals * n_keypoints, - ).reshape(position_array.shape[:-1]), - dims=dim_names[:-1], + ).reshape(position_array.shape[:1] + position_array.shape[2:]), + dims=dim_names[:1] + dim_names[2:], # exclude "space" ), }, coords={ @@ -408,7 +414,7 @@ def valid_poses_dataset(valid_position_array, request): "source_file": "test.h5", "ds_type": "poses", }, - ).transpose("time", "space", "keypoints", "individuals") + ) @pytest.fixture @@ -490,7 +496,8 @@ def valid_poses_dataset_uniform_linear_motion( """Return a valid poses dataset for two individuals moving in uniform linear motion, with 5 frames with low confidence values and time in frames. """ - dim_names = ValidPosesDataset.DIM_NAMES + dim_names = ("time", "individuals", "keypoints", "space") + # ValidPosesDataset.DIM_NAMES position_array = valid_poses_array_uniform_linear_motion["position"] confidence_array = valid_poses_array_uniform_linear_motion["confidence"] diff --git a/tests/test_unit/test_load_bboxes.py b/tests/test_unit/test_load_bboxes.py index f2d944992..4200b1c90 100644 --- a/tests/test_unit/test_load_bboxes.py +++ b/tests/test_unit/test_load_bboxes.py @@ -32,8 +32,8 @@ def valid_from_numpy_inputs_required_arrays(): rng = np.random.default_rng(seed=42) return { - "position_array": rng.random((n_frames, n_individuals, n_space)), - "shape_array": rng.random((n_frames, n_individuals, n_space)), + "position_array": rng.random((n_frames, n_space, n_individuals)), + "shape_array": rng.random((n_frames, n_space, n_individuals)), "confidence_array": rng.random((n_frames, n_individuals)), "individual_names": [ f"id_{id}" for id in individual_names_array.squeeze() @@ -128,13 +128,13 @@ def assert_dataset( # Confidence has the same shape as position, except for the space dim assert dataset.confidence.shape == position_shape[:1] + position_shape[2:] # Check the dims and coords - DIM_NAMES = ValidBboxesDataset.DIM_NAMES - expected_dim_length_dict = { - DIM_NAMES[idx]: position_shape[i] for i, idx in enumerate([0, 2, 1]) - } + dim_names = ValidBboxesDataset.DIM_NAMES + expected_dim_length_dict = dict( + zip(dim_names, position_shape, strict=True) + ) assert expected_dim_length_dict == dataset.sizes # Check the coords - for dim in DIM_NAMES[1:]: + for dim in dim_names[1:]: assert all(isinstance(s, str) for s in dataset.coords[dim].values) assert all(coord in dataset.coords["space"] for coord in ["x", "y"]) # Check the metadata attributes @@ -470,6 +470,6 @@ def test_position_numpy_array_from_via_tracks_file(via_tracks_file): # Compare to extracted position array assert np.allclose( - bboxes_arrays["position_array"], # frames, individuals, xy - np.stack(list_derived_centroids, axis=1), + bboxes_arrays["position_array"], # frames, xy, individuals + np.stack(list_derived_centroids, axis=-1), ) diff --git a/tests/test_unit/test_load_poses.py b/tests/test_unit/test_load_poses.py index df3bed2ad..d51e282ce 100644 --- a/tests/test_unit/test_load_poses.py +++ b/tests/test_unit/test_load_poses.py @@ -82,14 +82,13 @@ def assert_dataset( dataset.confidence.shape == position_shape[:1] + position_shape[2:] ) # Check the dims - DIM_NAMES = ValidPosesDataset.DIM_NAMES - expected_dim_length_dict = { - DIM_NAMES[idx]: position_shape[i] - for i, idx in enumerate([0, 3, 2, 1]) - } + dim_names = ValidPosesDataset.DIM_NAMES + expected_dim_length_dict = dict( + zip(dim_names, position_shape, strict=True) + ) assert expected_dim_length_dict == dataset.sizes # Check the coords - for dim in DIM_NAMES[1:]: + for dim in dim_names[1:]: assert all(isinstance(s, str) for s in dataset.coords[dim].values) assert all(coord in dataset.coords["space"] for coord in ["x", "y"]) # Check the metadata attributes diff --git a/tests/test_unit/test_validators/test_datasets_validators.py b/tests/test_unit/test_validators/test_datasets_validators.py index e41331f77..74ec18b1f 100644 --- a/tests/test_unit/test_validators/test_datasets_validators.py +++ b/tests/test_unit/test_validators/test_datasets_validators.py @@ -75,10 +75,10 @@ def position_array_params(request): f"Expected a numpy array, but got {type(list())}.", ), # not an ndarray ( - np.zeros((10, 2, 3)), + np.zeros((10, 3, 2)), f"Expected '{key}_array' to have 2 spatial " "coordinates, but got 3.", - ), # last dim not 2 + ), # `space` dim (at idx 1) not 2 ] for key in ["position", "shape"] } @@ -101,10 +101,10 @@ def position_array_params(request): "Expected 'position_array' to have 4 dimensions, but got 3.", ), # not 4d ( - np.zeros((10, 2, 3, 4)), + np.zeros((10, 4, 3, 2)), "Expected 'position_array' to have 2 or 3 " "spatial dimensions, but got 4.", - ), # last dim not 2 or 3 + ), # `space` dim (at idx 1) not 2 or 3 ], ) def test_poses_dataset_validator_with_invalid_position_array(