疲劳驾驶检测和识别4：C++实现疲劳驾驶检测和识别(含源码，可实时检测)

python libs/convert/convert_torch_to_onnx.py

"""
This code is used to convert the pytorch model into an onnx format model.
"""
import sys
import ossys.path.insert(0, os.getcwd())
import torch.onnx
import onnx
from classifier.models.build_models import get_models
from basetrainer.utils import torch_toolsdef build_net(model_file, net_type, input_size, num_classes, width_mult=1.0):""":param model_file: 模型文件:param net_type: 模型名称:param input_size: 模型输入大小:param num_classes: 类别数:param width_mult::return:"""model = get_models(net_type, input_size, num_classes, width_mult=width_mult, is_train=False, pretrained=False)state_dict = torch_tools.load_state_dict(model_file)model.load_state_dict(state_dict)return modeldef convert2onnx(model_file, net_type, input_size, num_classes, width_mult=1.0, device="cpu", onnx_type="default"):model = build_net(model_file, net_type, input_size, num_classes, width_mult=width_mult)model = model.to(device)model.eval()model_name = os.path.basename(model_file)[:-len(".pth")] + ".onnx"onnx_path = os.path.join(os.path.dirname(model_file), model_name)# dummy_input = torch.randn(1, 3, 240, 320).to("cuda")dummy_input = torch.randn(1, 3, input_size[1], input_size[0]).to(device)# torch.onnx.export(model, dummy_input, onnx_path, verbose=False,#                   input_names=['input'],output_names=['scores', 'boxes'])do_constant_folding = Trueif onnx_type == "default":torch.onnx.export(model, dummy_input, onnx_path, verbose=False, export_params=True,do_constant_folding=do_constant_folding,input_names=['input'],output_names=['output'])elif onnx_type == "det":torch.onnx.export(model,dummy_input,onnx_path,do_constant_folding=do_constant_folding,export_params=True,verbose=False,input_names=['input'],output_names=['scores', 'boxes', 'ldmks'])elif onnx_type == "kp":torch.onnx.export(model,dummy_input,onnx_path,do_constant_folding=do_constant_folding,export_params=True,verbose=False,input_names=['input'],output_names=['output'])onnx_model = onnx.load(onnx_path)onnx.checker.check_model(onnx_model)print(onnx_path)if __name__ == "__main__":net_type = "mobilenet_v2"width_mult = 1.0input_size = [128, 128]num_classes = 2model_file = "work_space/mobilenet_v2_1.0_CrossEntropyLoss/model/best_model_022_98.1848.pth"convert2onnx(model_file, net_type, input_size, num_classes, width_mult=width_mult)

