使用现有数据集和自己的数据训练yolov3&yolov3_tiny模型

        因为项目中需要检测定制的种类，所以需要使用yolov3重新训练新的数据。yolov3-tiny以及yolov4、yolov4-tiny也都顺便训练了一下，训练过程大同小异，主要是数据的收集、分类、标注、标签转换。

目录

一：开源数据集的下载和标注转换(从COCO转换成XML)

        1、使用fiftyone下载所需要的子集。

        a、下载google的open image(open-images-v6)

        b、下载COCO2017数据集子集

        2、直接去COCO官网下载完整的数据集(此处使用的是COCO-2017)。

二、自己的数据集收集和标注。

 1、数据集收集

2、标注

三：将所有图片和标注转换成YOLO所需要的格式

1、下载darknet源码。

2、为了方便后面的脚本使用，在darknet目录下面建立以下文件夹：

3、coco数据集的重新整理。

4、得到yolo格式的label文件

四、修改darknet的配置开始训练

1、修改cfg/voc.data

2、修改 cfg/yolov3-tiny.cfg

3、开始训练

4、测试效果

结束：

---

一：开源数据集的下载和标注转换(从COCO转换成XML)

        open-images-v6包含600多个种类的图片，coco-2017数据集包括了80个种类。其中很多种类是我们不需要的，所以只需要使用其中的子集就可以。

        1、使用fiftyone下载所需要的子集。

官网链接：FiftyOne — FiftyOne 0.13.2 documentation

        安装教程在链接中都有。按照教程中安装完成之后，提供两个python脚本下载数据集的子集。

        a、下载google的open image(open-images-v6)

try:
    import fiftyone as fo
except:
    print("0.0")
import fiftyone.zoo as foz


dataset = fo.zoo.load_zoo_dataset(
    "open-images-v6",
    split="train",
    label_types=["detections"],
    #classes=["Person"],
    #classes=["Car"],
    #classes=["Truck"],
    classes=["Fireplace"],

    max_samples=2000,
    dataset_dir="/home/easen/wyc/data/v6data/fire",
    dataset_name="fire_data",
)
session = fo.launch_app(dataset)
session.wait()

        此处dataset_dir设置一次之后不要轻易更改，更改后每次运行脚本需要重新下载4.9G的文件，网速快的就不存在这个问题了，哈哈。 

        使用try: import 是因为会报个错误，但是实际使用中并不影响。所以加了try过滤掉错误。

        b、下载COCO2017数据集子集

try:
    import fiftyone as fo
except:
    print("except error")
dataset = fo.zoo.load_zoo_dataset(
    "coco-2017",
#    split="validation",
    split="train",
    label_types=["detections"],
    classes=["truck"],
    max_samples=2000,
    dataset_dir="/home/easen/wyc/data/coco_data",
)
session = fo.launch_app(dataset)
session.wait()

        classes中的truck是小写，是因为COCO数据集的class文件中是这样的，在Google中首字母是大写。

        2、直接去COCO官网下载完整的数据集(此处使用的是COCO-2017)。

官网链接：COCO - Common Objects in Context

        下载完成之后包含以下几个.zip，图中红色的是下载的原始文件，annotations test2017 train2017 val2017是unzip之后的文件夹。

        得到完整的数据集后需要提取出其中的子集和XML格式的标注文件。此处提供个脚本，亲测没有问题。

        其中pycocotools的安装参考git clone https://github.com/cocodataset/cocoapi.git，下载后make ，却啥补啥。我是python3 -m venv coco了个虚拟环境折腾的。

from pycocotools.coco import COCO
import os
import shutil
from tqdm import tqdm
import skimage.io as io
import matplotlib.pyplot as plt
import cv2
from PIL import Image, ImageDraw
 
# 需要设置的路径 写自己的文件夹就行
savepath=" " 
img_dir="xx/images"
anno_dir="xx/annotations"
datasets_list=['train2017', 'val2017']

#coco有80类，这里写要提取类的名字，以person为例 
classes_names = ['person'] 
#包含所有类别的原coco数据集路径
'''
目录格式如下：
$COCO_PATH
----|annotations
----|train2017
----|val2017
----|test2017
'''
dataDir= '/path/to/coco_orgi/' 
 
headstr = """\

    VOC
    %s
    
        My Database
        COCO
        flickr
        NULL
    
    
        NULL
        company
    
    
        %d
        %d
        %d
    
    0
"""
objstr = """\
    
"""
 
tailstr = '''\

'''
 
