import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
classes = ['car', 'watcher', 'base','ignore','armor']
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
if w>=1:
w=0.99
if h>=1:
h=0.99
return (x,y,w,h)
def convert_annotation(rootpath,xmlname):
xmlpath = rootpath + '/image_annotation'
xmlfile = os.path.join(xmlpath,xmlname)
with open(xmlfile, "r", encoding='UTF-8') as in_file:
txtname = xmlname[:-4]+'.txt'
print(txtname)
txtpath = rootpath + '/worktxt'#生成的.txt文件会被保存在worktxt目录下
if not os.path.exists(txtpath):
os.makedirs(txtpath)
txtfile = os.path.join(txtpath,txtname)
with open(txtfile, "w+" ,encoding='UTF-8') as out_file:
tree=ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
out_file.truncate()
for obj in root.iter('object'):
#difficult = obj.find('difficulty').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')
if __name__ == "__main__":
rootpath='D:/DJI ROCO/robomaster_South China Regional Competition/'
xmlpath=rootpath+'/image_annotation'
list=os.listdir(xmlpath)
for i in range(0,len(list)) :
path = os.path.join(xmlpath,list[i])
if ('.xml' in path)or('.XML' in path):
convert_annotation(rootpath,list[i])
print('done', i)
else:
print('not xml file',i)
制作标签文件,将xml文件转化为yolov5需要的txt文件
转化时记得将某些异常数据进行处理。比如某些小于0的数据要让它变为大于0的数据,python中不能出现中文路径。
创建环境时特别慢
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes
这样可以切换到国内源,创建环境就会快很多
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
conda config --set show_channel_urls yes
这样下载python包的时候会快很多
创建conda环境:
conda create --name your conda name python=3.7
最后是激活
conda activate yolov5
就可以安装包了
最后就是配置train.py文件,
parser.add_argument(’–weights’, type=str, default=’./weights/yolov5s.pt’, help=‘initial weights path’)
这个是预训练模型地址,可以使用5s,5m,5l分别对应的是small,mideate,large,可以根据需要选择不同的模型
parser.add_argument(’–cfg’, type=str, default=’./models/yolov5s.yaml’, help=‘model.yaml path’)
这个为配置文件,需要自己手动修改。
#parameters
nc: 35 # number of classes
这里只需要修改训练个数
parser.add_argument(’–epochs’, type=int, default=2)
parser.add_argument(’–batch-size’, type=int, default=4, help=‘total batch size for all GPUs’)
第一个为训练轮数,表示训练几轮,第二个为每批处理的个数,这两个数根据自己显卡的性能来选择,可能耗时不一样。
parser.add_argument(’–logdir’, type=str, default=‘runs/’, help=‘logging directory’)
这个为数据可视化路径,可以使用tensorboard来查看训练过程
训练成功后会生成last.pt和best.pt两个模型,可以调用模型来进行检测,可以使用自己的摄像头。
(yolov5) PS D:\yolov5> python .\detect.py --weight .\runs\exp0\weights\best.pt --source .\test.mp4
emmmmmm,检测效果不太好。。。
安装pycocotools,在官网上先下载源代码
https://github.com/philferriere/cocoapi
下载之后解压并进入cocoapi-master/PythonAPI进入后用命令行来运行python程序, python .\setup.py build_ext --inplace
python .\setup.py build_ext install
这样就安装上了这个包。
pip 安装pygame
pip install pygame -i https://pypi.tuna.tsinghua.edu.cn/simple
切换为国内源来安装pygame
pip install --user -i https://pypi.tuna.tsinghua.edu.cn/simple opencv-python
下载时加入清华源可以加快下载速度