yolov4-自定义函数

这是一个在 TensorFlow、TFLite 和 TensorRT 中实现的 YOLOv4、YOLOv4-tiny、YOLOv3 和 YOLOv3-tiny 的多种自定义函数库。

免责声明：本仓库与我的另一个仓库 tensorflow-yolov4-tflite 非常相似。我创建这个仓库是为了探索使用 YOLOv4 实现自定义函数，但这些自定义函数可能会降低整体运行速度，并且在时间复杂度方面可能不够优化。因此，如果您希望运行最优化的 TensorFlow 版 YOLOv4 代码，请前往我的另一个仓库。而这个仓库则是为了探索利用 YOLOv4 可以实现的酷炫自定义功能和应用场景！

对象计数自定义函数的实际演示！

当前支持的自定义函数及标志

• 对象计数（总对象数及各类别对象数）
• 打印每检测结果的相关信息（类别、置信度、边界框坐标）
• 裁剪检测结果并保存为新图像
• 使用 Tesseract OCR 进行车牌识别
• 对检测结果应用 Tesseract OCR 提取文本

如果您希望看到某个特定的自定义功能被实现，请在 Issues 标签页中提交 issue 并提出建议！如果足够多的人提出相同的自定义需求，我会尽快将其添加到项目中！

快速入门

Conda（推荐）

# TensorFlow CPU
conda env create -f conda-cpu.yml
conda activate yolov4-cpu

# TensorFlow GPU
conda env create -f conda-gpu.yml
conda activate yolov4-gpu

Pip

# TensorFlow CPU
pip install -r requirements.txt

# TensorFlow GPU
pip install -r requirements-gpu.txt

Nvidia 驱动程序（适用于 GPU，若您未使用 Conda 环境且尚未配置 CUDA）

请确保使用 CUDA Toolkit 10.1 版本，因为这是本仓库所用 TensorFlow 版本的适配版本。
https://developer.nvidia.com/cuda-10.1-download-archive-update2

下载官方预训练权重

YOLOv4 已经预训练完成，能够检测 80 类物体。为了便于演示，我们将使用预训练权重。
下载预训练的 yolov4.weights 文件：https://drive.google.com/open?id=1cewMfusmPjYWbrnuJRuKhPMwRe_b9PaT

将下载的 yolov4.weights 文件从您的下载文件夹复制并粘贴到本仓库的 data 文件夹中。

如果您想使用 yolov4-tiny.weights——一个运行检测更快但精度稍低的小型模型——请从这里下载：https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v4_pre/yolov4-tiny.weights

使用自定义训练的 YOLOv4 权重

在此学习如何训练自定义 YOLOv4 权重：https://www.youtube.com/watch?v=mmj3nxGT2YQ

观看我在 TensorFlow 中使用自定义模型的详细讲解：https://www.youtube.com/watch?v=nOIVxi5yurE

使用我的车牌训练的自定义权重：https://drive.google.com/file/d/1EUPtbtdF0bjRtNjGv436vDY28EN5DXDH/view?usp=sharing

将您的自定义 .weights 文件复制并粘贴到 data 文件夹中，同时将自定义 .names 文件复制并粘贴到 data/classes/ 文件夹中。

为了让您的自定义模型正常工作，您需要在代码中仅做一处修改：将 core/config.py 文件第 14 行的路径更新为指向您的自定义 .names 文件，如下所示。（我的自定义 .names 文件名为 custom.names，但您的文件名可能不同）

注意：如果您使用的是预训练的 yolov4，则务必保持第 14 行为 coco.names。

使用 TensorFlow 的 YOLOv4（tf、.pb 模型）

要使用 TensorFlow 实现 YOLOv4，首先需要将 .weights 转换为对应的 TensorFlow 模型文件，然后运行该模型。

# 将 Darknet 权重转换为 TensorFlow
## yolov4
python save_model.py --weights ./data/yolov4.weights --output ./checkpoints/yolov4-416 --input_size 416 --model yolov4 

# 运行 YOLOv4 TensorFlow 模型
python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/kite.jpg

# 在视频上运行 YOLOv4
python detect_video.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --video ./data/video/video.mp4 --output ./detections/results.avi

