使用google colab + yolov3_tiny(pytorch) 训练自己的目标检测数据集

ref:

#1 google colab及基本使用    https://blog.csdn.net/nominior/article/details/105197850

#2 YOLOV3训练自己的数据集（PyTorch版本）    https://www.pianshen.com/article/1287380494/

#3 DARKNET YOLOV3-TINY 训练自己的数据集步骤   https://www.freesion.com/article/4196394355/

#4 KeyError: 'module_list.85.Conv2d.weight' #657    https://www.cnblogs.com/dgwblog/p/12225767.html

#5 深度学习中的遥感影像数据集   https://blog.csdn.net/nominior/article/details/105247990

 

 

 

1 数据准备：

a) 数据下载

下载代码：https://github.com/ultralytics/yolov3

下载权重文件：下载yolov3-tiny.pt权重文件，https://drive.google.com/drive/folders/1LezFG5g3BCW6iYaV89B2i64cqEUZD7e0

下载数据集：#2中有自己制作数据集的教程，本文使用的是#5中  RSOD-Dataset  的部分数据集

• aircraft dataset：  https://pan.baidu.com/s/1eRWFV5C
• playground：  https://pan.baidu.com/s/1nuD4KLb

b) 文件格式整理

解压代码：得到yolov3-master项目

 

移动权重：将下载的yolov3-tiny.pt文件移动到weights文件夹下

 

数据集整理：

在data下

• 新建Annotations、images、ImageSets、labels文件夹，
• 新建文本文件DIY_yolo.data、DIY_yolo.names，使用记事本打开，分别写入如下内容（name中为自己的类别名）

classes=2
train=data/train.txt
valid=data/test.txt
names=data/DIY_yolo.names

aircraft
playground

 

修改cfg/yolov3-tiny.cfg配置

找到[yolo]下的classes，改为类别数；找到[yolo]上的filters，改为(类别数+5)*3

注意！有2处yolo，都需要修改（修改2个classes，2个filters）

                          

 

解压数据集

• 将各类数据子集中的图片（JPEGImages文件夹中）复制到data/images中；
• 将各类数据子集中的xml标记（Annotation/xml文件夹中）复制到data/Annotations中；

 

c) 数据集预处理

在yolov3-master中，新建makeTxt.py文件，写入如下代码

import os
import random

trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'data/Annotations'
txtsavepath = 'data/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('data/ImageSets/trainval.txt', 'w')
ftest = open('data/ImageSets/test.txt', 'w')
ftrain = open('data/ImageSets/train.txt', 'w')
fval = open('data/ImageSets/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()

在yolov3-master中，新建voc_label.py文件，写入如下代码

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

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

classes = ["aircraft","playground"]  # 类别名


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/Annotations/%s.xml' % (image_id))
    out_file = open('data/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')
        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/labels/'):
        os.makedirs('data/labels/')
    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()
    list_file = open('data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

 

2 数据上传及运行：

*本地也可以运行，但可能显存不足，报cuda memory错误

a）上传数据

压缩yolov3-master文件夹

上传yolov3-master压缩包，到google drive中（时间较长）

上传完成后，打开google colab，新建notebook，使用GPU运行时，挂载google drive并切换到压缩包所在的路径（参考#1）

 

b）运行代码

环境准备：

根据个人路径，在colab notebook中依次运行如下代码，

import os
path = r'/content/drive/My Drive'
os.chdir(path)

!ls
!nvidia-smi

!unzip yolov3-master.zip

path = r'/content/drive/My Drive/yolov3-master'
os.chdir(path)

!ls

!pip install -r requirements.txt

 

 

数据准备：

运行如下代码

#数据集划分，在data/ImageSets下生成txt
!python makeTxt.py

#label生成，在data/labels生成可训练的txt标记,并在data下生成train、val、test.txt
!python voc_label.py

 

train：

运行代码

!python train.py --data data/DIY_yolo.data --cfg cfg/yolov3-tiny.cfg --weights weights/yolov3-tiny.pt --epochs 20

生成weigths/last.pt、weigths/best.pt，train_batch0.jpg，result.png，result.txt

test：

运行代码

!python test.py --data data/DIY_yolo.data --cfg cfg/yolov3-tiny.cfg --weights weights/best.pt --save-json 

生成test_batch0_gt.jpg，test_batch0_pred.jpg，results.json

 

infer:

新建source_dir、output_dir，将需要推理的图像移入source_dir（本文随机选择1张aircraft、1张playground复制到source_dir）

运行代码

!python detect.py  --cfg cfg/yolov3-tiny.cfg --names data/DIY_yolo.names --weights weights/best.pt --source source_dir --output output_dir 

在output_dir生成同名预测结果