retinaface update

Author: wang-xinyu

README.md

@@ -40,6 +40,7 @@ Following models are implemented, each one also has a readme inside.
 |[vgg](./vgg)| VGG 11-layer model |
 |[yolov3](./yolov3)| darknet-53, weights from yolov3 authors |
 |[yolov3-spp](./yolov3-spp)| darknet-53, weights from [ultralytics/yolov3](https://github.com/ultralytics/yolov3) |
+|[retinaface](./retinaface)| resnet-50, weights from [biubug6/Pytorch_Retinaface](https://github.com/biubug6/Pytorch_Retinaface) |
 
 ## Tricky Operations
 
@@ -59,13 +60,15 @@ Some tricky operations encountered in these models, already solved, but might ha
 |yolo layer v2| three yolo layers implemented in one plugin, see yolov3-spp. |
 |upsample| replaced by a deconvolution layer, see yolov3. |
 |hsigmoid| hard sigmoid is implemented as a plugin, hsigmoid and hswish are used in mobilenetv3 |
+|retinaface output decode| implement a plugin to decode bbox, confidence and landmarks, see retinaface. |
 
 ## Speed Benchmark
 
 | Models | Device | BatchSize | Mode | Input Shape(HxW) | FPS |
 |-|-|:-:|:-:|:-:|:-:|
-| yolov3(darknet53) | Xavier | 1 | FP16 | 320x320 | 55 |
-| yolov3-spp(darknet53) | GTX1080 | 1 | FP32 | 256x416 | 94 |
+| YOLOv3(darknet53) | Xavier | 1 | FP16 | 320x320 | 55 |
+| YOLOv3-spp(darknet53) | GTX1080 | 1 | FP32 | 256x416 | 94 |
+| RetinaFace(resnet50) | TX2 | 1 | FP16 | 384x640 | 15 |
 
 Help wanted, if you got speed results, please add an issue or PR.
 

retinaface/README.md

@@ -1,3 +1,7 @@
 # RetinaFace
 
-still working in progress
+## Notice
+
+- Only tested on TensorRT4.1.3
+
+

retinaface/decode.cu

@@ -34,22 +34,22 @@ namespace nvinfer1
     {
         //output the result to channel
         int totalCount = 0;
-        totalCount += input_h_ / 8 * input_w_ / 8 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
-        totalCount += input_h_ / 16 * input_w_ / 16 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
-        totalCount += input_h_ / 32 * input_w_ / 32 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
+        totalCount += decodeplugin::INPUT_H / 8 * decodeplugin::INPUT_W / 8 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
+        totalCount += decodeplugin::INPUT_H / 16 * decodeplugin::INPUT_W / 16 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
+        totalCount += decodeplugin::INPUT_H / 32 * decodeplugin::INPUT_W / 32 * 2 * sizeof(decodeplugin::Detection) / sizeof(float);
 
         return Dims3(totalCount + 1, 1, 1);
     }
 
     __device__ float Logist(float data){ return 1./(1. + exp(-data)); };
 
-    __global__ void CalDetection(const float *input, float *output, int num_elem, int input_h, int input_w, int step, int anchor) {
+    __global__ void CalDetection(const float *input, float *output, int num_elem, int step, int anchor) {
 
         int idx = threadIdx.x + blockDim.x * blockIdx.x;
         if (idx >= num_elem) return;
 
-        int h = input_h / step;
-        int w = input_w / step;
+        int h = decodeplugin::INPUT_H / step;
+        int w = decodeplugin::INPUT_W / step;
         int y = idx / w;
         int x = idx % w;
         const float *bbox_reg = &input[0];
@@ -60,7 +60,7 @@ namespace nvinfer1
             float conf1 = cls_reg[idx + k * num_elem * 2];
             float conf2 = cls_reg[idx + k * num_elem * 2 + num_elem];
             conf2 = exp(conf2) / (exp(conf1) + exp(conf2));
-            if (conf2 <= 0.002) continue;
+            if (conf2 <= 0.02) continue;
 
             float *res_count = output;
             int count = (int)atomicAdd(res_count, 1);
@@ -70,10 +70,8 @@ namespace nvinfer1
             float prior[4];
             prior[0] = ((float)x + 0.5) / w;
             prior[1] = ((float)y + 0.5) / h;
-            prior[2] = (float)anchor / input_w;
-            prior[3] = (float)anchor / input_h;
-            printf("prior0, %f\n", prior[0]);
-            printf("bbox0, %f\n", bbox_reg[idx + k * num_elem * 4]);
+            prior[2] = (float)anchor * (k + 1) / decodeplugin::INPUT_W;
+            prior[3] = (float)anchor * (k + 1) / decodeplugin::INPUT_H;
 
             //Location
             det->bbox[0] = prior[0] + bbox_reg[idx + k * num_elem * 4] * 0.1 * prior[2];
@@ -84,12 +82,17 @@ namespace nvinfer1
             det->bbox[1] -= det->bbox[3] / 2;
             det->bbox[2] += det->bbox[0];
             det->bbox[3] += det->bbox[1];
-            det->bbox[0] *= input_w;
-            det->bbox[1] *= input_h;
-            det->bbox[2] *= input_w;
-            det->bbox[3] *= input_h;
+            det->bbox[0] *= decodeplugin::INPUT_W;
+            det->bbox[1] *= decodeplugin::INPUT_H;
+            det->bbox[2] *= decodeplugin::INPUT_W;
+            det->bbox[3] *= decodeplugin::INPUT_H;
             det->class_confidence = conf2;
-            anchor *= 2;
+            for (int i = 0; i < 10; i += 2) {
+                det->landmark[i] = prior[0] + lmk_reg[idx + k * num_elem * 10 + num_elem * i] * 0.1 * prior[2];
+                det->landmark[i+1] = prior[1] + lmk_reg[idx + k * num_elem * 10 + num_elem * (i + 1)] * 0.1 * prior[3];
+                det->landmark[i] *= decodeplugin::INPUT_W;
+                det->landmark[i+1] *= decodeplugin::INPUT_H;
+            }
         }
     }
 
@@ -101,17 +104,15 @@ namespace nvinfer1
         int thread_count;
         for (unsigned int i = 0; i < 3; ++i)
         {
-            num_elem = input_h_ / base_step * input_w_ / base_step;
+            num_elem = decodeplugin::INPUT_H / base_step * decodeplugin::INPUT_W / base_step;
             thread_count = (num_elem < thread_count_) ? num_elem : thread_count_;
             CalDetection<<< (num_elem + thread_count - 1) / thread_count, thread_count>>>
-                (inputs[i], output, num_elem, input_h_, input_w_, base_step, base_anchor);
+                (inputs[i], output, num_elem, base_step, base_anchor);
             base_step *= 2;
             base_anchor *= 4;
         }
-
     }
 
-
     int DecodePlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream)
     {
         //assert(batchSize == 1);

retinaface/decode.h

@@ -6,11 +6,12 @@
 namespace decodeplugin
 {
     struct alignas(float) Detection{
-        //x y w h
-        float bbox[4];
+        float bbox[4];  //x1 y1 x2 y2
         float class_confidence;
         float landmark[10];
     };
+    static const int INPUT_H = 384;
+    static const int INPUT_W = 640;
 }
 
 
@@ -52,8 +53,6 @@ namespace nvinfer1
         void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1);
 
     private:
-        const int input_h_ = 384;
-        const int input_w_ = 640;
         int thread_count_ = 256;
     };
 };