poseNet.h

Go to the documentation of this file.

/*
 * Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
 *
 * 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 __POSE_NET_H__
#define __POSE_NET_H__
 
 
#include "tensorNet.h"
 
#include <array>
#include <vector>
 
 
#define POSENET_DEFAULT_INPUT   "input"
 
#define POSENET_DEFAULT_CMAP  "cmap"
 
#define POSENET_DEFAULT_PAF  "paf"
 
#define POSENET_DEFAULT_THRESHOLD 0.15f
 
#define POSENET_DEFAULT_KEYPOINT_SCALE 0.0052f
 
#define POSENET_DEFAULT_LINK_SCALE 0.0013f
 
#define POSENET_MODEL_TYPE "pose"
 
#define POSENET_USAGE_STRING  "poseNet arguments: \n"                                                           \
                  "  --network=NETWORK     pre-trained model to load, one of the following:\n"  \
                  "                            * resnet18-body (default)\n"                             \
                  "                            * resnet18-hand\n"                                               \
                  "                            * densenet121-body\n"                                            \
                  "  --model=MODEL         path to custom model to load (caffemodel, uff, or onnx)\n"                                   \
                  "  --prototxt=PROTOTXT   path to custom prototxt to load (for .caffemodel only)\n"                                    \
                  "  --labels=LABELS       path to text file containing the labels for each class\n"                                    \
                  "  --input-blob=INPUT    name of the input layer (default is '" POSENET_DEFAULT_INPUT "')\n"                  \
                  "  --output-cvg=COVERAGE name of the coverge output layer (default is '" POSENET_DEFAULT_CMAP "')\n"  \
                  "  --output-bbox=BOXES   name of the bounding output layer (default is '" POSENET_DEFAULT_PAF "')\n"  \
                  "  --mean-pixel=PIXEL    mean pixel value to subtract from input (default is 0.0)\n"                                  \
                  "  --threshold=THRESHOLD minimum threshold for detection (default is 0.5)\n"                                                  \
                  "  --overlay=OVERLAY     detection overlay flags (e.g. --overlay=links,keypoints)\n"                                  \
                  "                        valid combinations are:  'box', 'links', 'keypoints', 'none'\n"                              \
                  "  --keypoint-scale=X    radius scale for keypoints, relative to image (default: 0.0052)\n"                           \
                  "  --link-scale=X        line width scale for links, relative to image (default: 0.0013)\n"                           \
                  "  --profile             enable layer profiling in TensorRT\n\n"
 
 
class poseNet : public tensorNet
{
public: 
        struct ObjectPose
        {               
                uint32_t ID;    
                float Left;     
                float Right;    
                float Top;      
                float Bottom;   
                struct Keypoint
                {
                        uint32_t ID;    
                        float x;                
                        float y;                
                };
                
                std::vector<Keypoint> Keypoints;                        
                std::vector<std::array<uint32_t, 2>> Links;     
                inline int FindKeypoint(uint32_t id) const;         
                
                inline int FindLink(uint32_t a, uint32_t b) const;  
        };
        
        enum OverlayFlags
        {
                OVERLAY_NONE      = 0,         
                OVERLAY_BOX       = (1 << 0),  
                OVERLAY_LINKS     = (1 << 1),  
                OVERLAY_KEYPOINTS = (1 << 2),  
                OVERLAY_DEFAULT   = OVERLAY_LINKS|OVERLAY_KEYPOINTS,
        };
 
        static uint32_t OverlayFlagsFromStr( const char* flags );
 
        static poseNet* Create( const char* network="resnet18-body", float threshold=POSENET_DEFAULT_THRESHOLD, 
                                            uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, precisionType precision=TYPE_FASTEST, 
                                            deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static poseNet* Create( const char* model_path, const char* topology, const char* colors,
                                            float threshold=POSENET_DEFAULT_THRESHOLD, 
                                            const char* input = POSENET_DEFAULT_INPUT, 
                                            const char* cmap = POSENET_DEFAULT_CMAP, 
                                            const char* paf = POSENET_DEFAULT_PAF,
                                            uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                            precisionType precision=TYPE_FASTEST,
                                            deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
 
        static poseNet* Create( int argc, char** argv );
        
        static poseNet* Create( const commandLine& cmdLine );
        
        static inline const char* Usage()               { return POSENET_USAGE_STRING; }
 
        virtual ~poseNet();
        
        template<typename T> bool Process( T* image, uint32_t width, uint32_t height, std::vector<ObjectPose>& poses, uint32_t overlay=OVERLAY_DEFAULT )                                        { return Process((void*)image, width, height, imageFormatFromType<T>(), poses, overlay); }
 
        bool Process( void* image, uint32_t width, uint32_t height, imageFormat format, std::vector<ObjectPose>& poses, uint32_t overlay=OVERLAY_DEFAULT );
        
        template<typename T> bool Process( T* image, uint32_t width, uint32_t height, uint32_t overlay=OVERLAY_DEFAULT )                                                                                        { return Process((void*)image, width, height, imageFormatFromType<T>(), overlay); }
 
        bool Process( void* image, uint32_t width, uint32_t height, imageFormat format, uint32_t overlay=OVERLAY_DEFAULT );
 
        template<typename T> bool Overlay( T* input, T* output, uint32_t width, uint32_t height, const std::vector<ObjectPose>& poses, uint32_t overlay=OVERLAY_DEFAULT )               { return Overlay((void*)input, (void*)output, width, height, imageFormatFromType<T>(), overlay); }
        
        bool Overlay( void* input, void* output, uint32_t width, uint32_t height, imageFormat format, const std::vector<ObjectPose>& poses, uint32_t overlay=OVERLAY_DEFAULT );
        
        inline float GetThreshold() const                                                               { return mThreshold; }
 
        inline void SetThreshold( float threshold )                                             { mThreshold = threshold; }
 
        inline const char* GetCategory() const                                                  { return mTopology.category.c_str(); }
        
        inline uint32_t GetNumKeypoints() const                                                 { return mTopology.keypoints.size(); }
        
        inline const char* GetKeypointName( uint32_t index ) const                      { return mTopology.keypoints[index].c_str(); }
        
        inline int FindKeypointID( const char* name ) const;
 
        inline float4 GetKeypointColor( uint32_t index ) const                          { return mKeypointColors[index]; }
        
        inline void SetKeypointColor( uint32_t index, const float4& color )     { mKeypointColors[index] = color; }
        
        inline void SetKeypointAlpha( uint32_t index, float alpha )                     { mKeypointColors[index].w = alpha; }
        
        inline void SetKeypointAlpha( float alpha );
        
        inline float GetKeypointScale() const                                                   { return mKeypointScale; }
        
        inline void SetKeypointScale( float scale )                                             { mKeypointScale = scale; }
        
        inline float GetLinkScale() const                                                               { return mLinkScale; }
        
        inline void SetLinkScale( float scale )                                                 { mLinkScale = scale; }
        
protected:
 
        static const int CMAP_WINDOW_SIZE=5;
        static const int PAF_INTEGRAL_SAMPLES=7;
        static const int MAX_LINKS=100;
        static const int MAX_OBJECTS=100;
        
        struct Topology
        {
                std::string category;
                std::vector<std::string> keypoints;
                int links[MAX_LINKS * 4];
                int numLinks;
        };
        
        // constructor
        poseNet();
 
        bool init( const char* model_path, const char* topology, const char* colors, float threshold, 
                         const char* input, const char* cmap, const char* paf, uint32_t maxBatchSize, 
                         precisionType precision, deviceType device, bool allowGPUFallback );
 
        bool postProcess(std::vector<ObjectPose>& poses, uint32_t width, uint32_t height);
        
        bool loadTopology( const char* json_path, Topology* topology );
        bool loadKeypointColors( const char* filename );
        
        Topology mTopology;
        
        float mThreshold;
        float mLinkScale;
        float mKeypointScale;
 
        float4* mKeypointColors;
        
        // post-processing buffers
        int* mPeaks;
        int* mPeakCounts;
        int* mConnections;
        int* mObjects;
        int  mNumObjects;
        
        float* mRefinedPeaks;
        float* mScoreGraph;
        
        void* mAssignmentWorkspace;
        void* mConnectionWorkspace;
};
 
 
// FindKeypointID
inline int poseNet::FindKeypointID( const char* name ) const
{
        if( !name )
                return -1;
        
        const uint32_t numKeypoints = GetNumKeypoints();
        
        for( uint32_t n=0; n < numKeypoints; n++ )
        {
                if( strcasecmp(GetKeypointName(n), name) == 0 )
                        return n;
        }
        
        return -1;
}
 
// FindKeypoint
inline int poseNet::ObjectPose::FindKeypoint( uint32_t id ) const
{
        const uint32_t numKeypoints = Keypoints.size();
        
        for( uint32_t n=0; n < numKeypoints; n++ )
        {
                if( id == Keypoints[n].ID )
                        return n;
        }
        
        return -1;
}
 
// FindLink
inline int poseNet::ObjectPose::FindLink( uint32_t a, uint32_t b ) const
{
        const uint32_t numLinks = Links.size();
        
        for( uint32_t n=0; n < numLinks; n++ )
        {
                if( a == Keypoints[Links[n][0]].ID && b == Keypoints[Links[n][1]].ID )
                        return n;
        }
        
        return -1;
}
 
// SetKeypointAlpha
inline void poseNet::SetKeypointAlpha( float alpha )
{
        const uint32_t numKeypoints = GetNumKeypoints();
        
        for( uint32_t n=0; n < numKeypoints; n++ )
                mKeypointColors[n].w = alpha;
}
 
 
#endif