YOLOv3代码实现

学习资料：

https://www.bilibili.com/video/BV1t54y1C7ra/?spm_id_from=333.999.0.0

https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

1.安装环境

按照readme文件安装好环境，以下是我安装时出现的问题:

（1）在安装pytorch版本时，要注意自己的显卡适合的版本，比如NVIDIA GeForce RTX 3090，需要安装1.7.1版本的pytorch。

https://blog.csdn.net/a563562675/article/details/121656894

备注：查看cuda版本，打开cmd，命令行输入nvidia-smi

（2）安装requirements文件，在终端输入命令pip install -r requirements.txt，有时候安装不上，检查自己的网络，换个网络则会安装成功

2.准备自己的数据集

在下载自己的数据集，它已经是yolo格式，但只有图片和标签，并没有划分训练集和验证集。

https://blog.csdn.net/freezing_00/article/details/129097738

划分好自己的数据集之后，在data/pascal_voc_classes.json文件里修改自己数据集的类别。

在data文件夹下，创建my_data_label.names文件（txt），一个类别对应一行

然后在calculate_dataset.py文件下，修改相对应的路径，生成my_train_data.txt文件、my_val_data.txt文件以及my_data.data文件，并生成新的my_yolov3.cfg文件

数据集制作完成，然后创建weights文件夹，选择一个预训练权重下载到weights文件夹下

3.运行train文件

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--epochs', type=int, default=30)
    parser.add_argument('--batch-size', type=int, default=4) # 根据GPU的显存和内存进行设置
    parser.add_argument('--cfg', type=str, default='cfg/my_yolov3.cfg', help="*.cfg path")
    parser.add_argument('--data', type=str, default='data/my_data.data', help='*.data path')
    parser.add_argument('--hyp', type=str, default='cfg/hyp.yaml', help='hyperparameters path') # 超参数
    parser.add_argument('--multi-scale', type=bool, default=True,
                        help='adjust (67%% - 150%%) img_size every 10 batches') # 是否使用多尺度训练，多尺度训练是指在训练过程中使用不同尺度（大小、分辨率等）的数据输入来训练机器学习模型。
    parser.add_argument('--img-size', type=int, default=512, help='test size') # 指定网络输入的大小，在测试集
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--savebest', type=bool, default=True, help='only save best checkpoint') #是否保存map最高的权重
    parser.add_argument('--notest', action='store_true', help='only test final epoch') # 如果为True，则在训练的最后一个epoch进行验证，默认是关闭的
    parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') # 把图片缓存到内存中，action=store_false,默认值为True；action=store_true，默认值是False
    parser.add_argument('--weights', type=str, default='weights/yolov3-spp-ultralytics-512.pt',
                        help='initial weights path')
    parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
    parser.add_argument('--device', default='cuda:0', help='device id (i.e. 0 or 0,1 or cpu)')
    parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
    parser.add_argument('--freeze-layers', type=bool, default=False, help='Freeze non-output layers') # 是否冻结网络权重，只会预测三个预测器，如果为True，默认冻结Darknet53结构，训练其余
    # 是否使用混合精度训练(需要GPU支持混合精度)
    parser.add_argument("--amp", default=False, help="Use torch.cuda.amp for mixed precision training")
    opt = parser.parse_args()

    # 检查文件是否存在
    opt.cfg = check_file(opt.cfg) #配置文件
    opt.data = check_file(opt.data) #数据文件
    opt.hyp = check_file(opt.hyp) #超参数文件
    print(opt)

    with open(opt.hyp) as f:
        hyp = yaml.load(f, Loader=yaml.FullLoader)

    print('Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/')
    tb_writer = SummaryWriter(comment=opt.name)
    train(hyp)

先看train.py的267行作用这一部分，根据自己需求看是否需要修改；看55行作用，将 80改为自己数据集的类别个数，然后运行train.py文件。

hyp["cls"] *= nc / 80  # update coco-tuned hyp['cls'] to current dataset，对损失的增益进行调整，根据类别

（wz作者认为先训练三个预测器，然后训练除Darknet53之外的网络，训练结果较好。关于train文件的其他代码可以视频和wz作者的代码）

4.结果分析

打开runs文件夹下的文件，在空白处右键，选择在终端打开，然后输入tensorboard.exe --logdir "./"，回车

5.测试文件

（1）测试单张图片，在当前文件下载一张测试图片。打开predict_test，在16行左右按照自己的需要进行修改，在当前文件下载一张测试图片，即可运行。

def main():
    img_size = 512  # 必须是32的整数倍 [416, 512, 608]
    cfg = "cfg/my_yolov3.cfg"  # 改成生成的.cfg文件
    weights_path = "weights/best.pt"  # 改成自己训练好的权重文件
    json_path = "data/pascal_voc_classes.json"  # json标签文件
    img_path = "test.jpg"
    assert os.path.exists(cfg), "cfg file {} dose not exist.".format(cfg)
    assert os.path.exists(weights_path), "weights file {} dose not exist.".format(weights_path)
    assert os.path.exists(json_path), "json file {} dose not exist.".format(json_path)
    assert os.path.exists(img_path), "image file {} dose not exist.".format(img_path)

