pytorch训练自己的数据集记录

参考：https://github.com/tanluren/yolov3-channel-and-layer-pruning

1.测试集的单独制作：

首先从航线获取的视频中截取，放入data\photos_test_or目录下

cmd命令行运行下面命令将图片重命名：

python rename_photo.py

# -*- coding:utf-8 -*-
import os

def test():
	# 源地址和输出地址
	cwd = os.getcwd()
	photo_test_or_Path = cwd + '\\photos_test_or\\'
	photo_test_Path = cwd + '\\photos_test\\'
	if not os.path.exists(photo_test_Path):
		os.mkdir(photo_test_Path)
	filelist = os.listdir(photo_test_or_Path)
	label_i = 1
	for item in filelist:
		if item.endswith('.jpg'):
			src = os.path.join(os.path.abspath(photo_test_or_Path), item)  #原图的地址
			dst = os.path.join(os.path.abspath(photo_test_Path), str(label_i) + '.jpg')   #新图的地址
			try:
				os.rename(src, dst)
				print('converting %s to %s ...' % (src, dst))
				label_i += 1
			except:
				continue

if __name__ == '__main__':
    test()

创建labels_test文件夹（用于存放标签），用labelImg（我的软件在E:\software\win10\exe\labelImg-master\）进行标注，注意将该文件夹下的data\predefined_classes.txt改成你要检测的类别

测试集标注好后cmd命令行运行下面命令将图片和标签放入images和labels文件夹：

python rename_test_photo_label.py

# -*- coding:utf-8 -*-
import os

def test():
	# 源地址和输出地址
	cwd = os.getcwd()
	photo_test_Path = cwd + '\\photos_test\\'
	print(photo_test_Path)
	photos_Path = cwd + '\\images\\'
	print(photos_Path)
	filelist = os.listdir(photo_test_Path)
	photo_i = 1
	for item in filelist:
		if item.endswith('.jpg'):
			src = os.path.join(os.path.abspath(photo_test_Path), item)  #原图的地址
			dst = os.path.join(os.path.abspath(photos_Path), '0' + str(photo_i) + '.jpg')   #新图的地址
			try:
				os.rename(src, dst)
				print('converting %s to %s ...' % (src, dst))
				photo_i += 1
			except:
				continue
				
	label_test_Path = cwd + '\\labels_test\\'
	print(label_test_Path)
	labels_Path = cwd + '\\labels\\'
	print(labels_Path)
	filelist = os.listdir(label_test_Path)
	label_i = 1
	for item in filelist:
		if item.endswith('.txt'):
			src = os.path.join(os.path.abspath(label_test_Path), item)  #原图的地址
			dst = os.path.join(os.path.abspath(labels_Path), '0' + str(label_i) + '.txt')   #新图的地址
			try:
				os.rename(src, dst)
				print('converting %s to %s ...' % (src, dst))
				label_i += 1
			except:
				continue

if __name__ == '__main__':
    test()

2.训练集和测试集生成脚本的修改

python get_train_test_data.py

#-*- coding:utf-8 -*-
import os
import random

train_list = range(1,235) #训练集(234张图片)
#print(train_list)
test_list = range(1,61) #测试集(60张图片)
#print(test_list)

train = random.sample(train_list, 234)
#print(train)
test = random.sample(test_list, 60) 
#print(test) 

ftrain = open('ImageSets\\train.txt', 'w')
ftest = open('ImageSets\\test.txt', 'w') #只生成了训练集和测试集

for train_i in train:
    train_name = 'data\\images\\' + str(train_i) + '.jpg' + '\n'
    print(train_name)
    ftrain.write(train_name)

for test_i in test:
    test_name = 'data\\images\\' + '0' + str(test_i) + '.jpg' + '\n'
    print(test_name)
    ftest.write(test_name)
 
ftrain.close()
ftest.close()

生成后的train.txt和test.txt文件如下：

3.用yolov3-tiny进行训练

python train.py --cfg cfg/yolov3-tiny_landmark.cfg --data data/landmark.data --weights weights/yolov3-tiny.weights --epochs 400 --batch-size 32

Namespace(accumulate=2, adam=False, arc='defaultpw', batch_size=32, bucket='', cache_images=False, cfg='cfg/yolov3-tiny_landmark.cfg', data='data/landmark.data', device='', epochs=400, evolve=False, img_size=416, img_weights=False, multi_scale=False, name='', nosave=False, notest=False, prebias=False, prune=1, rect=False, resume=False, s=0.001, sr=False, t_cfg='', t_weights='', transfer=False, var=None, weights='weights/yolov3-tiny.weights')

4.训练完后做视频检测（检测的结果会保存到output，目录下）

python detect.py --cfg cfg/yolov3-tiny_landmark.