cmake_minimum_required(VERSION 3.5)
project(Detector)add_compile_options(-fPIC) # fix Bug: can not be used when making a shared object
set(CMAKE_CXX_FLAGS "-Wall -std=c++11 -pthread")
#set(CMAKE_CXX_FLAGS_RELEASE "-O2 -DNDEBUG")
#set(CMAKE_CXX_FLAGS_DEBUG "-g")if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)# -DCMAKE_BUILD_TYPE=Debug# -DCMAKE_BUILD_TYPE=Releasemessage(STATUS "No build type selected, default to Release")set(CMAKE_BUILD_TYPE "Release" CACHE STRING "Build type (default Debug)" FORCE)
endif ()# opencv set
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS} ./src/)
#MESSAGE(STATUS "OpenCV_INCLUDE_DIRS = ${OpenCV_INCLUDE_DIRS}")# base_utils
set(BASE_ROOT 3rdparty/base-utils) # 设置base-utils所在的根目录
add_subdirectory(${BASE_ROOT}/base_utils/ base_build) # 添加子目录到build中
include_directories(${BASE_ROOT}/base_utils/include)
include_directories(${BASE_ROOT}/base_utils/src)
MESSAGE(STATUS "BASE_ROOT = ${BASE_ROOT}")# TNN set
# Creates and names a library, sets it as either STATIC
# or SHARED, and provides the relative paths to its source code.
# You can define multiple libraries, and CMake buil ds it for you.
# Gradle automatically packages shared libraries with your APK.
# build for platform
# set(TNN_BUILD_SHARED OFF CACHE BOOL "" FORCE)
if (CMAKE_SYSTEM_NAME MATCHES "Android")set(TNN_OPENCL_ENABLE ON CACHE BOOL "" FORCE)set(TNN_ARM_ENABLE ON CACHE BOOL "" FORCE)set(TNN_BUILD_SHARED OFF CACHE BOOL "" FORCE)set(TNN_OPENMP_ENABLE ON CACHE BOOL "" FORCE)  # Multi-Thread#set(TNN_HUAWEI_NPU_ENABLE OFF CACHE BOOL "" FORCE)add_definitions(-DTNN_OPENCL_ENABLE)           # for OpenCL GPUadd_definitions(-DTNN_ARM_ENABLE)              # for Android CPUadd_definitions(-DDEBUG_ANDROID_ON)            # for Android Logadd_definitions(-DPLATFORM_ANDROID)
elseif (CMAKE_SYSTEM_NAME MATCHES "Linux")set(TNN_OPENCL_ENABLE ON CACHE BOOL "" FORCE)set(TNN_CPU_ENABLE ON CACHE BOOL "" FORCE)set(TNN_X86_ENABLE OFF CACHE BOOL "" FORCE)set(TNN_QUANTIZATION_ENABLE OFF CACHE BOOL "" FORCE)set(TNN_OPENMP_ENABLE ON CACHE BOOL "" FORCE)  # Multi-Threadadd_definitions(-DTNN_OPENCL_ENABLE)           # for OpenCL GPUadd_definitions(-DDEBUG_ON)                    # for WIN/Linux Logadd_definitions(-DDEBUG_LOG_ON)                # for WIN/Linux Logadd_definitions(-DDEBUG_IMSHOW_OFF)            # for OpenCV showadd_definitions(-DPLATFORM_LINUX)
elseif (CMAKE_SYSTEM_NAME MATCHES "Windows")set(TNN_OPENCL_ENABLE ON CACHE BOOL "" FORCE)set(TNN_CPU_ENABLE ON CACHE BOOL "" FORCE)set(TNN_X86_ENABLE ON CACHE BOOL "" FORCE)set(TNN_QUANTIZATION_ENABLE OFF CACHE BOOL "" FORCE)set(TNN_OPENMP_ENABLE ON CACHE BOOL "" FORCE)  # Multi-Threadadd_definitions(-DTNN_OPENCL_ENABLE)           # for OpenCL GPUadd_definitions(-DDEBUG_ON)                    # for WIN/Linux Logadd_definitions(-DDEBUG_LOG_ON)                # for WIN/Linux Logadd_definitions(-DDEBUG_IMSHOW_OFF)            # for OpenCV showadd_definitions(-DPLATFORM_WINDOWS)
endif ()
set(TNN_ROOT 3rdparty/TNN)
include_directories(${TNN_ROOT}/include)
include_directories(${TNN_ROOT}/third_party/opencl/include)
add_subdirectory(${TNN_ROOT}) # 添加外部项目文件夹
set(TNN -Wl,--whole-archive TNN -Wl,--no-whole-archive)# set TNN library
MESSAGE(STATUS "TNN_ROOT = ${TNN_ROOT}")# Detector
include_directories(src)
set(SRC_LISTsrc/object_detection.cppsrc/classification.cppsrc/Interpreter.cpp)
add_library(dmcv SHARED ${SRC_LIST})
target_link_libraries(dmcv ${OpenCV_LIBS} base_utils)
MESSAGE(STATUS "DIR_SRCS = ${SRC_LIST}")add_executable(Detector src/main.cpp)
#add_executable(Detector src/main_for_detect.cpp)
#add_executable(Detector src/main_for_yolov5.cpp)
target_link_libraries(Detector dmcv ${TNN} -lpthread)

