YOLOv5-模型训练与检测 - YOLOv5 - 人工智能 | Aayu Yain = 学无止境 = 世界上大部分事，都没太大意义。真理与热爱除外

# 下载 OIDv4_ToolKit 数据集

	git clone https://github.com/EscVM/OIDv4_ToolKit.git
	# 安装依赖
	pip install -r .\requirements.txt

python .\main.py downloader --classes Apple Orange --type_csv all

编写 python 脚本如下：

	from genericpath import isfile
	import os
	import cv2 as cv

	root_dir = "E:\\code\\github\\OIDv4_ToolKit\\OID\Dataset"
	valid_label_dir = "E:\\code\\python\\openvino_utils\\data\\labels\\valid"

	def read_annotation(data_type = "train", class_type = "Apple"):
	current_dir = os.path.join(root_dir, data_type, class_type)
	print(current_dir)
	files = os.listdir(current_dir)
	for f in files:
	if (os.path.isfile(os.path.join(current_dir, f))):
	image = cv.imread(os.path.join(current_dir, f))
	label_file = os.path.join(current_dir, "label", f.replace(".jpg", ".txt"))
	yolo_label = f.replace(".jpg", ".txt")
	data_label_text_f = os.path.join(valid_label_dir, yolo_label)
	file_write_obj = open(data_label_text_f, 'w')

	with open(label_file) as f:
	boxes = [line.strip() for line in f.readlines()]

	class_index = -1
	for box in boxes:
	anno_info = box.split(" ")
	print("class name: ", anno_info[0])
	x1 = float(anno_info[1])
	y1 = float(anno_info[2])
	x2 = float(anno_info[3])
	y2 = float(anno_info[4])

	if anno_info[0] == "Apple":
	class_index = 0
	else:
	class_index = 1

	h, w, c = image.shape
	cx = (x1 + x2) / (2 * w)
	cy = (y1 + y2) / (2 * h)
	sw = (x2 - x1) / w
	sh = (y2 - y1) / h
	file_write_obj.write("%d %f %f %f %f\n" % (class_index, cx, cy, sw, sh))

	file_write_obj.close()

	if __name__ == "__main__":
	read_annotation(data_type="validation", class_type="Orange")

转换完毕后，将标签和图片数据放在 yolov5 目录下的 fruit_training 下的 data 目录下，没有目录就创建

参考 yolov5\data\VOC.yaml 文件，在 yolov5\fruit_training\ 目录下创建 dataset.yaml 文件，编写内容如下：

	train: fruit_training/data/images/train/
	val: fruit_training/data/images/valid

	nc: 2

	names: ['Apple', 'Orange']

参考 yolov5\models\yolov5s.yaml 文件，拷贝到 yolov5\fruit_training\ 目录下，修改 nc 为 2。文件内如下：

	# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

	# Parameters
	nc: 2 # number of classes
	depth_multiple: 0.33 # model depth multiple
	width_multiple: 0.50 # layer channel multiple
	anchors:
	- [10,13, 16,30, 33,23] # P3/8
	- [30,61, 62,45, 59,119] # P4/16
	- [116,90, 156,198, 373,326] # P5/32

	# YOLOv5 backbone
	backbone:
	# [from, number, module, args]
	[[-1, 1, Focus, [64, 3]], # 0-P1/2
	[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
	[-1, 3, C3, [128]],
	[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
	[-1, 9, C3, [256]],
	[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
	[-1, 9, C3, [512]],
	[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
	[-1, 1, SPP, [1024, [5, 9, 13]]],
	[-1, 3, C3, [1024, False]], # 9
	]

	# YOLOv5 head
	head:
	[[-1, 1, Conv, [512, 1, 1]],
	[-1, 1, nn.Upsample, [None, 2, 'nearest']],
	[[-1, 6], 1, Concat, [1]], # cat backbone P4
	[-1, 3, C3, [512, False]], # 13

	[-1, 1, Conv, [256, 1, 1]],
	[-1, 1, nn.Upsample, [None, 2, 'nearest']],
	[[-1, 4], 1, Concat, [1]], # cat backbone P3
	[-1, 3, C3, [256, False]], # 17 (P3/8-small)

	[-1, 1, Conv, [256, 3, 2]],
	[[-1, 14], 1, Concat, [1]], # cat head P4
	[-1, 3, C3, [512, False]], # 20 (P4/16-medium)

	[-1, 1, Conv, [512, 3, 2]],
	[[-1, 10], 1, Concat, [1]], # cat head P5
	[-1, 3, C3, [1024, False]], # 23 (P5/32-large)

	[[17, 20, 23], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)
	]

运行命令： python .\train.py --data fruit_training/dataset.yaml --cfg fruit_training/yolov5s.yaml --weights yolov5s.pt --batch-size 1 --epochs 2

运行命令： tensorboard --logdir=E:\code\python\yolov5\runs\train\exp

下载几个苹果和橘子的图片放到 fruit_training 目录下，运行检测命令：

python .\detect.py --source .\fruit_training\apple.jpeg --weights .\fruit_training\best.pt --conf 0.25

哦对，还要把训练好的 best.pt 模型也放到 fruit_training 目录下

看 requirements.txt 文件里是否有安装 coremltools 和 onnx 这两个第三方库，没有的话就安装：

	pip install coremltools>=4.1 -i https://pypi.tuna.tsinghua.edu.cn/simple
	pip install onnx>=1.9.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成后，运行下面命令进行模型转换：

python .\export.py --weights .\fruit_training\best.pt --img 640 --batch 1