动手学习深度学习pytorch版——从零开始实现YOLOv1

最新推荐文章于 2025-07-15 19:50:25 发布

原创

最新推荐文章于 2025-07-15 19:50:25 发布 · 3.7w 阅读

127

572 ·

CC 4.0 BY-SA版权

文章标签：

#计算机视觉 #深度学习 #python

从零开始实现YOLOv1

1. 数据处理
2. YOLOv1网络实现
- 2.1 Loss函数设计
- 2.2 网络结构实现
3. 开始训练啦
4. 网络预测(Inference)
5. 结果展示

以下是YOLOv1网络的实践操作部分，具体理论可以参考我的博客“ 经典论文解析——YOLOv1——目标检测”。文中代码均使用了GPU，如果要用CPU尝试，可以将代码中的.cuda()都去掉，但是不建议这么做，因为这个网络还挺大的，CPU很可能跑不动。
完整代码文件有空我会放在github上，想参考的读者目前就只能复制粘贴了(但是文中提供的代码也是完整的了)。如有错误，请评论指出，共同进步，谢谢。

2021年3月27日更新：

介于很多人要求，将代码修正后放在github上，链接点这里。其中主要代码使用了我自己的pytorch项目模板，如果有需要可以留言我再传上去。主体部分代码和文章里一样，做了几个小修改：1). 计算loss时的m,n反了，进行修正；2). labels2bbox用矩阵运算完成，速度更快一点；
我自己又跑了一遍，没发现nan的问题，初始配置可以参见代码，我的显存不够，所以优化器只能用SGD，Adam我跑不起来。
注意，github里的代码和文章里的代码在bbox的格式上有点区别，github里bbox[0:4]是坐标，bbox[4]是置信度，bbox[5]是类别。所以相应代码略有不同。直接使用github里的代码就可以了，文章里的代码可用作为参考。

1. 数据处理

1.1 数据集介绍

搭建网络前首先要将原始数据集的数据格式转变为我们方便使用的格式。这里使用的是VOC2012数据集，下载下来有如下五个文件夹：
在这里插入图片描述
Annotations文件夹是存放图片对应的xml文件，比如“2007_000027.xml"存放的是图片2007_000027.jpg对应的信息，用记事本打开可以看到，这是xml格式的数据，里面除了图片的基本信息以外，还有一项< object >类，里面分别存放了类别的名称(< name >)，识别的难易程度(< difficult >)，以及bounding box的坐标信息< bndbox >，这里存放的box信息是以两点式存放，也就是左上角点和右下角点。当然，VOC数据集不只是用于目标检测任务，所以还存放了一些其他信息，比如人体的具体部分(< part >)等，这些就不用关注了。
在这里插入图片描述
ImageSets文件夹里存放了官方为我们划分好的训练集和验证集的txt文件。我们主要使用“ImageSets/Main/"文件夹下的train.txt和val.txt文件，train.txt文件存放了官方划分的训练集的图片名称，val.txt文件存放了验证集图片的名称。
还有一个需要关注的文件夹就是JEPGImages，里面存放了对应图片名称的原始图片。剩下的两个文件夹我们就不需要特别关注了。

1.2 数据提取和转换

了解数据集后，很明显，我们需要将图片对应的xml文件中bounding box的信息提取出来，并转换为我们需要的 $cls,x_c,y_c,w,h)$ 格式，其中cls是根据物体类别的序号决定的，物体类别排序储存在全局变量CLASSES中。

CLASSES = ['person', 'bird', 'cat', 'cow', 'dog', 'horse', 'sheep',
           'aeroplane', 'bicycle', 'boat', 'bus', 'car', 'motorbike', 'train',
           'bottle', 'chair', 'dining table', 'potted plant', 'sofa', 'tvmonitor']

使用下面三个子函数，调用make_label_txt()函数，就可以在当前项目文件夹的labesl文件夹下创造出与图片对应的txt文件，比如图片2007_000027.jpg，就有对应的2007_000027.txt文件，里面储存着图片2007_000027.jpg的所有bbox信息，每行一个。注意，要从.xml中提取信息，需要有xml库。