# 在网络摄像头上运行 YOLOv4
python detect_video.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --video 0 --output ./detections/results.avi

如果您想运行 yolov3 或 yolov3-tiny，请在上述命令中将 --model yolov3 和 .weights 文件替换为相应内容。

注意：您还可以通过修改 --images 标志一次性处理多张图片，例如：--images "./data/images/kite.jpg, ./data/images/dog.jpg"

输出图像（常规 TensorFlow）

您可以在 detections 文件夹中找到包含检测结果的输出图像。

预训练 YOLOv4 模型示例

输出视频

视频会保存到您指定的 --output 标志所指向的路径。如果您未设置该标志，则视频不会保存检测结果。

使用 TensorFlow 的 YOLOv4-Tiny

以下命令可用于运行 yolov4-tiny 模型。

# yolov4-tiny
python save_model.py --weights ./data/yolov4-tiny.weights --output ./checkpoints/yolov4-tiny-416 --input_size 416 --model yolov4 --tiny

# 运行 YOLOv4-Tiny TensorFlow 模型
python detect.py --weights ./checkpoints/yolov4-tiny-416 --size 416 --model yolov4 --images ./data/images/kite.jpg --tiny

# 自定义 YOLOv4
python save_model.py --weights ./data/custom.weights --output ./checkpoints/custom-416 --input_size 416 --model yolov4 

# 运行自定义 YOLOv4 TensorFlow 模型
python detect.py --weights ./checkpoints/custom-416 --size 416 --model yolov4 --images ./data/images/car.jpg

我在 core/functions.py 文件中创建了一个自定义函数，可用于统计并跟踪在每张图像或视频中某一时刻检测到的物体数量。该函数既可统计检测到的物体总数，也可按类别统计检测到的物体数量。

统计总物体数

要统计总物体数，只需在 detect.py 或 detect_video.py 命令中添加自定义标志 --count 即可。

# 运行 YOLOv4 模型并统计检测到的总物体数
python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --count

运行上述命令后，将统计检测到的物体总数，并将其输出到您的命令提示符或终端，同时也会保存在检测结果文件中，如下所示：

按类别统计物体数

要统计目标检测器中每个类别的物体数量，需添加自定义标志 --count，并修改 detect.py 或 detect_video.py 脚本中的某一行代码。默认情况下，count_objects 函数有一个名为 by_class 的参数，其值为 False。若将该参数改为 True，则会按类别进行统计。

要按类别统计，需将 detect.py 或 detect_video.py 修改为如下形式：

然后运行与前述相同的命令：

# 运行 YOLOv4 模型并按类别统计物体数
python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --count

运行上述命令后，将统计每个类别的检测物体数量，并将其输出到您的命令提示符或终端，同时也会保存在检测结果文件中，如下所示：

注意： 您也可以将 --count 标志添加到 detect_video.py 命令中！

如果您想修改输出的具体信息，可以编辑位于 core/utils.py 文件中的 draw_bbox 函数。打印信息的代码行如下：

添加 --info 标志后的示例命令：

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --info

终端或命令行中的输出结果如下：

注意： 您也可以将 --info 标志添加到 detect_video.py 命令中！

我在 core/functions.py 文件中创建了一个自定义函数，可应用于任何 detect.py 或 detect_video.py 命令，用于裁剪 YOLOv4 的检测结果，并将每个检测结果保存为单独的新图像。要裁剪检测结果，只需在任意命令中添加 --crop 标志即可。裁剪后的图像将保存在 detections/crop/ 文件夹中。

添加 --crop 标志后的示例命令：

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --crop

以下是上述命令生成的其中一个裁剪后的检测结果示例。

我创建了一个自定义函数，可将由我的自定义 YOLOv4 模型检测到的车牌边界框区域输入 Tesseract OCR，以读取并提取车牌号码。为了从图像中准确提取车牌号码，会对车牌进行充分的预处理。负责执行预处理和文本提取的函数名为 recognize_plate，位于 core/utils.py 文件中。

免责声明： 若要运行 Tesseract OCR，您必须先下载二进制文件并在本地机器上完成安装。请务必在继续操作前完成此步骤，否则命令将无法正常运行！

