yolov5 c++ 识别工件

记录一下训练yolov5识别工件，并用c++调用模型的过程
操作系统：Ubuntu18.04

训练过程

准备数据

使用网站https://app.roboflow.com/cv-rbynj进行数据标注。
创建一个新的工程，新建好的工程如下图所示

然后上传文件,点击上传好的图片，开始标注，框选出物体，并在左上角的框里写入物体的类别，这里为了识别工件，输入workpiece。

当数据都标注好了，点击左上角的白色返回按钮

然后点击右上角finish uploading

然后可以按照自己的需求进行设置，yolov5一般默认图片尺寸为640*640，所以可以把resize修改一下。这里全部当做训练集到出，后面再从训练集中取出一部分作为测试集

设置完成后点generate

然后点击export导出数据，导出时选择 Pascal VOC格式。

VOC格式下，jpg和xml文件会混在一起，将jpg和xml分到不同的文件夹下，可以使用如下代码：

import os
# from PIL import Image
import shutil
folder = "/home/victor/yolov5/workpiece.v1i.voc/train" #原始的路径
new_path_image = "/home/victor/yolov5/workpiece.v1i.voc/image" #图片路径
new_path_xml = "/home/victor/yolov5/workpiece.v1i.voc/xml" #xml文件路径
for index, i in enumerate(sorted(os.listdir(folder))):
    filename1 = os.path.splitext(i)[1]  # 读取文件后缀名
    filename0 = os.path.splitext(i)[0]  # 读取文件名
    print(filename1)
    # m = filename1 == '.xml'
    # print(m)
    if filename1 == '.xml':
        full_path = os.path.join(folder, i)
        despath = os.path.join(new_path_xml, i)
        shutil.move(full_path, despath)
        
    elif filename1 == '.jpg':
        full_path = os.path.join(folder, i)
        despath = os.path.join(new_path_image, i)
        shutil.move(full_path, despath)
    else:
        pass

数据处理主要参考：链接
训练时会报错：AssertionError: train: No labels in /home/robot410/code/yolov5/data/ImageSets/Main/train.cache. Can not train without labels.类似的错误
然后修改Main.py文件

"""
2021.01.19
author:alian
function: create train.txt and test.txt in ImageSets\Main
"""
import os
import random
 
trainval_percent = 0.2   # 可自行进行调节(设置训练和测试的比例是8：2)
train_percent = 1
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)
 
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
 
#ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
#fval = open('ImageSets/Main/val.txt', 'w')
 
for i in list:
    name = total_xml[i][:-4] + '.jpg'+'\n'
    if i in trainval:
        #ftrainval.write(name)
        if i in train:
            ftest.write('/home/robot410/code/yolov5/data/images/' + name)
        #else:
            #fval.write(name)
    else:
        ftrain.write('/home/robot410/code/yolov5/data/images/' + name)
 
#ftrainval.close()
ftrain.close()
#fval.close()
ftest.close()

主要是name那得加上+‘.jpg‘，以及ftest.write和ftrain.write要加上绝对路径，使其可以找到正确图像训练和测试路径，再运行一下:

python Main.py

训练命令参考：

python train.py --data data/workpiece.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 100 --device 0

workpiece.yaml中的内容如下：

# train and val datasets (image directory or *.txt file with image paths)
train: /home/robot410/code/yolov5/data/ImageSets/Main/train.txt  # 上一步我们生成的train.txt，根据自己的路径进行更改（建议用全路径）
val: /home/robot410/code/yolov5/data/ImageSets/Main/test.txt  # 上一步我们生成的test.txt
#test: ../coco/test-dev2017.txt  # 20k images for submission to https://competitions.codalab.org/competitions/20794
 
# number of classes
nc: 1   # 训练的类别
 
# class namesdata
names: ['workpiece']

weights选择yolo提供的几种模型中的一个即可。
训练好之后。测试命令参考:

