一、说明

本文学习资源来自Github开源项目：
https://github.com/breadbread1984/YOLOv3-tf2.0

由于条件限制，本文部分内容我没有实际运行测试。
本文摘选注释部分功能代码，完整项目代码可到开源地址获取。

二、代码实现

1. 训练COCO数据集

下载COCO数据集download_dataset.py

#!/usr/bin/python3

import tensorflow_datasets as tfds;

def main():
    
    # load dataset
    coco2014_builder = tfds.builder("coco");
    coco2014_builder.download_and_prepare();
    # try to load the dataset once
    coco2014_train = tfds.load(name = "coco", split = tfds.Split.TRAIN, download = False);
    coco2014_test = tfds.load(name = "coco", split = tfds.Split.TEST, download = False);

if __name__ == "__main__":
    
    assert tf.executing_eagerly();
    main();

数据集37.57GiB比较大，在google colab上超出了磁盘使用量，不适合在colab上复现。

tf eager训练coco数据集 train_eager.py

#!/usr/bin/python3

import os;
import numpy as np;
import cv2;
import tensorflow as tf;
import tensorflow_datasets as tfds;
from YOLOv3 import YOLOv3, Loss;
from Predictor import Predictor;
from preprocess import map_function;

os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1';
#os.environ['TF_CPP_MIN_VLOG_LEVEL'] = '3';
#os.environ['CUDA_VISIBLE_DEVICES'] = '';
batch_size = 8; # images of different sizes can't be stack into a batch

def main():

    # yolov3 model  导入yolov3模型
    yolov3 = YOLOv3((416,416,3), 80);
    yolov3_loss = Loss((416,416,3), 80);
    # load downloaded dataset 加载数据集
    trainset = tfds.load(name = "coco", split = tfds.Split.TRAIN, download = False);
    trainset = trainset.map(map_function).repeat(100).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE);
    validationset = tfds.load(name = "coco", split = tfds.Split.VALIDATION, download = False);
    validationset_iter = validationset.map(map_function).repeat(100).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE).__iter__();
    testset = tfds.load(name = "coco", split = tfds.Split.TEST, download = False); # without label
    testset = testset.repeat(100).prefetch(tf.data.experimental.AUTOTUNE);
    testset_iter = testset.__iter__();
    # restore from existing checkpoint 加载checkpoint
    optimizer = tf.keras.optimizers.Adam(1e-4);
    if False == os.path.exists('checkpoints'): os.mkdir('checkpoints');
    checkpoint = tf.train.Checkpoint(model = yolov3, optimizer = optimizer, optimizer_step = optimizer.iterations);
    checkpoint.restore(tf.train.latest_checkpoint('checkpoints'));
    # tensorboard summary
    log = tf.summary.create_file_writer('checkpoints');
    # train model 训练模型
    print("training...");
    train_loss = tf.keras.metrics.Mean(name = 'train loss', dtype = tf.float32);
    validation_loss = tf.keras.metrics.Mean(name = 'validation loss', dtype = tf.float32);
    for images, labels in trainset:
        with tf.GradientTape() as tape:
            outputs = yolov3(images);
            loss = yolov3_loss([*outputs, *labels]);
        # check whether the loss numberic is correct
        try:
            loss_check = tf.debugging.check_numerics(loss, 'the loss is not correct! cancel train_loss update!');
            with tf.control_dependencies([loss_check]):
                train_loss.update_state(loss);
                print('Step #%d Loss: %.6f' % (optimizer.iterations, loss));
        except BaseException as e:
            print(e.message);
        # write log
        if tf.equal(optimizer.iterations % 10, 0):
            with log.as_default():
                tf.summary.scalar('train loss',train_loss.result(), step = optimizer.iterations);
            train_loss.reset_states();
        grads = tape.gradient(loss, yolov3.trainable_variables);
        # check whether the grad numerics is correct
        try:
            grads_check = [tf.debugging.check_numerics(grad, 'the grad is not correct! cancel gradient apply!') for grad in grads];
            with tf.control_dependencies(grads_check):
                optimizer.apply_gradients(zip(grads, yolov3.trainable_variables));
        except BaseException as e:
            print(e.message);
        # save model
        if tf.equal(optimizer.iterations % 2000, 0):
            # save checkpoint every 1000 steps
            checkpoint.save(os.path.join('checkpoints','ckpt'));
            yolov3.save('yolov3.h5');
        # eval on testset 测试集
        if tf.equal(optimizer.iterations % 100, 0):
            # validate with latest model
            print("validating on validation set...");
            for i in range(10):
                images, labels = next(validationset_iter);
                outputs = yolov3(images);
                loss = yolov3_loss([*outputs, *labels]);
                # NOTE: validation loss is not important, numeric validity is not checked
                validation_loss.update_state(loss);
            with log.as_default():
                tf.summary.scalar('validation loss', validation_loss.result(), step = optimizer.iterations);
            validation_loss.reset_states();
            # evaluate every 1000 steps
            print("testing on test set...");
            features = next(testset_iter);
            img = features["image"].numpy().astype('uint8');
            predictor = Predictor(yolov3 = yolov3);
            boundings = predictor.predict(img);
            color_map = dict();
            for bounding in boundings:
                if bounding[5].numpy().astype('int32') in color_map:
                    clr = color_map[bounding[5].numpy().astype('int32')];
                else:
                    color_map[bounding[5].numpy().astype('int32')] = tuple(np.random.randint(low=0, high=256,size=(3,)).tolist());
                    clr = color_map[bounding[5].numpy().astype('int32')];
                cv2.rectangle(img, tuple(bounding[0:2].numpy().astype('int32')), tuple(bounding[2:4].numpy().astype('int32')), clr, 5);
            img = tf.expand_dims(img, axis = 0);
            with log.as_default():
                tf.summary.image('detect', img, step = optimizer.iterations);
    yolov3.save('yolov3.h5'); # 保存模型