# 检查目录是否存在，如果存在，先删除再创建，否则，直接创建
def mkr(path):
    if not os.path.exists(path):
        os.makedirs(path)  # 可以创建多级目录

def id2name(coco):
    classes=dict()
    for cls in coco.dataset['categories']:
        classes[cls['id']]=cls['name']
    return classes
 
def write_xml(anno_path,head, objs, tail):
    f = open(anno_path, "w")
    f.write(head)
    for obj in objs:
        f.write(objstr%(obj[0],obj[1],obj[2],obj[3],obj[4]))
    f.write(tail)
 
 
def save_annotations_and_imgs(coco,dataset,filename,objs):
    #将图片转为xml，例:COCO_train2017_000000196610.jpg-->COCO_train2017_000000196610.xml
    dst_anno_dir = os.path.join(anno_dir, dataset)
    mkr(dst_anno_dir)
    anno_path=dst_anno_dir + '/' + filename[:-3]+'xml'
    img_path=dataDir+dataset+'/'+filename
    print("img_path: ", img_path)
    dst_img_dir = os.path.join(img_dir, dataset)
    mkr(dst_img_dir)
    dst_imgpath=dst_img_dir+ '/' + filename
    print("dst_imgpath: ", dst_imgpath)
    img=cv2.imread(img_path)
    #if (img.shape[2] == 1):
    #    print(filename + " not a RGB image")
     #   return
    shutil.copy(img_path, dst_imgpath)
 
    head=headstr % (filename, img.shape[1], img.shape[0], img.shape[2])
    tail = tailstr
    write_xml(anno_path,head, objs, tail)
 
 
def showimg(coco,dataset,img,classes,cls_id,show=True):
    global dataDir
    I=Image.open('%s/%s/%s'%(dataDir,dataset,img['file_name']))
    #通过id，得到注释的信息
    annIds = coco.getAnnIds(imgIds=img['id'], catIds=cls_id, iscrowd=None)
    # print(annIds)
    anns = coco.loadAnns(annIds)
    # print(anns)
    # coco.showAnns(anns)
    objs = []
    for ann in anns:
        class_name=classes[ann['category_id']]
        if class_name in classes_names:
            print(class_name)
            if 'bbox' in ann:
                bbox=ann['bbox']
                xmin = int(bbox[0])
                ymin = int(bbox[1])
                xmax = int(bbox[2] + bbox[0])
                ymax = int(bbox[3] + bbox[1])
                obj = [class_name, xmin, ymin, xmax, ymax]
                objs.append(obj)
                draw = ImageDraw.Draw(I)
                draw.rectangle([xmin, ymin, xmax, ymax])
    if show:
        plt.figure()
        plt.axis('off')
        plt.imshow(I)
        plt.show()
 
    return objs
 
for dataset in datasets_list:
    #./COCO/annotations/instances_train2017.json
    annFile='{}/annotations/instances_{}.json'.format(dataDir,dataset)
 
    #使用COCO API用来初始化注释数据
    coco = COCO(annFile)
 
    #获取COCO数据集中的所有类别
    classes = id2name(coco)
    print(classes)
    #[1, 2, 3, 4, 6, 8]
    classes_ids = coco.getCatIds(catNms=classes_names)
    print(classes_ids)
    for cls in classes_names:
        #获取该类的id
        cls_id=coco.getCatIds(catNms=[cls])
        img_ids=coco.getImgIds(catIds=cls_id)
        print(cls,len(img_ids))
        # imgIds=img_ids[0:10]
        for imgId in tqdm(img_ids):
            img = coco.loadImgs(imgId)[0]
            filename = img['file_name']
            # print(filename)
            objs=showimg(coco, dataset, img, classes,classes_ids,show=False)
            print(objs)
            save_annotations_and_imgs(coco, dataset, filename, objs)

        脚本运行完之后就得到了annotations 和images两个文件夹。分别包含了xml格式的标注和图片。        

        找了个脚本看一下得到的xml文件中的坐标信息绘制在图片上的效果。这一步可以省略，脚本运行一会儿后建议ctrl+c退出  要不然会等所有文件绘制完之后才会退出。(感谢原作者csy，时间过去比较久，原文链接找不到了抱歉)

# coding: utf-8
# author:csy
# 2021-04-23


"""
将VOC数据标签映射到原图上
"""

import cv2
import os
import time
import xml.etree.cElementTree as ET


