Yolov5 + Deepsort 重新训练自己的数据(保姆级超详细)
从下面github库中拿代码:
https://github.com/mikel-brostrom/Yolov5_DeepSort_Pytorchhttps://github.com/mikel-brostrom/Yolov5_DeepSort_PytorchGitHub – Sharpiless/Yolov5-Deepsort: 最新版本yolov5+deepsort目标检测和追踪,能够显示目标类别,支持5.0版本可训练自己数据集最新版本yolov5+deepsort目标检测和追踪,能够显示目标类别,支持5.0版本可训练自己数据集 – GitHub – Sharpiless/Yolov5-Deepsort: 最新版本yolov5+deepsort目标检测和追踪,能够显示目标类别,支持5.0版本可训练自己数据集https://github.com/Sharpiless/Yolov5-Deepsort
下载好匹配的deeosort和yolov5代码很重要,题主折腾了一天,坑在版本上了!!
题主用的deeosort v3.0和yolov5 5.0版本,master似乎还不完善,没跑通,要是跑通了的读者希望可以交流一下。
直接进入正题:
一.目标追踪整体代码
分别主体是yolov5和deep_sort。
二.训练自己的数据集
yolov5和deep_sort分开训练。首先训练yolov5,这个不难,超链接如下。
三.训练deep_sort
准备deep_sort的数据集,和yolov5不一样,这是一个分类的数据集。
我们用代码把图像中的检测目标扣出来,作为我们的数据集。
代码如下:
import cv2
import xml.etree.ElementTree as ET
import numpy as np
import xml.dom.minidom
import os
import argparse
def main():
# JPG文件的地址
img_path = '/home/zqy/Desktop/yolov5-master/nxm_data/images_all/'
# XML文件的地址
anno_path = '/home/zqy/Desktop/yolov5-master/nxm_data/labels_xml/'
# 存结果的文件夹
cut_path = '/home/zqy/Desktop/yolov5-master/nxm_data/crops/'
if not os.path.exists(cut_path):
os.makedirs(cut_path)
# 获取文件夹中的文件
imagelist = os.listdir(img_path)
# print(imagelist
for image in imagelist:
image_pre, ext = os.path.splitext(image)
img_file = img_path + image
img = cv2.imread(img_file)
xml_file = anno_path + image_pre + '.xml'
# DOMTree = xml.dom.minidom.parse(xml_file)
# collection = DOMTree.documentElement
# objects = collection.getElementsByTagName("object")
tree = ET.parse(xml_file)
root = tree.getroot()
# if root.find('object') == None:
# return
obj_i = 0
for obj in root.iter('object'):
obj_i += 1
print(obj_i)
cls = obj.find('name').text
xmlbox = obj.find('bndbox')
b = [int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)),
int(float(xmlbox.find('xmax').text)),
int(float(xmlbox.find('ymax').text))]
img_cut = img[b[1]:b[3], b[0]:b[2], :]
path = os.path.join(cut_path, cls)
# 目录是否存在,不存在则创建
mkdirlambda = lambda x: os.makedirs(x) if not os.path.exists(x) else True
mkdirlambda(path)
try:
cv2.imwrite(os.path.join(cut_path, cls, '{}_{:0>2d}.jpg'.format(image_pre, obj_i)), img_cut)
except:
continue
print("&&&&")
if __name__ == '__main__':
main()
注意:这里数据集可能会存在负样本,导致img_cut为空,我在这里修改了代码,加了try判断,只算入了正样本。
上述代码在自己的数据集上生成了crops文件夹,目录如下:
接着要把这些数据分为训练集和验证集,跟类别有关系,注意类别和目标是两个概念。
我在这检测的目标只有一个,但是可以有不同的类别,在这里的类别有112个,因此训练集和测试集下边的类别就应该有112个。可以自己整理,也可以用代码分,代码如下:
import os
from PIL import Image
from shutil import copyfile, copytree, rmtree, move
PATH_DATASET = '/home/zqy/Desktop/yolov5-master/nxm_data/crops' # 需要处理的文件夹
PATH_NEW_DATASET = '/home/zqy/Desktop/yolov5-master/nxm_data/stitches' # 处理后的文件夹
PATH_ALL_IMAGES = PATH_NEW_DATASET + '/all_images'
PATH_TRAIN = PATH_NEW_DATASET + '/train'
PATH_TEST = PATH_NEW_DATASET + '/test'
# 定义创建目录函数
def mymkdir(path):
path = path.strip() # 去除首位空格
path = path.rstrip("\\") # 去除尾部 \ 符号
isExists = os.path.exists(path) # 判断路径是否存在
if not isExists:
os.makedirs(path) # 如果不存在则创建目录
print(path + ' 创建成功')
return True
else:
# 如果目录存在则不创建,并提示目录已存在
print(path + ' 目录已存在')
return False
class BatchRename():
'''
批量重命名文件夹中的图片文件
'''
def __init__(self):
self.path = PATH_DATASET # 表示需要命名处理的文件夹
# 修改图像尺寸
def resize(self):
for aroot, dirs, files in os.walk(self.path):
# aroot是self.path目录下的所有子目录(含self.path),dir是self.path下所有的文件夹的列表.
filelist = files # 注意此处仅是该路径下的其中一个列表
# print('list', list)
# filelist = os.listdir(self.path) #获取文件路径
total_num = len(filelist) # 获取文件长度(个数)
for item in filelist:
if item.endswith('.jpg'): # 初始的图片的格式为jpg格式的(或者源文件是png格式及其他格式,后面的转换格式就可以调整为自己需要的格式即可)
src = os.path.join(os.path.abspath(aroot), item)
# 修改图片尺寸到128宽*256高
im = Image.open(src)
out = im.resize((128, 256), Image.ANTIALIAS) # resize image with high-quality
out.save(src) # 原路径保存
def rename(self):
for aroot, dirs, files in os.walk(self.path):
# aroot是self.path目录下的所有子目录(含self.path),dir是self.path下所有的文件夹的列表.
filelist = files # 注意此处仅是该路径下的其中一个列表
# print('list', list)
# filelist = os.listdir(self.path) #获取文件路径
total_num = len(filelist) # 获取文件长度(个数)
i = 1 # 表示文件的命名是从1开始的
for item in filelist:
if item.endswith('.jpg'): # 初始的图片的格式为jpg格式的(或者源文件是png格式及其他格式,后面的转换格式就可以调整为自己需要的格式即可)
src = os.path.join(os.path.abspath(aroot), item)
# 根据图片名创建图片目录
dirname = str(item.split('_')[0])
# 为相同车辆创建目录
# new_dir = os.path.join(self.path, '..', 'bbox_all', dirname)
new_dir = os.path.join(PATH_ALL_IMAGES, dirname)
if not os.path.isdir(new_dir):
mymkdir(new_dir)
# 获得new_dir中的图片数
num_pic = len(os.listdir(new_dir))
dst = os.path.join(os.path.abspath(new_dir),
dirname + 'C1T0001F' + str(num_pic + 1) + '.jpg')
# 处理后的格式也为jpg格式的,当然这里可以改成png格式 C1T0001F见mars.py filenames 相机ID,跟踪指数
# dst = os.path.join(os.path.abspath(self.path), '0000' + format(str(i), '0>3s') + '.jpg') 这种情况下的命名格式为0000000.jpg形式,可以自主定义想要的格式
try:
copyfile(src, dst) # os.rename(src, dst)
print('converting %s to %s ...' % (src, dst))
i = i + 1
except:
continue
print('total %d to rename & converted %d jpgs' % (total_num, i))
def split(self):
# ---------------------------------------
# train_test
images_path = PATH_ALL_IMAGES
train_save_path = PATH_TRAIN
test_save_path = PATH_TEST
if not os.path.isdir(train_save_path):
os.mkdir(train_save_path)
os.mkdir(test_save_path)
for _, dirs, _ in os.walk(images_path, topdown=True):
for i, dir in enumerate(dirs):
for root, _, files in os.walk(images_path + '/' + dir, topdown=True):
for j, file in enumerate(files):
if (j == 0): # test dataset;每个车辆的第一幅图片
print("序号:%s 文件夹: %s 图片:%s 归为测试集" % (i + 1, root, file))
src_path = root + '/' + file
dst_dir = test_save_path + '/' + dir
if not os.path.isdir(dst_dir):
os.mkdir(dst_dir)
dst_path = dst_dir + '/' + file
move(src_path, dst_path)
else:
src_path = root + '/' + file
dst_dir = train_save_path + '/' + dir
if not os.path.isdir(dst_dir):
os.mkdir(dst_dir)
dst_path = dst_dir + '/' + file
move(src_path, dst_path)
rmtree(PATH_ALL_IMAGES)
if __name__ == '__main__':
demo = BatchRename()
demo.resize()
demo.rename()
demo.split()
分好后train和test下各有112个文件夹,代表着112个类别。
将train和test移动到deep_sort/deep目录下。
修改train.py中train dataset的预处理如下:
transform_train = torchvision.transforms.Compose([
torchvision.transforms.Resize((128, 64)),
torchvision.transforms.RandomCrop((128, 64), padding=4),
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
[0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
接着修改147行,以免权重保存覆盖原始权重:
torch.save(checkpoint, './checkpoint/ckpt1.t7')
接着在model.py中修改类别,这类是112个类别:
class Net(nn.Module):
def __init__(self, num_classes= 112 ,reid=False):
super(Net,self).__init__()
# 3 128 64
self.conv = nn.Sequential(
nn.Conv2d(3,64,3,stride=1,padding=1),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
# nn.Conv2d(32,32,3,stride=1,padding=1),
# nn.BatchNorm2d(32),
# nn.ReLU(inplace=True),
nn.MaxPool2d(3,2,padding=1),
)
然后在deep_sort/deep目录下打开终端,运行:
python train.py --data-dir data/
得到结果如下:
权重结果保存在deep/checkpoint中。
四.测试结果
python track.py --yolo_weights 你的权重 --source 你的视频 --deep_sort_weights 你的权重 --device 0 --save-vid
–save-vid要调用,否则不会保存结果。
结束!
ps:如果报错
File "/home/zqy/Desktop/Yolov5_DeepSort_Pytorch-3.0/deep_sort_pytorch/deep_sort/deep/feature_extractor.py", line 37, in _resize
return cv2.resize(im.astype(np.float32)/255., size)
cv2.error: OpenCV(4.5.5) /io/opencv/modules/imgproc/src/resize.cpp:4052: error: (-215:Assertion failed) !ssize.empty() in function 'resize'
原因,track.py里的iou和nms调太低了,如果仍然报错,建议加个try跳过这些空的im。
来源:武大人民泌外I科人工智能团队