if __name__ == "__main__":
    
    assert tf.executing_eagerly();
    main();

keras训练数据集 train_keras.py

#!/usr/bin/python3

import os;
import numpy as np;
import tensorflow as tf;
import tensorflow_datasets as tfds;
from YOLOv3 import YOLOv3, Loss;
from preprocess import map_function;

os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1';
#os.environ['TF_CPP_MIN_VLOG_LEVEL'] = '3';
#os.environ['CUDA_VISIBLE_DEVICES'] = '';
batch_size = 8; # images of different sizes can't be stack into a batch

def main():

    # yolov3 model
    yolov3 = YOLOv3((416,416,3,), 80);
    @tf.function
    def loss(outputs, labels):
        return Loss((416,416,3,),80)([outputs[0], outputs[1], outputs[2], labels[0], labels[1], labels[2]]);
    yolov3.compile(optimizer = tf.keras.optimizers.Adam(1e-4), loss = loss);
    # load downloaded dataset
    trainset = tfds.load(name = "coco2014", split = tfds.Split.TRAIN, download = False);
    trainset = trainset.map(map_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE);
    validationset = tfds.load(name = "coco2014", split = tfds.Split.VALIDATION, download = False);
    validationset = validationset.map(map_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE);
    yolov3.fit(trainset, epochs = 100, validation_data = validationset);
    yolov3.save('yolov3.h5');

if __name__ == "__main__":
    
    assert tf.executing_eagerly();
    main();

训练过程截图：
训练集损失函数：

测试集损失函数：

检测结果：

2. 使用训练好的模型预测Predictor.py

调用方式：

python3 Predictor.py <path/to/image>

#!/usr/bin/python3

import sys;
from os.path import exists;
import numpy as np;
import cv2;
import tensorflow as tf;
from YOLOv3 import YOLOv3, OutputParser;

