detectNet.h

Go to the documentation of this file.

/*
 * Copyright (c) 2017, 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 __DETECT_NET_H__
#define __DETECT_NET_H__


#include "tensorNet.h"


#define DETECTNET_DEFAULT_INPUT   "data"

#define DETECTNET_DEFAULT_COVERAGE  "coverage"

#define DETECTNET_DEFAULT_BBOX  "bboxes"

#define DETECTNET_DEFAULT_THRESHOLD 0.5f

#define DETECTNET_USAGE_STRING  "detectNet arguments: \n"                                                               \
                  "  --network NETWORK     pre-trained model to load, one of the following:\n"          \
                  "                            * pednet (default)\n"                                                    \
                  "                            * multiped\n"                                                            \
                  "                            * facenet\n"                                                                     \
                  "                            * ssd-mobilenet-v1\n"                                                    \
                  "                            * ssd-mobilenet-v2\n"                                                    \
                  "                            * ssd-inception-v2\n"                                                    \
                  "                            * coco-airplane\n"                                                       \
                  "                            * coco-bottle\n"                                                                 \
                  "                            * coco-chair\n"                                                          \
                  "                            * coco-dog\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"                                    \
                  "  --class_labels LABELS path to text file containing the labels for each class\n"                                    \
                  "  --threshold THRESHOLD minimum threshold for detection (default is 0.5)\n"                                                  \
                  "  --input_blob INPUT    name of the input layer (default is '" DETECTNET_DEFAULT_INPUT "')\n"                        \
                  "  --output_cvg COVERAGE name of the coverge output layer (default is '" DETECTNET_DEFAULT_COVERAGE "')\n"    \
                  "  --output_bbox BOXES   name of the bounding output layer (default is '" DETECTNET_DEFAULT_BBOX "')\n"       \
                  "  --mean_pixel PIXEL    mean pixel value to subtract from input (default is 0.0)\n"                                  \
                  "  --batch_size BATCH    maximum batch size (default is 1)\n"


class detectNet : public tensorNet
{
public:
        struct Detection
        {
                // Object Info
                uint32_t Instance;      
                uint32_t ClassID;       
                float Confidence;       
                // Bounding Box Coordinates
                float Left;             
                float Right;            
                float Top;              
                float Bottom;           
                inline float Width() const                                              { return Right - Left; }

                inline float Height() const                                             { return Bottom - Top; }

                inline float Area() const                                               { return Width() * Height(); }

                inline void Center( float* x, float* y ) const          { if(x) *x = Left + Width() * 0.5f; if(y) *y = Top + Height() * 0.5f; }

                inline bool Contains( float x, float y ) const          { return x >= Left && x <= Right && y >= Top && y <= Bottom; }

                inline bool Overlaps( const Detection& det ) const      { return !(det.Left > Right || det.Right < Left || det.Top > Bottom || det.Bottom < Top); }
                
                inline bool Overlaps( float x1, float y1, float x2, float y2 ) const    { return !(x1 > Right || x2 < Left || y1 > Bottom || y2 < Top); }
                
                inline bool Expand( float x1, float y1, float x2, float y2 )         { if(!Overlaps(x1, y1, x2, y2)) return false; Left = fminf(x1, Left); Top = fminf(y1, Top); Right = fmaxf(x2, Right); Bottom = fmaxf(y2, Bottom); return true; }
                
                inline bool Expand( const Detection& det )                      { if(!Overlaps(det)) return false; Left = fminf(det.Left, Left); Top = fminf(det.Top, Top); Right = fmaxf(det.Right, Right); Bottom = fmaxf(det.Bottom, Bottom); return true; }
                
                inline void Reset()                                                             { Instance = 0; ClassID = 0; Confidence = 0; Left = 0; Right = 0; Top = 0; Bottom = 0; }                                                                
                
                inline Detection()                                                              { Reset(); }
        };

        enum OverlayFlags
        {
                OVERLAY_NONE  = 0,                      
                OVERLAY_BOX   = (1 << 0),       
                OVERLAY_LABEL = (1 << 1)                
        };
        
        enum NetworkType
        {
                CUSTOM = 0,             
                COCO_AIRPLANE,          
                COCO_BOTTLE,            
                COCO_CHAIR,             
                COCO_DOG,                       
                FACENET,                        
                PEDNET,                 
                PEDNET_MULTI,           
#if NV_TENSORRT_MAJOR > 4
                SSD_MOBILENET_V1,       
                SSD_MOBILENET_V2,       
                SSD_INCEPTION_V2        
#endif
        };

        static NetworkType NetworkTypeFromStr( const char* model_name );

        static detectNet* Create( NetworkType networkType=PEDNET_MULTI, float threshold=DETECTNET_DEFAULT_THRESHOLD, 
                                                 uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, precisionType precision=TYPE_FASTEST, 
                                                 deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static detectNet* Create( const char* prototxt_path, const char* model_path, const char* mean_binary, 
                                                 const char* class_labels, float threshold=DETECTNET_DEFAULT_THRESHOLD, 
                                                 const char* input = DETECTNET_DEFAULT_INPUT, 
                                                 const char* coverage = DETECTNET_DEFAULT_COVERAGE, 
                                                 const char* bboxes = DETECTNET_DEFAULT_BBOX,
                                                 uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                 precisionType precision=TYPE_FASTEST,
                                                 deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
                                                          
        static detectNet* Create( const char* prototxt_path, const char* model_path, float mean_pixel=0.0f, 
                                                 const char* class_labels=NULL, float threshold=DETECTNET_DEFAULT_THRESHOLD, 
                                                 const char* input = DETECTNET_DEFAULT_INPUT, 
                                                 const char* coverage = DETECTNET_DEFAULT_COVERAGE, 
                                                 const char* bboxes = DETECTNET_DEFAULT_BBOX,
                                                 uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                 precisionType precision=TYPE_FASTEST,
                                                 deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static detectNet* Create( const char* model_path, const char* class_labels, float threshold, 
                                                 const char* input, const Dims3& inputDims, 
                                                 const char* output, const char* numDetections,
                                                 uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                 precisionType precision=TYPE_FASTEST,
                                                 deviceType device=DEVICE_GPU, bool allowGPUFallback=true );

        static detectNet* Create( int argc, char** argv );
        
        static inline const char* Usage()               { return DETECTNET_USAGE_STRING; }

        virtual ~detectNet();
        
        int Detect( float* input, uint32_t width, uint32_t height, Detection** detections, uint32_t overlay=OVERLAY_BOX );

        int Detect( float* input, uint32_t width, uint32_t height, Detection* detections, uint32_t overlay=OVERLAY_BOX );
        
        bool Overlay( float* input, float* output, uint32_t width, uint32_t height, Detection* detections, uint32_t numDetections, uint32_t flags=OVERLAY_BOX );
        
        inline float GetThreshold() const                                                       { return mCoverageThreshold; }

        inline void SetThreshold( float threshold )                                     { mCoverageThreshold = threshold; }

        inline uint32_t GetMaxDetections() const                                        { return mMaxDetections; } 
                
        inline uint32_t GetNumClasses() const                                           { return mNumClasses; }

        inline const char* GetClassDesc( uint32_t index )       const           { return mClassDesc[index].c_str(); }
        
        inline const char* GetClassSynset( uint32_t index ) const               { return mClassSynset[index].c_str(); }
        
        inline const char* GetClassPath() const                                         { return mClassPath.c_str(); }

        void SetClassColor( uint32_t classIndex, float r, float g, float b, float a=255.0f );
        
        
protected:

        // constructor
        detectNet( float meanPixel=0.0f );

        bool allocDetections();
        bool defaultColors();
        void defaultClassDesc();
        bool loadClassDesc( const char* filename );

        bool init( const char* prototxt_path, const char* model_path, const char* mean_binary, const char* class_labels, 
                         float threshold, const char* input, const char* coverage, const char* bboxes, uint32_t maxBatchSize, 
                         precisionType precision, deviceType device, bool allowGPUFallback );
        
        int clusterDetections( Detection* detections, uint32_t width, uint32_t height );

        float  mCoverageThreshold;
        float* mClassColors[2];
        float  mMeanPixel;

        std::vector<std::string> mClassDesc;
        std::vector<std::string> mClassSynset;

        std::string mClassPath;
        uint32_t    mCustomClasses;
        uint32_t          mNumClasses;

        Detection* mDetectionSets[2];   // list of detections, mNumDetectionSets * mMaxDetections
        uint32_t   mDetectionSet;       // index of next detection set to use
        uint32_t         mMaxDetections;        // number of raw detections in the grid

        static const uint32_t mNumDetectionSets = 16; // size of detection ringbuffer
};


#endif