pytorch 实现yolo3详细理解（五）训练自己数据集和csv数据集标签处理

摘要

前面基本已经将yolo3的大致细节都分析了，那么现在就要训练自己下载的数据集来看下效果，进行迁移学习，首先我会先对github本身自定义的custom数据集进行训练，只有一张照片，一个标签签，之后训练自己的数据集是要从xml文件先提取标签，完全按照custom中的格式进行布局，然后修改一下cfg文件就可以运行。dataset源码是对txt文件的处理，在实际运行中对数据进行分析是利用panda，所以直接利用panda生成csv文件进行读取更加方便。代码https://github.com/eriklindernoren/PyTorch-YOLOv3

custom数据集文件

在data文件夹下，这里一种有二个文件夹，三个文件，这里主要是一个模板，用于将自己的数据集处理成这种形式方便训练，这里也可以直接训练custom数据集，现在先来操作一下训练custom

训练custom数据集

在config文件夹下，如果是linux系统就可以直接运行.sh文件产生custom的cfg文件，这个文件的作用是darknet的网络部分，由于数据集的类别变化，输出的信息是由变化的，window系统可以直接复制yolo3.cfg文件 然后修改里面的参数，从新命名为yolo3-custom.cfg文件。具体修改如下（打开文件在最下面）

custom 中只有一个分类的train，所以类别为1，还有filters=3*（5+1）=18，3是yolo3中每个预测大小都有三种，5是框的坐标和置信度，1是我们需要训练的类别，
这里只写个一处位置的修改，上面还有二处，找到【yolo】的位置，修改classes，和上一层的filters=18就行，

开始训练

打开train.py文件，修改这二个地方，改成custom的数据集即可。到这里是可以运行的，当然你的电脑配置可能到不到设置的要求，在按照下面的修改一下batch 和num，主要看你的电脑配置了，

到目前为止，应该可以正常训练，如果不行，就是包的问题，

注意这里的tensorflow是cpu版本，cpu和gpu不能同时安装，主要是利用tensorboard，如果是这里有问题的话，可以将train.py利用tensorboard的代码全部注释掉，如下

自定义数据处理成custom格式

首先看下我下载数据

二个文件夹，一个放xml文件，一个是照片，


这里我们需要读取每一个xml文件来生成label，一个xml文件生成一个label，名字也是对应照片上的名字。这里还有一个小细节，就是在生成之前label，我们需要全部的分类，比如这里有五个分类目标，这五个分类目标是遍历每一个xml物体的名称，去掉重复的部分，就是分类了，这部分代码我没有展示，可以需要自己写（比较简单），下面放的这个代码是遍历全部的xml文件生成每一个对应的label，而且这里的框的位置是xywh形式。

生成lable的代码

import os,shutil
import numpy as np
import xml.etree.ElementTree as ET
import os #   海胆        全息      扇贝              海星          海草
_classes = ('echinus', 'holothurian', 'scallop', 'starfish', 'waterweeds')
from PIL import Image
import torchvision.transforms as transforms
_class_to_ind = dict(zip(_classes, range(5)))
_ind_to_class = dict(zip(range(5), _classes))

a='E:\pytorch\JPEGImages/train/box/'
e='E:\pytorch\JPEGImages/train/label/'
d = 'E:\pytorch\JPEGImages/train/image/'
b=os.listdir(a)
xml = []
txt = []
img = []
for i in range(len(b)):
    xml.append(a + b[i])
    path,filename = os.path.split(xml[i])
    name = os.path.splitext(filename)
    txt.append(e+name[0]+'.txt')
    img.append(d+name[0]+'.jpg')

for j in range(len(b)):
    xml_list = []
    nn = []
    d = img[j]
    image = transforms.ToTensor()(Image.open(d).convert('RGB'))
    _, height, weight = image.shape
    a = xml[j]
    tree = ET.parse(a)
    root = tree.getroot()
    for member in root.findall('object'):
        a = member[0].text
        b = _class_to_ind[a]
        value = (b,
                 int(member[1][0].text),
                 int(member[1][1].text),
                 int(member[1][2].text),
                 int(member[1][3].text)
                 )
        xml_list.append(value)
    for ii in range(len(xml_list)):
        tt = int(xml_list[ii][0])
        x1 = float(xml_list[ii][1])
        y1 = float(xml_list[ii][2])
        x2 = float(xml_list[ii][3])
        y2 = float(xml_list[ii][4])
        x = ((x2+x1)/2)/weight
        y = ((y2+y1)/2)/height
        w = (x2-x1)/weight
        h = (y2-y1)/height
        x = round(x,8)
        y = round(y,8)
        w = round(w,8)
        h = round(h,8)
        sum = (tt,x,y,w,h)
        nn.append(sum)
    e = txt[j]
    file = open(e, 'w')
    for p in range(len(nn)):
        label = str(nn[p][0])
        x = str(nn[p][1])
        y = str(nn[p][2])
        w = str(nn[p][3])
        h = str(nn[p][4])
        file.write(label + ' ')
        file.write(x + ' ')
        file.write(y + ' ')
        file.write(w + ' ')
        file.write(h + '\n')
    file.close()

代码对应每一个不同的xml可能不同，主要是你需要学会读取xml文件的信息，在我之前的博客有讲解ET的使用。生成效果如下


这里边处理好全部的label了，还需要在整理二个txt文件保存路径即可，自己创建train.txt文件，保存训练照片的路径，在创建valid.txt保存验证数据集的路径，还需要穿件一个classes.names文件，保存类整体如下，对比custom还是很容易实现这一步的。