Tesseract OCR 官方 GitHub 仓库：tesseract-ocr/tessdoc

关于如何在 Mac 或 Linux 机器上安装 Tesseract 的优秀文章：PyImageSearch 文章

对于 Windows 系统，推荐使用：Windows 安装指南

成功安装 Tesseract 后，即可继续操作。若您尚未训练出用于检测车牌的 YOLOv4 模型，也可直接使用我已训练好的模型。虽然该模型并不完美，但效果尚佳。在此下载车牌检测模型，并学习如何使用 TensorFlow 保存和运行它

在图像上运行车牌识别（下方为视频示例）

车牌识别在图像上的表现非常出色。您只需在运行自定义 YOLOv4 模型的命令中添加 --plate 标志即可。

试试在仓库中的这张图片上运行吧！

# 运行车牌识别
python detect.py --weights ./checkpoints/custom-416 --size 416 --model yolov4 --images ./https://oss.gittoolsai.com/images/theAIGuysCode_yolov4-custom-functions_readme_22c1b2608d4b.jpg --plate

示例结果图

上述命令的输出将在您的命令终端中显示所有检测到的车牌号码，并将以下图像输出并保存到 detections 文件夹中。

您应该能够在屏幕上看到车牌号码，显示在 YOLOv4 检测到的边界框上方。

幕后揭秘

本节将重点介绍我为实现基于 YOLOv4 的车牌识别所采取的步骤，以及未来可进一步优化的方向。

本次演示将以如下原始图像为例，展示整个流程的每一步操作。

第一步是从 YOLOv4 中获取边界框坐标，并简单地截取该边界框内的子图像区域。由于这张图非常小，我们通常会使用 cv2.resize() 将其放大至原图尺寸的 3 倍。

随后，我们将图像转换为灰度图，并应用一个小的高斯模糊以使其更加平滑。

接着，对图像进行二值化处理，使白字黑底的文本更清晰，并同时采用 Otsu 方法进行自动阈值分割。这种白字黑底的图像有助于后续提取轮廓。

然后，利用 OpenCV 对图像进行膨胀操作，以增强轮廓的可见性，便于后续步骤中更好地检测到这些轮廓。

接下来，我们使用 OpenCV 在图像中查找所有矩形形状的轮廓，并按从左到右的顺序对其进行排序。

如您所见，这样做会检测到许多非车牌数字字符的轮廓。为了过滤掉这些不必要的区域，我们设置了若干参数来判定一个轮廓是否有效。这些参数主要包括高度与宽度的比例（例如，区域的高度至少应为图像总高度的 1/6），此外还设置了区域面积等其他条件。具体细节请参阅代码。经过这一系列筛选后，最终得到的结果如下。

现在，车牌数字中的每个字符都成为了我们关注的唯一区域。我们将每个子图像单独分割出来，并应用 bitwise_not 掩码，将图像反转为黑字白底，这样 Tesseract 的识别准确率会更高。最后一步是对图像进行小幅中值模糊处理，然后将其送入 Tesseract 进行文字识别，从而提取出对应的字母或数字。以下是字母输入 Tesseract 后的示例效果。

随后，我们将每个字母或数字依次拼接成字符串，最终便得到了完整的车牌识别结果！砰！

在视频上运行车牌识别

直接在视频上实时运行车牌识别，并同时进行 YOLOv4 的目标检测，会遇到一些问题。Tesseract OCR 的时间复杂度较高，会导致视频处理速度大幅下降，几乎慢得像蜗牛一样。不过，通过在任何 detect_video.py 命令中添加 --plate 命令行参数，仍然可以实现这一功能。

然而，我认为最佳方案是先不加 --plate 参数运行视频检测，而是使用 --crop 参数，将屏幕上检测到的目标裁剪下来并保存为新图像。点击此处查看具体操作。待视频以较高帧率完成处理后，所有车牌图像都会被裁剪并保存到 detections/crop 文件夹中。我在仓库中新增了一个名为 license_plate_recognizer.py 的简易脚本，您可以直接运行它来识别车牌。此外，您还可以轻松自定义该脚本，以进一步提升识别效果。未来我将在仓库的后续提交中实现这一功能的自动化集成。

