diff --git a/dart/biomechanics/SubjectOnDisk.cpp b/dart/biomechanics/SubjectOnDisk.cpp index cfc6829c5..077ceab40 100644 --- a/dart/biomechanics/SubjectOnDisk.cpp +++ b/dart/biomechanics/SubjectOnDisk.cpp @@ -219,6 +219,40 @@ proto::BasicTrialType basicTrialTypeToProto(BasicTrialType type) return proto::BasicTrialType::other; } +DataQuality dataQualityFromProto(proto::DataQuality quality) +{ + switch (quality) + { + case proto::DataQuality::pilotData: + return pilotData; + case proto::DataQuality::experimentalData: + return experimentalData; + case proto::DataQuality::internetData: + return internetData; + case proto::DataQuality_INT_MIN_SENTINEL_DO_NOT_USE_: + return internetData; + break; + case proto::DataQuality_INT_MAX_SENTINEL_DO_NOT_USE_: + return internetData; + break; + } + return internetData; +} + +proto::DataQuality dataQualityToProto(DataQuality quality) +{ + switch (quality) + { + case pilotData: + return proto::DataQuality::pilotData; + case experimentalData: + return proto::DataQuality::experimentalData; + case internetData: + return proto::DataQuality::internetData; + } + return proto::DataQuality::pilotData; +} + DetectedTrialFeature detectedTrialFeatureFromProto( proto::DetectedTrialFeature feature) { @@ -1568,6 +1602,12 @@ std::vector SubjectOnDisk::getMissingGRF(int trial) return mHeader->mTrials[trial]->mMissingGRFReason; } +/// This returns the user supplied enum of type 'DataQuality' +DataQuality SubjectOnDisk::getQuality() +{ + return mHeader->getQuality(); +} + int SubjectOnDisk::getNumProcessingPasses() { return mHeader->mPasses.size(); @@ -3680,6 +3720,17 @@ SubjectOnDiskHeader& SubjectOnDiskHeader::setNotes(const std::string& notes) return *this; } +SubjectOnDiskHeader& SubjectOnDiskHeader::setQuality(DataQuality quality) +{ + mDataQuality = quality; + return *this; +} + +DataQuality SubjectOnDiskHeader::getQuality() +{ + return mDataQuality; +} + std::shared_ptr SubjectOnDiskHeader::addProcessingPass() { @@ -3915,6 +3966,7 @@ void SubjectOnDiskHeader::write(dart::proto::SubjectOnDiskHeader* header) { header->add_exo_dof_index(index); } + header->set_data_quality(dataQualityToProto(mDataQuality)); if (!header->IsInitialized()) { @@ -4059,6 +4111,8 @@ void SubjectOnDiskHeader::read(const dart::proto::SubjectOnDiskHeader& proto) { mExoDofIndices.push_back(proto.exo_dof_index(i)); } + + mDataQuality = dataQualityFromProto(proto.data_quality()); } void SubjectOnDiskHeader::writeSensorsFrame( diff --git a/dart/biomechanics/SubjectOnDisk.hpp b/dart/biomechanics/SubjectOnDisk.hpp index fae66f9b0..3bcb4619b 100644 --- a/dart/biomechanics/SubjectOnDisk.hpp +++ b/dart/biomechanics/SubjectOnDisk.hpp @@ -463,6 +463,8 @@ class SubjectOnDiskHeader SubjectOnDiskHeader& setSubjectTags(std::vector subjectTags); SubjectOnDiskHeader& setHref(const std::string& sourceHref); SubjectOnDiskHeader& setNotes(const std::string& notes); + SubjectOnDiskHeader& setQuality(DataQuality quality); + DataQuality getQuality(); std::shared_ptr addProcessingPass(); std::vector> getProcessingPasses(); std::shared_ptr addTrial(); @@ -522,6 +524,9 @@ class SubjectOnDiskHeader // This is exoskeleton data std::vector mExoDofIndices; + // This is the user supplied quality of the data + DataQuality mDataQuality; + friend class SubjectOnDisk; friend struct Frame; friend struct FramePass; @@ -637,6 +642,9 @@ class SubjectOnDisk /// include `notMissingGRF`. std::vector getMissingGRF(int trial); + /// This returns the user supplied enum of type 'DataQuality' + DataQuality getQuality(); + int getNumProcessingPasses(); ProcessingPassType getProcessingPassType(int processingPass); diff --git a/dart/biomechanics/enums.hpp b/dart/biomechanics/enums.hpp index 6a66ac0f8..5b8ad9b47 100644 --- a/dart/biomechanics/enums.hpp +++ b/dart/biomechanics/enums.hpp @@ -45,6 +45,13 @@ enum DetectedTrialFeature flatTerrain }; +enum DataQuality +{ + pilotData, + experimentalData, + internetData +}; + enum MissingGRFStatus { no = 0, // no will cast to `false` diff --git a/dart/proto/SubjectOnDisk.proto b/dart/proto/SubjectOnDisk.proto index f82103276..80463e40b 100644 --- a/dart/proto/SubjectOnDisk.proto +++ b/dart/proto/SubjectOnDisk.proto @@ -44,6 +44,12 @@ enum DetectedTrialFeature { flatTerrain = 3; } +enum DataQuality { + pilotData = 0; + experimentalData = 1; + internetData = 2; +} + // Many of the ML tasks we want to support from SubjectOnDisk data include // effectively predicting the results of a downstream processing task from // an upstream processing task. Trivially, that's predicting physics from @@ -166,6 +172,8 @@ message SubjectOnDiskHeader { repeated int32 exo_dof_index = 22; // Details about the subject tags provided on the AddBiomechanics platform repeated string subject_tag = 23; + // This is what the user has tagged this subject as, in terms of data quality + DataQuality data_quality = 25; } message SubjectOnDiskProcessingPassFrame { diff --git a/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp b/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp index 1a82aa0aa..6e13db8e8 100644 --- a/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp +++ b/python/_nimblephysics/biomechanics/SubjectOnDisk.cpp @@ -73,6 +73,21 @@ void SubjectOnDisk(py::module& m) "This is a trial that doesn't fit into any of the " "other categories."); + auto dataQuality + = ::py::enum_(m, "DataQuality") + .value( + "PILOT_DATA", + dart::biomechanics::DataQuality::pilotData, + "This is data that was collected as part of a pilot study.") + .value( + "EXPERIMENTAL_DATA", + dart::biomechanics::DataQuality::experimentalData, + "This is data that was collected as part of an experiment.") + .value( + "INTERNET_DATA", + dart::biomechanics::DataQuality::internetData, + "This is data that was collected from the internet."); + auto detectedTrialFeature = ::py::enum_( m, "DetectedTrialFeature") @@ -1089,6 +1104,13 @@ Note that these are specified in the local body frame, acting on the body at its "setNotes", &dart::biomechanics::SubjectOnDiskHeader::setNotes, ::py::arg("notes")) + .def( + "setQuality", + &dart::biomechanics::SubjectOnDiskHeader::setQuality, + ::py::arg("quality")) + .def( + "getQuality", + &dart::biomechanics::SubjectOnDiskHeader::getQuality) .def( "addProcessingPass", &dart::biomechanics::SubjectOnDiskHeader::addProcessingPass) @@ -1323,6 +1345,11 @@ Note that these are specified in the local body frame, acting on the body at its This method is provided to give a cheaper way to filter out frames we want to ignore for training, without having to call the more expensive :code:`loadFrames()` and examine frames individually. )doc") + .def( + "getQuality", + &dart::biomechanics::SubjectOnDisk::getQuality, + "This returns the user-supplied quality of the data in this " + "subject") // int getNumProcessingPasses(); .def( "getNumProcessingPasses",