在这里,先给出原博主的代码:https://blog.csdn.net/qq_36756866/article/details/106299705
这个代码有一些bug,并且生成数据的数量并不能任意改变,于是我自己修改了一部分代码,实现生成任意数据量的功能。使用只需修改文件路径即可,非常方便,不会有任何bug。代码如下:
# coding:utf-8
import cv2
import os
import random
import numpy as np
import xml.etree.ElementTree as ET
import xml.dom.minidom
img_path = '/root/wu_datasets/ReCurrentPapper/data/DOTA/DOTAship/VOCdevkit/VOC2012/JPEGImages/' # 原始图片文件夹路径
save_path = '/root/wu_datasets/ReCurrentPapper/data/DOTA_mixup/VOCdevkit/VOC2012/JPEGImages/' # mixup的图片文件夹路径
xml_path = '/root/wu_datasets/ReCurrentPapper/data/DOTA/DOTAship/VOCdevkit/VOC2012/Annotations/' # 原始图片对应的标注文件xml文件夹的路径
save_xml = '/root/wu_datasets/ReCurrentPapper/data/DOTA_mixup/VOCdevkit/VOC2012//Annotations/' # mixup的图片对应的标注文件xml的文件夹路径
img_names = os.listdir(img_path)
img_num = len(img_names)
#print(img_names)
# sum 代表生成数据的总和
sum = []
# 这里表示比原先数据集多10倍
for i in range(10):
for i in range(0, img_num):
sum.append(i)
# print(sum)
cnt = 0
for k in sum:
cnt += 1
print(cnt)
imgname = img_names[k]
imgpath = img_path + imgname
img = cv2.imread(imgpath)
img_h, img_w = img.shape[0], img.shape[1]
# print(img_h,img_w)
i = random.randint(0, img_num - 1)
add_path = img_path + img_names[i]
addimg = cv2.imread(add_path)
add_h, add_w = addimg.shape[0], addimg.shape[1]
if add_h != img_h or add_w != img_w:
print('resize!')
addimg = cv2.resize(addimg, (img_w, img_h), interpolation=cv2.INTER_LINEAR)
scale_h, scale_w = img_h / add_h, img_w / add_w
lam = np.random.beta(1.5, 1.5)
# print(lam)
mixed_img = lam * img + (1 - lam) * addimg
save_img = save_path + str(cnt) + '.jpg'
cv2.imwrite(save_img, mixed_img)
# print(save_img)
# print(imgname, img_names[i])
if imgname != img_names[i]:
xmlfile1 = xml_path + imgname[:-4] + '.xml'
xmlfile2 = xml_path + img_names[i][:-4] + '.xml'
# print(xmlfile1,xmlfile2)
tree1 = ET.parse(xmlfile1)
tree2 = ET.parse(xmlfile2)
doc = xml.dom.minidom.Document()
root = doc.createElement("annotation")
doc.appendChild(root)
for folds in tree1.findall("folder"):
folder = doc.createElement("folder")
folder.appendChild(doc.createTextNode(str(folds.text)))
root.appendChild(folder)
for filenames in tree1.findall("filename"):
filename = doc.createElement("filename")
filename.appendChild(doc.createTextNode("{}.jpg".format(cnt)))
root.appendChild(filename)
for paths in tree1.findall("path"):
path = doc.createElement("path")
path.appendChild(doc.createTextNode(str(paths.text)))
root.appendChild(path)
for sources in tree1.findall("source"):
source = doc.createElement("source")
database = doc.createElement("database")
database.appendChild(doc.createTextNode(str("Unknow")))
source.appendChild(database)
root.appendChild(source)
for sizes in tree1.findall("size"):
size = doc.createElement("size")
width = doc.createElement("width")
height = doc.createElement("height")
depth = doc.createElement("depth")
width.appendChild(doc.createTextNode(str(img_w)))
height.appendChild(doc.createTextNode(str(img_h)))
depth.appendChild(doc.createTextNode(str(3)))
size.appendChild(width)
size.appendChild(height)
size.appendChild(depth)
root.appendChild(size)
nodeframe = doc.createElement("frame")
nodeframe.appendChild(doc.createTextNode(imgname[:-4]))
objects = []
for obj in tree1.findall("object"):
obj_struct = {}
obj_struct["name"] = obj.find("name").text
obj_struct["pose"] = obj.find("pose").text
obj_struct["truncated"] = obj.find("truncated").text
obj_struct["difficult"] = obj.find("difficult").text
bbox = obj.find("bndbox")
obj_struct["bbox"] = [int(bbox.find("xmin").text),
int(bbox.find("ymin").text),
int(bbox.find("xmax").text),
int(bbox.find("ymax").text)]
objects.append(obj_struct)
for obj in tree2.findall("object"):
obj_struct = {}
obj_struct["name"] = obj.find("name").text
obj_struct["pose"] = obj.find("pose").text
obj_struct["truncated"] = obj.find("truncated").text
obj_struct["difficult"] = obj.find("difficult").text # 有的版本的labelImg改参数为小写difficult
bbox = obj.find("bndbox")
obj_struct["bbox"] = [int(int(bbox.find("xmin").text) * scale_w),
int(int(bbox.find("ymin").text) * scale_h),
int(int(bbox.find("xmax").text) * scale_w),
int(int(bbox.find("ymax").text) * scale_h)]
objects.append(obj_struct)
for obj in objects:
nodeobject = doc.createElement("object")
nodename = doc.createElement("name")
nodepose = doc.createElement("pose")
nodetruncated = doc.createElement("truncated")
nodedifficult = doc.createElement("difficult")
nodebndbox = doc.createElement("bndbox")
nodexmin = doc.createElement("xmin")
nodeymin = doc.createElement("ymin")
nodexmax = doc.createElement("xmax")
nodeymax = doc.createElement("ymax")
nodename.appendChild(doc.createTextNode(obj["name"]))
nodepose.appendChild(doc.createTextNode(obj["pose"]))
nodetruncated.appendChild(doc.createTextNode(obj["truncated"]))
nodedifficult.appendChild(doc.createTextNode(obj["difficult"]))
nodexmin.appendChild(doc.createTextNode(str(obj["bbox"][0])))
nodeymin.appendChild(doc.createTextNode(str(obj["bbox"][1])))
nodexmax.appendChild(doc.createTextNode(str(obj["bbox"][2])))
nodeymax.appendChild(doc.createTextNode(str(obj["bbox"][3])))
nodebndbox.appendChild(nodexmin)
nodebndbox.appendChild(nodeymin)
nodebndbox.appendChild(nodexmax)
nodebndbox.appendChild(nodeymax)
nodeobject.appendChild(nodename)
nodeobject.appendChild(nodepose)
nodeobject.appendChild(nodetruncated)
nodeobject.appendChild(nodedifficult)
nodeobject.appendChild(nodebndbox)
root.appendChild(nodeobject)
fp = open(save_xml + str(cnt) + ".xml", "w")
doc.writexml(fp, indent='\t', addindent='\t', newl='\n', encoding="utf-8")
fp.close()
else:
xmlfile1 = xml_path + imgname[:-4] + '.xml'
# print(xmlfile1)
tree1 = ET.parse(xmlfile1)
doc = xml.dom.minidom.Document()
root = doc.createElement("annotation")
doc.appendChild(root)
for folds in tree1.findall("folder"):
folder=doc.createElement("folder")
folder.appendChild(doc.createTextNode(str(folds.text)))
root.appendChild(folder)
for filenames in tree1.findall("filename"):
filename=doc.createElement("filename")
filename.appendChild(doc.createTextNode("{}.jpg".format(cnt)))
root.appendChild(filename)
for paths in tree1.findall("path"):
path = doc.createElement("path")
path.appendChild(doc.createTextNode(str(paths.text)))
root.appendChild(path)
for sources in tree1.findall("source"):
source = doc.createElement("source")
database = doc.createElement("database")
database.appendChild(doc.createTextNode(str("Unknow")))
source.appendChild(database)
root.appendChild(source)
for sizes in tree1.findall("size"):
size = doc.createElement("size")
width = doc.createElement("width")
height = doc.createElement("height")
depth = doc.createElement("depth")
width.appendChild(doc.createTextNode(str(img_w)))
height.appendChild(doc.createTextNode(str(img_h)))
depth.appendChild(doc.createTextNode(str(3)))
size.appendChild(width)
size.appendChild(height)
size.appendChild(depth)
root.appendChild(size)
nodeframe = doc.createElement("frame")
nodeframe.appendChild(doc.createTextNode(imgname[:-4]))
objects = []
for obj in tree1.findall("object"):
obj_struct = {}
obj_struct["name"] = obj.find("name").text
obj_struct["pose"] = obj.find("pose").text
obj_struct["truncated"] = obj.find("truncated").text
obj_struct["difficult"] = obj.find("difficult").text
bbox = obj.find("bndbox")
obj_struct["bbox"] = [int(bbox.find("xmin").text),
int(bbox.find("ymin").text),
int(bbox.find("xmax").text),
int(bbox.find("ymax").text)]
objects.append(obj_struct)
for obj in objects:
nodeobject = doc.createElement("object")
nodename = doc.createElement("name")
nodepose = doc.createElement("pose")
nodetruncated = doc.createElement("truncated")
nodedifficult = doc.createElement("difficult")
nodebndbox = doc.createElement("bndbox")
nodexmin = doc.createElement("xmin")
nodeymin = doc.createElement("ymin")
nodexmax = doc.createElement("xmax")
nodeymax = doc.createElement("ymax")
nodename.appendChild(doc.createTextNode(obj["name"]))
nodepose.appendChild(doc.createTextNode(obj["pose"]))
nodetruncated.appendChild(doc.createTextNode(obj["truncated"]))
nodedifficult.appendChild(doc.createTextNode(obj["difficult"]))
nodexmin.appendChild(doc.createTextNode(str(obj["bbox"][0])))
nodeymin.appendChild(doc.createTextNode(str(obj["bbox"][1])))
nodexmax.appendChild(doc.createTextNode(str(obj["bbox"][2])))
nodeymax.appendChild(doc.createTextNode(str(obj["bbox"][3])))
nodebndbox.appendChild(nodexmin)
nodebndbox.appendChild(nodeymin)
nodebndbox.appendChild(nodexmax)
nodebndbox.appendChild(nodeymax)
nodeobject.appendChild(nodename)
nodeobject.appendChild(nodepose)
nodeobject.appendChild(nodetruncated)
nodeobject.appendChild(nodedifficult)
nodeobject.appendChild(nodebndbox)
root.appendChild(nodeobject)
fp = open(save_xml + str(cnt) + ".xml", "w")
doc.writexml(fp, indent='\t', addindent='\t', newl='\n', encoding="utf-8")
fp.close()