RK3588 环境配置第二弹 训练yolo权重文件

笨比(我)训练yolo

在上一次配置好ubuntu端的rknn环境后,我们现在需要训练一个检测器,由于项目需要,我们选择YOLO

数据集就更简单了,常见的数据集都可以,为了方便后续,选择:红外航拍人车检测数据集,光电红外开源平台

下载数据集

数据集结构如下:

datasets
	---dataset 
		----00000.jpg
		----00001.jpg
		-----…………
		-----00001.xml
		-----00002.xml
		………………
	---readme.txt

我们需要将数据集分开 ,并移动到代码的data_dir中

classes

根据数据集中的readme,数据集分为6类,

classes = ["person","car" ,"bus","cyclist","bike","truck"]

Stupid Methods

为了让小白也可以开train,我们对数据集的处理尽量使用笨方法:

第一步 将所有的jpg的路径打包为txt

create_txt.py


# -*- coding:utf-8 -*-
import glob
 
imageList = glob.glob(r"reddata\data\dataset\*.jpg")  # 图片所在文件夹的路径
f = open('train.txt', 'a')  # 创建标签文件存放图片文件名
for item in imageList:
    print(item)  
    img_name = item.split('/')[-1]  # 图片文件名018.jpg
    # img_name = item.split('\\')[-1]  # 图片文件名018.jpg
    f.write(img_name + '\n')  # 将图片文件名逐行写入txt
 
f.close()
print('OK')
#所有的jpg文件
将所有的xml文件转为labels

xml2labels.py

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
 
sets = []
classes = ["person","car" ,"bus","cyclist","bike","truck"]  ##修改为自己的类别
 
#原样保留。size为图片大小
# 将ROI的坐标转换为yolo需要的坐标
# size是图片的w和h
# box里保存的是ROI的坐标(x,y的最大值和最小值)
# 返回值为ROI中心点相对于图片大小的比例坐标,和ROI的w、h相对于图片大小的比例
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(image_add):
    #image_add进来的是带地址的.jpg
    image_add = os.path.split(image_add)[1]        #截取文件名带后缀
    image_add = image_add[0:image_add.find('.',1)] #删除后缀,现在只有文件名没有后缀
    #现在传进来的只有图片名没有后缀
   
    in_file = open('reddata/data/xml/' + image_add + '.xml',encoding='utf8', errors="ignore")        #修改为你自己的输入目录
    out_file = open('reddata/data/labels/%s.txt'%(image_add), 'w',encoding='utf8', errors="ignore")  #修改为你自己的输出目录
 
    tree=ET.parse(in_file)
    root = tree.getroot()
    
    if root.find('size'):
 
        size = root.find('size')
        w = int(size.find('width').text)    #偶尔xml标记出错,width或height设置为0了
        h = int(size.find('height').text)   #需要标记出来,便于单独处理
        if w==0:
            print("出错! width或height为0:  "+image_add)
            return
        #在一个XML中每个Object的迭代
        for obj in root.iter('object'):
            #iter()方法可以递归遍历元素/树的所有子元素
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            #如果训练标签中的品种不在程序预定品种,或者difficult = 1,跳过此object
            if cls not in classes or int(difficult)==1:
                continue
            #cls_id 只等于1
            cls_id = classes.index(cls)
            xmlbox = obj.find('bndbox')
            #b是每个Object中,一个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')
    else:
        print("出错!xml缺少size:  "+image_add)      #偶尔xml缺少size,需要标记出来,便于单独处理
    
 
 
image_adds = open("train.txt",encoding='utf8')        #修改为你自己的训练数据集目录
for image_add in image_adds:
    print(image_add)
    image_add = image_add.strip()
    print (image_add)
    convert_annotation(image_add)
 

翻看yolo代码,发现他的数据集格式是这样的

datasets
	---test 
		----images
			----00056.jpg
		----labels
			----00056.txt
			……
	---train
	----images
			----00073.jpg
		----labels
			----00073.txt
			……
	----val
	----images
			----02752.jpg
		----labels
			----02752.txt
			……
			
	

我们仿照他的结构

