1.准备图片
建立VOCData文件夹
里面建立annotations和images文件夹
补充一个视频转图片的代码:
import cv2
import os
#此删除文件夹内容的函数来源于网上
def del_file(filepath):
"""
删除某一目录下的所有文件或文件夹
:param filepath: 路径
:return:
"""
del_list = os.listdir(filepath)
for f in del_list:
file_path = os.path.join(filepath, f)
if os.path.isfile(file_path):
os.remove(file_path)
def video_to_images(fps,path):
cv = cv2.VideoCapture(path)
if(not cv.isOpened()):
print("\n打开视频失败!请检查视频路径是否正确\n")
exit(0)
if not os.path.exists("images/"):
os.mkdir("images/") # 创建文件夹
else:
del_file('images/') # 清空文件夹
order = 0 #序号
h = 0
while True:
h=h+1
rval, frame = cv.read()
if h == fps:
h = 0
order = order + 1
if rval:
cv2.imwrite('./images/' + str(order) + '.jpg', frame) #图片保存位置以及命名方式
cv2.waitKey(1)
else:
break
cv.release()
print('\nsave success!\n')
# 参数设置
fps = 1 # 隔多少帧取一张图 1表示全部取
path="tram_result.mp4" # 视频路径 比如 D:\\images\\tram_result.mp4 或者 D:/images/tram_result.mp4
if __name__ == '__main__':
video_to_images(fps,path)
# 会在代码的当前文件夹下 生成images文件夹 用于保存图片
# 如果有images文件夹,会清空文件夹!
2.划分数据集
划分训练集,测试集和验证集
# coding:utf-8
import os
import random
import argparse
parser = argparse.ArgumentParser()
#xml文件的地址,根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='Annotations', type=str, help='VOCData/Annotations')
#数据集的划分,地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 0.8 # 训练集和验证集所占比例。
train_percent = 0.7 # 训练集所占比例,可自己进行调整
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
for i in list_index:
name = total_xml[i][:-4] + '\n'
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
运行完毕后 会生成 ImagesSets\Main 文件夹,且在其下生成 测试集、训练集、验证集,存放图片的名字(无后缀.jpg)
xml_to_txt文件 必须是绝对路径
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd
sets = ['train', 'val', 'test']
classes = ["light", "post"] # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)
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_id):
in_file = open('D:/Yolov5/yolov5/VOCData/Annotations/%s.xml' % (image_id), encoding='UTF-8')
out_file = open('D:/Yolov5/yolov5/VOCData/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
#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))
b1, b2, b3, b4 = b
# 标注越界修正
if b2 > w:
b2 = w
if b4 > h:
b4 = h
b = (b1, b2, b3, b4)
bb = convert((w, h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
for image_set in sets:
if not os.path.exists('D:/Yolov5/yolov5/VOCData/labels/'):
os.makedirs('D:/Yolov5/yolov5/VOCData/labels/')
image_ids = open('D:/Yolov5/yolov5/VOCData/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
if not os.path.exists('D:/Yolov5/yolov5/VOCData/dataSet_path/'):
os.makedirs('D:/Yolov5/yolov5/VOCData/dataSet_path/')
list_file = open('dataSet_path/%s.txt' % (image_set), 'w')
# 这行路径不需更改,这是相对路径
for image_id in image_ids:
list_file.write('D:/Yolov5/yolov5/VOCData/images/%s.jpg\n' % (image_id))
convert_annotation(image_id)
list_file.close()
运行后会生成如下 labels 文件夹和 dataSet_path 文件夹。
其中 labels 中为不同图像的标注文件。每个图像对应一个txt文件,文件每一行为一个目标的信息,包括class, x_center, y_center, width, height格式,这种即为 yolo_txt格式
dataSet_path文件夹包含三个数据集的txt文件,train.txt等txt文件为划分后图像所在位置的路径,如train.txt就含有所有训练集图像的路径。
3.配置文件
在 yolov5 目录下的 data 文件夹下 新建一个 myvoc.yaml文件(可以自定义命名),用记事本打开。
内容是:
训练集以及验证集(train.txt和val.txt)的路径(可以改为相对路径)
以及 目标的类别数目和类别名称。
train: D:/Yolov5/yolov5/VOCData/dataSet_path/train.txt
val: D:/Yolov5/yolov5/VOCData/dataSet_path/val.txt
# number of classes
nc: 2
# class names
names: ["light", "post"]
选用yolov5s.yaml
4.开始训练
train.py 更改weights,cfg改成yolov5s.yaml,data改成自己的myvoc.yaml路径 ,batchsize,patience,epoch
python train.py --weights weights/yolov5s.pt --cfg models/yolov5s.yaml --data data/myvoc.yaml --epoch 300 --batch-size 8 --img 640 --patience 20
训练时或者训练后可以利用 tensorboard 查看训练可视化
tensorboard --logdir=runs
与训练一样,进入pytorch环境,进入yolov5文件夹执行。
5.检测
python detect.py --source 0 # webcam 自带摄像头
file.jpg # image 图片
file.mp4 # video 视频
path/ # directory
path/*.jpg # glob
'https://youtu.be/NUsoVlDFqZg' # YouTube
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
python detect.py --weights runs/train/exp/weights/best.pt --source data/video/tram.mp4
遇到的一些问题
2.JPEGImages
3.torch2.0问题