参考文献:
将图片和标签随机到画布中_xukobe97的博客-CSDN博客
YOLOv5训练自己的数据集_xukobe97的博客-CSDN博客_yolov5训练自己数据集
需求背景:
想要识别抽烟和玩手机,但是数据集不够,所以找了一些抽烟玩手机的数据集,但是单独训练效果不理想,所以想着如何把抽烟玩手机的图片粘贴到食堂图片中,再进行训练。
1.找到空白图片
由于我们添加的图片最好不要把食堂图片中原有的人遮挡住,所以我们尽量选择食堂数据集中没有人的图片,即没有标签存在的图片。
对于原始数据集,我们找到其标签大小为0kb的图片,通过如下脚本找到对应的图片,作为我们粘贴的背景;
# -*- coding: utf-8 -*-
import os
import shutil
def file_name(path1,path2,path3):
path1_list = []
path2_list = []
for root, dirs, files in os.walk(path1):
for file in files:
path1_list.append(os.path.splitext(file)[0])
for root, dirs, files in os.walk(path2):
for file in files:
path2_list.append(os.path.splitext(file)[0])
same = list(set(path1_list)&set(path2_list))
for root, dirs, files in os.walk(path2):
for file in files:
if os.path.splitext(file)[0] in same:
shutil.copy(os.path.join(root,file),path3)
path1 = ''
path2 = ''
path3 = ''
file_name(path1,path2,path3)
通过上述脚本找到无人的食堂图片作为背景;
2.进行抽烟玩手机数据集的粘贴
扫描二维码关注公众号,回复:
14715387 查看本文章
采用之间随机画布的脚本,只是之前的标签变成整个粘贴上的图片,本次仍然还是原来的标签位置。所以对xml更改时需要在随机的位置的基础上,再在原来的xmax、xmin、ymax、ymin乘以原图片缩放相同的比例,这样标签的位置才不会变。更改代码如下:
for objects in objectlist:
bndbox = objects.getElementsByTagName('bndbox')
for box in bndbox:
x1_list = box.getElementsByTagName('xmin')
y1_list = box.getElementsByTagName('ymin')
x2_list = box.getElementsByTagName('xmax')
y2_list = box.getElementsByTagName('ymax')
xmin = int(x1_list[0].childNodes[0].data)
ymin = int(y1_list[0].childNodes[0].data)
xmax = int(x2_list[0].childNodes[0].data)
ymax = int(y2_list[0].childNodes[0].data)
"下列四行代码是调整后坐标值与未调整前的坐标关系"
if w1 > 1280 or h1 > 720:
if 1280 < w1 < 1600 or 720 < h1 < 900:
x1min = 0.65 * xmin + list_all[i][0]
y1min = 0.65 * ymin + list_all[i][1]
x1max = 0.65 * xmax + list_all[i][0]
y1max = 0.65 * ymax + list_all[i][1]
elif 1280 < w1 < 2560 or 720 < h1 < 1440:
x1min = 0.25 * xmin + list_all[i][0]
y1min = 0.25 * ymin + list_all[i][1]
x1max = 0.25 * xmax + list_all[i][0]
y1max = 0.25 * ymax + list_all[i][1]
elif 2560 < w1 < 5120 or 1440 < h1 < 2880:
x1min = 0.20 * xmin + list_all[i][0]
y1min = 0.20 * ymin + list_all[i][1]
x1max = 0.20 * xmax + list_all[i][0]
y1max = 0.20 * ymax + list_all[i][1]
else:
x1min = 0.10 * xmin + list_all[i][0]
y1min = 0.10 * ymin + list_all[i][1]
x1max = 0.10 * xmax + list_all[i][0]
y1max = 0.10 * ymax + list_all[i][1]
else:
x1min = xmin + list_all[i][0]
y1min = ymin + list_all[i][1]
x1max = xmax + list_all[i][0]
y1max = ymax + list_all[i][1]
x1_list[0].childNodes[0].data = int(x1min)
y1_list[0].childNodes[0].data = int(y1min)
x2_list[0].childNodes[0].data = int(x1max)
y2_list[0].childNodes[0].data = int(y1max)
通过更改后的随机图片代码即可将抽烟玩手机随机到食堂图片中;
3.将xml转化为txt
通过YOLOv5训练自己的数据集中的voc_label.py和makeTxt.py脚本将xml转化为txt;注意地是,voc_label.py中的classes记得更改,不然产生的txt无标签。
4.更改txt中标签的数字
txt中的标签数字只有一类一般是0,但是我们想要它作为第8类或者第9类,采用下述脚本进行标签数字的更改
import os
import random
import numpy as np
from numpy import *
txtfilepath = "" #原始txt文件所存文件夹,文件夹可以有一个或多个txt文件
savefilepath = "" #更改后txt文件存放的文件夹
total_txt = os.listdir(txtfilepath) # 返回指定的文件夹包含的文件或文件夹的名字的列表
files = os.listdir(savefilepath)
num = len(total_txt)
list = range(num) #创建从0到num的整数列表
for i in list: #遍历每一个文件
name = total_txt[i]
readfile = open(txtfilepath+"/"+name, 'r') #读取文件
fline = readfile.readlines() #读取txt文件中每一行
savetxt = open(savefilepath+"/"+name,'w')
for j in fline:
b ='8'+ j[1:] #替换固定行内容
savetxt.write(b) #写入新的文件中
最后得到粘贴好的图片和更改好的txt,即可得到新的数据集。