【YOLOV4】(14) 视频目标检测，预测阶段 TensorFlow 完整代码

大家好，研究了两天终于把 YOLOV4 预测阶段的完整代码复现出来了。本文只用函数方法，最直观的向大家展示代码。强烈建议大家在阅读本文之前，先看以下文章：

（1）YOLOV4特征提取网络：https://blog.csdn.net/dgvv4/article/details/123818580

（2）先验框解码调整：https://blog.csdn.net/dgvv4/article/details/124076352

（3）NMS非极大值抑制：https://blog.csdn.net/dgvv4/article/details/124062839

由于本文篇幅有点长，可以用右侧的目录栏跳转。理论部分我已经在上面几篇博文中写的很清楚了，这里就不再赘述了。

完整代码在我的Gitee中，有需要的自取：

https://gitee.com/dgvv4/yolo-target-detection/tree/master/YOLOV4%E9%A2%84%E6%B5%8B%E4%BB%A3%E7%A0%81

YOLOV4的VOC数据集的权重文件下载地址：

链接：https://pan.baidu.com/s/1WEq4eZ6UFsByWGZQf0AE1w   提取码：snlw

先放张图看效果：

---

1. 特征提取网络

这个部分在之前的文章中已经构建过，理论方面不懂的可以看上面的链接。这里就简单提一下。

1.1 骨干特征提取网络

CSPDarkNet53 骨干特征提取网络在 YOLOV3 的 DarkNet53网络 的基础上引入了 CSP结构。该结构增强了卷积神经网络的学习能力；移除了计算瓶颈；降低了显存的使用；加快了网络的推理速度。

输入图像的shape为[416,416,3]，网络不断进行下采样来获得更高的语义信息，输出三个有效特征层，feat1的shape为 [52,52,256] ，feat2的shape为[26,26,512] ，feat3的shape为[13,13,1024]

代码如下，文件命名为 CSPDarknet53.py

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, regularizers

#（1）Mish激活函数
def mish(x):
    # x*tanh(ln(1+ex))
    x = x * tf.math.tanh(tf.math.softplus(x))
    return x

#（2）标准卷积块
def conv_block(inputs, filters, kernel_size, strides):
    # 卷积+BN+Mish
    x = layers.Conv2D(filters, kernel_size, strides, 
                      padding='same', use_bias=False,  # 有BN不要偏置
                      kernel_regularizer=regularizers.l2(5e-4))(inputs)  # l2正则化

    x = layers.BatchNormalization()(x)
    x = mish(x)

    return x

