segNet.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 __SEGMENTATION_NET_H__
#define __SEGMENTATION_NET_H__


#include "tensorNet.h"


#define SEGNET_DEFAULT_INPUT   "data"

#define SEGNET_DEFAULT_OUTPUT  "score_fr_21classes"


class segNet : public tensorNet
{
public:
        enum NetworkType
        {
                FCN_ALEXNET_PASCAL_VOC,             
                FCN_ALEXNET_SYNTHIA_CVPR16,         
                FCN_ALEXNET_SYNTHIA_SUMMER_HD,    
                FCN_ALEXNET_SYNTHIA_SUMMER_SD,    
                FCN_ALEXNET_CITYSCAPES_HD,          
                FCN_ALEXNET_CITYSCAPES_SD,          
                FCN_ALEXNET_AERIAL_FPV_720p,        
                /* add new models here */
                SEGNET_CUSTOM
        };

        enum FilterMode
        {
                FILTER_POINT,           
                FILTER_LINEAR           
        };

        static NetworkType NetworkTypeFromStr( const char* model_name );

        static FilterMode FilterModeFromStr( const char* str, FilterMode default_value=FILTER_LINEAR );

        static segNet* Create( NetworkType networkType=FCN_ALEXNET_CITYSCAPES_SD, uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE,
                                           precisionType precision=TYPE_FASTEST, deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        
        static segNet* Create( const char* prototxt_path, const char* model_path, 
                                           const char* class_labels, const char* class_colors=NULL,
                                           const char* input = SEGNET_DEFAULT_INPUT, 
                                           const char* output = SEGNET_DEFAULT_OUTPUT,
                                           uint32_t maxBatchSize=DEFAULT_MAX_BATCH_SIZE, 
                                           precisionType precision=TYPE_FASTEST, 
                                           deviceType device=DEVICE_GPU, bool allowGPUFallback=true );
        

        static segNet* Create( int argc, char** argv );
        
        virtual ~segNet();
        
        bool Process( float* input, uint32_t width, uint32_t height, const char* ignore_class="void" );

        bool Mask( uint8_t* output, uint32_t width, uint32_t height );

        bool Mask( float* output, uint32_t width, uint32_t height, FilterMode filter=FILTER_LINEAR );

        bool Overlay( float* output, uint32_t width, uint32_t height, FilterMode filter=FILTER_LINEAR );

        int FindClassID( const char* label_name );

        inline uint32_t GetNumClasses() const                                           { return DIMS_C(mOutputs[0].dims); }

        inline const char* GetClassLabel( uint32_t id ) const           { return id < mClassLabels.size() ? mClassLabels[id].c_str() : NULL; }
        
        inline float* GetClassColor( uint32_t id ) const                                { return mClassColors[0] + (id*4); }

        void SetClassColor( uint32_t classIndex, float r, float g, float b, float a=255.0f );
        
        void SetGlobalAlpha( float alpha, bool explicit_exempt=true );

        inline const char* GetClassPath() const                                         { return mClassPath.c_str(); }

        inline uint32_t GetGridWidth() const                                            { return DIMS_W(mOutputs[0].dims); }

        inline uint32_t GetGridHeight() const                                           { return DIMS_H(mOutputs[0].dims); }

        inline NetworkType GetNetworkType() const                                       { return mNetworkType; }

        inline const char* GetNetworkName() const                                       { return (mNetworkType != SEGNET_CUSTOM ? "FCN_Alexnet" : "segNet"); }

protected:
        segNet();
        
        bool classify( const char* ignore_class );

        bool overlayPoint( float* input, uint32_t in_width, uint32_t in_height, float* output, uint32_t out_width, uint32_t out_height, bool mask_only );
        bool overlayLinear( float* input, uint32_t in_width, uint32_t in_height, float* output, uint32_t out_width, uint32_t out_height, bool mask_only );
        
        bool loadClassColors( const char* filename );
        bool loadClassLabels( const char* filename );
        
        std::vector<std::string> mClassLabels;
        std::string mClassPath;

        float*   mClassColors[2];       
        uint8_t* mClassMap[2];          
        float*   mLastInputImg;         
        uint32_t mLastInputWidth;       
        uint32_t mLastInputHeight;      
        NetworkType mNetworkType;       
};


#endif