YOLOv3 训练自己的数据集

 

关于YOLOv3的介绍不再多言，可以看一下论文。本文旨在尝试使用利用YOLO训练自己的数据集。

一、 环境

建议使用anaconda搭建环境。

建立一个名为yolov3（自定义）的环境，并安装pytorch和pillow。

conda create -n yolov3 python=3.7 pytorch=1.4 pillow

激活yolov3环境。

activate yolov3

安装opencv用pip比较方便。

pip install opencv-python

至此，关键的包就装完了，其余包跑的时候发现错误随便用pip或conda装一下就好。

二、下载并测试YOLO

这里选择github上ultralytics的基于pytorch的实现。找一个路径，

git clone https://github.com/ultralytics/yolov3.git

或者直接到https://github.com/ultralytics/yolov3下载。

其中weights文件在google云盘，不能**的同学可以在在这里下载自己需要的weights文件。下载后放到weights文件夹下

链接：https://pan.baidu.com/s/18Q0T5OSsYt_ubfP0QmzlLg 
提取码：w05u 

有兴趣的同学可以仔细阅读以下README文件。

使用摄像头进行测试：

python detect.py --cfg cfg/yolov3-tiny.cfg --weights weights/yolov3-tiny.weights --source 0

正常情况下你会看到一个被persion框圈起来的自己。

当然你也可以选择不用摄像头，把--source后面的0改成文件（图片or视频）就好，那样你就可以在output文件夹下找到对应的处理完的图像or视频。

三、训练自己的数据集

关于训练自己的数据集也可以看看原文档https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data

1. 收集照片

提供一个自己拍摄照片的脚本：

import cv2
import os

data_path = './data'  # 数据储存的文件夹
dataset_name = 'your_dataset'  # 数据集名称
image_path = data_path+'/'+dataset_name+'/images/'
MAX_image = 1000
video_source = 0

if __name__ == "__main__":
    cap = cv2.VideoCapture(video_source)
    rat, frame = cap.read()
    image_count = 0
    if not os.path.exists(data_path):
        os.mkdir(data_path)
    if not os.path.exists(data_path+'/'+dataset_name):
        os.mkdir(data_path+'/'+dataset_name)
    if not os.path.exists(image_path):
        os.mkdir(image_path)
    while rat:
        rat, frame = cap.read()
        cv2.imshow("create custom dataset", frame)
        keyPress = cv2.waitKey(1) & 0xff
        if keyPress == 27:
            break
        if keyPress == ord('s'):
            full_path = image_path+dataset_name+'{:04.0f}'.format(image_count)+'.jpg'
            cv2.imwrite(full_path, frame)
            print("save "+full_path)
            image_count += 1
            if image_count == MAX_image:
                print(" you have almost saved %d samples!" % MAX_image)
                break
    cv2.destroyAllWindows()
    cap.release()

使用方法：

• 修改数据集名称
• 设置预期数据数
• 设置视频源，0表示摄像头
• 运行程序后，按s键储存图像，按esc退出
• 若按帧直接存储，自行修改

2. 标注

可以使用labelimg进行标注，标注方法也很简单，框起来输入label，设置下difficult（遮挡严重就选，否则不用管），保存即可。

如有和我一样的懒人，可以试试用跟踪算法写一个辅助标注的工具。这里给一个参考https://github.com/Carey-cc/Label_image_with_medianflow。但效果并不是特别理想，局限性也很大。要求高的还是建议手动标注，更准确。

3. 生成coco格式数据集

标注完成后，文件结构应该是这样的：

|
|-.github
|-cfg
|-data
  |
  |-your_dataset
    |
    |-images
      |-xxxx.jpg
      |-xxxx.xml
  |-....
|-...

在images文件夹下，搜索xml，选中所有xml文件，然后剪切。

在your_dataset文件夹下建立如下文件结构，并把xml文件粘贴到指定位置

|
|-your_dataset
  |-Annotations
    |-xxxx.xml
  |-images
    |-xxxx.jpg
  |-ImageSets
  |-labels

然后划分数据集：data_split.py