python detect.py --weight /home/robot410/code/yolov5/runs/train/exp10/weights/best.pt --source /home/robot410/code/yolov5/2021-12-1/2021-12-10-16-35/rgb

weight后是训练好的模型的路径，source是测试图片的根目录，文件默认保存在runs/detect下。

C++下的部署主要参考https://github.com/UNeedCryDear/yolov5-opencv-dnn-cpp
适配于yolov5-6.0，这个代码是Windows系统下的代码，环境只依赖opencv，需要4.5及以上版本，这里附上自己的CMakeLists:

cmake_minimum_required(VERSION 2.8)
project(yolo)

IF(NOT CMAKE_BUILD_TYPE)
    SET(CMAKE_BUILD_TYPE Release)
ENDIF()

MESSAGE("Build type: " ${CMAKE_BUILD_TYPE})

set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}  -Wall  -O3 -march=native ")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall   -O3 -march=native")

# Check C++11 or C++0x support
include(CheckCXXCompilerFlag)
CHECK_CXX_COMPILER_FLAG("-std=c++11" COMPILER_SUPPORTS_CXX11)
CHECK_CXX_COMPILER_FLAG("-std=c++0x" COMPILER_SUPPORTS_CXX0X)
if(COMPILER_SUPPORTS_CXX11)
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11")
    add_definitions(-DCOMPILEDWITHC11)
    message(STATUS "Using flag -std=c++11.")
elseif(COMPILER_SUPPORTS_CXX0X)
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++0x")
    add_definitions(-DCOMPILEDWITHC0X)
    message(STATUS "Using flag -std=c++0x.")
else()
    message(FATAL_ERROR "The compiler ${CMAKE_CXX_COMPILER} has no C++11 support. Please use a different C++ compiler.")
endif()

LIST(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake_modules)

set(OpenCV_DIR "/home/robot410/software/opencv-4.5.4/build")
find_package(OpenCV REQUIRED)
#if(NOT OpenCV_FOUND)
#    find_package(OpenCV 2.4.3 QUIET)
#    if(NOT OpenCV_FOUND)
#        message(FATAL_ERROR "OpenCV > 2.4.3 not found.")
#    endif()
#endif()

message( ${OpenCV_INCLUDE_DIRS})
message( ${OpenCV_LIBS})
#find_package(Eigen3 3.1.0 REQUIRED)
#find_package(Pangolin REQUIRED)

include_directories(
        ${PROJECT_SOURCE_DIR}
#        ${PROJECT_SOURCE_DIR}/include
#        ${EIGEN3_INCLUDE_DIR}
#        ${Pangolin_INCLUDE_DIRS}
)
add_library(${PROJECT_NAME} SHARED
        yolo.cpp
        )

target_link_libraries(${PROJECT_NAME}
        ${OpenCV_LIBS}
        )
add_executable(main
        main.cpp)
target_link_libraries(main ${PROJECT_NAME})

需要将yolo训练好的pt模型转换成onnx模型才可以用c++的opencv调用，yolov5中的export.py文件可以直接拿来用，其中的opset需要设置成12（原来默认的是13，在opencv下面会报错），运行代码参考：

python export.py --weights /home/robot410/code/yolov5/runs/train/exp10/weights/best.pt --img 640 --include onnx

onnx模型默认会在pt模型的相同目录下
Ubuntu没有stdafx.h文件，直接把含有#include"stdafx.h"的都注释掉，不影响
修改main.cpp中的模型路径和测试图片路径，这里的测试路径最好用绝对路径，使用原本默认的相对路径会有问题

string img_path = "/home/robot410/code/yolov5/2021-12-1/2021-12-10-16-35/rgb/000066.png";
	string model_path = "/home/robot410/yolov5-opencv-dnn-cpp/best.onnx";

yolo.h中的类型需要做相对应的修改：

std::vector<std::string> className = { "workpiece" };

然后对c++文件进行编译

mkdir build
cd build
cmake ..
make

最后运行：

./main

运行效果如下：