使用 detect_video.py 运行车牌识别的命令如下：

python detect_video.py --weights ./checkpoints/custom-416 --size 416 --model yolov4 --video ./data/video/license_plate.mp4 --output ./detections/recognition.avi --plate

而我认为更为高效的推荐方案是使用以下命令：在代码中自定义检测结果的裁剪频率。

python detect_video.py --weights ./checkpoints/custom-416 --size 416 --model yolov4 --video ./data/video/license_plate.mp4 --output ./detections/recognition.avi --crop

现在就动手玩一玩 license_plate_recognizer.py，尽情享受吧！

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --ocr

# 保存 TF 模型以便转换为 TFLite
python save_model.py --weights ./data/yolov4.weights --output ./checkpoints/yolov4-416 --input_size 416 --model yolov4 --framework tflite

# YOLOv4 转换为 TFLite
python convert_tflite.py --weights ./checkpoints/yolov4-416 --output ./checkpoints/yolov4-416.tflite

# YOLOv4 量化为 FP16
python convert_tflite.py --weights ./checkpoints/yolov4-416 --output ./checkpoints/yolov4-416-fp16.tflite --quantize_mode float16

# YOLOv4 量化为 INT8
python convert_tflite.py --weights ./checkpoints/yolov4-416 --output ./checkpoints/yolov4-416-int8.tflite --quantize_mode int8 --dataset ./coco_dataset/coco/val207.txt

# 运行 TFLite 模型
python detect.py --weights ./checkpoints/yolov4-416.tflite --size 416 --model yolov4 --images ./data/images/kite.jpg --framework tflite

python save_model.py --weights ./data/yolov3.weights --output ./checkpoints/yolov3.tf --input_size 416 --model yolov3
python convert_trt.py --weights ./checkpoints/yolov3.tf --quantize_mode float16 --output ./checkpoints/yolov3-trt-fp16-416

# yolov3-tiny
python save_model.py --weights ./data/yolov3-tiny.weights --output ./checkpoints/yolov3-tiny.tf --input_size 416 --tiny
python convert_trt.py --weights ./checkpoints/yolov3-tiny.tf --quantize_mode float16 --output ./checkpoints/yolov3-tiny-trt-fp16-416

# yolov4
python save_model.py --weights ./data/yolov4.weights --output ./checkpoints/yolov4.tf --input_size 416 --model yolov4
python convert_trt.py --weights ./checkpoints/yolov4.tf --quantize_mode float16 --output ./checkpoints/yolov4-trt-fp16-416
python detect.py --weights ./checkpoints/yolov4-trt-fp16-416 --model yolov4 --images ./data/images/kite.jpg --framework trt

save_model.py:
  --weights: 权重文件路径
    （默认：'./data/yolov4.weights'）
  --output: 输出路径
    （默认：'./checkpoints/yolov4-416'）
  --[no]tiny: yolov4 或 yolov4-tiny
    （默认：'False'）
  --input_size: 定义导出模型的输入尺寸
    （默认：416）
  --framework: 使用何种框架（tf、trt、tflite）
    （默认：tf）
  --model: yolov3 或 yolov4
    （默认：yolov4）

detect.py:
  --images: 输入图像路径，以逗号分隔的字符串形式
    （默认：'./data/images/kite.jpg'）
  --output: 输出文件夹路径
    （默认：'./detections/'）
  --[no]tiny: yolov4 或 yolov4-tiny
    （默认：'False'）
  --weights: 权重文件路径
    （默认：'./checkpoints/yolov4-416'）
  --framework: 使用何种框架（tf、trt、tflite）
    （默认：tf）
  --model: yolov3 或 yolov4
    （默认：yolov4）
  --size: 将图像调整为
    （默认：416）
  --iou: IoU 阈值
    （默认：0.45）
  --score: 置信度阈值
    （默认：0.25）
  --count: 统计图像中的目标数量
    （默认：False）
  --dont_show: 不显示图像输出
    （默认：False）
  --info: 打印检测信息
    （默认：False）
  --crop: 裁剪检测结果并保存为新图像
    （默认：False）
    
