YOLOv1的pytorch复现版本，博主亲自测试完整复现。

在学习目标检测时候发现网上很多只讲原理流程的没有现场复现成功的，以下就是我在网上找了许久成功运行的YOLOv1   pytorch复现版本.

由于我对模型进行训练了，里面采用的前半部分卷积网络是使用resnet50，在官方训练好的resnet50网络参数上在进行的训练，这就意味着模型有点大，截图如下

仅仅这个模型训练好的参数就900M，所有资料如下。

链接：https://pan.baidu.com/s/1O31dOYinjjD5wawIhSxtIQ 
提取码：2022

以下是整个模块的大小，说白了就是这个模型参数占空间。

整个模型部署的文件：

 成果展示：

里面代码我都进行了注释，主要的模块如下：

find_classes.py

from dataset.data import xml2dict

import xml.etree.ElementTree as ET
from tqdm import tqdm
import json
import os

# 本模块主要检索voc文件中存放的类，并以json字典的方式存放

json_path = './classes.json'  # 存放对应的类和编号

root = r'G:\数据集\voc\VOCtest_06-Nov-2007\VOCdevkit\VOC2007'
annotation_root = os.path.join(root, 'Annotations')  # 进入标签文件夹
annotation_list = os.listdir(annotation_root)  # 标签文件夹的名称
annotation_list = [os.path.join(annotation_root, a) for a in annotation_list]  # 打开每一个标签文件

s = set()  # 存放无序不重复元素，若重复则不添加
for annotation in tqdm(annotation_list):
    xml = ET.parse(os.path.join(annotation)).getroot()
    data = xml2dict(xml)['object']  # 读取一个文件中的所有对象，将存在的对象都存放s中
    if isinstance(data, list):  # 添加容错率，防止出现格式问题
        for d in data:
            s.add(d['name'])
    else:
        s.add(data['name'])


s = list(s)  # s由字典形式变成列表形式
s.sort()  # 排序
data = {value: i for i, value in enumerate(s)}  # 将列表中的元素按照标号存放在字典中
json_str = json.dumps(data)   # 转换成json格式

with open(json_path, 'w') as f:
    f.write(json_str)  # 将字典写入json文件中

transform.py

import torch
import torchvision
import random

# 图像处理函数
class Compose:
    def __init__(self, transforms):
        self.transforms = transforms  # 转换函数

    def __call__(self, image, label):
        for t in self.transforms:  # 轮流使用图片转换函数
            image, label = t(image, label)
        return image, label


class ToTensor: #转换类型
    def __init__(self):
        self.totensor = torchvision.transforms.ToTensor()

    def __call__(self, image, label):
        image = self.totensor(image)
        label = torch.tensor(label)
        return image, label


class RandomHorizontalFlip:
    def __init__(self, p=0.5):
        self.p = p

    def __call__(self, image, label):
        """
        :param label: xmin, ymin, xmax, ymax
        如果图片被水平翻转,那么label的xmin与xmax会互换，变成 xmax, ymin, xmin, ymax
        由于YOLO的输出是(center_x, center_y, w, h) ,因此label的xmin与xmax换位不会影响损失计算与训练
        但是需要注意w,h计算时使用abs
        """
        if random.random() < self.p:
            #print("测试1",image.shape)  torch.Size([3, 375, 500])
            height, width = image.shape[-2:]   # height:375  width:500
            # print("测试2",label)      tensor([[300., 167., 397., 268.,  10.]])
            image = image.flip(-1)      # 水平翻转
            bbox = label[:, :4]
            # print("测试3",bbox)      tensor([[300., 167., 397., 268.]])
            # bbox: xmin, ymin, xmax, ymax
            bbox[:, [0, 2]] = width - bbox[:, [0, 2]]
            label[:, :4] = bbox
            # print("测试4",label) tensor([[200., 167., 103., 268.,  10.]]}
        return image, label


class Resize:
    def __init__(self, image_size, keep_ratio=True):
        """
        :param image_size: int
        keep_ratio = True  保留宽高比
        keep_ratio = False 填充成正方形
        """
        self.image_size = image_size
        self.keep_ratio = keep_ratio

    def __call__(self, image, label):
        """
        :param in_image: tensor [3, h, w]
        :param label: xmin, ymin, xmax, ymax
        :return:
        """
        # 将所有图片左上角对齐构成448*448tensor的Transform

        h, w = tuple(image.size()[1:])
        label[:, [0, 2]] = label[:, [0, 2]] / w
        label[:, [1, 3]] = label[:, [1, 3]] / h

        if self.keep_ratio:
            r_h = min(self.image_size / h, self.image_size / w)
            r_w = r_h
        else:
            r_h = self.image_size / h
            r_w = self.image_size / w

        h, w = int(r_h * h), int(r_w * w)
        h, w = min(h, self.image_size), min(w, self.image_size)
        label[:, [0, 2]] = label[:, [0, 2]] * w
        label[:, [1, 3]] = label[:, [1, 3]] * h

        T = torchvision.transforms.Resize([h, w])

        Padding = torch.nn.ZeroPad2d((0, self.image_size - w, 0, self.image_size - h))
        image = Padding(T(image))

        assert list(image.size()) == [3, self.image_size, self.image_size]
        # print("测试5",label)  tensor([[268.8000, 149.6320, 355.7120, 240.1280,  10.0000]]}
        return image, label