（2）测试多张图片

首先在my_yolo_dataset文件下有test文件，打开calculate_dataset.py文件，在所有关于train、val代码下添加test代码。#+++是我修改的。然后运行该文件。

"""
该脚本有3个功能：
1.统计训练集和验证集的数据并生成相应.txt文件
2.创建data.data文件，记录classes个数, train以及val数据集文件(.txt)路径和label.names文件路径
3.根据yolov3-spp.cfg创建my_yolov3.cfg文件修改其中的predictor filters以及yolo classes参数(这两个参数是根据类别数改变的)
"""
import os

train_annotation_dir = "./my_yolo_dataset1/train/labels"
val_annotation_dir = "./my_yolo_dataset1/val/labels"
test_annotation_dir = "./my_yolo_dataset1/test/labels" #+++++
classes_label = "./data/my_data_label.names"
cfg_path = "./cfg/yolov3-spp.cfg"

assert os.path.exists(train_annotation_dir), "train_annotation_dir not exist!"
assert os.path.exists(val_annotation_dir), "val_annotation_dir not exist!"
assert os.path.exists(test_annotation_dir), "val_annotation_dir not exist!"
assert os.path.exists(classes_label), "classes_label not exist!"
assert os.path.exists(cfg_path), "cfg_path not exist!"


def calculate_data_txt(txt_path, dataset_dir):
    # create my_data.txt file that record image list
    with open(txt_path, "w") as w:
        for file_name in os.listdir(dataset_dir):
            if file_name == "classes.txt":
                continue

            img_path = os.path.join(dataset_dir.replace("labels", "images"),
                                    file_name.split(".")[0]) + ".jpg"
            line = img_path + "\n"
            assert os.path.exists(img_path), "file:{} not exist!".format(img_path)
            w.write(line)


def create_data_data(create_data_path, label_path, train_path, val_path, test_path,classes_info):
    # create my_data.data file that record classes, train, valid and names info.
    # shutil.copyfile(label_path, "./data/my_data_label.names")
    with open(create_data_path, "w") as w:
        w.write("classes={}".format(len(classes_info)) + "\n")  # 记录类别个数
        w.write("train={}".format(train_path) + "\n")           # 记录训练集对应txt文件路径
        w.write("valid={}".format(val_path) + "\n")             # 记录验证集对应txt文件路径
        w.write("test={}".format(test_path) + "\n")       #+++++
        w.write("names=data/my_data_label.names" + "\n")        # 记录label.names文件路径


def change_and_create_cfg_file(classes_info, save_cfg_path="./cfg/my_yolov3.cfg"):
    # create my_yolov3.cfg file changed predictor filters and yolo classes param.
    # this operation only deal with yolov3-spp.cfg
    filters_lines = [636, 722, 809]
    classes_lines = [643, 729, 816]
    cfg_lines = open(cfg_path, "r").readlines()

    for i in filters_lines:
        assert "filters" in cfg_lines[i-1], "filters param is not in line:{}".format(i-1)
        output_num = (5 + len(classes_info)) * 3
        cfg_lines[i-1] = "filters={}\n".format(output_num)

    for i in classes_lines:
        assert "classes" in cfg_lines[i-1], "classes param is not in line:{}".format(i-1)
        cfg_lines[i-1] = "classes={}\n".format(len(classes_info))

    with open(save_cfg_path, "w") as w:
        w.writelines(cfg_lines)


def main():
    # 统计训练集和验证集的数据并生成相应txt文件
    train_txt_path = "data/my_train_data.txt"
    val_txt_path = "data/my_val_data.txt"
    test_txt_path = "data/my_test_data.txt"   #+++++
    calculate_data_txt(train_txt_path, train_annotation_dir)
    calculate_data_txt(val_txt_path, val_annotation_dir)
    calculate_data_txt(test_txt_path, test_annotation_dir)

    classes_info = [line.strip() for line in open(classes_label, "r").readlines() if len(line.strip()) > 0]
    # 创建data.data文件，记录classes个数, train以及val数据集文件(.txt)路径和label.names文件路径
    create_data_data("./data/my_data.data", classes_label, train_txt_path, val_txt_path, test_txt_path,classes_info)    #+++++

    # 根据yolov3-spp.cfg创建my_yolov3.cfg文件修改其中的predictor filters以及yolo classes参数(这两个参数是根据类别数改变的)
    change_and_create_cfg_file(classes_info)


if __name__ == '__main__':
    main()

打开validation文件，修改：

98： test_path = data_dict["valid"]----test_path = data_dict["test"]

132： for imgs, targets, paths, shapes, img_index in tqdm(val_dataset_loader, desc="validation..."):

for imgs, targets, paths, shapes, img_index in tqdm(val_dataset_loader, desc="test..."):

207： parser.add_argument('--weights', default='./weights/yolov3spp-voc-512.pt', type=str, help='training weights')

parser.add_argument('--weights', default='./weights/best.pt', type=str, help='training weights')

然后运行，报错

voc_map_info_list.append(" {:15}: {}".format(category_index[i + 1], stats[1])) KeyError: 4

把这一行的i+1改成i，运行成功，但是如何运行出来像测试单张图片一样的结果，还不会