imageNet.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 __IMAGE_NET_H__
#define __IMAGE_NET_H__


#include "tensorNet.h"


#define IMAGENET_DEFAULT_INPUT   "data"

#define IMAGENET_DEFAULT_OUTPUT  "prob"


#define IMAGENET_USAGE_STRING  "imageNet arguments: \n"                                                         \
                  "  --network NETWORK    pre-trained model to load, one of the following:\n"   \
                  "                           * alexnet\n"                                                              \
                  "                           * googlenet (default)\n"                                  \
                  "                           * googlenet-12\n"                                                         \
                  "                           * resnet-18\n"                                                    \
                  "                           * resnet-50\n"                                                    \
                  "                           * resnet-101\n"                                                   \
                  "                           * resnet-152\n"                                                   \
                  "                           * vgg-16\n"                                                               \
                  "                           * vgg-19\n"                                                               \
                  "                           * inception-v4\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 '" IMAGENET_DEFAULT_INPUT "')\n"  \
                  "  --output_blob OUTPUT name of the output layer (default is '" IMAGENET_DEFAULT_OUTPUT "')\n"        \
                  "  --batch_size BATCH   maximum batch size (default is 1)\n"


class imageNet : public tensorNet
{
public:
        enum NetworkType
        {
                CUSTOM,        
                ALEXNET,                
                GOOGLENET,      
                GOOGLENET_12,   
                RESNET_18,      
                RESNET_50,      
                RESNET_101,     
                RESNET_152,     
                VGG_16,         
                VGG_19,         
                INCEPTION_V4,   
        };

        static NetworkType NetworkTypeFromStr( const char* model_name );

        static const char* NetworkTypeToStr( NetworkType network );

        static imageNet* Create( NetworkType networkType=GOOGLENET, uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                precisionType precision=TYPE_FASTEST,
                                                deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static imageNet* Create( const char* prototxt_path, const char* model_path, 
                                                const char* mean_binary, const char* class_labels, 
                                                const char* input=IMAGENET_DEFAULT_INPUT, 
                                                const char* output=IMAGENET_DEFAULT_OUTPUT, 
                                                uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                precisionType precision=TYPE_FASTEST,
                                                deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static imageNet* Create( int argc, char** argv );

        static inline const char* Usage()               { return IMAGENET_USAGE_STRING; }

        virtual ~imageNet();
        
        int Classify( float* rgba, uint32_t width, uint32_t height, float* confidence=NULL );

        int Classify( float* confidence=NULL );

        bool PreProcess( float* rgba, uint32_t width, uint32_t height );

        bool Process();

        inline uint32_t GetNumClasses() const                                           { return mOutputClasses; }
        
        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(); }

        inline NetworkType GetNetworkType() const                                       { return mNetworkType; }

        inline const char* GetNetworkName() const                                       { NetworkTypeToStr(mNetworkType); }

protected:
        imageNet();
        
        bool init( NetworkType networkType, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
        bool init(const char* prototxt_path, const char* model_path, const char* mean_binary, const char* class_path, const char* input, const char* output, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
        bool loadClassInfo( const char* filename );
        
        uint32_t mCustomClasses;
        uint32_t mOutputClasses;
        
        std::vector<std::string> mClassSynset;  // 1000 class ID's (ie n01580077, n04325704)
        std::vector<std::string> mClassDesc;

        std::string mClassPath;
        NetworkType mNetworkType;
};


#endif