//
// Created by Pan on 2020/6/24.
//#include "object_detection.h"
#include "classification.h"
#include <iostream>
#include <string>
#include <vector>
#include "file_utils.h"
#include "image_utils.h"using namespace dl;
using namespace vision;
using namespace std;const int num_thread = 1; // 开启CPU线程数目
DeviceType device = GPU;  // 选择运行设备CPU/GPU
// 人脸检测模型
const char *det_model_file = (char *) "../data/tnn/face/rfb-face-mask-320-320_sim.opt.tnnmodel";
const char *det_proto_file = (char *) "../data/tnn/face/rfb-face-mask-320-320_sim.opt.tnnproto";
ObjectDetectionParam model_param = FACE_MODEL;//模型参数
// 疲劳驾驶分类模型
const char *cls_model_file = (char *) "../data/tnn/drowsy/mobilenet_v2_112_112.tnnmodel";
const char *cls_proto_file = (char *) "../data/tnn/drowsy/mobilenet_v2_112_112.tnnproto";
ClassificationParam ClassParam = DROWSY_MODEL;//模型参数// 设置检测阈值
const float scoreThresh = 0.5;
const float iouThresh = 0.3;
ObjectDetection *detector = new ObjectDetection(det_model_file,det_proto_file,model_param,num_thread,device);Classification *classifier = new Classification(cls_model_file,cls_proto_file,ClassParam,num_thread,device);/**** 测试图片文件*/
void test_image_file() {//测试图片的目录string image_dir = "../data/test_image";std::vector<string> image_list = get_files_list(image_dir);for (string image_path:image_list) {cv::Mat bgr_image = cv::imread(image_path);if (bgr_image.empty()) continue;FrameInfo resultInfo;// 进行人脸检测detector->detect(bgr_image, &resultInfo, scoreThresh, iouThresh);// 进行图像分类classifier->detect(bgr_image, &resultInfo);// 可视化检测结果classifier->visualizeResult(bgr_image, &resultInfo);}delete detector;detector = nullptr;delete classifier;classifier = nullptr;printf("FINISHED.\n");
}/**** 测试视频文件* @return*/
int test_video_file() {//测试视频文件string video_file = "../data/video/video-test.mp4";cv::VideoCapture cap;bool ret = get_video_capture(video_file, cap);cv::Mat frame;while (ret) {cap >> frame;if (frame.empty()) break;FrameInfo resultInfo;// 进行人脸检测detector->detect(frame, &resultInfo, scoreThresh, iouThresh);// 进行图像分类classifier->detect(frame, &resultInfo);// 可视化检测结果classifier->visualizeResult(frame, &resultInfo, 20);}cap.release();delete detector;detector = nullptr;delete classifier;classifier = nullptr;printf("FINISHED.\n");return 0;}/**** 测试摄像头* @return*/
int test_camera() {int camera = 0; //摄像头ID号(请修改成自己摄像头ID号)cv::VideoCapture cap;bool ret = get_video_capture(camera, cap);cv::Mat frame;while (ret) {cap >> frame;if (frame.empty()) break;FrameInfo resultInfo;// 进行人脸检测detector->detect(frame, &resultInfo, scoreThresh, iouThresh);// 进行图像分类classifier->detect(frame, &resultInfo);// 可视化检测结果classifier->visualizeResult(frame, &resultInfo, 20);}cap.release();delete detector;detector = nullptr;delete classifier;classifier = nullptr;printf("FINISHED.\n");return 0;}int main() {test_image_file();   // 测试图片文件//test_video_file();   // 测试视频文件//test_camera();       //测试摄像头return 0;
}

#!/usr/bin/env bash
if [ ! -d "build/" ];thenmkdir "build"
elseecho "exist build"
fi
cd build
cmake ..
make -j4
sleep 1
./demo