2022年6月23日更新:
- 完整源码部分,修改一些以前出现的bug,增加中文路径的读取和保存。在多个检测项中使用了,对于YOLOV4来说其实调整对比度在cfg文件中就有了,所以这里的gamma变换一般不用了,另外对于真实场景来说水平或者数值镜像翻转后的状态可能是不会出现的,这是一种寿星对称,中心旋转是有用的,这是观察视角转了180度。
2021年8月2日更新:
- 完整源码部分,重命名部分填充位数保证一定可以重命名。
2020年4月15日更新:
原代码中图片都是jpg的,更新了可支持.png,.bmp,.jpg不同格式混合在一起的情况。
完整demo上传百度盘了,有兴趣可以看一下。
链接: https://pan.baidu.com/s/1mgqUsJmn6ynFHqKVAmYPSA 提取码: xgnt
简介:本篇文章把darknet训练YOLO模型的过程中需要用到的一些预处理函数打包封装成一个类,方便应用,图片标注工具是Labelme。封装的类中包含labelme标注的标签文件转为darknetYOLO所需的txt格式;数据增强(包括水平、竖直翻转,平滑滤波,gamma变换),在数据增强的时候也会把对应的txt文件做相同的处理;从json标注文件复查标签;从txt文件复查标签。
造轮子。
由于图片可能有尺度缩放所以YOLO里边的位置用的是图片的百分比坐标,如下
labelme标注的文件是json格式,保存的坐标是绝对坐标,需要转换为2.中的txt标签格式。
参考:https://www.jianshu.com/p/61436306b479
4.1假设已经标注好了图片,且图片和标签文件保存如下:原始图片保存在srcImg,标注的json文件保存在jsonLabel。
4.2 把json标注格式转为txt格式,并加一些数据增强操作的完整流程,类代码有点长附在文末。
if __name__ == "__main__":
srcImageFolder = "K:\imageData\polarity\image1"# 存放原始图片的文件夹,有标注过的和没标注过的图片
jsonFolder = "K:\imageData\polarity\jsonLabel1"# labelme标注的json文件保存的文件夹
labeledImageFolder = "K:/imageData/polarity/temp/labeledImage"# 筛选出所有标注过的图片
txtFolder = "K:/imageData/polarity/temp/txtLabel"# txt文件保存的路径
objFolder = "K:/imageData/polarity/temp/obj"# 把图片和对应的txt文件放在一起,可以直接拿去训练的那种
trainAndTestFileSaveFolder = "K:/imageData/polarity/temp"
# 创建文件夹
try:
if not os.path.exists(labeledImageFolder):
os.makedirs(labeledImageFolder)
if not os.path.exists(txtFolder):
os.makedirs(txtFolder)
if not os.path.exists(objFolder):
os.makedirs(objFolder)
except:
print("os.makedirs failed!")
# 初始化类
tool = darknetTool()
# 获取所有原始标注文件
jsonPathList = tool.getAllJsonFilePath(jsonFolder)
# 从json标签文件复查标注
# for item in jsonPathList:
# tool.showLabelFromJson(item,imageFolder)
# 从json标注文件查找图片,并保存到labeledImageFolder
srcLabeledImagePathList = tool.fromJsonFileToSearchImage(jsonPathList, srcImageFolder)
tool.copyAllLabeledImageToNewFolder(srcLabeledImagePathList, labeledImageFolder)
# json格式转为txt格式
for item in jsonPathList:
tool.changeLabelFromJsonToTxt(item, txtFolder)
txtPathList = tool.getAllTxtPath(txtFolder)
# 获取在labeledImageFolder中的图片的路径
dstLabeledImagePathList = tool.getAllImagePath(labeledImageFolder)
# 从txt标签文件查看标注结果
# 所有标注的图片路径都在dstLabeledImagePathList,通过图片路径去找对应的txt文件
# 所有的txt文件都在txtFolder文件夹下
# for imagePath in dstLabeledImagePathList:
# tool.showLabelFromTxt(imagePath,txtFolder)
# 数据增强
# 均值滤波
for imagePath in dstLabeledImagePathList:
tool.imageAugment_smooth(imagePath, txtFolder, labeledImageFolder,n=3) # n=9是滤波核大小
# 水平翻转
for imagePath in dstLabeledImagePathList:
tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=0) # flag=0 水平翻转
# 竖直翻转
for imagePath in dstLabeledImagePathList:
tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=1) # flag=1 竖直翻转
# 中心对称
for imagePath in dstLabeledImagePathList:
tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=2) # flag=2 中心对称
# gamma变换
for imagePath in dstLabeledImagePathList:
tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=0.5)
tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=0.8)
tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=1.5)
#tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=2)
#图片重命名
# 图片和标签重命名为00001.jpg 00001.txt.......(可选)
# n=5表示zfill()的位数,n=5则为00001.jpg...,n=3则为001.jpg....
tool.renameAllLabeledImageAndTxt(txtFolder, labeledImageFolder, n=6, startN=0)
#检查转换的最终结果
dstLabeledImagePathList = tool.getAllImagePath(labeledImageFolder)
for imagePath in dstLabeledImagePathList:
tool.showLabelFromTxt(imagePath,txtFolder)
#划分训练集和测试集
imagePathList = tool.getAllImagePath(labeledImageFolder)
# prePath = "data/obj/"则train.txt文件中的路径就会是data/obj/00001.jpg,data/obj/00002.jpg......
# prop=0.8表示80%的数据用来训练,其余的做验证集
tool.devideTrainSetAndTestSet(imagePathList, trainAndTestFileSaveFolder, prePath="data/obj/", prop=0.8)
4.4复查结果
#从txt文件复查数据增强后的结果
txtPathList_2 = tool.getAllTxtPath(txtFolder)
for item in txtPathList_2:
tool.showLabelFromTxt(item,labeledImageFolder)
5.总结
这个类就是一些函数的封装,单独拿去用也是没有什么问题的。
6.代码
2021年8月2日更新重命名部分填充位数保证一定可以重命名。
2022年6月23日更新一些以前出现的bug,增加中文路径的读取和保存。
import json
import os
import random
import shutil
import cv2
import numpy as np
class darknetTool():
def __init__(self, help_=False):
# opencv保存图片时会根据图片的类型(后缀名)对图片进行处理,bmp格式保存时
#常会因为内存占用太大而导致内存不足,可换成png或jpg格式,总体来说bmp>png>jpg
self.externName = "png"
if help_:
self.getHelp()
def getAllImagePath(self, folder, recursion=False):
expandName = ["jpg", "JPG", "jpeg", "JPEG", "png", "PNG", "bmp", "BMP"]
imagePathList = []
if not recursion:
for item in os.listdir(folder):
if os.path.isdir(os.path.join(folder, item)):
continue
if item.rsplit(".",1)[-1] in expandName:
imagePathList.append(os.path.join(folder, item))
return imagePathList
else:
for item in os.listdir(folder):
if os.path.isdir(os.path.join(folder, item)):
subPathList = self.getAllImagePath(os.path.join(folder, item), True)
imagePathList.extend(item for item in subPathList)
else:
if item.rsplit(".",1)[-1] in expandName:
imagePathList.append(os.path.join(folder, item))
return imagePathList
def getAllJsonFilePath(self, folder):
jsonPathList = []
for item in os.listdir(folder):
if item.rsplit(".",1)[-1] == 'json':
jsonPathList.append(os.path.join(folder, item))
return jsonPathList
def getAllTxtPath(self, folder):
txtPathList = []
for item in os.listdir(folder):
if item.rsplit(".",1)[-1] == 'txt':
txtPathList.append(os.path.join(folder, item))
return txtPathList
def fromJsonFileToSearchImage(self, jsonPathList, imageFolder):
imagePathList = []
for item in jsonPathList:
with open(item, "r", encoding='utf-8') as f:
jsonData = json.load(f)
imageName = jsonData["imagePath"].rsplit("\\",1)[-1]
jsonName = item.rsplit("\\",1)[-1]
imagePath = os.path.join(imageFolder, imageName)
# print(imagePath)
imagePathList.append(imagePath)
return imagePathList
def copyAllLabeledImageToNewFolder(self, imagePathList, saveFolder):
for item in imagePathList:
imageName = item.rsplit("\\",1)[-1]
if "//" in imageName:
imageName = imageName.rsplit("//",1)[-1]# 有的文件名是..//image//0001.bmp,这样的
savePath = os.path.join(saveFolder, imageName)
# item = os.path.abspath(item)
# savePath = os.path.abspath(savePath)
print("copyAllLabeledImageToNewFolder: srcPath", item)
print("--->copyAllLabeledImageToNewFolder: dstPath", savePath)
#check the image exits or not
if not os.path.exists(item): continue
shutil.copyfile(item, savePath)
def renameAllLabeledImageAndTxt(self, txtFolder, labeledImageFolder, n=5, startN=0):
# extendName是图像后缀名,n用于zfill(n),00001.jpg
i = startN
imagePathList = self.getAllImagePath(labeledImageFolder)
if (len(imagePathList[0].split("\\")[-1])-4) == n:
n += 1
for imagePath in imagePathList:
i += 1
imageName_dst = str(i).zfill(n) + "." + imagePath.rsplit(".",1)[-1]
image_srcPath = imagePath
image_dstPath = os.path.join(labeledImageFolder, imageName_dst)
txtName_src = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] +".txt"
txtName_dst = str(i).zfill(n) + ".txt"
txt_srcPath = os.path.join(txtFolder, txtName_src)
txt_dstPath = os.path.join(txtFolder, txtName_dst)
try:
if not os.path.exists(image_dstPath):
os.rename(image_srcPath, image_dstPath)
if not os.path.exists(txt_dstPath):
os.rename(txt_srcPath, txt_dstPath)
print("renameAllLabeledImageAndTxt")
print("image_srcPath:", image_srcPath)
print("image_dstPath:", image_dstPath)
print("txt_srcPath:", txt_srcPath)
print("txt_dstPath:", txt_dstPath)
except:
print("ERROR:renameAllLabeledImageAndTxt error")
print("image_srcPath:", image_srcPath)
print("image_dstPath:", image_dstPath)
print("txt_srcPath:", txt_srcPath)
print("txt_dstPath:", txt_dstPath)
def changeLabelFromJsonToTxt(self, jsonPath, txtSaveFolder):
with open(jsonPath, "r", encoding='utf-8') as f:
jsonData = json.load(f)
img_h = jsonData["imageHeight"]
img_w = jsonData["imageWidth"]
txtName = jsonPath.rsplit("\\",1)[-1].rsplit(".",1)[0] + ".txt"
txtPath = os.path.join(txtSaveFolder, txtName)
with open(txtPath, "w") as f:
for item in jsonData["shapes"]:
label = item["label"]
pt1 = item["points"][0]
pt2 = item["points"][1]
xCenter = (pt1[0] + pt2[0]) / 2
yCenter = (pt1[1] + pt2[1]) / 2
obj_h = pt2[1] - pt1[1]
obj_w = pt2[0] - pt1[0]
f.write(" {} ".format(label))
f.write(" {} ".format(xCenter / img_w))
f.write(" {} ".format(yCenter / img_h))
f.write(" {} ".format(obj_w / img_w))
f.write(" {} ".format(obj_h / img_h))
f.write(" \n")
def showLabelFromJson(self, jsonPath, imageFolder):
cv2.namedWindow("img", 0)
font = cv2.FONT_HERSHEY_SIMPLEX
with open(jsonPath, "r") as f:
jsonData = json.load(f)
imageName = jsonData["imagePath"].rsplit("\\",1)[-1]
imagePath = os.path.join(imageFolder, imageName)
img = cv2.imread(imagePath)
print("showLabelFromJson: image path:", imagePath)
for item in jsonData["shapes"]:
label = item["label"]
p1 = (int(item["points"][0][0]), int(item["points"][0][1]))
p2 = (int(item["points"][1][0]), int(item["points"][1][1]))
cv2.putText(img, label, (p1[0], p1[1] - 10), font, 1.2, (0, 0, 255), 2)
cv2.rectangle(img, p1, p2, (0, 255, 0), 2)
cv2.imshow("img", img)
cv2.waitKey(0)
cv2.destroyWindow("img")
def showLabelFromTxt(self, imagePath, txtFolder):
# label xcenter ycenter w h
txtName = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + ".txt"
txtPath = os.path.join(txtFolder, txtName)
print("showLabelFromTxt: image path:", imagePath)
img = cv2.imread(imagePath)
h, w = img.shape[:2]
cv2.namedWindow("img", 0)
font = cv2.FONT_HERSHEY_SIMPLEX
with open(txtPath, "r") as f:
lines = f.readlines()
for line in lines:
tempL = line.split(" ")
label = tempL[1]
obj_w = float(tempL[7])
obj_h = float(tempL[9])
topLeftx = (float(tempL[3]) - obj_w / 2) * w
topLefty = (float(tempL[5]) - obj_h / 2) * h
bottomRightx = (float(tempL[3]) + obj_w / 2) * w
bottomRighty = (float(tempL[5]) + obj_h / 2) * h
p1 = (int(topLeftx), int(topLefty))
p2 = (int(bottomRightx), int(bottomRighty))
cv2.putText(img, label, (p1[0], p1[1] - 10), font, 1.2, (0, 255, 255), 2)
cv2.rectangle(img, p1, p2, (0, 255, 0), 2)
cv2.imshow("img", img)
cv2.waitKey(0)
cv2.destroyWindow("img")
def imageAugment_smooth(self, imagePath, txtFolder, labeledImageFolder,n=5):
txtName = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + ".txt"
txtPath = os.path.join(txtFolder, txtName)
print("smooth: image path:", imagePath)
img = cv2.imread(imagePath)
imageName_smooth = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_smooth." + self.externName
imagePath_smooth = os.path.join(labeledImageFolder, imageName_smooth)#平滑图像保存路径
txtName_smooth = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_smooth.txt"
txtPath_smooth = os.path.join(txtFolder, txtName_smooth)#标签保存路径
shutil.copyfile(txtPath, txtPath_smooth)
print("imageAugment_smooth: source txtPath", txtPath)
print("imageAugment_smooth: dst txtPath", txtPath_smooth)
print("imageAugment_smooth: source imagePath", imagePath)
print("imageAugment_smooth: dst imagePath", imagePath_smooth)
dst = cv2.blur(img, (n, n))
cv2.imwrite(imagePath_smooth, dst)
def imageAugment_flip(self, imagePath, txtFolder, labeledImageFolder,flag=0):
# flag = 0水平翻转 flag = 1,竖直翻转, flag=2 水平翻转+竖直翻转
txtName = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + ".txt"
txtPath = os.path.join(txtFolder, txtName)
print("imageAugment_flip:",flag,": image path:", imagePath)
# img = cv2.imread(imagePath)
img = cv2.imdecode(np.fromfile(imagePath,dtype=np.uint8),-1)
if flag == 0:
imageName_flip = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + "_flipx." + self.externName
imagePath_flip = os.path.join(labeledImageFolder, imageName_flip)
txtName_flip = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_flipx.txt"
txtPath_flip = os.path.join(txtPath.rsplit("\\",1)[0:-1][0], txtName_flip)
elif flag == 1:
imageName_flip = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + "_flipy." + self.externName
imagePath_flip = os.path.join(labeledImageFolder, imageName_flip)
txtName_flip = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_flipy.txt"
txtPath_flip = os.path.join(txtPath.rsplit("\\",1)[0:-1][0], txtName_flip)
elif flag == 2:
imageName_flip = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + "_flipxy." + self.externName
imagePath_flip = os.path.join(labeledImageFolder, imageName_flip)
txtName_flip = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_flipxy.txt"
txtPath_flip = os.path.join(txtPath.rsplit("\\",1)[0:-1][0], txtName_flip)
# 打开原来的txt标签文件,修改原来的x坐标为 xcenter’= 1 - xcenter
with open(txtPath, "r") as fsrc:
with open(txtPath_flip, "w") as f:
lines = fsrc.readlines()
for line in lines:
temp = line.split(" ")
label, xcenter, ycenter, objw, objh = temp[1], temp[3], temp[5], temp[7], temp[9]
if flag == 0:
xcenter = 1 - float(xcenter)
elif flag == 1:
ycenter = 1 - float(ycenter)
elif flag == 2:
xcenter = 1 - float(xcenter)
ycenter = 1 - float(ycenter)
f.write(" {} ".format(label))
f.write(" {} ".format(xcenter))
f.write(" {} ".format(ycenter))
f.write(" {} ".format(objw))
f.write(" {} ".format(objh))
f.write(" \n")
if flag == 0:
dst = cv2.flip(img, 1)
elif flag == 1:
dst = cv2.flip(img, 0)
elif flag == 2:
dst = cv2.flip(img, 1)
dst = cv2.flip(dst, 0)
# cv2.imwrite(imagePath_flip, dst)
print("imagePath_flip: \t ", imagePath_flip)
cv2.imencode(".jpg",dst)[1].tofile(imagePath_flip)
def imageAugment_gamma(self, imagePath, txtFolder, labeledImageFolder,gamma=2):
txtName = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + ".txt"
txtPath = os.path.join(txtFolder, txtName)
print("imageAugment_gamma: image path:", imagePath)
img = cv2.imread(imagePath)
imageName_gamma = imagePath.rsplit("\\",1)[-1].rsplit(".",1)[0] + "_gamma{}.".format(gamma) + self.externName
imagePath_gamma = os.path.join(labeledImageFolder, imageName_gamma)
txtName_gamma = txtPath.rsplit("\\",1)[-1].split(".txt")[0] + "_gamma{}.txt".format(gamma)
txtPath_gamma = os.path.join(txtPath.rsplit("\\",1)[0:-1][0], txtName_gamma)
shutil.copyfile(txtPath, txtPath_gamma)
print("source txtPath", txtPath)
print("dst txtPath", txtPath_gamma)
print("source imagePath", imagePath)
print("dst imagePath", imagePath_gamma)
table = []
for i in range(256):
table.append(((i / 255.0) ** gamma) * 255)
table = np.array(table).astype("uint8")
dst = cv2.LUT(img, table)
cv2.imwrite(imagePath_gamma, dst)
def devideTrainSetAndTestSet(self, imagePathList, saveFolder, prePath="", prop=0.8):
trainTxt = saveFolder + "/train.txt"
testTxt = saveFolder + "/test.txt"
with open(trainTxt, "w") as ftrain:
with open(testTxt, "w") as ftest:
for item in imagePathList:
imagePath = prePath + item.rsplit("\\",1)[-1]
if random.random() < prop:
ftrain.write(imagePath)
ftrain.write("\n")
else:
ftest.write(imagePath)
ftest.write("\n")
def getHelp(self):
pass
if __name__ == "__main__":
rootPath = r"\\192.168.1.247\Pictures\自动生成元件标注数据\光板图标注\0610\small"
dataSavePath = r"D:\myAPP\imageData\autoPad\subset" #生成的数据集的保存路劲
srcImageFolder = rootPath + "\\image" # 存放原始图片的文件夹,有标注过的和没标注过的图片
jsonFolder = rootPath + "\\label" # labelme标注的json文件保存的文件夹
labeledImageFolder = dataSavePath + "/temp/obj"# 筛选出所有标注过的图片然后copy到这个文件夹
txtFolder = dataSavePath +"/temp/obj"# txt文件保存的路径
objFolder = dataSavePath +"/temp/obj"# 把图片和对应的txt文件放在一起,可以直接拿去训练的那种
trainAndTestFileSaveFolder = dataSavePath +"/temp"
prePath = "Pad/obj/"#保存train.txt和test.txt的相对路径eg:total_board/obj/00001.png
trainRatio = 0.9 #训练集和测试集的划分比例
zfillnum = 6 #零填充数,保存文件时填充例如,"1"-->"00001"
startnum = 1 #重命名文件时的起始数字
# 创建文件夹
try:
if not os.path.exists(labeledImageFolder):
os.makedirs(labeledImageFolder)
if not os.path.exists(txtFolder):
os.makedirs(txtFolder)
if not os.path.exists(objFolder):
os.makedirs(objFolder)
except:
print("os.makedirs failed!")
# 初始化类
tool = darknetTool()
# 获取所有原始标注文件
jsonPathList = tool.getAllJsonFilePath(jsonFolder)
# 从json标签文件复查标注
# for item in jsonPathList:
# tool.showLabelFromJson(item,imageFolder)
# 从json标注文件查找图片,并保存到labeledImageFolder
srcLabeledImagePathList = tool.fromJsonFileToSearchImage(jsonPathList, srcImageFolder)
tool.copyAllLabeledImageToNewFolder(srcLabeledImagePathList, labeledImageFolder)
# json格式转为txt格式
for item in jsonPathList:
tool.changeLabelFromJsonToTxt(item, txtFolder)
txtPathList = tool.getAllTxtPath(txtFolder)
# 获取在labeledImageFolder中的图片的路径
dstLabeledImagePathList = tool.getAllImagePath(labeledImageFolder)
# 从txt标签文件查看标注结果
# 所有标注的图片路径都在dstLabeledImagePathList,通过图片路径去找对应的txt文件
# 所有的txt文件都在txtFolder文件夹下
# for imagePath in dstLabeledImagePathList:
# tool.showLabelFromTxt(imagePath,txtFolder)
# 数据增强
# 均值滤波
# for imagePath in dstLabeledImagePathList:
# tool.imageAugment_smooth(imagePath, txtFolder, labeledImageFolder,n=5) # n=9是滤波核大小
# # 水平翻转
# for imagePath in dstLabeledImagePathList:
# tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=0) # flag=0 水平翻转
# # 竖直翻转
# for imagePath in dstLabeledImagePathList:
# tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=1) # flag=1 竖直翻转
# 中心对称
# for imagePath in dstLabeledImagePathList:
# tool.imageAugment_flip(imagePath, txtFolder, labeledImageFolder, flag=2) # flag=2 中心对称
# gamma变换
# for imagePath in dstLabeledImagePathList:
# #tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=0.5)
# tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=0.8)
# tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=1.2)
# #tool.imageAugment_gamma(imagePath, txtFolder, labeledImageFolder, gamma=2)
#图片重命名
# 图片和标签重命名为00001.jpg 00001.txt.......(可选)
# n=5表示zfill()的位数,n=5则为00001.jpg...,n=3则为001.jpg....
tool.renameAllLabeledImageAndTxt(txtFolder, labeledImageFolder, n=zfillnum, startN=startnum)
dstLabeledImagePathList = tool.getAllImagePath(labeledImageFolder)
# 检查转换的最终结果
# for imagePath in dstLabeledImagePathList:
# tool.showLabelFromTxt(imagePath,txtFolder)
#划分训练集和测试集
imagePathList = tool.getAllImagePath(labeledImageFolder)
# prePath = "data/obj/"则train.txt文件中的路径就会是data/obj/00001.jpg,data/obj/00002.jpg......
# prop=0.8表示80%的数据用来训练,其余的做验证集
tool.devideTrainSetAndTestSet(imagePathList, trainAndTestFileSaveFolder, prePath=prePath, prop=trainRatio)