init commit nanodet

nanodet/CMakeLists.txt

@@ -0,0 +1,39 @@
+cmake_minimum_required(VERSION 3.5)
+
+project(nanodet_trt)
+
+set(CMAKE_CXX_STANDARD 14)
+
+# CUDA
+find_package(CUDA REQUIRED)
+message(STATUS "Find CUDA include at ${CUDA_INCLUDE_DIRS}")
+message(STATUS "Find CUDA libraries: ${CUDA_LIBRARIES}")
+
+# TensorRT
+set(TENSORRT_ROOT /usr/src/tensorrt/)
+find_path(TENSORRT_INCLUDE_DIR NvInfer.h
+        HINTS ${TENSORRT_ROOT} PATH_SUFFIXES include/)
+message(STATUS "Found TensorRT headers at ${TENSORRT_INCLUDE_DIR}")
+find_library(TENSORRT_LIBRARY_INFER nvinfer
+        HINTS ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
+        PATH_SUFFIXES lib lib64 lib/x64)
+find_library(TENSORRT_LIBRARY_ONNXPARSER nvonnxparser
+        HINTS  ${TENSORRT_ROOT} ${TENSORRT_BUILD} ${CUDA_TOOLKIT_ROOT_DIR}
+        PATH_SUFFIXES lib lib64 lib/x64)
+set(TENSORRT_LIBRARY ${TENSORRT_LIBRARY_INFER} ${TENSORRT_LIBRARY_ONNXPARSER})
+message(STATUS "Find TensorRT libs: ${TENSORRT_LIBRARY}")
+
+# OpenCV
+find_package(OpenCV REQUIRED)
+message(STATUS "Find OpenCV include at ${OpenCV_INCLUDE_DIRS}")
+message(STATUS "Find OpenCV libraries: ${OpenCV_LIBRARIES}")
+
+set(COMMON_INCLUDE ../includes/common)
+set(YAML_INCLUDE ../includes/yaml-cpp/include)
+set(YAML_LIB_DIR ../includes/yaml-cpp/libs)
+
+include_directories(${CUDA_INCLUDE_DIRS} ${TENSORRT_INCLUDE_DIR} ${OpenCV_INCLUDE_DIRS} ${COMMON_INCLUDE} ${YAML_INCLUDE})
+link_directories(${YAML_LIB_DIR})
+
+add_executable(nanodet_trt main.cpp nanodet.cpp)
+target_link_libraries(nanodet_trt ${OpenCV_LIBRARIES} ${CUDA_LIBRARIES} ${TENSORRT_LIBRARY} yaml-cpp)

nanodet/coco.names

@@ -0,0 +1,80 @@
+person
+bicycle
+car
+motorbike
+aeroplane
+bus
+train
+truck
+boat
+traffic light
+fire hydrant
+stop sign
+parking meter
+bench
+bird
+cat
+dog
+horse
+sheep
+cow
+elephant
+bear
+zebra
+giraffe
+backpack
+umbrella
+handbag
+tie
+suitcase
+frisbee
+skis
+snowboard
+sports ball
+kite
+baseball bat
+baseball glove
+skateboard
+surfboard
+tennis racket
+bottle
+wine glass
+cup
+fork
+knife
+spoon
+bowl
+banana
+apple
+sandwich
+orange
+broccoli
+carrot
+hot dog
+pizza
+donut
+cake
+chair
+sofa
+pottedplant
+bed
+diningtable
+toilet
+tvmonitor
+laptop
+mouse
+remote
+keyboard
+cell phone
+microwave
+oven
+toaster
+sink
+refrigerator
+book
+clock
+vase
+scissors
+teddy bear
+hair drier
+toothbrush

nanodet/config.yaml

@@ -0,0 +1,14 @@
+nanodet:
+  onnx_file:     "../nanodet-m.onnx"
+  engine_file:   "../nanodet-m.trt"
+  labels_file:   "../coco.names"
+  BATCH_SIZE:    1
+  INPUT_CHANNEL: 3
+  IMAGE_WIDTH:   320
+  IMAGE_HEIGHT:  320
+  img_mean:      [ 103.53, 116.28, 123.675 ]
+  img_std:       [ 57.375,  57.12,  58.395 ]
+  obj_threshold: 0.35
+  nms_threshold: 0.6
+  strides:       [8, 16, 32]
+  reg_max:       7

nanodet/export_onnx.py

@@ -0,0 +1,32 @@
+import onnx
+import torch
+from nanodet.model.arch import build_model
+from nanodet.util import Logger, cfg, load_config, load_model_weight
+from onnxsim import simplify
+import argparse
+
+
+def main(config, model_path, output_path, input_shape=(320, 320), batch_size=1):
+    logger = Logger(-1, config.save_dir, False)
+    model = build_model(config.model)
+    checkpoint = torch.load(model_path, map_location=lambda storage, loc: storage)
+    load_model_weight(model, checkpoint, logger)
+    dummy_input = torch.autograd.Variable(torch.randn(batch_size, 3, input_shape[0], input_shape[1]))
+    torch.onnx.export(model, dummy_input, output_path, verbose=True, keep_initializers_as_inputs=True, opset_version=10)
+    onnx_model = onnx.load(output_path)  # load onnx model
+    model_simp, check = simplify(onnx_model)
+    assert check, "Simplified ONNX model could not be validated"
+    onnx.save(model_simp, output_path)
+    print('finished exporting onnx ')
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cfg_file', type=str, default='./config/nanodet-m.yml', help='config file path')
+    parser.add_argument('--weights_file', type=str, default='./weights/nanodet_m_oldversion.pth', help='weights file path')
+    parser.add_argument('--output_file', type=str, default='./nanodet-m.onnx', help='onnx file path')
+    parser.add_argument('--img_size', nargs='+', type=int, default=[320, 320], help='image size')
+    parser.add_argument('--batch_size', type=int, default=1, help='batch size')
+    opt = parser.parse_args()
+    load_config(cfg, opt.cfg_file)
+    main(cfg, opt.weights_file, opt.output_file, input_shape=opt.img_size, batch_size=opt.batch_size)

nanodet/main.cpp

@@ -0,0 +1,16 @@
+#include "nanodet.h"
+
+int main(int argc, char **argv)
+{
+    if (argc < 3)
+    {
+        std::cout << "Please design config file and folder name!" << std::endl;
+        return -1;
+    }
+    std::string config_file = argv[1];
+    std::string folder_name = argv[2];
+    nanodet nanodet(config_file);
+    nanodet.LoadEngine();
+    nanodet.InferenceFolder(folder_name);
+    return 0;
+}