data.py

from dataset.transform import *

from torch.utils.data import Dataset
import xml.etree.ElementTree as ET
from PIL import Image
import numpy as np
import json
import os


def get_file_name(root, layout_txt):
    # 读取root/layout_txt文件，得到一个字符串，以\n回车符为分隔符分割成list，[:-1]去除末尾的分隔符
    with open(os.path.join(root, layout_txt)) as layout_txt:
        file_name = layout_txt.read().split('\n')[:-1]
    return file_name


def xml2dict(xml):
    # 产生一个字典框架：data：{'folder': None, 'filename': None, 'source': None, 'owner': None, 'size': None, 'segmented': None, 'object': None}
    # 将xml文件中的信息慢慢加载到字典data中去
    data = {c.tag: None for c in xml}
    # print("测试",data)
    # exit()
    for c in xml:
        def add(data, tag, text):
            if data[tag] is None:
                data[tag] = text
            elif isinstance(data[tag], list):
                data[tag].append(text)
            else:
                data[tag] = [data[tag], text]
            return data

        if len(c) == 0:
            data = add(data, c.tag, c.text)
        else:
            data = add(data, c.tag, xml2dict(c))
    return data


class VOC0712Dataset(Dataset):
    def __init__(self, root, class_path, transforms, mode, data_range=None, get_info=False):
        # label: xmin, ymin, xmax, ymax, class
        # root:根目录  class_path:包含所有类的字典路径  transforms:图片处理   mode：train or test
        with open(class_path, 'r') as f:
            json_str = f.read()
            self.classes = json.loads(json_str)  # 将类都加在到classes变量中
        layout_txt = None
        if mode == 'train':
            root = [root[0], root[0], root[1], root[1]]  # 主要目的就是同时将voc2007和voc2012联合起来
            layout_txt = [r'ImageSets\Main\train.txt', r'ImageSets\Main\val.txt',
                          r'ImageSets\Main\train.txt', r'ImageSets\Main\val.txt']
        elif mode == 'test':
            if not isinstance(root, list):  # 防止输入为单元素报错，并将其改成列表形式
                root = [root]
            layout_txt = [r'ImageSets\Main\test.txt']
        assert layout_txt is not None, 'Unknown mode'  # 如果mode不是train or test，抛出异常 Unknown mode

        self.transforms = transforms
        self.get_info = get_info

        self.image_list = []
        self.annotation_list = []
        for r, txt in zip(root, layout_txt):
            self.image_list += [os.path.join(r, 'JPEGImages', t + '.jpg') for t in get_file_name(r, txt)]
            # 将每个图片的地址保存在image_list中
            self.annotation_list += [os.path.join(r, 'Annotations', t + '.xml') for t in get_file_name(r, txt)]
            # 将每个图片的xml文件信息保存在annotation_list中
        # data_range是个二元数组，如data_range = [200,1000]表示数据集取200到1000区间的数据
        if data_range is not None:
            self.image_list = self.image_list[data_range[0]: data_range[1]]
            self.annotation_list = self.annotation_list[data_range[0]: data_range[1]]

    def __len__(self):
        # 返回数据集的长度
        return len(self.annotation_list)

    def __getitem__(self, idx):
        image = Image.open(self.image_list[idx])  # 获取图片
        image_size = image.size  # 获取图片的尺寸
        label = self.label_process(self.annotation_list[idx])
        # 获取图片标签label  label: xmin, ymin, xmax, ymax, class
        # label = [[156.  97. 351. 270.   6.]]
        if self.transforms is not None: #是否转换图片
            image, label = self.transforms(image, label)
        if self.get_info:  # 表示是否需要图像名称以及图像大小信息
            return image, label, os.path.basename(self.image_list[idx]).split('.')[0], image_size
        else:
            return image, label

    def label_process(self, annotation):
        xml = ET.parse(os.path.join(annotation)).getroot()
        data = xml2dict(xml)['object']

        '''
        此为xml文件中object中的格式