228 lines
7.8 KiB
C++
228 lines
7.8 KiB
C++
/*###############################################################################
|
|
#
|
|
# Copyright 2020 NVIDIA Corporation
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy of
|
|
# this software and associated documentation files (the "Software"), to deal in
|
|
# the Software without restriction, including without limitation the rights to
|
|
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
|
|
# the Software, and to permit persons to whom the Software is furnished to do so,
|
|
# subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
|
|
# FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
|
|
# COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
|
|
# IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
|
|
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
|
|
#
|
|
###############################################################################*/
|
|
#ifndef __BODY_ENGINE__
|
|
#define __BODY_ENGINE__
|
|
|
|
#include <random>
|
|
#include <chrono>
|
|
#include "nvAR.h"
|
|
#include "nvCVOpenCV.h"
|
|
#include "opencv2/opencv.hpp"
|
|
// #include "FeatureVertexName.h"
|
|
#define FITBODY_PRIVATE
|
|
|
|
class KalmanFilter1D {
|
|
private:
|
|
float Q_; // Covariance of the process noise
|
|
float xhat_; // Current prediction
|
|
float xhatminus_; // Previous prediction
|
|
float P_; // Estimated accuracy of xhat_
|
|
float Pminus_; // Previous P_
|
|
float K_; // Kalman gain
|
|
float R_; // Covariance of the observation noise
|
|
bool bFirstUse_;
|
|
|
|
public:
|
|
KalmanFilter1D() { reset(); }
|
|
|
|
KalmanFilter1D(float Q, float R) { reset(Q, R); }
|
|
|
|
void reset() {
|
|
R_ = 0.005f * 0.005f;
|
|
Q_ = 1e-5f;
|
|
xhat_ = 0.0f;
|
|
xhatminus_ = 0.0f;
|
|
P_ = 1;
|
|
bFirstUse_ = true;
|
|
Pminus_ = 0.0f;
|
|
K_ = 0.0f;
|
|
}
|
|
|
|
void reset(float Q, float R) {
|
|
reset();
|
|
Q_ = Q;
|
|
R_ = R;
|
|
}
|
|
|
|
float update(float val) {
|
|
if (bFirstUse_) {
|
|
xhat_ = val;
|
|
bFirstUse_ = false;
|
|
}
|
|
|
|
xhatminus_ = xhat_;
|
|
Pminus_ = P_ + Q_;
|
|
K_ = Pminus_ / (Pminus_ + R_);
|
|
xhat_ = xhatminus_ + K_ * (val - xhatminus_);
|
|
P_ = (1 - K_) * Pminus_;
|
|
|
|
return xhat_;
|
|
}
|
|
};
|
|
|
|
bool CheckResult(NvCV_Status nvErr, unsigned line);
|
|
|
|
#define BAIL_IF_ERR(err) \
|
|
do { \
|
|
if (0 != (err)) { \
|
|
goto bail; \
|
|
} \
|
|
} while (0)
|
|
|
|
#define BAIL_IF_CVERR(nvErr, err, code) \
|
|
do { \
|
|
if (!CheckResult(nvErr, __LINE__)) { \
|
|
err = code; \
|
|
goto bail; \
|
|
} \
|
|
} while (0)
|
|
|
|
typedef struct KeyPointsProperties {
|
|
int numPoints;
|
|
float confidence_threshold;
|
|
}KeyPointsProperties;
|
|
|
|
// This default focal length matches a logitech webcam
|
|
static const float FOCAL_LENGTH_DEFAULT = 800.f;
|
|
|
|
/********************************************************************************
|
|
* BodyEngine
|
|
********************************************************************************/
|
|
|
|
class BodyEngine {
|
|
public:
|
|
enum Err { errNone, errGeneral, errRun, errInitialization, errRead, errEffect, errParameter };
|
|
int input_image_width, input_image_height, input_image_pitch;
|
|
|
|
void setInputImageWidth(int width) { input_image_width = width; }
|
|
void setInputImageHeight(int height) { input_image_height = height; }
|
|
int getInputImageWidth() { return input_image_width; }
|
|
int getInputImageHeight() { return input_image_height; }
|
|
int getInputImagePitch() { return input_image_pitch = input_image_width * 3 * sizeof(unsigned char); }
|
|
void setBodyModel(const char *bodyModel) { body_model = bodyModel; }
|
|
|
|
Err createFeatures(const char* modelPath, unsigned int _batchSize = 1);
|
|
Err createBodyDetectionFeature(const char* modelPath, CUstream stream);
|
|
Err createKeyPointDetectionFeature(const char* modelPath, unsigned int batchSize, CUstream stream);
|
|
void destroyFeatures();
|
|
void destroyBodyDetectionFeature();
|
|
void destroyKeyPointDetectionFeature();
|
|
Err initFeatureIOParams();
|
|
Err initBodyDetectionIOParams(NvCVImage* _inputImageBuffer);
|
|
Err initKeyPointDetectionIOParams(NvCVImage* _inputImageBuffer);
|
|
void releaseFeatureIOParams();
|
|
void releaseBodyDetectionIOParams();
|
|
void releaseKeyPointDetectionIOParams();
|
|
|
|
unsigned findBodyBoxes();
|
|
NvAR_Rect* getLargestBox();
|
|
NvCV_Status findKeyPoints();
|
|
NvAR_BBoxes* getBoundingBoxes();
|
|
NvAR_Point2f* getKeyPoints();
|
|
NvAR_Point3f* getKeyPoints3D();
|
|
NvAR_Quaternion* getJointAngles();
|
|
NvAR_BBoxes* getBBoxes();
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
NvAR_TrackingBBoxes* getTrackingBBoxes();
|
|
#endif
|
|
float* getKeyPointsConfidence();
|
|
float getAverageKeyPointsConfidence();
|
|
void enlargeAndSquarifyImageBox(float enlarge, NvAR_Rect& box, int FLAG_variant);
|
|
unsigned findLargestBodyBox(NvAR_Rect& bodyBox, int variant = 0);
|
|
unsigned acquireBodyBox(cv::Mat& src, NvAR_Rect& bodyBox, int variant = 0);
|
|
unsigned acquireBodyBoxAndKeyPoints(cv::Mat& src, NvAR_Point2f* refMarks, NvAR_Point3f* refKeyPoints3D,
|
|
NvAR_Quaternion* refJointAngles, NvAR_BBoxes* refBodyBoxes, int variant = 0);
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
unsigned acquireBodyBoxAndKeyPoints(cv::Mat& src, NvAR_Point2f* refMarks, NvAR_Point3f* refKeyPoints3D,
|
|
NvAR_Quaternion* refJointAngles, NvAR_TrackingBBoxes* refBodyBoxes, int variant = 0);
|
|
#endif
|
|
void setBodyStabilization(bool);
|
|
void setMode(int);
|
|
BodyEngine::Err setFocalLength(float);
|
|
void useCudaGraph(bool); // Using cuda graph improves model latency
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
void enablePeopleTracking(bool _bEnablePeopleTracking, unsigned int _shadowTrackingAge = 90, unsigned int _probationAge = 10, unsigned int _maxTargetsTracked = 30);
|
|
#endif
|
|
int getNumKeyPoints() { return numKeyPoints; }
|
|
std::vector<NvAR_Point3f> getReferencePose() { return referencePose; }
|
|
|
|
NvCVImage inputImageBuffer{}, tmpImage{};
|
|
NvAR_FeatureHandle bodyDetectHandle{}, keyPointDetectHandle{};
|
|
std::vector<NvAR_Point2f> keypoints;
|
|
std::vector<float> keypoints_confidence;
|
|
std::vector<NvAR_Point3f> keypoints3D;
|
|
std::vector<NvAR_Quaternion> jointAngles;
|
|
CUstream stream{};
|
|
std::vector<NvAR_Rect> output_bbox_data;
|
|
std::vector<float> output_bbox_conf_data;
|
|
NvAR_BBoxes output_bboxes{};
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
NvAR_TrackingBBoxes output_tracking_bboxes{};
|
|
std::vector<NvAR_TrackingBBox> output_tracking_bbox_data;
|
|
#endif
|
|
|
|
int batchSize;
|
|
int nvARMode;
|
|
std::mt19937 ran;
|
|
unsigned int numKeyPoints;
|
|
std::vector<NvAR_Point3f> referencePose;
|
|
float confidenceThreshold;
|
|
std::string body_model;
|
|
|
|
bool bStabilizeBody;
|
|
bool bUseOTAU;
|
|
char *bdOTAModelPath, *ldOTAModelPath;
|
|
float bFocalLength;
|
|
bool bUseCudaGraph;
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
bool bEnablePeopleTracking;
|
|
unsigned int shadowTrackingAge;
|
|
unsigned int probationAge;
|
|
unsigned int maxTargetsTracked;
|
|
#endif
|
|
BodyEngine() {
|
|
batchSize = 1;
|
|
nvARMode = 1;
|
|
bStabilizeBody = true;
|
|
bUseCudaGraph = true;
|
|
#if NV_MULTI_OBJECT_TRACKER
|
|
bEnablePeopleTracking = false;
|
|
shadowTrackingAge = 90;
|
|
probationAge = 10;
|
|
maxTargetsTracked = 30;
|
|
#endif
|
|
bFocalLength = FOCAL_LENGTH_DEFAULT;
|
|
confidenceThreshold = 0.f;
|
|
appMode = keyPointDetection;
|
|
input_image_width = 960;
|
|
input_image_height = 544;
|
|
input_image_pitch = 3 * input_image_width * sizeof(unsigned char); // RGB
|
|
bUseOTAU = false;
|
|
bdOTAModelPath = NULL;
|
|
ldOTAModelPath = NULL;
|
|
}
|
|
enum mode { bodyDetection = 0, keyPointDetection } appMode;
|
|
void setAppMode(BodyEngine::mode _mAppMode);
|
|
};
|
|
#endif
|