def get_bbox(xml_path):
    tree = ET.ElementTree(file=xml_path)
    root = tree.getroot()
    object_set = root.findall('object')
    object_bbox = list()

    for Object in object_set:
        bbox = Object.find('bndbox')
        x1 = int(bbox.find('xmin').text.split('.')[0])
        y1 = int(bbox.find('ymin').text.split('.')[0])
        x2 = int(bbox.find('xmax').text.split('.')[0])
        y2 = int(bbox.find('ymax').text.split('.')[0])
        obj_bbox = [x1, y1, x2, y2]  
        object_bbox.append(obj_bbox)
    return object_bbox


def drow_object(img_file, bndboxes):

    img = cv2.imread(img_file)

    for i in range(len(bndboxes)):
        xmin = bndboxes[i][0]
        ymin = bndboxes[i][1]
        xmax = bndboxes[i][2]
        ymax = bndboxes[i][3]

        cv2.rectangle(img, (xmin, ymax), (xmax, ymin), (0, 0, 255), 2)

    cv2.imwrite('result' + str(time.time()) + '.jpg', img)


if __name__ == '__main__':
    xml_dir = 'annotations/train2017/'
    img_dir = 'images/train2017/'
    for file in os.listdir(xml_dir):
        file_name = file.split('.')[0]
        xml = file_name + '.xml'
        pic = file_name + '.jpg'
        bndboxes = get_bbox(xml_path=os.path.join(xml_dir, xml))

        drow_object(img_file=os.path.join(img_dir, pic), bndboxes=bndboxes)

        到此开源数据集的子集获取完毕，先不着急转换成YOLO格式的label。

二、自己的数据集收集和标注。

 1、数据集收集

        收集自己的数据集，通常就是在网上爬一些图片以及在实际使用的环境下取照片。收集完所有照片后，要保证自己的文件名称和之前的文件名称不一样，此处记录一个重命名脚本：

rename_dataset.py

使用方法：python3 rename_dataset.py  images_path  start_number

import sys
import os

i = len(sys.argv)
print("argv is %d"%(i))
if len(sys.argv) != 3:
    print("***********************************\nerror: please input %s image_path  start_number\n***************************" %(sys.argv[0]))
    exit()

print("start rename %s/images"%(sys.argv[1]))
path = sys.argv[1] #得到要修改的文件夹路径

filelist = os.listdir(path) 
count=int(sys.argv[2])#设置图片编号从传入的参数开始
for file in filelist:#打印出所有图片原始的文件名
    print(file)
for file in filelist:   #遍历所有文件
    Olddir=os.path.join(path,file)   #原来的文件路径
    if os.path.isdir(Olddir):   #如果是文件夹则跳过
        continue
    filename=os.path.splitext(file)[0]   #文件名
    filetype=os.path.splitext(file)[1]   #文件扩展名
    Newdir=os.path.join(path,str(count).zfill(6)+filetype)  #用字符串函数zfill 以0补全所需位数
    os.rename(Olddir,Newdir)#重命名
    count+=1

2、标注

        此处使用labelimg进行标注，标注好的格式为XML。

        下载源文件 ： git clone https://github.com/tzutalin/labelImg.git  

sudo apt-get install pyqt5-dev-tools
sudo pip3 install -r requirements/requirements-linux-python3.txt
make qt5py3
python3 labelImg.py
点击change save dir 指定保存的XML文件夹。

labelimg快捷方式:w画框、d下一张、a上一张、空格：图中没有。

所有图片标注完成以后得到XML格式文件。

        至此所有的图片文件和标注文件都已完成，将所有图片放到一个文件夹，XML文件放到一个文件夹。

三：将所有图片和标注转换成YOLO所需要的格式

1、下载darknet源码。

git clone https://github.com/AlexeyAB/darknet

        修改MAKEFILE：

        GPU=1

        CUDNN=1

        OPENCV=1

        make编译生成darknet可执行文件，

        下载模型权重文件

wget https://pjreddie.com/media/files/yolov3-tiny.weights

        测试一下：

./darknet detect cfg/yolov3-tiny.cfg yolov3-tiny.weights data/dog.jpg

2、为了方便后面的脚本使用，在darknet目录下面建立以下文件夹：

        把自己数据的所有xml文件放入annotations文件夹，所有图片放入JPEGImages文件夹，运行脚本在Main目录下生成train.txt和val.txt。train.txt中是所有train数据的名称(不带后缀)。val是所有验证数据的名称。

3、coco数据集的重新整理。

        编写脚本获取数据集的train.txt和val.txt。

import os
import random

train_xmlfilepath = 'annotations/train2017/'
val_xmlfilepath = 'annotations/val2017/'
total_train_xml = os.listdir(train_xmlfilepath)
total_val_xml = os.listdir(val_xmlfilepath)