detect_video.py:
  --video: 输入视频路径（使用 0 表示摄像头）
    （默认：'./data/video/video.mp4'）
  --output: 输出视频路径（请确保为给定格式设置正确的编解码器，例如 .avi 格式使用 XVID）
    （默认：None）
  --output_format: 保存视频时 VideoWriter 使用的编解码器
    （默认：'XVID'）
  --[no]tiny: yolov4 或 yolov4-tiny
    （默认：'false'）
  --weights: 权重文件路径
    （默认：'./checkpoints/yolov4-416'）
  --framework: 使用何种框架（tf、trt、tflite）
    （默认：tf）
  --model: yolov3 或 yolov4
    （默认：yolov4）
  --size: 将图像调整为
    （默认：416）
  --iou: IoU 阈值
    （默认：0.45）
  --score: 置信度阈值
    （默认：0.25）
  --count: 统计视频中的目标数量
    （默认：False）
  --dont_show: 不显示视频输出
    （默认：False）
  --info: 打印检测信息
    （默认：False）
  --crop: 裁剪检测结果并保存为新图像
    （默认：False）

• tensorflow-yolov4-tflite

yolov4-custom-functions 快速上手指南

环境准备

系统要求

• 操作系统：支持 Linux、macOS 或 Windows（推荐使用 Linux 或 macOS）
• Python 版本：3.6+（建议使用 3.8）
• TensorFlow：2.x（根据是否使用 GPU，选择 CPU 或 GPU 版本）

前置依赖

• 安装 Conda 或 pip（推荐使用 Conda 管理环境）
• 如果使用 GPU：
  • 安装 NVIDIA 驱动
  • 安装 CUDA Toolkit 10.1（与 TensorFlow 兼容）
  • 安装 cuDNN（与 CUDA 10.1 对应的版本）

安装步骤

使用 Conda（推荐）

# TensorFlow CPU
conda env create -f conda-cpu.yml
conda activate yolov4-cpu

# TensorFlow GPU
conda env create -f conda-gpu.yml
conda activate yolov4-gpu

使用 Pip

# TensorFlow CPU
pip install -r requirements.txt

# TensorFlow GPU
pip install -r requirements-gpu.txt

注意：如果使用国内网络，可以考虑使用清华源加速安装：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt

下载预训练权重

YOLOv4 默认支持检测 80 个类别。为了方便演示，我们使用官方提供的预训练权重。

• 下载 YOLOv4 预训练权重：yolov4.weights
• 下载 YOLOv4-tiny 预训练权重（更小更快）：yolov4-tiny.weights

将下载的 .weights 文件复制到项目目录下的 data 文件夹中。

基本使用

使用预训练模型进行目标检测

以下命令将使用预训练的 YOLOv4 模型对图片进行检测：

# 将 .weights 转换为 TensorFlow 模型
python save_model.py --weights ./data/yolov4.weights --output ./checkpoints/yolov4-416 --input_size 416 --model yolov4

# 对单张图片进行检测
python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/kite.jpg

使用自定义模型

如果你有自己训练的 .weights 和 .names 文件，请将它们分别放入 data 和 data/classes/ 目录中，并修改 core/config.py 中第 14 行指向你的 .names 文件。

示例命令如下：

# 将自定义 .weights 转换为 TensorFlow 模型
python save_model.py --weights ./data/custom.weights --output ./checkpoints/custom-416 --input_size 416 --model yolov4

# 对图片进行检测
python detect.py --weights ./checkpoints/custom-416 --size 416 --model yolov4 --images ./data/images/car.jpg

使用自定义功能（以对象计数为例）

要启用对象计数功能，只需在命令中添加 --count 标志：

python detect.py --weights ./checkpoints/yolov4-416 --size 416 --model yolov4 --images ./data/images/dog.jpg --count

如需按类别统计对象数量，还需在 detect.py 中设置参数 by_class=True。

使用 Tesseract OCR 提取文本

确保已正确安装 Tesseract OCR 并配置好环境变量后，可使用 OCR 功能提取检测到的文本内容。

以上就是 yolov4-custom-functions 的快速上手指南，你可以根据需要进一步探索更多自定义功能！