#（3）残差块
def res_block(inputs, filters, num):

    residual = inputs  # 残差边
    # 1*1卷积调整通道
    x = conv_block(inputs, filters//2, kernel_size=(1,1), strides=1)
    # 3*3卷积提取特征
    x = conv_block(x, filters//2 if num!=1 else filters, kernel_size=(3,3), strides=1)
    # 残差连接输入和输出
    x = layers.Add()([x, residual])
    
    return x

#（4）CSP结构
def csp_bolck(inputs, filters, num):
    
    # 卷积下采样
    x = conv_block(inputs, filters, kernel_size=(3,3), strides=2)
    
    # num!=1时1*1卷积在通道维度上会下降一半
    shortcut = conv_block(x, filters//2 if num!=1 else filters, (1,1), strides=1)  # 残差边
    mainconv = conv_block(x, filters//2 if num!=1 else filters, (1,1), strides=1)  # 主干卷积

    # 重复执行残差结构
    for _ in range(num):
        mainconv = res_block(inputs=mainconv, filters=filters, num=num)

    # 1*1卷积调整通道
    mainconv = conv_block(mainconv, filters//2 if num!=1 else filters, (1,1), strides=1)

    # 输入和输出在通道维度堆叠
    x = layers.concatenate([mainconv, shortcut])

    # 1*1卷积整合通道
    x = conv_block(x, filters, (1,1), strides=1)

    return x

#（5）主干网络
def cspdarknet(inputs):

    # [416,416,3]==>[416,416,32]
    x = conv_block(inputs, filters=32, kernel_size=(3,3), strides=1)
    # [416,416,32]==>[208,208,64]
    x = csp_bolck(x, filters=64, num=1)
    # [208,208,64]==>[104,104,128]
    x = csp_bolck(x, filters=128, num=2)
    
    # [104,104,128]==>[52,52,256]
    x = csp_bolck(x, filters=256, num=8)
    feat1 = x

    # [52,52,256]==>[26,26,512]
    x = csp_bolck(x, filters=512, num=8)
    feat2 = x

    # [26,26,512]==>[13,13,1024]
    x = csp_bolck(x, filters=1024, num=4)
    feat3 = x

    return feat1, feat2, feat3

---

1.2 SPP结构

对网络模型输出的 feat3 先经过三个卷积层调整通道数，然后分别使用池化核size为 5*5，9*9，13*13 的最大池化，通过padding='same' 使得池化前后的特征图的shape是完全相同的。然后将原始输入和三种池化的结果特征图在通道维度上堆叠。最后经过三次卷积融合通道信息。

代码如下，文件命名为 SPP.py

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from CSPDarknet53 import cspdarknet  # 导入网络模型
from CSPDarknet53 import conv_block  # 导入标准卷积块

def spp(inputs):

    # 获取网络的三个输出特征层
    feat1, feat2, feat3 = cspdarknet(inputs)

    # 对最后一个输出特征层进行3次卷积
    # [13,13,1024]==>[13,13,512]
    p5 = conv_block(feat3, filters=512, kernel_size=(1,1), strides=1)
    # [13,13,512]==>[13,13,1024]
    p5 = conv_block(p5, filters=1024, kernel_size=(3,3), strides=1)
    # [13,13,1024]==>[13,13,512]
    p5 = conv_block(p5, filters=512, kernel_size=(1,1), strides=1)

    # 经过不同尺度的最大池化后相堆叠
    maxpool1 = layers.MaxPooling2D(pool_size=(13,13), strides=1, padding='same')(p5)
    maxpool2 = layers.MaxPooling2D(pool_size=(9,9), strides=1, padding='same')(p5)
    maxpool3 = layers.MaxPooling2D(pool_size=(5,5), strides=1, padding='same')(p5)
    
    # 四种尺度在通道维度上堆叠[13,13,2048]
    p5 = layers.concatenate([maxpool1, maxpool2, maxpool3, p5])

    # 三次卷积调整通道数
    # [13,13,2048]==>[13,13,512]
    p5 = conv_block(p5, filters=512, kernel_size=(1,1), strides=1)
    # [13,13,512]==>[13,13,1024]
    p5 = conv_block(p5, filters=1024, kernel_size=(3,3), strides=1)
    # [13,13,1024]==>[13,13,512]
    p5 = conv_block(p5, filters=512, kernel_size=(1,1), strides=1)

    return feat1, feat2, p5

---

1.3 PANet 金字塔

PANet 将网络输出的有效特征层和SPP结构的输出进行特征融合，它是由两个特征金字塔组成，一个是将低层的语义信息向高层融合，另一个是将高层的语义信息向低层融合。

代码如下，文件命名为 PANet.py

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from CSPDarknet53 import conv_block  # 网络模型和标准卷积块
from SPP import spp  # 导入spp加强特征提取模块

# 5次卷积操作提取特征减少参数量
def five_conv(x, filters):
    x = conv_block(x, filters, (1,1), strides=1)
    x = conv_block(x, filters*2, (3,3), strides=1)
    x = conv_block(x, filters, (1,1), strides=1)    
    x = conv_block(x, filters*2, (3,3), strides=1)
    x = conv_block(x, filters, (1,1), strides=1)
    return x

def panet(inputs):

    # 获得网络的三个有效输出特征层
    feat1, feat2, p5 = spp(inputs)

    #（1）
    # 对spp结构的输出进行卷积和上采样
    # [13,13,512]==>[13,13,256]==>[26,26,256]
    p5_upsample = conv_block(p5, filters=256, kernel_size=(1,1), strides=1)
    p5_upsample = layers.UpSampling2D(size=(2,2))(p5_upsample)

    # 对feat2特征层卷积后再与p5_upsample堆叠
    # [26,26,512]==>[26,26,256]==>[26,26,512]
    p4 = conv_block(feat2, filters=256, kernel_size=(1,1), strides=1)
    p4 = layers.concatenate([p4, p5_upsample])

    # 堆叠后进行5次卷积[26,26,512]==>[26,26,256]
    p4 = five_conv(p4, filters=256)

    #（2）
    # 对p4卷积上采样
    # [26,26,256]==>[26,26,512]==>[52,52,512]
    p4_upsample = conv_block(p4, filters=128, kernel_size=(1,1), strides=1)
    p4_upsample = layers.UpSampling2D(size=(2,2))(p4_upsample)

    # feat1层卷积后与p4_upsample堆叠
    # [52,52,256]==>[52,52,128]==>[52,52,256]
    p3 = conv_block(feat1, filters=128, kernel_size=(1,1), strides=1)
    p3 = layers.concatenate([p3, p4_upsample])

    # 堆叠后进行5次卷积[52,52,256]==>[52,52,128]
    p3 = five_conv(p3, filters=128)

    # 存放第一个特征层的输出
    p3_output = p3

    #（3）
    # p3卷积下采样和p4堆叠
    # [52,52,128]==>[26,26,256]==>[26,26,512]
    p3_downsample = conv_block(p3, filters=256, kernel_size=(3,3), strides=2)
    p4 = layers.concatenate([p3_downsample, p4])

    # 堆叠后的结果进行5次卷积[26,26,512]==>[26,26,256]
    p4 = five_conv(p4, filters=256)

    # 存放第二个有效特征层的输出
    p4_output = p4

    #（4）
    # p4卷积下采样和p5堆叠
    # [26,26,256]==>[13,13,512]==>[13,13,1024]
    p4_downsample = conv_block(p4, filters=512, kernel_size=(3,3), strides=2)
    p5 = layers.concatenate([p4_downsample, p5])

    # 堆叠后进行5次卷积[13,13,1024]==>[13,13,512]
    p5 = five_conv(p5, filters=512)

    # 存放第三个有效特征层的输出
    p5_output = p5

    # 返回输出层结果
    return p3_output, p4_output, p5_output

---

1.4 YOLO预测头

YOLOHead 由一个3*3卷积层和一个1*1卷积层构成，3*3卷积整合之前获得的所有特征信息，1*1卷积获得三个有效特征层的输出结果。

代码如下，命名为 YOLOHEAD.py

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, Model
from PANet import panet  # 导入panet加强特征提取方法
from CSPDarknet53 import conv_block  # 导入标准卷积快

# 对PANet的特征输出层处理获得最终的预测结果
def yoloHead(inputs, num_anchors, num_classes):
    '''
    num_anchors每个网格包含先验框的数量, num_classes分类数
    num_anchors(5+num_classes)代表: 每个先验框有5+num_classes个参数, 即(x,y,w,h,c)和20个类别的条件概率
    每一个特征层的输出代表: 每一个网格上每一个先验框内部是否包含物体, 以及包含物体的种类, 和先验框的调整参数
    '''

    # 获得三个有效特征层
    p3_output, p4_output, p5_output = panet(inputs)

    # 3*3卷积[52,52,128]==>[52,52,256]
    p3_output = conv_block(p3_output, filters=256, kernel_size=(3,3), strides=1)
    # [52,52,256]==>[52,52,num_anchors(5+num_classes)]
    p3_output = conv_block(p3_output, filters=num_anchors*(5+num_classes),
                           kernel_size=(1,1), strides=1)

    # [26,26,256]==>[26,26,516]
    p4_output = conv_block(p4_output, filters=512, kernel_size=(3,3), strides=1)
    # [26,26,512]==>[26,26,num_anchors(5+num_classes)]
    p4_output = conv_block(p4_output, filters=num_anchors*(5+num_classes),
                           kernel_size=(1,1), strides=1)

    # [13,13,512]==>[13,13,1024]
    p5_output = conv_block(p5_output, filters=1024, kernel_size=(3,3), strides=1)
    # [13,13,1024]==>[13,13,num_anchors(5+num_classes)]
    p5_output = conv_block(p5_output, filters=num_anchors*(5+num_classes),
                           kernel_size=(1,1), strides=1)
    
    # 返回三个有效特征层
    return p5_output, p4_output, p3_output

---

1.5 模型构建

输入是416*416的图像，输出是三个有效特征层，13*13检测层负责预测大目标，26*26检测层负责预测中目标，52*52的特征层负责预测小目标。

代码如下，文件命名为 model.py

import tensorflow as tf
from tensorflow import keras
from YOLOHEAD import yoloHead  # 加载网络的预测头

# 构造模型
def MODEL(inputs, num_anchors, num_classes, summary=False):

    # 接收输出的三个有效特征层
    p5_output, p4_output, p3_output = yoloHead(inputs, num_anchors, num_classes)
    # 构造模型
    model = keras.Model(inputs, [p5_output, p4_output, p3_output])
    # 是否查看模型结构
    if summary:
        model.summary()
    # 返回模型
    return model

---

2. 先验框解码，获得预测框

首先要对网络的输出特征层解码，它们的shape分别是 [52, 52, (3*(5+num_classes))]，[26, 26, (3*(5+num_classes))]，[13, 13, (3*(5+num_classes))]。

其中通道数 3*(5+num_classes) 可以理解为：每个网格生成 3 个预测框，每个预测框包含了预测框相较于先验框的偏移量，坐标偏移量(tx, ty)，宽高偏移量(tw, th)，预测框中是否包含目标物体 c，预测框内的物体属于每个类别的条件概率num_classes，在VOC数据集中num_classes=20。其中 (tx, ty) 是负无穷到正无穷的任何数，(tw, th) 是归一化后的宽高。

以某个网格的先验框的调整为例，如下图所示，虚线框代表：和物体的真实标签框 iou 值最大的那个先验框，该先验框的宽高为(pw, ph)；蓝色框代表微调先验框后生成的预测框。

（cx,cy）是先验框中心点所在的网格的左上坐标（归一化后的坐标），由于坐标偏移量 (tx,ty) 可以是从负无穷到正无穷的任何数，为了防止坐标调整偏移过大，给偏移量添加sigmoid函数。将坐标偏移量限制在0-1之间，将预测框的中心点限制在它所在的网格内。

接下来将预测框的坐标和高度宽度归一化处理，再将归一化后的预测框信息映射到原始图像的大小，例如，在 13*13 的图像上得到的归一化后的预测框，需要映射到 1280*720 的原始图片大小上

代码如下，命名为 anchors.py

# 先验框解码调整预测框
import tensorflow as tf

#（一）对某一个输出特征层解码
def anchors_decode(feats, anchors, num_classes, inputs_shape):
    '''
    feats是某一个特征层的输出结果, 如shape=[b, 13, 13, 3*(5+num_classes)]
    anchors代表每个特征层, 每个网格会生成三个先验框[3,2]
    num_classes代表分类类别的数量
    inputs_shape代表输入图像的高宽[416,416]
    '''
    # 计算每个网格先验框的个数=3
    num_anchors = len(anchors)
    # 获得网格的高宽(h,w)=(13,13)
    grid_shape = feats.shape[1:3]

    #（1）构造网格，将图像划分成13*13个网格
    # 获得网格的x坐标 (13)==>[1,13,1,1]
    grid_x = tf.reshape(range(grid_shape[1]), shape=[1,-1,1,1])
    # 在y和anchors维度上复制扩张[1,13,1,1]==>[13,13,3,1]
    grid_x = tf.tile(grid_x, multiples=[grid_shape[0], 1, num_anchors, 1])

    # 同理获得网格的y坐标(13)==>[13,1,1,1]==>[13,13,3,1]
    grid_y = tf.tile(tf.reshape(range(grid_shape[0]), shape=[-1,1,1,1]), multiples=[1, grid_shape[1], num_anchors, 1])

    # 在通道维度上合并[13,13,3,2], 并且转换成tf.float32类型
    grid = tf.cast(tf.concat([grid_x, grid_y], axis=-1), dtype=tf.float32)
    
    #（2）调整先验框宽高shape, 13*13个网格，每个网格3个先验框，每个先验框(x,y)
    # [3,2]==>[1,1,3,2]
    anchors_tensor = tf.reshape(anchors, shape=[1,1,num_anchors,2])
    # 扩充[1,1,3,2]==>[13,13,3,2]
    anchors_tensor = tf.tile(anchors_tensor, multiples=[grid_shape[0], grid_shape[1], 1, 1])
    # 转换数据类型
    anchors_tensor = tf.cast(anchors_tensor, dtype=tf.float32)

    #（3）调整网络输出层的shape
    # [b,13,13,num_anchors*(5+num_classes)]==>[b,13,13,num_anchors,5+num_classes]
    feats = tf.reshape(feats, [-1, grid_shape[0], grid_shape[1], num_anchors, 5+num_classes])
    '''
    代表13*13个网格, 每个网格有3个先验框, 每个先验框有(5+num_classes)项信息
    其中, 5代表: 中心点坐标(x,y), 宽高(w,h), 置信度c
    num_classes: 检测框属于某个类别的条件概率, VOC数据集中等于20
    '''

    #（4）调整网格的先验框中心坐标和宽高
    # 归一化中心点坐标偏移量，把预测框中心点限制在该网格中
    anchors_xy = tf.nn.sigmoid(feats[..., :2])
    box_xy = anchors_xy + grid
    # 坐标归一化
    grid_shape_wh = tf.cast(grid_shape[::-1], feats.dtype)  #
    box_xy = box_xy / grid_shape_wh
    
    # 调整先验框宽高，得到预测框
    anchors_wh = tf.exp(feats[..., 2:4])
    box_wh = anchors_wh * anchors_tensor
    # 宽高度归一化
    inputs_shape_wh = tf.cast(inputs_shape[::-1], feats.dtype)
    box_wh = box_wh / inputs_shape_wh
    
    # 获得预测框的置信度和每个类别的条件概率
    box_conf = tf.nn.sigmoid(feats[..., 4:5])
    box_prob = tf.nn.sigmoid(feats[..., 5:])

    # 返回预测框的信息
    return box_xy, box_wh, box_conf, box_prob


#（二）将归一化后的预测框信息转换为对应的原图坐标
def correct_box(box_xy, box_wh, inputs_shape, image_shape):
    # 调整shape, y轴放前面方便宽高相乘
    box_yx = box_xy[..., ::-1]
    box_hw = box_wh[..., ::-1]
    
    # 调整输入图像的数据类型
    inputs_shape = tf.cast(inputs_shape, box_yx.dtype)
    image_shape = tf.cast(image_shape, box_hw.dtype)

    # 计算预测框在网格上的左上角和右下角坐标
    box_min = box_yx - (box_hw / 2.0)
    box_max = box_yx + (box_hw / 2.0)

    # 保存网格的左上和右下坐标
    boxes  = tf.concat([box_min[..., 0:1], box_min[..., 1:2], box_max[..., 0:1], box_max[..., 1:2]], axis=-1)
    # 获得原图上的左上和右下坐标
    boxes = boxes * tf.concat([image_shape, image_shape], axis=-1)
    
    return boxes


#（三）计算一个特征层的检测框坐标和类别的概率
def anchors_prediction(feats, anchors, num_classes, inputs_shape, image_shape):
    '''
    feats是某一个特征层的输出结果, 如shape=[b, 13, 13, 3*(5+num_classes)]
    anchors代表每个特征层, 每个网格会生成三个先验框[3,2]
    num_classes代表分类类别的数量
    image_shape:原始图像大小
    inputs_shape:网络输入的图片大小(h,w)=416*416
    '''
    # 获得某个特征层输出的预测框信息
    box_xy, box_wh, box_conf, box_prob = anchors_decode(feats, anchors, num_classes, inputs_shape)

    # 获得在原图像上的预测框
    boxes = correct_box(box_xy, box_wh, inputs_shape, image_shape)
    # [4]==>[n,4]
    boxes = tf.reshape(boxes, shape=[-1, 4])

    # 获得每个类别的真实概率
    box_scores = box_conf * box_prob
    # [num_classes]==>[n,num_classes]
    box_scores = tf.reshape(box_scores, shape=[-1,num_classes])

    # 返回每个框的坐标和每个类别的概率
    return boxes, box_scores

---

3. NMS非极大值抑制剔除冗余框

图像经过神经网络模型运算，会得到多个目标预测框。然而，每个目标只有一个真实标签框，算法却会生成多个预测框，这样会造成检测框冗余或预测错误的情况。因此，深度神经网络通过非极大值抑制 NMS 来筛选预测框，去掉交并比更大并且置信度较低的冗余检测框，剩下置信度较高，定位较准确的预测框。

代码如下，命名为 NMS.py

import tensorflow as tf
import numpy as np
from anchors import anchors_prediction  # 导入对网络每一输出层的解码方法

# 对所有输出层解码，并使用非极大值抑制
def predict(feats, image_shape, num_classes, pre_anchors, conf_thresh, nms_thresh, max_boxes=20):
    '''
    feats:代表模型的三个有效输出特征层
    image_shape:原始图像的(h,w)
    max_boxes:最大预测框数量
    conf_thresh:分类概率小于这个值的框被删除
    nms_thresh:两个框计算交并比, iou小于这个值被保留, 删除重复的框
    '''
    # 每个特征层对应的三个先验框
    anchor_mask = [[6,7,8], [3,4,5], [0,1,2]]

    boxes = []  # 存放每张图的预测框坐标
    box_scores = []  # 每个预测框的类别概率

    # 获得网络输入的图像的shape=416*416
    inputs_shape = tf.constant([416,416])
    
    # 获取三个有效输出特征层的预测框坐标和类别概率
    for i in range(len(feats)):
        #获取每个输出层的预测框坐标和概率
        feat_boxes, feat_box_scores = anchors_prediction(feats=feats[i], 
                                                         anchors=pre_anchors[anchor_mask[i]],  # 每个输出层对应的3个先验框
                                                         num_classes=num_classes,  # 分类数
                                                         inputs_shape=inputs_shape,  # 网络的输出图像大小[416,416]
                                                         image_shape=image_shape)  # 原始图像大小
        # 保存每一层的预测框坐标和概率
        # feat_boxes.shape=[2028,4], feat_box_scores.shape=[2028,20]
        boxes.append(feat_boxes)
        box_scores.append(feat_box_scores)

    # 调整排序方式，放在一行中 [[], [], []]==>[ , , ]
    boxes = tf.concat(boxes, axis=0)
    box_scores = tf.concat(box_scores, axis=0)

    # 设置一个阈值，保留概率高于该阈值的预测框
    # mask.shape=[3*2028, 20]
    mask = (box_scores) > conf_thresh
    # 设置每张图片最多出现几个预测框
    max_boxes = tf.constant(max_boxes, tf.int32)

    '''
    box_score是每个检测框对应的20个类别的概率
    mask是每个检测框的20个类别的概率是否满足要求
    '''

    # 保存每个预测框的坐标, 概率, 类别名称
    predict_boxes = []
    predict_score = []
    predict_classes = []

    # 对每一个类别删除重复的框
    for c in range(num_classes):
        # 取出所有类为c且满足最小概率要求的box
        class_boxes = tf.boolean_mask(boxes, mask[:, c])
        # 取出所有类为c的分数
        class_box_scores = tf.boolean_mask(box_scores[:, c], mask[:, c])
        # 非极大抑制
        nms_index = tf.image.non_max_suppression(class_boxes, class_box_scores, max_boxes, iou_threshold = nms_thresh)

        # 取出筛选后的结果
        class_boxes = tf.gather(class_boxes, nms_index)  # 检测框坐标
        class_box_scores = tf.gather(class_box_scores, nms_index)  # 每个框的类别概率
        classes = tf.ones_like(class_box_scores, dtype=tf.int32) * c  # 1*c获得类别索引

        # 保存筛选后的预测框坐标、概率、类别
        predict_boxes.append(class_boxes)
        predict_score.append(class_box_scores)
        predict_classes.append(classes)

    # 重新排序预测框, 放在同一行
    predict_boxes = tf.concat(predict_boxes, axis=0)
    predict_score = tf.concat(predict_score, axis=0)
    predict_classes = tf.concat(predict_classes, axis=0)

    return predict_boxes, predict_score, predict_classes

---

4. 视频目标检测

首先对读入的帧图像预处理，opencv读入的图片是BGR类型，需要转换成RGB，然后再转换成 Image 类型，调整图像尺寸从1280*720调整为416*416，然后再将像素值归一化，并为单张图像扩充batchsize维度  [h, w, c] ==> [b, h, w, c]。将预处理后的图像传入网络预测 yolo_model.predict(img)。

这里使用的 VOC数据集的权重和分类，以及聚类后的先验框的宽高比

代码如下，命名为 results.py 

import cv2
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow import keras
from model import MODEL  # 导入网络模型结构
from NMS import predict  # 导入预测框处理方法
# 调用GPU加速
gpus = tf.config.experimental.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# -------------------------------------------- #
# yolov4权重文件路径, 视频文件路径
# -------------------------------------------- #
yolo_weights = 'yolo4_voc_weights.h5'
video_path = 'D:/deeplearning/video/car.mp4'
cap = cv2.VideoCapture(video_path)  # 视频捕获

# -------------------------------------------- #
# class_names: VOC数据集的分类名
# anchors: 先验框的长宽
# num_anchors: 每个网格生成几个先验框
# num_classes: 一共有几个类别
# inputs_shape: 输入图像的尺寸
# inputs: 网络输入层
# conf_thresh:分类概率小于这个值的框被删除
# nms_thresh:两个框计算交并比, iou小于这个值被保留, 删除重复的框
# -------------------------------------------- #
class_names = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
              'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']

anchors = np.array([[12, 16],  [19, 36],  [40, 28],  [36, 75],  [76, 55],  [72, 146],  [142, 110],  [192, 243],  [459, 401]])
num_anchors = 3
num_classes = len(class_names)
input_shape = [416,416,3]
inputs = keras.Input(shape=input_shape)
conf_thresh = 0.6
nms_thresh = 0.4

# ----------------------------------------------------- #
# 模型构造, 加载权重, 自定义的网络层名报错就用by_name=True
# ----------------------------------------------------- #
yolo_model = MODEL(inputs, num_anchors, num_classes, summary=False)
yolo_model.load_weights(yolo_weights, by_name=True) 

# ----------------------------------------------------- #
# 图像预处理
# ----------------------------------------------------- # 
def preprocessing(image, inputs_shape):
    
    # opencv读入的图像是BGR格式，转换成RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    # 转换成Image类型
    image = Image.fromarray(np.uint8(image)) 
    # 调整图像尺寸
    image_data = image.resize(size=(inputs_shape[1], inputs_shape[0]))
    # 改变数据类型
    image_data = np.array(image_data, dtype=np.float32)
    # 归一化
    image_data = image_data / 255.0
    # 添加batch维度 [416,416,3]==>[1,416,416,3]
    image_data = np.expand_dims(image_data, axis=0)

    return image_data

# ----------------------------------------------------- #
# 处理视频帧图像
# ----------------------------------------------------- # 
while True:
    # 返回图像是否读取成功success，以及读取的帧图像img
    success, img = cap.read()
    # 将读入的图像复制一份
    frame = img.copy()
    # 原始输入图像的(h,w)
    image_shape = img.shape[0:2]
    # 图像预处理
    img = preprocessing(img, input_shape)

    # 返回模型输出的三个有效特征层
    p5_output, p4_output, p3_output = yolo_model.predict(img)
    
    # 整合一下输出特征层
    feats = [p5_output, p4_output, p3_output]
    
    # ----------------------------------------------------- #
    # 将输出的预测框信息解码、调整先验框、非极大值抑制
    # feats:代表模型的三个有效输出特征层
    # image_shape:原始图像的(h,w)
    # num_classes:分类类别数
    # anchors:先验框的高宽
    # conf_thresh:分类概率小于这个值的框被删除
    # nms_thresh:两个框计算交并比, iou小于这个值被保留, 删除重复的框
    # max_boxes:最大预测框数量
    # ----------------------------------------------------- #
    # predict_boxes预测框左上坐标和宽高
    # predict_score预测框的类别概率
    # predict_classes预测框所属类别的索引
    # ----------------------------------------------------- #
    predict_boxes, predict_score, predict_classes = predict(feats, image_shape, num_classes, anchors, conf_thresh, nms_thresh, max_boxes=100)
    
    # ----------------------------------------------------- #
    # 绘制预测框
    # ----------------------------------------------------- # 
    for box, score, class_index in zip(predict_boxes, predict_score, predict_classes):

        # 获取预测框的左上坐标和宽高
        box_y1, box_x1, box_y2, box_x2 = box[0], box[1], box[2], box[3]
        # 每个预测框的类别名称
        class_name = class_names[class_index]
        # 将名称和分数组合在一起, 字符串
        class_name_score = class_name + ': ' + str(round(score.numpy(),2))

        # 绘制预测框
        cv2.rectangle(frame, (box_x1, box_y1), (box_x2, box_y2), color=(0,255,0), thickness=2)
        # 显示类别和概率
        cv2.putText(frame, class_name_score, (box_x1, box_y1-5), cv2.FONT_HERSHEY_COMPLEX, 1, (0,0,255), 2)
    
    # ----------------------------------------------------- #
    # 显示图像
    # ----------------------------------------------------- # 
    cv2.imshow('frame', frame)  # 传入窗口名和帧图像
    # 每帧图像滞留10毫秒后消失，按下键盘上的ESC键退出程序
    if cv2.waitKey(10) & 0xFF == 27:
        break

# 释放视频资源
cap.release()
cv2.destroyAllWindows()

---

视频效果展示：