根据训练集,测试集,验证集的比例将数据集划分

data_split.py


import os, random, shutil
def moveFile(fileDir):
        pathDir = os.listdir(fileDir)    #取图片的原始路径
        filenumber=len(pathDir)
        rate=1    #自定义抽取图片的比例,比方说100张抽10张,那就是0.1
        picknumber=int(filenumber*rate) #按照rate比例从文件夹中取一定数量图片
        sample = random.sample(pathDir, picknumber)  #随机选取picknumber数量的样本图片
        print (sample)
        for name in sample:
                shutil.move(fileDir+name, tarDir+name)
        return

if __name__ == '__main__':
	fileDir = "./newdatastes/reddata/data/dataset/"    #源图片文件夹路径
	tarDir = './datasets/val/images/'    #移动到新的文件夹路径
	moveFile(fileDir)


这里可以根据自己想要的比例抽取数据集,需要注意的是:运行代码后已经移动的jpg无法恢复

根据已经存在的train,test,val中的 jpg 将相应的labels 移动过来

由于我们是随机抽取的jpg数据集,所以我们不能直接运行上面的代码继续移动labels
所以我们先让数据集产生txt目录,然后根据这些索引去移动对应的labels:
name2txt.py

import os 
 
file_path = "./val/images/"  # Path
path_list = os.listdir(file_path)
 
path_name = []
 
for i in path_list:
    path_name.append(i.split(".")[0])
 
#path_name.sort()
 
for file_name in path_name:
    with open("val.txt", "a") as file:
        file.write(file_name +".txt" + "\n")
        print(file_name)
    file.close()

分别将三个文件夹下的jpg目录产出:
train.txt
test.txt
val.txt

根据txt文件移动对应的labels

movetxt.py

import shutil

if __name__ == '__main__':
    file_object = open('./datasets/val.txt')
    try:
        for line in file_object:
            # print(line)
            shutil.move('./newdatastes/reddata/data/labels/'+line.rstrip('\n'), "./datasets/val/labels")
    finally:
        file_object.close()

注意不要有中文路径

成功制作数据集!!

其实用另一种方法可以很快速的制作出这种数据集,但是在前两天有个学弟问过我一些制作数据集的问题后,我意识到,一些没做过这方面的人可能更倾向于笨方法去处理,我也很笨hh,所以还是写出来了

更改yolo 代码

更改data/xxxx.yaml
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ./datasets # dataset root dir
train: ./datasets/train/images  # train images (relative to 'path') 128 images
val:   ./datasets/val/images #images/train2017  # val images (relative to 'path') 128 images
test:  ./datasets/test/images # test images (optional)

# number of classes
nc: 6  #更改为你自己的classesnum

# Classes
names: [ 'person', 'car', 'bus', 'cyclist', 'bike', 'truck']#更改为你自己的classes name

其他的就没什么大的修改点了,有什么bug就去查好了(doge)

记得install requirements

我是用的yolo 代码依赖于1.8.0的torch,但是torch版本
但是有GPU 不用是笨蛋hh
所以我安装了1.8.1touch+cu111

记得去看看你们的cuda是否可用!!!

遇事不决 train一发

先conda 激活环境

conda activate env

然后既可以开始train,根据你自己的设备情况选择epochs和batch-size

python train.py  --epochs 300 --batch-sie 32 --device 0

在这里插入图片描述
经过十个小时的训练,已经成功训练完毕,爱训练过程中可以利用tensorboard 去看看 训练的收敛情况,但是我懒hh

如果不出意外的话,训练好的pt文件会保存在runs/train/exp_x/weights/best.pt
在weights 中还保留有训练中产生的各种参数:
such as :

在这里插入图片描述
还有其他的就不一一展示了

对于笨比(我)来说,看曲线不如detect

所以:

python detect.py --weights ./runs/train/weights/best.pt  --source ./reference/images -devie 0

在这里插入图片描述

可以看到速度还是很快的:

在这里插入图片描述
效果也很不错

敬请期待! pt --> onnx -->rknn

  • 0
    点赞
  • 4
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值