import os
import random

dataset_name = "your_dataset"
trainval_percent = 0.1
train_percent = 0.9

xmlfilepath = './data/'+dataset_name+'/Annotations'
txtsavepath = './data/'+dataset_name+'/ImageSets'
total_xml = os.listdir(xmlfilepath)

num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

ftrainval = open(txtsavepath+'/trainval.txt', 'w')
ftest = open(txtsavepath+'/test.txt', 'w')
ftrain = open(txtsavepath+'/train.txt', 'w')
fval = open(txtsavepath+'/val.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

coco格式label：voc_label.py

import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'test', 'val']

classes = ["your class"]  # 设置你的类名，要求和标注的时候输入的一样
data_path = './data/your_dataset/'  #设置你的数据集路径


def convert(size, box):
    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):
    in_file = open(data_path+'Annotations/%s.xml' % (image_id))
    out_file = open(data_path+'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'):
        try:
            difficult = obj.find('difficulty').text
        except AttributeError:
            difficult = '0'
        cls = obj.find('name').text
        if cls not in classes or int(difficult) >= 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists(data_path+'labels'):
        os.makedirs(data_path+'labels')
    image_ids = open(data_path+'ImageSets/%s.txt' % (image_set)).read().strip().split('\n')
    list_file = open(data_path+'%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(data_path+'images/%s.jpg\n' % (image_id))
        print(image_id)
        convert_annotation(image_id)
    list_file.close()

然后新建.names和.data文件

文件内容如下：

your_dataset.names （你标注的类名）

your class

your_dataset.data

classes=1
train=data/your_dataset/train.txt
valid=data/your_dataset/test.txt
names=data/your_dataset/your_dataset.names
backup=backup/
eval=coco

其中classes为类别数

至此，你的数据集文件夹下的文件结构为：

|
|-your_dataset
  |-Annotations
    |-xxxx.xml
  |-images
    |-xxxx.jpg
  |-ImageSets
    |-test.txt
    |-train.txt
    |-trainval.txt
    |-val.txt
  |-labels
    |-xxxx.txt
  |-your_dataset.names
  |-your_dataset.data
  |-text.txt
  |-train.txt
  |-val.txt

 4. 训练

以yolov3-tiny为例，打开cfg/yolov3-tiny.cfg

找到文件中的两处yolo

[convolutional]
size=1
stride=1
pad=1
filters=255  -> 3*(1+5) = 18
activation=linear



[yolo]
mask = 3,4,5
anchors = 10,14,  23,27,  37,58,  81,82,  135,169,  344,319
classes=80 -> 1
num=6
jitter=.3
ignore_thresh = .7
truth_thresh = 1
random=1

修改其中 [convolutional]里的（只修改每个[yolo]前的这个）filters=3*(class数+5），下面[yolo]里的classes改成class数。

然后开始训练

python train.py --cfg cfg/yolov3-tiny.cfg --weights weights/yolov3-tiny.pt --data data/your_dataset/your_dataset.data

如果想指定epochs或batch-size，加上--epochs xxx --batch-size xxx.

训练默认都是用GPU 的，因此不用管这个，如果你有多块GPU要选择，加上--device xxx。

然后等待......

训练结束后可以看到训练结果：

5. 测试

训练好的模型会保存在weights文件夹下，bets.pt是最佳，last.pt是最后训练的，这里提醒一下，可以利用这个中断和继续训练，即训练的时候加载预训练模型选last.pt.

测试模型和开始的时候差不多，改一下weights文件就好。

python detect.py --cfg cfg/yolov3-tiny.cfg --weights weights/best.weights --source 0

 

总结

yolo的用户体验真的很不错，十分“丝滑”，精度也挺高。

另外，解决有关中文路径的问题：

在utils/datasets.py中，第514行做些修改，否则对中文路径解码会出现错误，提示Image Not Found，如下注释处

    if img is None:  # not cached
        img_path = self.img_files[index]
        # img = cv2.imread(img_path)  # BGR
        img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)  # 解决中文路径问题
        assert img is not None, 'Image Not Found ' + img_path