backgroundNet.h

Go to the documentation of this file.

/*
 * Copyright (c) 2022, 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 __BACKGROUND_NET_H__
#define __BACKGROUND_NET_H__
 
 
#include "tensorNet.h"
#include <jetson-utils/cudaFilterMode.h>
 
 
#define BACKGROUNDNET_DEFAULT_INPUT   "input_0"
 
#define BACKGROUNDNET_DEFAULT_OUTPUT  "output_0"
 
#define BACKGROUNDNET_MODEL_TYPE "background"
 
#define BACKGROUNDNET_USAGE_STRING  "backgroundNet arguments: \n"                                       \
                  "  --network=NETWORK    pre-trained model to load, one of the following:\n"   \
                  "                           * u2net (default)\n"                                              \
                  "  --model=MODEL        path to custom model to load (caffemodel, uff, or onnx)\n"                            \
                  "  --input-blob=INPUT   name of the input layer (default is '" BACKGROUNDNET_DEFAULT_INPUT "')\n"     \
                  "  --output-blob=OUTPUT name of the output layer (default is '" BACKGROUNDNET_DEFAULT_OUTPUT "')\n"   \
                  "  --profile            enable layer profiling in TensorRT\n\n"
 
 
class backgroundNet : public tensorNet
{
public:
        static backgroundNet* Create( const char* network="u2net", uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                        precisionType precision=TYPE_FASTEST, deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static backgroundNet* Create( const char* model_path, 
                                                        const char* input=BACKGROUNDNET_DEFAULT_INPUT, 
                                                        const char* output=BACKGROUNDNET_DEFAULT_OUTPUT, 
                                                        uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                                        precisionType precision=TYPE_FASTEST,
                                                        deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static backgroundNet* Create( int argc, char** argv );
 
        static backgroundNet* Create( const commandLine& cmdLine );
 
        static inline const char* Usage()               { return BACKGROUNDNET_USAGE_STRING; }
 
        virtual ~backgroundNet();
        
        template<typename T> int Process( T* image, uint32_t width, uint32_t height,
                                                            cudaFilterMode filter=FILTER_LINEAR, bool maskAlpha=true )  { return Process((void*)image, width, height, imageFormatFromType<T>(), filter, maskAlpha); }
        
        template<typename T> int Process( T* input, T* output, uint32_t width, uint32_t height,
                                                            cudaFilterMode filter=FILTER_LINEAR, bool maskAlpha=true )  { return Process((void*)input, (void*)output, width, height, imageFormatFromType<T>(), filter, maskAlpha); }
        
        inline bool Process( void* image, uint32_t width, uint32_t height, imageFormat format,
                                         cudaFilterMode filter=FILTER_LINEAR, bool maskAlpha=true )                                     { return Process(image, image, width, height, format, filter, maskAlpha); }
                            
        bool Process( void* input, void* output, uint32_t width, uint32_t height, imageFormat format,
                            cudaFilterMode filter=FILTER_LINEAR, bool maskAlpha=true );
 
protected:
        backgroundNet();
        
        bool init(const char* model_path, const char* input, const char* output, uint32_t maxBatchSize, precisionType precision, deviceType device, bool allowGPUFallback );
 
};
 
 
#endif