BodyPose2D.h

Go to the documentation of this file.

/*
 * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 */
 
#ifndef NV_BODYPOSE_2D_H_
#define NV_BODYPOSE_2D_H_
 
#include <map>
#include <memory>
 
#include <cuda_runtime.h>
#include <cv/core/Array.h>
#include <cv/core/BBox.h>
#include <cv/core/Core.h>
#include <cv/core/Image.h>
#include <cv/core/MathTypes.h>
#include <cv/core/Model.h>
#include <cv/core/Tensor.h>
 
namespace cvcore { namespace bodypose2d {
 
struct BodyPart
{
    int part_idx;       
    Vector2i loc;       
    float score = 0.0f; 
};
 
struct Human
{
    std::map<int, BodyPart> body_parts; 
    BBox boundingBox;                   
    float score = 0.0f;                 
};
 
struct BodyPose2DPostProcessorParams
{
    size_t numJoints;            
    size_t nmsWindowSize;        
    size_t featUpsamplingFactor; 
    float threshHeat;            
    float threshVectorScore;     
    int threshVectorCnt1;        
    int threshPartCnt;           
    float threshHumanScore;      
    std::vector<std::pair<uint32_t, uint32_t>>
        jointEdges; 
};
 
enum class HandType : uint32_t
{
    LEFT_HAND = 0,
    RIGHT_HAND
};
 
struct HandBoundingBoxParams
{
    HandType handType;                         
    bool enableHandBoxDynamicScaling;          
    float defaultHandBoxScalingFactor;         
    float handBoxScalingFactorMinRange;        
    float handBoxScalingFactorMaxRange;        
    Vector2d handBoxScalingFactorTrainDims;    
    Vector2f handBoxScalingFactorNeckHipCoeff; 
    float handBoxNoseNeckHipRatio;             
    /* Wrist and elbow indices are necessary for computing the hand bounding box*/
    uint32_t wristIndex; 
    uint32_t elbowIndex; 
    /* Hip, nose and neck joint indices are optional.
     * If these indices are not available, the index can be set to any
     * invalid index value i.e index > numJoints */
    uint32_t leftHipIndex;  
    uint32_t rightHipIndex; 
    uint32_t noseIndex;     
    uint32_t neckIndex;     
};
 
CVCORE_API extern const HandBoundingBoxParams defaultHandParams;
 
CVCORE_API extern const ImagePreProcessingParams defaultPreProcessorParams;
 
CVCORE_API extern const ModelInputParams defaultModelInputParams;
 
CVCORE_API extern const ModelInferenceParams defaultInferenceParams;
 
CVCORE_API extern const BodyPose2DPostProcessorParams defaultPostProcessorParams;
 
class CVCORE_API BodyPose2DPreProcessor
{
public:
    BodyPose2DPreProcessor() = delete;
 
    BodyPose2DPreProcessor(const ImagePreProcessingParams &preProcessorParams,
                           const ModelInputParams &modelInputParams);
 
    ~BodyPose2DPreProcessor();
 
    void execute(Tensor<NHWC, C3, F32> &output, const Tensor<NHWC, C3, U8> &input, cudaStream_t stream = 0);
 
    void execute(Tensor<NCHW, C3, F32> &output, const Tensor<NHWC, C3, U8> &input, cudaStream_t stream = 0);
 
private:
    struct BodyPose2DPreProcessorImpl;
    std::unique_ptr<BodyPose2DPreProcessorImpl> m_pImpl;
};
 
class CVCORE_API BodyPose2D
{
public:
    static constexpr int MAX_HUMAN_COUNT = 200;
 
    BodyPose2D() = delete;
 
    BodyPose2D(const ImagePreProcessingParams &preProcessorParams, const ModelInputParams &modelInputParams,
               const ModelInferenceParams &inferenceParams, const BodyPose2DPostProcessorParams &postProcessorParams);
 
    ~BodyPose2D();
 
    void execute(Array<ArrayN<Human, MAX_HUMAN_COUNT>> &output, const Tensor<NHWC, C3, U8> &input,
                 cudaStream_t stream = 0);
 
private:
    struct BodyPose2DImpl;
    std::unique_ptr<BodyPose2DImpl> m_pImpl;
};
 
class CVCORE_API BodyPose2DPostProcessor
{
public:
    BodyPose2DPostProcessor() = delete;
 
    BodyPose2DPostProcessor(const BodyPose2DPostProcessorParams &params, const ModelInputParams &modelInputParams);
 
    ~BodyPose2DPostProcessor();
 
    void allocateStagingBuffers();
 
    void execute(Array<ArrayN<Human, BodyPose2D::MAX_HUMAN_COUNT>> &output, const Tensor<NCHW, CX, F32> &pafMap,
                 const Tensor<NCHW, CX, F32> &heatMap, int imageWidth, int imageHeight, cudaStream_t stream = 0);
 
private:
    struct BodyPose2DPostProcessorImpl;
    std::unique_ptr<BodyPose2DPostProcessorImpl> m_pImpl;
};
 
class CVCORE_API HandBoundingBoxGenerator
{
public:
    static constexpr size_t MOVING_AVG_WINDOW_SIZE = 100;
    HandBoundingBoxGenerator() = delete;
    HandBoundingBoxGenerator(const HandBoundingBoxParams &params);
    bool execute(BBox &box, const Human &pose, size_t imageWidth, size_t imageHeight);
 
    ~HandBoundingBoxGenerator();
 
private:
    struct HandBoundingBoxGeneratorImpl;
    std::unique_ptr<HandBoundingBoxGeneratorImpl> m_pImpl;
};
 
}} // namespace cvcore::bodypose2d
#endif