Face-Recognition/pluginImplement.h at master · AastaNV/Face-Recognition · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
#ifndef __PLUGIN_LAYER_H__
#define __PLUGIN_LAYER_H__

#include <cassert>
#include <iostream>
#include <cudnn.h>
#include <cstring>

#include "NvCaffeParser.h"
#include "NvInferPlugin.h"

#define CHECK(status)                                                                                           \
    {                                                                                                                           \
        if (status != 0)                                                                                                \
        {                                                                                                                               \
            std::cout << "Cuda failure: " << cudaGetErrorString(status) \
                      << " at line " << __LINE__                                                        \
                      << std::endl;                                                                     \
            abort();                                                                                                    \
        }                                                                                                                               \
    }

using namespace nvinfer1;
using namespace nvcaffeparser1;
using namespace plugin;

enum FunctionType
{
    SELECT=0,
    SUMMARY
};

class bboxProfile {
public:
    bboxProfile(float4& p, int idx): pos(p), bboxNum(idx) {}

    float4 pos;
    int bboxNum = -1;
    int labelID = -1;
};

class tagProfile {
public:
    tagProfile(int b, int l): bboxID(b), label(l) {}
    int bboxID;
    int label;
};

class BboxMergeLayer : public IPlugin
{
public:
    BboxMergeLayer() {};
    BboxMergeLayer(const void* buffer, size_t size);

    inline int getNbOutputs() const override { return 1; };
    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

    int initialize() override;
    inline void terminate() override { ; }

    inline size_t getWorkspaceSize(int) const override { return 0; }
    int enqueue(int batchSize, const void*const *inputs, void** outputs, void*, cudaStream_t stream) override;

    size_t getSerializationSize() override;
    void serialize(void* buffer) override;

    void configure(const Dims*inputs, int nbInputs, const Dims* outputs, int nbOutputs, int) override;

protected:
    void mergeRect(std::vector<float4>& rects, const float4& rect);

    DimsCHW dimsData;
    DimsCHW dimsConf;
    DimsCHW dimsBbox;

    int ow;
    int oh;
    int owh;
    int cls;
    float cell_width;
    float cell_height;
};

class RecognitionLayer : public IPlugin
{
public:
    RecognitionLayer(FunctionType t) { type = t; };
    RecognitionLayer(FunctionType t, const void* buffer, size_t size);

    int getNbOutputs() const override;
    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

    inline int initialize() override { return 0; }
    inline void terminate() override { ; }

    inline size_t getWorkspaceSize(int) const override { return 0; }
    int enqueue(int batchSize, const void*const *inputs, void** outputs, void*, cudaStream_t stream) override;

    size_t getSerializationSize() override;
    void serialize(void* buffer) override;

    inline void configure(const Dims*inputs, int nbInputs, const Dims* outputs, int nbOutputs, int) override { ; };

protected:
    int select(const void*const *inputs, void** outputs);
    int summary(const void*const *inputs, void** outputs);
    int bboxExist(const float4& pos, const int idx);
    int tagExist(int label, int idx);

    size_t classNum;
    FunctionType type;
    static std::vector<bboxProfile*> bboxTable;
    static std::vector<tagProfile*> tagTable;
};

class DataRoiLayer : public IPlugin
{
public:
    DataRoiLayer() {};
    DataRoiLayer(const void* buffer, size_t size);

    inline int getNbOutputs() const override { return 1; };
    Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

    int initialize() override;
    inline void terminate() override { ; }

    inline size_t getWorkspaceSize(int) const override { return 0; }
    int enqueue(int batchSize, const void*const *inputs, void** outputs, void*, cudaStream_t stream) override;

    size_t getSerializationSize() override;
    void serialize(void* buffer) override;

    void configure(const Dims*inputs, int nbInputs, const Dims* outputs, int nbOutputs, int) override;

protected:
    DimsCHW dimsData;
    DimsCHW dimsRoi;
};

class PluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactory
{
public:
    virtual nvinfer1::IPlugin* createPlugin(const char* layerName, const nvinfer1::Weights* weights, int nbWeights) override;
    IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override;

    bool isPlugin(const char* name) override;
    void destroyPlugin();

    std::unique_ptr<BboxMergeLayer> mBboxMergeLayer{ nullptr };
    std::unique_ptr<DataRoiLayer> mDataRoiLayer{ nullptr };
    std::unique_ptr<RecognitionLayer> mSelectLayer{ nullptr };
    std::unique_ptr<RecognitionLayer> mSummaryLayer{ nullptr };
};

#endif