From b5d91bb4afdf9659458134f42147178d6a7e5d19 Mon Sep 17 00:00:00 2001 From: Liezl Maree <38435167+roomrys@users.noreply.github.com> Date: Tue, 1 Aug 2023 14:18:21 -0700 Subject: [PATCH] Initial implementation of a PoseTraining Group (#9) --- spec/ndx-pose.extensions.yaml | 120 ++++++++++++++++++--- spec/ndx-pose.namespace.yaml | 4 + src/spec/create_extension_spec.py | 172 ++++++++++++++++++++++++++---- 3 files changed, 260 insertions(+), 36 deletions(-) diff --git a/spec/ndx-pose.extensions.yaml b/spec/ndx-pose.extensions.yaml index 0beec02..b67af9c 100644 --- a/spec/ndx-pose.extensions.yaml +++ b/spec/ndx-pose.extensions.yaml @@ -1,4 +1,33 @@ groups: +- neurodata_type_def: Skeleton + neurodata_type_inc: NWBDataInterface + default_name: Skeleton + doc: Group that holds node and edge data for defining parts of a pose and their + connections to one another. + attributes: + - name: id + dtype: text + doc: Unique ID associated with the skeleton. + datasets: + - name: nodes + dtype: text + dims: + - num_body_parts + shape: + - null + doc: Array of body part names corresponding to the names of the PoseEstimationSeries + objects or PoseTraining objects. + - name: edges + dtype: uint8 + dims: + - num_edges + - nodes_index, nodes_index + shape: + - null + - 2 + doc: Array of pairs of indices corresponding to edges between nodes. Index values + correspond to row indices of the 'nodes' dataset. Index values use 0-indexing. + quantity: '?' - neurodata_type_def: PoseEstimationSeries neurodata_type_inc: SpatialSeries doc: Estimated position (x, y) or (x, y, z) of a body part over time. @@ -92,27 +121,88 @@ groups: dtype: text doc: Version string of the software tool used. required: false - - name: nodes + groups: + - neurodata_type_inc: PoseEstimationSeries + doc: Estimated position data for each body part. + quantity: '*' + links: + - target_type: Skeleton + doc: Layout of body part locations and connections. +- neurodata_type_def: TrainingFrame + neurodata_type_inc: NWBDataInterface + default_name: TrainingFrame + doc: Group that holds ground-truth position data for all instances of a skeleton in a single frame. + attributes: + - name: annotator dtype: text - dims: - - num_body_parts - shape: - - null - doc: Array of body part names corresponding to the names of the PoseEstimationSeries - objects within this group. + doc: Name of annotator who labeled the TrainingFrame. + required: false + groups: + - neurodata_type_inc: Instance + doc: Position data for a single instance of a skeleton in a single training frame. + quantity: '*' + - name: source_video + neurodata_type_inc: ImageSeries + doc: Path to original video file and frame used. quantity: '?' - - name: edges + attributes: + - name: frame_index + dtype: uint8 + doc: Frame index of TrainingFrame in original video file. + required: true + - name: source_frame + neurodata_type_inc: Images + doc: Collection of images used for training (stored internally or externally). + quantity: 1 + attributes: + - name: images_index + dtype: uint8 + doc: Index of TrainingFrame in collection of frames stored in Images. + required: true +- neurodata_type_def: Instance + neurodata_type_inc: NWBDataInterface + default_name: Instance + doc: Group that holds ground-truth pose data for single subject in a single frame. + attributes: + - name: id dtype: uint8 + doc: ID used to differentiate instances. + required: false + datasets: + - name: node_locations + dtype: float dims: - - num_edges - - nodes_index, nodes_index + - - num_body_parts + - x, y + - - num_body_parts + - x, y, z + shape: + - - null + - 2 + - - null + - 3 + doc: Locations (x, y) or (x, y, z) of nodes for single skeleton instance in single frame. + - name: node_visibility + dtype: bool + dims: + - num_body_parts shape: - null - - 2 - doc: Array of pairs of indices corresponding to edges between nodes. Index values - correspond to row indices of the 'nodes' dataset. Index values use 0-indexing. + doc: Markers for node visibility where true corresponds to a visible node and + false corresponds to an occluded node. quantity: '?' + links: + - target_type: Skeleton + doc: Layout of body part locations and connections. +- neurodata_type_def: PoseTraining + neurodata_type_inc: NWBDataInterface + default_name: PoseTraining + doc: Group that holds images, ground-truth annotations, and metadata for training + a pose estimator. groups: - - neurodata_type_inc: PoseEstimationSeries - doc: Estimated position data for each body part. + - neurodata_type_inc: Skeleton + doc: Skeleton used in project where each skeleton corresponds to a unique morphology. + quantity: '*' + - neurodata_type_inc: TrainingFrame + doc: Frames and ground-truth annotations for training a pose estimator. quantity: '*' diff --git a/spec/ndx-pose.namespace.yaml b/spec/ndx-pose.namespace.yaml index 7146cd3..d249699 100644 --- a/spec/ndx-pose.namespace.yaml +++ b/spec/ndx-pose.namespace.yaml @@ -3,10 +3,14 @@ namespaces: - Ryan Ly - Ben Dichter - Alexander Mathis + - Liezl Maree + - Chris Brozdowski contact: - rly@lbl.gov - bdichter@lbl.gov - alexander.mathis@epfl.ch + - lmaree@salk.edu + - cbroz@datajoint.com doc: NWB extension to store pose estimation data name: ndx-pose schema: diff --git a/src/spec/create_extension_spec.py b/src/spec/create_extension_spec.py index e0f8b6d..ee086d5 100644 --- a/src/spec/create_extension_spec.py +++ b/src/spec/create_extension_spec.py @@ -10,13 +10,47 @@ def main(): doc='NWB extension to store pose estimation data', name='ndx-pose', version='0.1.1', - author=['Ryan Ly', 'Ben Dichter', 'Alexander Mathis'], - contact=['rly@lbl.gov', 'bdichter@lbl.gov', 'alexander.mathis@epfl.ch'], + author=['Ryan Ly', 'Ben Dichter', 'Alexander Mathis', 'Liezl Maree', 'Chris Brozdowski'], + contact=['rly@lbl.gov', 'bdichter@lbl.gov', 'alexander.mathis@epfl.ch', 'lmaree@salk.edu', 'cbroz@datajoint.com'], ) ns_builder.include_type('SpatialSeries', namespace='core') ns_builder.include_type('NWBDataInterface', namespace='core') + skeleton = NWBGroupSpec( + neurodata_type_def='Skeleton', + neurodata_type_inc='NWBDataInterface', + doc='Group that holds node and edge data for defining parts of a pose and their connections to one another.', + default_name='Skeleton', + attributes=[ + NWBAttributeSpec( + name='id', + doc='Unique ID associated with the skeleton.', + dtype='text', + ), + ], + datasets=[ + NWBDatasetSpec( + name='nodes', + doc=('Array of body part names corresponding to the names of the PoseEstimationSeries objects or ' + 'PoseTraining objects.'), + dtype='text', + dims=['num_body_parts'], + shape=[None], + quantity=1, + ), + NWBDatasetSpec( + name='edges', + doc=("Array of pairs of indices corresponding to edges between nodes. Index values correspond to row " + "indices of the 'nodes' dataset. Index values use 0-indexing."), + dtype='uint8', + dims=['num_edges', 'nodes_index, nodes_index'], + shape=[None, 2], + quantity='?', + ), + ] + ) + pose_estimation_series = NWBGroupSpec( neurodata_type_def='PoseEstimationSeries', neurodata_type_inc='SpatialSeries', @@ -72,6 +106,13 @@ def main(): quantity='*', ), ], + links=[ + NWBLinkSpec( + doc='Layout of body part locations and connections.', + target_type='Skeleton', + quantity=1 + ), + ], datasets=[ NWBDatasetSpec( name='description', @@ -123,24 +164,6 @@ def main(): ), ], ), - NWBDatasetSpec( - name='nodes', - doc=('Array of body part names corresponding to the names of the PoseEstimationSeries objects within ' - 'this group.'), - dtype='text', - dims=['num_body_parts'], - shape=[None], - quantity='?', - ), - NWBDatasetSpec( - name='edges', - doc=("Array of pairs of indices corresponding to edges between nodes. Index values correspond to row " - "indices of the 'nodes' dataset. Index values use 0-indexing."), - dtype='uint8', - dims=['num_edges', 'nodes_index, nodes_index'], - shape=[None, 2], - quantity='?', - ), ], # TODO: collections of multiple links is currently buggy in PyNWB/HDMF # links=[ @@ -152,7 +175,114 @@ def main(): # ], ) - new_data_types = [pose_estimation_series, pose_estimation] + training_frame = NWBGroupSpec( + neurodata_type_def='TrainingFrame', + neurodata_type_inc='NWBDataInterface', + doc='Group that holds ground-truth position data for all instances of a skeleton in a single frame.', + default_name='TrainingFrame', + groups=[ + NWBGroupSpec( + neurodata_type_inc='Instance', + doc='Position data for a single instance of a skeleton in a single training frame.', + quantity='*', + ), + NWBGroupSpec( + name='source_video', + doc='Path to original video file and frame used.', + quantity='?', + attributes=[ + NWBAttributeSpec( + name='path', + doc='Path to original video file.', + dtype='text', + required=False, + ), + NWBAttributeSpec( + name='frame_index', + doc='Frame index of TrainingFrame in original video file.', + dtype='uint8', + required=False, + ), + ], + ), + NWBGroupSpec( + neurodata_type_inc='Image', + name='source_frame', + doc='Image frame used for training (stored either internally or externally).', + quantity='1', + ), + ], + attributes=[ + NWBAttributeSpec( + name='annotator', + doc='Name of annotator who labeled the TrainingFrame.', + dtype='text', + required=False, + ), + ], + ) + + instance = NWBGroupSpec( + neurodata_type_def='Instance', + neurodata_type_inc='NWBDataInterface', + doc='Group that holds ground-truth pose data for a single instance of a skeleton in a single frame.', + default_name='Instance', + links=[ + NWBLinkSpec( + doc='Layout of body part locations and connections.', + target_type='Skeleton', + quantity=1 + ), + ], + attributes=[ + NWBAttributeSpec( + name='id', + doc='ID used to differentiate skeleton instances.', + dtype='uint8', + required=False, + ), + ], + datasets=[ + NWBDatasetSpec( + name='node_locations', + doc=('Locations (x, y) or (x, y, z) of nodes for single instance in single frame.'), + dtype='float', + dims=[['num_body_parts', 'x, y'], ['num_body_parts', 'x, y, z']], + shape=[[None, 2], [None, 3]], + quantity=1, + ), + NWBDatasetSpec( + name='node_visibility', + doc=('Markers for node visibility where true corresponds to a visible node and false corresponds to ' + 'an occluded node.'), + dtype='bool', + dims=['num_body_parts'], + shape=[None], + quantity='?', + ), + ], + ) + + pose_training = NWBGroupSpec( + neurodata_type_def='PoseTraining', + neurodata_type_inc='NWBDataInterface', + doc='Group that holds images, ground-truth annotations, and metadata for training a pose estimator.', + default_name='PoseTraining', + groups=[ + NWBGroupSpec( + neurodata_type_inc='Skeleton', + doc='Skeleton used in project where each skeleton corresponds to a unique morphology.', + quantity='*', + ), + NWBGroupSpec( + neurodata_type_inc='TrainingFrame', + doc='Frames and ground-truth annotations for training a pose estimator.', + quantity='*', + ), + ], + ) + + new_data_types = [skeleton, pose_estimation_series, pose_estimation, training_frame, instance, pose_training] # export the spec to yaml files in the spec folder output_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'spec'))