Add Mobilenetv2 SSD Lite post and preprocessing, add mobilenet demo

Signed-off-by: xavier <micaelaverucchi@gmail.com>

include/evaluation.h

@@ -1,5 +1,5 @@
 #ifndef EVALUATION_H
-#define EVALUATION_H_H
+#define EVALUATION_H
 
 #include <iostream>
 #include <vector>

include/tkDNN/Layer.h

@@ -17,6 +17,7 @@ enum layerType_t {
     LAYER_ACTIVATION_CRELU,
     LAYER_ACTIVATION_LEAKY,
     LAYER_FLATTEN,
+    LAYER_RESHAPE,
     LAYER_MULADD,
     LAYER_POOLING,
     LAYER_SOFTMAX,
@@ -62,6 +63,7 @@ public:
             case LAYER_ACTIVATION_CRELU:    return "ActivationCReLU";
             case LAYER_ACTIVATION_LEAKY:    return "ActivationLeaky";
             case LAYER_FLATTEN:             return "Flatten";
+            case LAYER_RESHAPE:             return "Reshape";
             case LAYER_MULADD:              return "MulAdd";
             case LAYER_POOLING:             return "Pooling";
             case LAYER_SOFTMAX:             return "Softmax";
@@ -265,6 +267,20 @@ public:
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 };
 
+/**
+    Reshape layer
+*/
+class Reshape : public Layer {
+
+public:
+    Reshape(Network *net, dataDim_t new_dim, bool final=false); 
+    virtual ~Reshape();
+    virtual layerType_t getLayerType() { return LAYER_RESHAPE; };
+
+    virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+
+};
+
 
 /**
     MulAdd layer
@@ -329,11 +345,13 @@ protected:
 class Softmax : public Layer {
 
 public:
-    Softmax(Network *net); 
+    Softmax(Network *net, const tk::dnn::dataDim_t* dim=nullptr, bool final=false, const cudnnSoftmaxMode_t mode=CUDNN_SOFTMAX_MODE_CHANNEL); 
     virtual ~Softmax();
     virtual layerType_t getLayerType() { return LAYER_SOFTMAX; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
+    dataDim_t dim;
+    cudnnSoftmaxMode_t mode;
 };
 
 /**
@@ -343,7 +361,7 @@ public:
 class Route : public Layer {
 
 public:
-    Route(Network *net, Layer **layers, int layers_n); 
+    Route(Network *net, Layer **layers, int layers_n, bool final=false); 
     virtual ~Route();
     virtual layerType_t getLayerType() { return LAYER_ROUTE; };
 
@@ -410,6 +428,11 @@ struct box {
     int cl;
     float x, y, w, h;
     float prob;
+
+    void print() 
+    {
+        std::cout<<"x: "<<x<<"\ty: "<<y<<"\tw: "<<w<<"\th: "<<h<<"\tcl: "<<cl<<"\tprob: "<<prob<<std::endl;
+    }
 };
 struct sortable_bbox {
     int index;

include/tkDNN/MobilenetDetection.h

@@ -0,0 +1,112 @@
+#ifndef MOBILENETDETECTION_H
+#define MOBILENETDETECTION_H
+
+#include <iostream>
+#include "tkdnn.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/videoio.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#define N_COORDS 4
+
+
+struct SSDSpec
+{
+    int feature_size = 0;
+    int shrinkage = 0;
+    int box_width = 0;
+    int box_height = 0;
+    int ratio1 = 0;
+    int ratio2 = 0;
+
+    SSDSpec() {}
+
+    SSDSpec(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2) : feature_size(feature_size), shrinkage(shrinkage), box_width(box_width), box_height(box_height),
+                                                                                                      ratio1(ratio1), ratio2(ratio2) {}
+
+    void setAll(int feature_size, int shrinkage, int box_width, int box_height, int ratio1, int ratio2)
+    {
+        this->feature_size = feature_size;
+        this->shrinkage = shrinkage;
+        this->box_width = box_width;
+        this->box_height = box_height;
+        this->ratio1 = ratio1;
+        this->ratio2 = ratio2;
+    }
+
+    void print()
+    {
+        std::cout << "fsize: " << feature_size << "\tshrinkage: " << shrinkage << "\t box W:" << box_width << "\tbox H: " << box_height << "\t x ratio:" << ratio1 << "\t y ratio:" << ratio2 << std::endl;
+    }
+};
+
+
+namespace tk
+{
+namespace dnn
+{
+class MobilenetDetection
+{
+
+private:
+    tk::dnn::NetworkRT *netRT = nullptr;
+
+    int classes = 21;
+    float iou_threshold = 0.45;
+    float center_variance = 0.1;
+    float size_variance = 0.2;
+    float conf_thresh = 0.4;
+    int input_h = 300;
+    int input_w = 300;
+    int image_size = 300;
+
+    float *priors = nullptr;
+    int n_priors = 0;
+
+    cv::Mat origImg;
+    cv::Mat bgr[3];
+
+    float *input, *input_d;
+    float *locations_h, *confidences_h;
+
+    tk::dnn::dataDim_t dim;
+
+    dnnType *conf;
+    dnnType *loc;
+
+    
+
+    float __colors[6][3] = {{1, 0, 1}, {0, 0, 1}, {0, 1, 1}, {0, 1, 0}, {1, 1, 0}, {1, 0, 0}};
+    int baseline = 0;
+    float fontScale = 0.5;
+    int thickness = 2;
+
+    void generate_ssd_priors(const SSDSpec *specs, const int n_specs, bool clamp = true);
+    void convert_locatios_to_boxes_and_center(float *priors, const int n_priors, float *locations, const float center_variance, const float size_variance);
+    float iou(const tk::dnn::box &a, const tk::dnn::box &b);
+    std::vector<tk::dnn::box> postprocess(float *locations, float *confidences, const int n_values, const float threshold, const int n_classes, const float iou_thresh, const int width, const int height);
+    float get_color2(int c, int x, int max);
+
+    cv::Scalar colors[256];
+    std::vector<std::string> voc_class_name;
+
+public:
+    // keep track of inference times (ms)
+    std::vector<double> stats;
+    std::vector<tk::dnn::box> detected;
+
+    MobilenetDetection() {}
+    ~MobilenetDetection() {}
+
+    void init(std::string tensor_path);
+    cv::Mat draw();
+    void update(cv::Mat &img);
+};
+
+} // namespace dnn
+} // namespace tk
+
+
+#endif /*MOBILENETDETECTION_H*/

