yolov5已经普及常用，最常规的步骤就是获取数据（图片）、标注图片（标注xml文件）、对标注的数据增强（椒盐模糊、高斯模糊、翻转等） 、训练数据。
在这里，假如客户新增一个目标，需要客户自己标注数据、增强数据和训练数据不合理，我们可以希望实现这样一个场景：建立一个界面，打开界面后：1.客户对新的产品从不同的角度拍七八张照片；2.客户在界面上对七八张图片的目标用长方形画框；3.界面系统自动对这几张标注好的图片进行数据扩增；4.将扩增后的数据自动移入训练的文件夹；5.训练模型。
其实，希望客户做以上的1、2、5三步就行，3和4由界面系统自动完成，这样就能很简单的完成新模型的标注和训练。下面分别是标注和数据扩增的代码。
1.图片标注
代码如下：
import cv2
import numpy as np
#1. 在图片中标注矩形框
def draw_circle(event, x, y, flags, param):
global ix, iy, drawing, px, py
if event == cv2.EVENT_LBUTTONDOWN:
drawing = True
ix, iy = x, y
elif event == cv2.EVENT_MOUSEMOVE:
if drawing == True:
# cv2.rectangle(img, (ix, iy), (px, py), (0, 0, 0), thickness=2) # 将刚刚拖拽的矩形涂黑
# cv2.rectangle(img, (ix, iy), (x, y), (0, 255, 0), 0)
px, py = x, y
elif event == cv2.EVENT_LBUTTONUP:
drawing = False
cv2.rectangle(img, (ix, iy), (x, y), (0, 255, 0), 0)
preds.append([classes, [iy, ix, y, x]])
# px, py = -1, -1
# 2. 将矩形框的坐标和名称写入xml文件
# 2.写入xml文件
def img_xml(img_path, xml_path, img_name, pred):
if len(pred) != 0:
# 1.读取图片(xml需要写入图片的长宽高)
img = cv2.imread(img_path)
# 2.写入xml文件
# (1)写入文件头部
files_path = img_path.split("\\")[-2]
print("..:", files_path)
xml_file = open((xml_path + img_name + '.xml'), 'w')
xml_file.write('<annotation>\n')
xml_file.write(' <folder>' + files_path + '</folder>\n')
xml_file.write(' <filename>' + img_name + '.jpg' + '</filename>\n')
xml_file.write(' <path>' + img_path + '</path>\n')
xml_file.write(' <source>\n')
xml_file.write(' <database>Unknown</database>\n')
xml_file.write(' </source>\n')
# （2）写入图片的长宽高信息
xml_file.write(' <size>\n')
xml_file.write(' <width>' + str(img.shape[1]) + '</width>\n')
xml_file.write(' <height>' + str(img.shape[0]) + '</height>\n')
xml_file.write(' <depth>' + str(img.shape[2]) + '</depth>\n')
xml_file.write(' </size>\n')
xml_file.write(' <segmented>0</segmented>\n')
# 3.写入字符串信息:[["water",[15,20,30,40]],["red",[12,13,14,15]]]
for item in pred:
xml_file.write(' <object>\n')
xml_file.write(' <name>' + str(item[0]) + '</name>\n')
xml_file.write(' <pose>Unspecified</pose>\n')
xml_file.write(' <truncated>0</truncated>\n')
xml_file.write(' <difficult>0</difficult>\n')
xml_file.write(' <bndbox>\n')
# 写入字符串信息
# [top, left, bottom, right]
xml_file.write(' <xmin>' + str(item[1][1]) + '</xmin>\n')
xml_file.write(' <ymin>' + str(item[1][0]) + '</ymin>\n')
xml_file.write(' <xmax>' + str(item[1][3]) + '</xmax>\n')
xml_file.write(' <ymax>' + str(item[1][2]) + '</ymax>\n')
xml_file.write(' </bndbox>\n')
xml_file.write(' </object>\n')
xml_file.write('</annotation>\n')
if __name__ == "__main__":
# 1.读取图片
img_path = "D:\AI\yq\datas\imgs\\4.jpg"
img = cv2.imread(img_path)
xml_path = "D:\AI\yq\datas\labels\\"
img_name = "4"
# pred = [["water", [15, 20, 30, 40]], ["red", [12, 13, 14, 15]]] #[top,left,bottom,right]
classes = "luoding"
preds = []
#2. 鼠标给图片画矩形，并将矩形的左上角和右下角的坐标记录下来
drawing = False # 鼠标按下为真
mode = True # 如果为真，画矩形，按m切换为曲线
ix, iy = -1, -1
px, py = -1, -1
cv2.namedWindow('image')
cv2.setMouseCallback('image', draw_circle)
while (1):
cv2.imshow('image', img)
k = cv2.waitKey(1) & 0xFF
if k == ord('q'):
print(preds)
img_xml(img_path, xml_path, img_name, preds)
break
elif k == 27:
break
cv2.destroyAllWindows()
以上是标注图片，并生成xml文件的代码，结果如下：
在以上代码中，有2个函数，函数draw_circle(event, x, y, flags, param)是在图片中用鼠标画矩形框，并将矩形框的左上角和右下角的坐标记录下来（preds用来记录每个目标的坐标和目标名称），preds的打印结果如下：
这个数据加上图片的路径和名称就可以写出图片目标的xml文件了。
函数img_xml(img_path, xml_path, img_name, pred)是将用鼠标标注的图片信息写入xml文件，函数中的参数为
img_path：标注图片的路径
xml_path：保存xml文件的路径
img_name:保存图片的名称（不带.后缀，如1.jpg的图片名称为1）
pred：图片中目标的信息，是函数draw_circle(event, x, y, flags, param)的返回值
输出结果如下：
2.数据扩增
常见的数据扩增的方式有：原数据成倍增加、椒盐模糊、高斯模糊、水平翻转、垂直翻转、按照某个角度翻转、平移、缩放、裁剪、亮度、直方图均衡化、调整白平衡、mixup、cutmix、cutout等。这些数据增强的方式，有的只需要处理图片，而图片中的xml文件中的坐标值不变，有的则会改变目标在图片中的具体位置，这时需要修改xml文件中的坐标值。下面分别介绍不同的扩增：
2.1图片的水平翻转
图片水平翻转后，其目标的坐标会发生变换。其xml文件需要修改扩增后的图片路径、图片名称、图片中目标的名称、图片中目标的坐标，其他信息可以根据自己的需要对该代码稍微修改即可。代码如下：
import cv2
import os
import random
import xml.etree.ElementTree as ET
from shutil import copyfile
# 2.写入xml文件
def img_xml(img_path, xml_path, img_name, pred):
if len(pred) != 0:
# 1.读取图片(xml需要写入图片的长宽高)
img = cv2.imread(img_path)
# 2.写入xml文件
# (1)写入文件头部
files_path = img_path.split("\\")[-2]
print("..:", files_path)
xml_file = open((xml_path + img_name + '.xml'), 'w')
xml_file.write('<annotation>\n')
xml_file.write(' <folder>' + files_path + '</folder>\n')
xml_file.write(' <filename>' + img_name + '.jpg' + '</filename>\n')
xml_file.write(' <path>' + img_path + '</path>\n')
xml_file.write(' <source>\n')
xml_file.write(' <database>Unknown</database>\n')
xml_file.write(' </source>\n')
# （2）写入图片的长宽高信息
xml_file.write(' <size>\n')
xml_file.write(' <width>' + str(img.shape[1]) + '</width>\n')
xml_file.write(' <height>' + str(img.shape[0]) + '</height>\n')
xml_file.write(' <depth>' + str(img.shape[2]) + '</depth>\n')
xml_file.write(' </size>\n')
xml_file.write(' <segmented>0</segmented>\n')
# 3.写入字符串信息:[["water",[15,20,30,40]],["red",[12,13,14,15]]]
for item in pred:
xml_file.write(' <object>\n')
xml_file.write(' <name>' + str(item[0]) + '</name>\n')
xml_file.write(' <pose>Unspecified</pose>\n')
xml_file.write(' <truncated>0</truncated>\n')
xml_file.write(' <difficult>0</difficult>\n')
xml_file.write(' <bndbox>\n')
# 写入字符串信息
# [top, left, bottom, right]
xml_file.write(' <xmin>' + str(item[1][1]) + '</xmin>\n')
xml_file.write(' <ymin>' + str(item[1][0]) + '</ymin>\n')
xml_file.write(' <xmax>' + str(item[1][3]) + '</xmax>\n')
xml_file.write(' <ymax>' + str(item[1][2]) + '</ymax>\n')
xml_file.write(' </bndbox>\n')
xml_file.write(' </object>\n')
xml_file.write('</annotation>\n')
# 3. 图片的水平翻转
def img_horizen_rotation(img, xmls,img3_path, xmls_path):
print("............... 图片翻转 ..................")
# 1. 图片翻转
imgs = cv2.flip(img, 1)
# 2. xml 文件的修改
tree = ET.parse(xml_path)
root = tree.getroot()
ss = []
modifys = []
width = 0
for element in root.iter():
#(1 按照[y_min, x_min, y_max, x_max]保存数据，用来后面处理数据
# print(element.tag, element.text)
if element.tag == "width":
width = int(element.text)
# print("width:", width)
if element.tag == "name":
name = element.text
ss.append(name)
if element.tag == "xmin":
xmin = int(element.text)
# print("xmin:", xmin)
ss.append(xmin)
if element.tag == "ymin":
ymin = int(element.text)
# print("ymin:", ymin)
ss.append(ymin)
if element.tag == "xmax":
xmax = int(element.text)
# print("xmax:", xmax)
ss.append(xmax)
if element.tag == "ymax":
ymax = int(element.text)
# print("ymax:", ymax)
ss.append(ymax)
if len(ss) == 5:
# print(ss)
x_min = width - ss[3]
y_min = ss[2]
x_max = width - ss[1]
y_max = ss[4]
modifys.append([ss[0], [y_min, x_min, y_max, x_max]])
ss = []
img_name = img3_path.split("\\")[-1].split(".")[0]
# （2 将翻转后的图片和xml文件保存在特点路径
img_xml(img_path, xmls_path, img_name, modifys)
cv2.imwrite(img3_path, imgs)
# 3. 画图
for item in modifys:
a = (item[1][1], item[1][0])
b = (item[1][3], item[1][2])
cv2.rectangle(imgs, a, b, (0, 255, 0), 2)
cv2.imshow("s2", imgs)
cv2.waitKey(0)
return img, xmls
if __name__ == "__main__":
# 要处理的图片和对应的xml文件路径
img_path = "D:\AI\yq\datas\strange\imgs\\1.jpg"
xml_path = "D:\AI\yq\datas\strange\labels\\1.xml"
# 保存处理后的图片和xml文件路径
img_saves = "D:\AI\yq\datas\strange\imgs_strange\\"
xml_saves = "D:\AI\yq\datas\strange\labels_strange\\"
# 水平翻转后的图片路径
img3_path = "D:\AI\yq\datas\imgs\\5.jpg"
# 3. 图片的水平翻转
img3, xml3 = img_horizen_rotation(img, xml_path, img3_path, xml_saves)
以上代码中进行了原图片中目标坐标的转换，并写入新的xml文件。结果如下：
后面一张是前面一张翻转后的图片，画的框是前面图片标注的框翻转后画在后面图片的目标上，可以看到转换后的坐标框是和目标物体吻合的。
2.2 椒盐算法
代码如下：
# 1.椒盐算法增强
# (1.img:要处理的图片；
# (2.xmls：要处理的图片对应的xml文件；
# (3. percentage： 椒盐算法加噪点的比率
def pepper_and_salt(img, xmls, img3_path, xml_saves, percentage):
# 1. 椒盐算法
num=int(percentage*img.shape[0]*img.shape[1])# 椒盐噪声点数量
random.randint(0, img.shape[0])
img2=img.copy()
for i in range(num):
X=random.randint(0,img2.shape[0]-1)#从0到图像长度之间的一个随机整数,因为是闭区间所以-1
Y=random.randint(0,img2.shape[1]-1)
if random.randint(0,1) ==0: #黑白色概率55开
img2[X,Y] = (255,255,255) # 白色
else:
img2[X,Y] = (0, 0, 0) # 黑色
# 保存图片
cv2.imwrite(img3_path, img2)
imgs = cv2.imread(img3_path)
# 2. xml 文件的修改
tree = ET.parse(xml_path)
root = tree.getroot()
ss = []
modifys = []
width = 0
for element in root.iter():
# (1 按照[y_min, x_min, y_max, x_max]保存数据，用来后面处理数据
# print(element.tag, element.text)
if element.tag == "width":
width = int(element.text)
# print("width:", width)
if element.tag == "name":
name = element.text
ss.append(name)
if element.tag == "xmin":
xmin = int(element.text)
# print("xmin:", xmin)
ss.append(xmin)
if element.tag == "ymin":
ymin = int(element.text)
# print("ymin:", ymin)
ss.append(ymin)
if element.tag == "xmax":
xmax = int(element.text)
# print("xmax:", xmax)
ss.append(xmax)
if element.tag == "ymax":
ymax = int(element.text)
# print("ymax:", ymax)
ss.append(ymax)
if len(ss) == 5:
# print(ss)
x_min = ss[1]
y_min = ss[2]
x_max = ss[3]
y_max = ss[4]
modifys.append([ss[0], [y_min, x_min, y_max, x_max]])
ss = []
img_name = img3_path.split("\\")[-1].split(".")[0]
# （2 将翻转后的图片和xml文件保存在特点路径
img_xml(img_path, xml_saves, img_name, modifys)
# 3. 画图
for item in modifys:
a = (item[1][1], item[1][0])
b = (item[1][3], item[1][2])
cv2.rectangle(imgs, a, b, (0, 255, 0), 2)
cv2.imshow("s1", img)
cv2.imshow("s2", imgs)
cv2.waitKey(0)
return img2, xmls这里只写了算法的函数，由于椒盐算法没有改变目标物在图片中的位置，所以不需要修改xml文件或者只需要修改新的图片路径即可。输出效果如下：
前面的图片是原图，后面的图片是对原图做了椒盐变换，并将标注的坐标画成了框，显然是吻合的。