之后就是修改cfg文件了，一共有二个位置修改网络yolo3中的【yolo】层calsses修改自己分类和fitlers=3*（5+分类数）和创建custom.data用来数据集读取，具体如下之后再train.py再一次修改一下就可以了。

csv保存标签，重写dataset类

import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET
_classes = ('echinus', 'holothurian', 'scallop', 'starfish', 'waterweeds')
from PIL import Image
import torchvision.transforms as transforms
_class_to_ind = dict(zip(_classes, range(5)))
_ind_to_class = dict(zip(range(5), _classes))
d = 'E:\pytorch\JPEGImages/train/image/'

def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        print(xml_file)
        path, filename = os.path.split(xml_file)
        name = os.path.splitext(filename)
        num = int(name[0])
        img=(d + name[0] + '.jpg')
        image = transforms.ToTensor()(Image.open(img).convert('RGB'))
        _, height, weight = image.shape
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            a = member[0].text
            classes = _class_to_ind[a]
            value = (num,
                     img,
                     weight,
                     height,
                     classes,
                     int(member[1][0].text),
                     int(member[1][1].text),
                     int(member[1][2].text),
                     int(member[1][3].text)
                     )
            xml_list.append(value)
    column_name = ['num','filename','width','height','class','xmin','ymin','xmax','ymax']
    xml_df = pd.DataFrame(xml_list,columns=column_name)
    return xml_df
def main():
    image_path = os.path.join(os.getcwd(),'E:\pytorch\JPEGImages/train/box')
    xml_df = xml_to_csv(image_path)
    xml_df.to_csv('E:\pytorch\JPEGImages/train/labes2.csv',index=None)  #带路径  'E://'
    print('finish')
main()

这一步是生成csv文件， 这里方便讲解就放了一点数据，如下图所示，我写num的目的是为了标记，比如全部是0的时候，说明这些信息全是一张照片的信息，一个照片可能有多个框，所以这里是为了一张照片的信息。

from PIL import Image
import pandas as pd
import numpy as np
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import torch
import torch.nn.functional as F
import random
# 定义读取文件的格式
def pad_to_square(img, pad_value):
    c, h, w = img.shape
    dim_diff = np.abs(h - w)
    # (upper / left) padding and (lower / right) padding
    pad1, pad2 = dim_diff // 2, dim_diff - dim_diff // 2
    # Determine padding
    pad = (0, 0, pad1, pad2) if h <= w else (pad1, pad2, 0, 0)
    # Add padding
    img = F.pad(img, pad, "constant", value=pad_value)

    return img, pad
def resize(image, size):
    image = F.interpolate(image.unsqueeze(0), size=size, mode="nearest").squeeze(0)
    return image


def random_resize(images, min_size=288, max_size=448):
    new_size = random.sample(list(range(min_size, max_size + 1, 32)), 1)[0]
    images = F.interpolate(images, size=new_size, mode="nearest")
    return images

def horisontal_flip(images, targets):
    images = torch.flip(images, [-1])
    targets[:, 2] = 1 - targets[:, 2]
    return images, targets
class Mydataset(Dataset):
    def __init__(self,csv,transform=None, target_transform=None,):
        super(Mydataset,self).__init__()
        self.path = []
        self.df=pd.read_csv(csv)
        for i in range(3):
            a = self.df[self.df['num'] == i]
            n = a['filename'].index[0]
            self.path.append(a['filename'][n])
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, index):
        path=self.path[index]
        img = transforms.ToTensor()(Image.open(path).convert('RGB'))
        a = self.df[self.df['num'] == index]
        d = a.iloc[:, 4:9]
        e = np.array(d)
        targets = np.zeros((len(a), 6))
        targets[:, 1:] = e
        img, pad = pad_to_square(img, 0)
        _, padded_h, padded_w = img.shape
        x1 = targets[:, 2]
        y1 = targets[:, 3]
        x2 = targets[:, 4]
        y2 = targets[:, 5]
        x1 += pad[0]
        y1 += pad[2]
        x2 += pad[1]
        y2 += pad[3]
        targets[:, 2] = ((x1 + x2) / 2) / padded_w
        targets[:, 3] = ((y1 + y2) / 2) / padded_h
        targets[:, 4] = (x2 - x1) / padded_w
        targets[:, 5] = (y2 - y1) / padded_h
        if np.random.random() < 0.5:
            img, targets = horisontal_flip(img, targets)
        return path,img,targets
    def collate_fn(self, batch):
        paths, imgs, targets = list(zip(*batch))
        targets = [boxes for boxes in targets if boxes is not None]
        # Add sample index to targets
        for i, boxes in enumerate(targets):
            boxes[:, 0] = i
        targets = torch.cat(targets, 0)
        # Selects new image size every tenth batch
        if self.multiscale and self.batch_count % 10 == 0:
            self.img_size = random.choice(range(self.min_size, self.max_size + 1, 32))
        # Resize images to input shape
        imgs = torch.stack([resize(img, self.img_size) for img in imgs])
        self.batch_count += 1
        return paths, imgs, targets
    def __len__(self):
        return len(self.path)
val_data=Mydataset(csv='E:\pytorch\JPEGImages/train/22.csv', transform=transforms.ToTensor())
print(len(val_data))
trainloader = DataLoader(val_data, batch_size=1,shuffle=True, num_workers=0)
for i,(p,img,e) in enumerate(trainloader):
    print(img.shape)
    print(e)

重写dataset的核心就是要让输出值有img和targets，代码细节不做过多讲解了，主要说明这里的 len(self):这个函数决定这__getitem__(self, index): index，这一步是读取csv文件的关键点。利用csv文件可以进行数据分析，比如各个大小形式的照片比例是多少，框的大小，都影响着yolo3的效果，