Merge pull request #167 from CoinCheung/dev

add ncnn demo

Author: CoinCheung

.gitignore

@@ -116,3 +116,4 @@ datasets/coco/val.txt
 pretrained/*
 dist_train.sh
 
+tvm/

ncnn/CMakeLists.txt

@@ -0,0 +1,15 @@
+CMAKE_MINIMUM_REQUIRED(VERSION 3.15)
+
+PROJECT(segment)
+
+set(CMAKE_CXX_FLAGS "-std=c++14 -O2")
+
+
+set (ncnn_DIR ${NCNN_ROOT}/lib/cmake/ncnn)
+find_package(OpenCV REQUIRED)
+find_package(ncnn REQUIRED)
+
+
+add_executable(segment segment.cpp)
+target_include_directories(segment PUBLIC ${OpenCV_INCLUDE_DIRS})
+target_link_libraries(segment ${OpenCV_LIBRARIES} ncnn)

ncnn/README.md

@@ -0,0 +1,67 @@
+
+### My platform
+
+* ubuntu 18.04
+* Intel(R) Xeon(R) Gold 6240 CPU @ 2.60GHz
+* cmake 3.17.1
+* opencv built from source
+
+### NOTE
+
+Though this demo runs on x86 platform, you can also use it on mobile platforms. NCNN is better optimized on mobile platforms.
+
+
+### Install ncnn
+
+1. dependencies
+```
+$ python -m pip install onnx-simplifier
+```
+
+2. build ncnn
+Just following the ncnn official tutoral: [build-for-linux](https://github.com/Tencent/ncnn/wiki/how-to-build#build-for-linux) to install ncnn:
+
+1) dependencies
+```
+# apt install build-essential git libprotobuf-dev protobuf-compiler 
+```
+
+2) (optional) install vulkan
+
+3) install opencv from source
+
+4) build 
+```
+## I am using commit 9391fae741a1fb8d58cdfdc92878a5e9800f8567, and I have not tested over newer commits
+$ git clone https://github.com/Tencent/ncnn.git
+$ $cd ncnn
+$ git submodule update --init
+$ mkdir -p build
+$ cmake -DCMAKE_TOOLCHAIN_FILE=../toolchains/host.gcc.toolchain.cmake ..
+$ make -j
+$ make install 
+```
+
+### convert model, build and run the demo
+
+1. convert pytorch model to ncnn model via onnx
+```
+$ cd BiSeNet/
+$ python tools/export_onnx.py --aux-mode eval --config configs/bisenetv2_city.py --weight-path /path/to/your/model.pth --outpath ./model_v2.onnx 
+$ python -m onnxsim model_v2.onnx model_v2_sim.onnx
+$ /path/to/ncnn/build/tools/onnx/onnx2ncnn model_v2_sim.onnx model_v2_sim.param model_v2_sim.bin
+$ mv model_v2_sim.param ncnn/models 
+$ mv model_v2_sim.bin ncnn/models 
+```
+
+2. compile demo code
+```
+cd ncnn/build
+cmake .. -DNCNN_ROOT=/path/to/ncnn/build/install
+make
+```
+
+3. run demo
+```
+./segment
+```

ncnn/segment.cpp

@@ -0,0 +1,117 @@
+
+#include "net.h"
+#include "mat.h"
+
+#include <opencv2/core/core.hpp>
+#include <opencv2/highgui/highgui.hpp>
+#include <opencv2/imgproc/imgproc.hpp>
+
+#include <iostream>
+#include <random>
+#include <algorithm>
+#include <stdio.h>
+#include <vector>
+
+
+using std::string;
+using std::vector;
+using cv::Mat;
+
+
+vector<vector<uint8_t>> get_color_map();
+void inference();
+
+
+int main(int argc, char** argv) {
+    inference();
+    return 0;
+}
+
+
+void inference() {
+    bool use_fp16 = false;
+    // load model
+    ncnn::Net mod;
+#if NCNN_VULKAN
+    int gpu_count = ncnn::get_gpu_count();
+    if (gpu_count <= 0) {
+        fprintf(stderr, "we do not have gpu device\n");
+        return;
+    }
+    mod.opt.use_vulkan_compute = 1;
+    mod.set_vulkan_device(1);
+#endif 
+    mod.load_param("../models/model_v2_sim.param");
+    mod.load_model("../models/model_v2_sim.bin");
+    mod.opt.use_fp16_packed = use_fp16;
+    mod.opt.use_fp16_storage = use_fp16;
+    mod.opt.use_fp16_arithmetic = use_fp16;
+
+    // load image, and copy to ncnn mat
+    int oH{1024}, oW{2048}, n_classes{19};
+    float mean[3] = {0.3257f, 0.3690f, 0.3223f};
+    float var[3] = {0.2112f, 0.2148f, 0.2115f};
+    cv::Mat im = cv::imread("../../example.png");
+    if (im.empty()) {
+        fprintf(stderr, "cv::imread failed\n");
+        return;
+    }
+    ncnn::Mat inp = ncnn::Mat::from_pixels_resize(
+            im.data, ncnn::Mat::PIXEL_BGR, im.cols, im.rows, oW, oH);
+    for (float &el : mean) el *= 255.;
+    for (float &el : var) el = 1. / (255. * el);
+    inp.substract_mean_normalize(mean, var);
+
+    // set input, run, get output
+    ncnn::Extractor ex = mod.create_extractor();
+    // ex.set_num_threads(1);
+#if NCNN_VULKAN
+    ex.set_vulkan_compute(true);
+#endif
+
+    ex.input("input_image", inp);
+    ncnn::Mat out;
+    ex.extract("preds", out); // output is nchw, as onnx, where here n=1
+
+    // generate colorful output, and dump
+    vector<vector<uint8_t>> color_map = get_color_map();
+    Mat pred(cv::Size(oW, oH), CV_8UC3);
+    for (int i{0}; i < oH; ++i) {
+        uint8_t *ptr = pred.ptr<uint8_t>(i);
+        for (int j{0}; j < oW; ++j) {
+            // compute argmax
+            int idx, offset, argmax{0}; 
+            float max;
+            idx = i * oW + j;
+            offset = oH * oW;
+            max = out[idx];
+            for (int k{1}; k < n_classes; ++k) {
+                idx += offset;
+                if (max < out[idx]) {
+                    max = out[idx];
+                    argmax = k;
+                }
+            }
+            // color the result
+            ptr[0] = color_map[argmax][0];
+            ptr[1] = color_map[argmax][1];
+            ptr[2] = color_map[argmax][2];
+            ptr += 3;
+        }
+    }
+    cv::imwrite("out.png", pred);
+
+}
+
+
+vector<vector<uint8_t>> get_color_map() {
+    vector<vector<uint8_t>> color_map(256, vector<uint8_t>(3));
+    std::minstd_rand rand_eng(123);
+    std::uniform_int_distribution<uint8_t> u(0, 255);
+    for (int i{0}; i < 256; ++i) {
+        for (int j{0}; j < 3; ++j) {
+            color_map[i][j] = u(rand_eng);
+        }
+    }
+    return color_map;
+}

tools/export_libtorch.py

@@ -25,7 +25,7 @@
 if cfg.use_sync_bn: cfg.use_sync_bn = False
 
 net = model_factory[cfg.model_type](cfg.n_cats, aux_mode='pred')
-net.load_state_dict(torch.load(args.weight_pth), strict=False)
+net.load_state_dict(torch.load(args.weight_pth, map_location='cpu'), strict=False)
 net.eval()
 
 

tools/export_onnx.py

@@ -18,14 +18,16 @@
         default='model_final.pth')
 parse.add_argument('--outpath', dest='out_pth', type=str,
         default='model.onnx')
+parse.add_argument('--aux-mode', dest='aux_mode', type=str,
+        default='pred')
 args = parse.parse_args()
 
 
 cfg = set_cfg_from_file(args.config)
 if cfg.use_sync_bn: cfg.use_sync_bn = False
 
-net = model_factory[cfg.model_type](cfg.n_cats, aux_mode='pred')
-net.load_state_dict(torch.load(args.weight_pth), strict=False)
+net = model_factory[cfg.model_type](cfg.n_cats, aux_mode=args.aux_mode)
+net.load_state_dict(torch.load(args.weight_pth, map_location='cpu'), strict=False)
 net.eval()
 
 

tools/train.py

@@ -200,7 +200,7 @@ def train():
 
     logger.info('\nevaluating the final model')
     torch.cuda.empty_cache()
-    heads, mious = eval_model(cfg, net)
+    heads, mious = eval_model(cfg, net.module)
     logger.info(tabulate([mious, ], headers=heads, tablefmt='orgtbl'))
 
     return