v0.7 Beta - NVIDIA GTC April 2021 Release

Update to v0.7 Beta

samples/BodyTrack/BodyEngine.h

@@ -0,0 +1,200 @@
+/*###############################################################################
+#
+# 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();
+  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_Rect& bodyBox, int variant = 0);
+  void setBodyStabilization(bool);
+  void setMode(int);
+  void setFocalLength(float);
+  void useCudaGraph(bool); // Using cuda graph improves model latency
+  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{};
+  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;
+
+  BodyEngine() {
+    batchSize = 1;
+    nvARMode = 1;
+    bStabilizeBody = true;
+    bUseCudaGraph = true;
+    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