基于YOLOv3-Tiny的验证码定位实践 (JavaScript 版本)

在深度学习的目标检测领域，常见的检测方法主要有两大类：one-stage 和 two-stage。one-stage 方法中，YOLO（You Only Look Once）模型因其速度和准确率而备受瞩目。YOLOv3-Tiny 版本尤其适用于需要快速检测的应用场景。在本文中，我们将介绍如何使用 YOLOv3-Tiny 模型来进行验证码的定位任务。

准备工作
系统与环境
操作系统：Ubuntu 16.04
Node.js 版本：14.x
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 的文件，这是我们验证出来的结果文件。

JavaScript 代码实现
我们使用 JavaScript 来实现上述过程中的部分关键步骤。以下是 JavaScript 代码示例：

安装依赖库
首先，我们需要安装必要的 JavaScript 库。我们将使用 tensorflow.js 来进行图像处理和目标检测。

bash

npm install @tensorflow/tfjs-node @tensorflow-models/coco-ssd
读取配置文件并进行推理
我们将使用 tensorflow.js 和 node-canvas 来进行图像处理和目标检测。

javascript

const fs = require('fs');
const path = require('path');
const canvas = require('canvas');
const cocoSsd = require('@tensorflow-models/coco-ssd');
const tf = require('@tensorflow/tfjs-node');

async function loadImage(imgPath) {
    const image = await canvas.loadImage(imgPath);
    const cnv = canvas.createCanvas(image.width, image.height);
    const ctx = cnv.getContext('2d');
    ctx.drawImage(image, 0, 0, image.width, image.height);
    return tf.browser.fromPixels(cnv);
}

async function detectObjects(model, imgPath) {
    const image = await loadImage(imgPath);
    const predictions = await model.detect(image);
    return predictions;
}

async function drawBoundingBoxes(imgPath, predictions, outputPath) {
    const image = await canvas.loadImage(imgPath);
    const cnv = canvas.createCanvas(image.width, image.height);
    const ctx = cnv.getContext('2d');
    ctx.drawImage(image, 0, 0, image.width, image.height);

    ctx.strokeStyle = 'red';
    ctx.lineWidth = 4;

    predictions.forEach(prediction => {
        ctx.beginPath();
        ctx.rect(...prediction.bbox);
        ctx.stroke();
        ctx.font = '18px Arial';
        ctx.fillStyle = 'red';
        ctx.fillText(prediction.class, prediction.bbox[0], prediction.bbox[1] > 10 ? prediction.bbox[1] - 5 : 10);
    });

    const buffer = cnv.toBuffer('image/jpeg');
    fs.writeFileSync(outputPath, buffer);
}

async function main() {
    const imgPath = 'test_imgs/1.jpg';
    const outputPath = 'predictions.jpg';

    const model = await cocoSsd.load();
    const predictions = await detectObjects(model, imgPath);
    await drawBoundingBoxes(imgPath, predictions, outputPath);

    console.log(`Result saved to ${outputPath}`);
}

main();