num_train = len(total_train_xml)
list_train = range(num_train)
num_val = len(total_val_xml)
list_val = range(num_val)

ftrain = open('annotations/train.txt', 'w')
fval = open('annotations/val.txt', 'w')

for i in list_train:
    name = total_train_xml[i][:-4] + '\n'
    ftrain.write(name)
ftrain.close()

for j in list_val:
    name = total_val_xml[j][:-4] + '\n'
    fval.write(name)
fval.close()

        将新生成的train.txt、val.txt内容分别追加到上一步生成的文件中。当然也可以直接用a+方式open文件直接生成。

        将Annotations中的train和val的xml文件全部复制到VOCdevkit/VOC2007/AnnotationsZ/中。

        将images中的train和val图片文件全部复制到VOCdevkit/VOC2007/JPEGImage中。

当文件过多的时候会报错：

        bash: /usr/bin/cp: Argument list too long

使用：

find source_folder -name "*.jpg" | xargs -i cp {} des_folder

需要将source_folder和des_folder改为自己的源文件夹和目的文件夹。 

在使用find + xargs方法的过程中发现复制xml文件时还是会报错。

写个脚本解决一下：

import os
import shutil

train_xmlfilepath = 'annotations/train2017/'
val_xmlfilepath = 'annotations/val2017/'
total_train_xml = os.listdir(train_xmlfilepath)
total_val_xml = os.listdir(val_xmlfilepath)

num_train = len(total_train_xml)
list_train = range(num_train)
num_val = len(total_val_xml)
list_val = range(num_val)


for i in list_train:
    shutil.copy(train_xmlfilepath+total_train_xml[i],"des_dir")//目的路径自己替换
for j in list_val:
    shutil.copy(val_xmlfilepath+total_val_xml[j],"des_dir")
print("ok")

4、得到yolo格式的label文件

修改darknet/scripts/voc_label.py。

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

sets=[ ('2007', 'train'), ('2007', 'val')]

classes = ["person","car","truck"]

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    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(year, image_id):
    in_file = open('VOCdevkit/VOC%s/Annotations/%s.xml'%(year, image_id))
    out_file = open('VOCdevkit/VOC%s/labels/%s.txt'%(year, 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')
        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()

for year, image_set in sets:
    if not os.path.exists('VOCdevkit/VOC%s/labels/'%(year)):
        os.makedirs('VOCdevkit/VOC%s/labels/'%(year))
    image_ids = open('VOCdevkit/VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n'%(wd, year, image_id))
        convert_annotation(year, image_id)
    list_file.close()

os.system("cat 2007_train.txt 2007_val.txt  > easen_train.txt")
os.system("cat 2007_train.txt 2007_val.txt  > easen_train.all.txt")

         运行脚本生成要用的文件。到这一步数据集和标注已经全部完成。

四、修改darknet的配置开始训练

1、修改cfg/voc.data

classes= 3
train  = /home/easen/wyc/test/darknet/easen/easen_train.txt
valid  = /home/easen/wyc/test/darknet/easen/2007_val.txt
names = data/voc.names
backup = /home/easen/wyc/test/darknet/easen/weights/  #保存权重的路径

修改data/voc.names

        修改成对应的名称。

2、修改 cfg/yolov3-tiny.cfg

# Testing
#batch=1
#subdivisions=1
# Training
batch=64                #根据电脑性能确定
subdivisions=8

burn_in=1000            
max_batches = 120000       #训练次数


filters=（5+class_num）*3   #这三项有两处要修改  在vim中 :/filters 按n查找下一处
anchors可改可不改
classes=3   #改成自己的class数量

其中anchors的修改使用kmeans生成：

git clone https://github.com/lars76/kmeans-anchor-boxes

    代码就不贴了。写的太累了。有兴趣的可以留言。

3、开始训练

/darknet partial ./cfg/yolov3-tiny.cfg ./yolov3-tiny.weights ./yolov3-tiny.conv.15 15

sudo ./darknet detector train cfg/voc.data cfg/yolov3-tiny.cfg yolov3-tiny.conv.15

4、测试效果

./darknet detector test cfg/voc.data cfg/yolov3-tiny.cfg easen/weights/yolov3-tiny_last.weights VOCdevkit/VOC2007/JPEGImages/000000035128.jpg

结束：

        记录了训练全过程，接下来把模型部署到arm上。因为yolov3只能跑到10fps左右，所以使用了yolov3-tiny，能跑到45fps左右，足够实时检测。