include/tkDNN/NetworkRT.h

@@ -34,6 +34,8 @@ using namespace nvinfer1;
 #include "pluginsRT/ResizeLayerRT.h"
 //#include "pluginsRT/Int8Calibrator.h"
 #include "pluginsRT/DeformableConvRT.h"
+#include "pluginsRT/FlattenConcatRT.h"
+#include "pluginsRT/ReshapeRT.h"
 
 class PluginFactory : IPluginFactory
 {
@@ -83,6 +85,8 @@ public:
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Pooling *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Softmax *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Route *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Flatten *l);
+    nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reshape *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reorg *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Region *l);
     nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l);

include/tkDNN/pluginsRT/FlattenConcatRT.h

@@ -0,0 +1,76 @@
+#include<cassert>
+
+class FlattenConcatRT : public IPlugin {
+
+public:
+	FlattenConcatRT() {
+		stat = cublasCreate(&handle);
+		if (stat != CUBLAS_STATUS_SUCCESS) {
+			printf ("CUBLAS initialization failed\n");
+			return;
+  		}
+	}
+
+	~FlattenConcatRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{ inputs[0].d[0] * inputs[0].d[1] * inputs[0].d[2], 1, 1};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+		assert(nbOutputs == 1 && nbInputs ==1);
+		rows = inputDims[0].d[0];
+		cols = inputDims[0].d[1] * inputDims[0].d[2];
+		c = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2];
+		h = 1;
+		w = 1;
+	}
+
+	int initialize() override {
+		return 0;
+	}
+
+	virtual void terminate() override {
+		checkERROR(cublasDestroy(handle));
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
+		checkCuda( cudaMemcpy(dstData, srcData, rows*cols*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+	
+		float const alpha(1.0);
+		float const beta(0.0);
+		checkERROR( cublasSgeam( handle, CUBLAS_OP_T, CUBLAS_OP_N, rows, cols, &alpha, srcData, cols, &beta, srcData, rows, dstData, rows ));
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 5*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, c);
+		tk::dnn::writeBUF(buf, h);
+		tk::dnn::writeBUF(buf, w);
+		tk::dnn::writeBUF(buf, rows);
+		tk::dnn::writeBUF(buf, cols);
+	}
+
+	int c, h, w;
+	int rows, cols;
+	cublasStatus_t stat; 
+	cublasHandle_t handle; 
+};

include/tkDNN/pluginsRT/ReshapeRT.h

@@ -0,0 +1,61 @@
+#include<cassert>
+
+class ReshapeRT : public IPlugin {
+
+public:
+	ReshapeRT(dataDim_t new_dim) {
+		n = new_dim.n;
+		c = new_dim.c;
+		h = new_dim.h;
+		w = new_dim.w;
+	}
+
+	~ReshapeRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsCHW{ c,h,w};
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
+	}
+
+	int initialize() override {
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
+
+		checkCuda( cudaMemcpy(dstData, srcData, c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice));
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 4*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, n);
+		tk::dnn::writeBUF(buf, c);
+		tk::dnn::writeBUF(buf, h);
+		tk::dnn::writeBUF(buf, w);
+	}
+
+	int n, c, h, w;
+};

include/tkDNN/pluginsRT/SoftmaxRT.h

@@ -0,0 +1,64 @@
+#include<cassert>
+
+class SoftmaxRT : public IPlugin {
+
+public:
+	SoftmaxRT(const tk::dnn::dataDim_t* dim) {
+		assert(dim != nullptr);
+		this->dim.n = dim->n;
+		this->dim.c = dim->c;
+		this->dim.h = dim->h;
+		this->dim.w = dim->w;
+		this->dim.l = dim->l;
+	}
+
+	~SoftmaxRT(){
+
+	}
+
+	int getNbOutputs() const override {
+		return 1;
+	}
+
+	Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
+		return DimsNCHW{this->dim.n,this->dim.c,this->dim.h,this->dim.w };
+	}
+
+	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {		
+	}
+
+	int initialize() override {
+		return 0;
+	}
+
+	virtual void terminate() override {
+	}
+
+	virtual size_t getWorkspaceSize(int maxBatchSize) const override {
+		return 0;
+	}
+
+	virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override {
+		dnnType *srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
+		dnnType *dstData = reinterpret_cast<dnnType*>(outputs[0]);
+		
+		
+		return 0;
+	}
+
+
+	virtual size_t getSerializationSize() override {
+		return 5*sizeof(int);
+	}
+
+	virtual void serialize(void* buffer) override {
+		char *buf = reinterpret_cast<char*>(buffer);
+		tk::dnn::writeBUF(buf, this->dim.n);
+		tk::dnn::writeBUF(buf, this->dim.c);
+		tk::dnn::writeBUF(buf, this->dim.h);
+		tk::dnn::writeBUF(buf, this->dim.w);
+		tk::dnn::writeBUF(buf, this->dim.l);
+	}
+
+	dataDim_t dim;
+};