在深度学习领域,YOLO(You Only Look Once)模型因其速度和准确率在目标检测任务中备受瞩目。YOLOv3-Tiny 版本尤其适用于需要快速检测的应用场景。在这篇文章中,我们将使用 C++ 实现基于 YOLOv3-Tiny 的验证码识别任务。
准备工作
系统与环境
操作系统:Ubuntu 16.04
C++ 编译器:GCC
CUDA 版本:10.0
GPU:GeForce GTX 1080
工具与资源
Darknet:深度学习框架,用于实现 YOLO 模型 GitHub - Darknet
labelImg:用于标注训练数据的工具 GitHub - labelImg
实践步骤
安装环境
首先,确保 CUDA 等环境已安装完毕。我们将使用 GPU 进行训练。然后,将 Darknet 克隆到本地:
bash
git clone https://github.com/pjreddie/darknet
cd darknet
编译 Darknet
进入 darknet 文件夹,修改 Makefile 文件:
Makefile
GPU=1 # 使用 GPU
CUDNN=1 # 使用 CUDNN
OPENCV=0 # 不使用 OpenCV
OPENMP=0
DEBUG=0
使用以下命令进行编译:
bash
make
编译完成后,会生成一个名为 darknet 的二进制文件。
修改配置文件
我们选择 cfg/yolov3-tiny.cfg 文件进行修改,以适应我们的任务需求:
ini
# Line 3
batch=24 # 每个训练步骤使用 24 张图片
# Line 4
subdivisions=8 # 将 batch 分成 8 个部分
# Line 127
filters=(classes + 5)*3 # 我们的任务中 filters=18
# Line 135
classes=1 # 我们要检测的类别数
# Line 171
filters=(classes + 5)*3 # 同上
# Line 177
classes=1 # 同上
关于 filters = (classes + 5) * 3 的计算方式,可以参考 YOLOv3 论文第 2.3 节内容 YOLOv3 Paper。
准备数据
我们使用 Selenium 的 Google Driver 下载了大量验证码图片,然后使用 labelImg 工具进行数据标注。
下载并编译 labelImg:
bash
git clone https://github.com/tzutalin/labelImg
cd labelImg
pip install pyqt5 lxml
make qt5py3
完成安装后,使用 labelImg 工具进行图片的标注工作。由于我们只有一类内容,因此标注的框只有一类。
生成训练数据
将标注好的数据整理成训练所需的格式。我们编写了 data.py 脚本生成训练数据配置文件,并将标注文件转换为 darknet 可识别的格式。
python
import os
import glob
def generate_txt_file(img_dir, output_file):
img_files = glob.glob(os.path.join(img_dir, '*.jpg'))
with open(output_file, 'w') as f:
for img in img_files:
f.write(img + '\n')
# 使用示例
generate_txt_file('imgs', 'train.txt')
generate_txt_file('imgs_val', 'val.txt')
下载预训练权重
在 weights 文件夹下使用 darknet53.sh 文件下载预训练的权重文件:
bash
cd weights
./darknet53.sh
配置文件准备
在 data 文件夹中准备如下文件:
train.txt - 训练数据文件,由 data.py 生成
char.names - 类别配置文件
val.txt - 验证数据文件
train.data - 训练配置表
开始训练
进入 darknet 文件夹,使用以下命令开始训练:
bash
./darknet detector train data/train.data cfg/yolov3-tiny.cfg weights/darknet53.conv.74 -gpus 0
通常,当训练的 loss 达到 0.6 左右时,就可以停止训练。
验证结果
使用以下命令验证训练好的模型:
bash
./darknet detector test data/train.data cfg/yolov3-tiny.cfg weights/yolov3-tiny_final.weights test_imgs/1.jpg -thresh 0.5 -gpus 0
在 darknet 文件夹中会生成一个名为 predictions.jpg 的文件,这是我们验证出来的结果文件。
C++ 代码实现
我们使用 C++ 语言来实现上述过程中的部分关键步骤。以下是 C++ 代码示例:
安装依赖库
首先,我们需要安装必要的 C++ 库。我们将使用 OpenCV 库来进行图像处理和目标检测。
bash
sudo apt-get update
sudo apt-get install libopencv-dev
读取配置文件并进行推理
我们将使用 OpenCV 库来进行图像处理和目标检测。
cpp
#include <opencv2/opencv.hpp>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
using namespace cv;
using namespace std;
void drawPred(float conf, int left, int top, int right, int bottom, Mat& frame) {
// 绘制边界框
rectangle(frame, Point(left, top), Point(right, bottom), Scalar(0, 255, 0));
// 在边界框上显示置信度
string label = format("%.2f", conf);
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = max(top, labelSize.height);
rectangle(frame, Point(left, top - round(1.5 * labelSize.height)), Point(left + round(1.5 * labelSize.width), top + baseLine), Scalar(255, 255, 255), FILLED);
putText(frame, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 0.75, Scalar(0, 0, 0), 1);
}
int main(int argc, char** argv) {
// 加载 YOLOv3-tiny 模型
String modelConfiguration = "cfg/yolov3-tiny.cfg";
String modelWeights = "yolov3-tiny.weights";
Net net = readNetFromDarknet(modelConfiguration, modelWeights);
// 设置为使用 GPU
net.setPreferableBackend(DNN_BACKEND_CUDA);
net.setPreferableTarget(DNN_TARGET_CUDA);
// 读取图片
Mat frame = imread("test_imgs/1.jpg");
// 创建 blob
Mat blob;
blobFromImage(frame, blob, 1/255.0, Size(416, 416), Scalar(), true, false);
// 设置输入 blob
net.setInput(blob);
// 推理检测
vector<Mat> outs;
vector<String> outNames = net.getUnconnectedOutLayersNames();
net.forward(outs, outNames);
// 解析检测结果
float confThreshold = 0.5;
for (size_t i = 0; i < outs.size(); ++i) {
float* data = (float*)outs[i].data;
for (int j = 0; j < outs[i].rows; ++j, data += outs[i].cols) {
Mat scores = outs[i].row(j).colRange(5, outs[i].cols);
Point classIdPoint;
double confidence;
minMaxLoc(scores, 0, &confidence, 0, &classIdPoint);
if (confidence > confThreshold) {
int centerX = (int)(data[0] * frame.cols);
int centerY = (int)(data[1] * frame.rows);
int width = (int)(data[2] * frame.cols);
int height = (int)(data[3] * frame.rows);
int left = centerX - width / 2;
int top = centerY - height / 2;
// 绘制边界框
drawPred((float)confidence, left, top, left + width, top + height, frame);
}
}
}
// 保存结果图片
imwrite("predictions.jpg", frame);
cout << "Result saved to predictions.jpg" << endl;
return 0;
}