import xml.etree.ElementTree as ET
import os
import cv2

def convert(size, box):
    """将bbox的左上角点、右下角点坐标的格式，转换为bbox中心点+bbox的w,h的格式
    并进行归一化"""
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def convert_annotation(image_id):
    """把图像image_id的xml文件转换为目标检测的label文件(txt)
    其中包含物体的类别，bbox的左上角点坐标以及bbox的宽、高
    并将四个物理量归一化"""
    in_file = open(DATASET_PATH + 'Annotations/%s' % (image_id))
    image_id = image_id.split('.')[0]
    out_file = open('./labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in CLASSES or int(difficult) == 1:
            continue
        cls_id = CLASSES.index(cls)
        xmlbox = obj.find('bndbox')
        points = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), points)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


def make_label_txt():
    """在labels文件夹下创建image_id.txt，对应每个image_id.xml提取出的bbox信息"""
    filenames = os.listdir(DATASET_PATH + 'Annotations')
    for file in filenames:
        convert_annotation(file)

生成好图片对应的.txt文件后，要注意进行检查，自己写项目的时候要养成习惯，做完一部分工作后，最好写代码验证一下是否能正常完成自己想要的任务。比如我们将xml文件的bounding box提取出来，就需要检验一下提出的数据经过转换后是否保持原样，最简单的检验方式就是在原图中，利用.txt里的数据绘出bbox并输出对应框的类别信息，看看是否合理。可以用以下代码检验。以图片2008_000007.jpg为例，可以看到如下图结果，确认信息提取无误。
在这里插入图片描述

def show_labels_img(imgname):
	"""imgname是输入图像的名称，无下标"""
    img = cv2.imread(DATASET_PATH + "JPEGImages/" + imgname + ".jpg")
    h, w = img.shape[:2]
    print(w,h)
    label = []
    with open("./labels/"+imgname+".txt",'r') as flabel:
        for label in flabel:
            label = label.split(' ')
            label = [float(x.strip()) for x in label]
            print(CLASSES[int(label[0])])
            pt1 = (int(label[1] * w - label[3] * w / 2), int(label[2] * h - label[4] * h / 2))
            pt2 = (int(label[1] * w + label[3] * w / 2), int(label[2] * h + label[4] * h / 2))
            cv2.putText(img,CLASSES[int(label[0])],pt1,cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,0,255))
            cv2.rectangle(img,pt1,pt2,(0,0,255,2))

    cv2.imshow("img",img)
    cv2.waitKey(0)

1.3 pytorch的Dataset类构造

准备好数据后，就要按老套路开始准备训练了。首先要准备训练使用的Dataset类。具体代码如下：

class VOC2012(Dataset):
    def __init__(self,is_train=True,is_aug=True):
        """
        :param is_train: 调用的是训练集(True)，还是验证集(False)
        :param is_aug:  是否进行数据增广
        """
        self.filenames = []  # 储存数据集的文件名称
        if is_train:
            with open(DATASET_PATH + "ImageSets/Main/train.txt", 'r') as f: # 调用包含训练集图像名称的txt文件
                self.filenames = [x.strip() for x in f]
        else:
            with open(DATASET_PATH + "ImageSets/Main/val.txt", 'r') as f:
                self.filenames = [x.strip() for x in f]
        self.imgpath = DATASET_PATH + "JPEGImages/"  # 原始图像所在的路径
        self.labelpath = "./labels/"  # 图像对应的label文件(.txt文件)的路径
        self.is_aug = is_aug

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, item):
        img = cv2.imread(self.imgpath+self.filenames[item]+".jpg")  # 读取原始图像
        h,w = img.shape[0:2]
        input_size = 448  # 输入YOLOv1网络的图像尺寸为448x448
        # 因为数据集内原始图像的尺寸是不定的，所以需要进行适当的padding，将原始图像padding成宽高一致的正方形
        # 然后再将Padding后的正方形图像缩放成448x448
        padw, padh = 0, 0