class Predictor(object):

    anchors = {2: [[10, 13], [16, 30], [33, 23]], 1: [[30, 61], [62, 45], [59, 119]], 0: [[116, 90], [156, 198], [373, 326]]};

    def __init__(self, input_shape = (416,416,3), class_num = 80, yolov3 = None):

        if yolov3 is None:
            self.input_shape = input_shape;
            if exists("yolov3.h5"):
                # restore from serialized file
                self.yolov3 = tf.keras.models.load_model('yolov3.h5', compile = False);
            else:
                # restore from checkpoint
                self.yolov3 = YOLOv3(input_shape, class_num);
                optimizer = tf.keras.optimizers.Adam(1e-4);
                checkpoint = tf.train.Checkpoint(model = self.yolov3, optimizer = optimizer, optimizer_step = optimizer.iterations);
                checkpoint.restore(tf.train.latest_checkpoint('checkpoints'));
        else:
            self.input_shape = tuple(yolov3.input.shape[1:]);
            self.yolov3 = yolov3;
        self.parsers = [OutputParser(tuple(self.yolov3.outputs[l].shape[1:]), self.input_shape, self.anchors[l]) for l in range(3)];

    def predict(self, image, conf_thres = 0.5, nms_thres = 0.5):

        images = tf.expand_dims(image, axis = 0);
        resize_images = tf.image.resize(images, self.input_shape[:2], method = tf.image.ResizeMethod.BICUBIC, preserve_aspect_ratio = True);
        resize_shape = resize_images.shape[1:3];
        top_pad = (self.input_shape[0] - resize_shape[0]) // 2;
        bottom_pad = self.input_shape[0] - resize_shape[0] - top_pad;
        left_pad = (self.input_shape[1] - resize_shape[1]) // 2;
        right_pad = self.input_shape[1] - resize_shape[1] - left_pad;
        resize_images = tf.pad(resize_images,[[0,0], [top_pad,bottom_pad], [left_pad,right_pad], [0,0]], constant_values = 128);
        deviation = tf.constant([left_pad / self.input_shape[1], top_pad / self.input_shape[0], 0, 0], dtype = tf.float32);
        scale = tf.constant([
            self.input_shape[1] / resize_shape[1], self.input_shape[0] / resize_shape[0],
            self.input_shape[1] / resize_shape[1], self.input_shape[0] / resize_shape[0]
        ], dtype = tf.float32);
        images_data = tf.cast(resize_images, tf.float32) / 255.;
        outputs = self.yolov3(images_data);
        whole_targets = tf.zeros((0,6), dtype = tf.float32);
        for i in range(3):
            pred_xy, pred_wh, pred_box_confidence, pred_class = self.parsers[i](outputs[i]);
            pred_box = tf.keras.layers.Concatenate(axis = -1)([pred_xy, pred_wh]);
            # target_mask.shape = (h, w, anchor num)
            target_mask = tf.greater(pred_box_confidence, conf_thres);
            # pred_box_confidence = (pred target num, 1)
            pred_box_confidence = tf.boolean_mask(pred_box_confidence, target_mask);
            pred_box_confidence = tf.expand_dims(pred_box_confidence, axis = -1);
            # pred_box.shape = (pred target num, 4)
            pred_box = tf.boolean_mask(pred_box, target_mask);
            pred_box = (pred_box - deviation) * scale * [image.shape[1], image.shape[0], image.shape[1], image.shape[0]];
            # pred_class.shape = (pred target num, 1)
            pred_class = tf.boolean_mask(pred_class, target_mask);
            pred_class = tf.math.argmax(pred_class, axis = -1);
            pred_class = tf.cast(tf.expand_dims(pred_class, axis = -1), dtype = tf.float32);
            # targets,sgaoe = (pred target num, 6)
            targets = tf.keras.layers.Concatenate(axis = -1)([pred_box, pred_box_confidence, pred_class]);
            whole_targets = tf.keras.layers.Concatenate(axis = 0)([whole_targets, targets]);
        # nms
        descend_idx = tf.argsort(whole_targets[..., 4], direction = 'DESCENDING');
        i = 0;
        while i < descend_idx.shape[0]:
            idx = descend_idx[i];
            cur_upper_left = whole_targets[idx, 0:2] - whole_targets[idx, 2:4] / 2;
            cur_down_right = cur_upper_left + whole_targets[idx, 2:4];
            wh = whole_targets[idx, 2:4];
            area = wh[..., 0] * wh[..., 1];
            following_idx = descend_idx[i+1:];
            following_targets = tf.gather(whole_targets, following_idx);
            following_upper_left = following_targets[..., 0:2] - following_targets[..., 2:4] / 2;
            following_down_right = following_upper_left + following_targets[..., 2:4];
            following_wh = following_targets[..., 2:4];
            following_area = following_wh[..., 0] * following_wh[..., 1];
            max_upper_left = tf.math.maximum(cur_upper_left, following_upper_left);
            min_down_right = tf.math.minimum(cur_down_right, following_down_right);
            intersect_wh = min_down_right - max_upper_left;
            intersect_wh = tf.where(tf.math.greater(intersect_wh, 0), intersect_wh, tf.zeros_like(intersect_wh));
            intersect_area = intersect_wh[..., 0] * intersect_wh[..., 1];
            overlap = intersect_area / (area + following_area - intersect_area);
            indices = tf.where(tf.less(overlap, nms_thres));
            following_idx = tf.gather_nd(following_idx, indices);
            descend_idx = tf.concat([descend_idx[:i + 1], following_idx], axis = 0);
            i += 1;
        whole_targets = tf.gather(whole_targets, descend_idx);
        upper_left = (whole_targets[..., 0:2] - whole_targets[..., 2:4] / 2)
        down_right = (upper_left + whole_targets[..., 2:4])
        boundings = tf.keras.layers.Concatenate(axis = -1)([upper_left, down_right, whole_targets[..., 4:]]);
        return boundings;

if __name__ == "__main__":

    assert tf.executing_eagerly() == True;
    if len(sys.argv) != 2:
        print("Usage: " + sys.argv[0] + " <image>");
        exit(0);
    predictor = Predictor();
    img = cv2.imread(sys.argv[1]);
    if img is None:
        print("invalid image!");
        exit(1);
    boundings = predictor.predict(img);
    color_map = dict();
    for bounding in boundings:
        if bounding[5].numpy().astype('int32') in color_map:
            clr = color_map[bounding[5].numpy().astype('int32')];
        else:
            color_map[bounding[5].numpy().astype('int32')] = tuple(np.random.randint(low=0, high=256,size=(3,)).tolist());
            clr = color_map[bounding[5].numpy().astype('int32')];
        cv2.rectangle(img, tuple(bounding[0:2].numpy().astype('int32')), tuple(bounding[2:4].numpy().astype('int32')), clr, 2);
    cv2.imshow('people', img);
    cv2.waitKey();

下面是运行效果：
python3 Predictor.py pics/detection1.png

3. 训练自己的数据集

标签文件格式

<path/to/image1> <target num>
<x> <y> <width> <height> <label>
<x> <y> <width> <height> <label>
...
<x> <y> <width> <height> <label>
<path/to/image2> <target num>
...

生成tfrecord文件

python3 create_dataset.py <path/to/annotation>

读取tfrecord可以使用：

from create_dataset import parse_function_generator;
trainset = tf.data.TFRecordDataset('trainset.tfrecord').map(parse_function_generator(num_classes = num_classes)).repeat(100).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE);