roi区域xml数据增强
保留roi区域内的完整或非完整目标(python实现)
代码中含有很多elif,包含了很多目标位置可能性!
"""
@Author : Laily
@Time : 2020/11/06 下午3:04
@Comment :目标检测roi区域数据增强
"""
import cv2
import os
import xml.dom.minidom
from xml.dom.minidom import parse
# 原始图片和xml地址
image_file = '/home/tensorsight/Documents/RoadAdministration/datasets/my/all/all_img'
xml_file = '/home/tensorsight/Documents/RoadAdministration/datasets/my/all/all_xml'
# 增强后图片和xml存放地址
save_file = '/home/tensorsight/Documents/RoadAdministration/datasets/my/all/roi/img'
save_xml = '/home/tensorsight/Documents/RoadAdministration/datasets/my/all/roi/xml'
def xml_solve(dir_xml,roi,x_width,y_height):
print("[ 当前增强 ] ",dir_xml)
doc = parse(dir_xml)
rootNode = doc.documentElement
xml_width = rootNode.getElementsByTagName("width")[0]
xml_width.childNodes[0].data = x_width
xml_height = rootNode.getElementsByTagName("height")[0]
xml_height.childNodes[0].data = y_height
xml_objects = rootNode.getElementsByTagName("object")
for xml_object in xml_objects:
xml_name = xml_object.getElementsByTagName("name")[0]
objectNode = xml_name.parentNode
print("roi:",roi)
xml_xmin = objectNode.getElementsByTagName("xmin")[0]
xml_ymin = objectNode.getElementsByTagName("ymin")[0]
xml_xmax = objectNode.getElementsByTagName("xmax")[0]
xml_ymax = objectNode.getElementsByTagName("ymax")[0]
# print(xml_ymin.childNodes[0].data)
# print(xml_ymax.childNodes[0].data)
# print(xml_xmin.childNodes[0].data)
# print(xml_xmax.childNodes[0].data)
ymin = int(xml_ymin.childNodes[0].data)
ymax = int(xml_ymax.childNodes[0].data)
xmin = int(xml_xmin.childNodes[0].data)
xmax = int(xml_xmax.childNodes[0].data)
#1.
if roi_ymax >= ymin >= roi_ymin and roi_ymin <= ymax <= roi_ymax \
and roi_xmax >= xmin >= roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************1-1")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = xmax-roi_xmin
#2.
elif ymin < roi_ymin and roi_ymin <= ymax <= roi_ymax \
and xmin < roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************2-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif ymin < roi_ymin and roi_ymin <= ymax <= roi_ymax \
and roi_xmax >= xmin >= roi_xmin and xmax > roi_xmax:
# print("******************************2-2")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = x_width
elif roi_ymax >= ymin >= roi_ymin and ymax > roi_ymax \
and xmin < roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************2-3")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif roi_ymax >= ymin >= roi_ymin and ymax > roi_ymax \
and roi_xmax >= xmin >= roi_xmin and xmax > roi_xmax:
# print("******************************2-4")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = x_width
#3.
elif ymin < roi_ymin and ymax > roi_ymax \
and roi_xmax >= xmin >= roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************3-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif roi_ymax >= ymin >= roi_ymin and roi_ymin <= ymax <= roi_ymax \
and xmin < roi_xmin and xmax > roi_xmax:
# print("******************************3-2")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = x_width
#4
elif ymin < roi_ymin and ymax > roi_ymax \
and xmin < roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************4-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif ymin < roi_ymin and ymax > roi_ymax \
and roi_xmax >= xmin >= roi_xmin and xmax > roi_xmax:
# print("******************************4-2")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = x_width
#5
elif ymin < roi_ymin and roi_ymin <= ymax <= roi_ymax \
and xmin < roi_xmin and xmax > roi_xmax:
# print("******************************5-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = x_width
elif roi_ymax >= ymin >= roi_ymin and ymax > roi_ymax \
and xmin < roi_xmin and xmax > roi_xmax:
# print("******************************5-2")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = x_width
#6
elif ymin < roi_ymin and roi_ymin <= ymax <= roi_ymax \
and roi_xmax >= xmin >= roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************6-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif roi_ymax >= ymin >= roi_ymin and ymax > roi_ymax \
and roi_xmax >= xmin >= roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************6-2")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif roi_ymax >= ymin >= roi_ymin and roi_ymin <= ymax <= roi_ymax \
and xmin < roi_xmin and roi_xmin <= xmax <= roi_xmax:
# print("******************************6-3")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = xmax-roi_xmin
elif roi_ymax >= ymin >= roi_ymin and roi_ymin <= ymax <= roi_ymax \
and roi_xmax >= xmin >= roi_xmin and xmax > roi_xmax:
# print("******************************6-4")
xml_ymin.childNodes[0].data = ymin-roi_ymin
xml_ymax.childNodes[0].data = ymax-roi_ymin
xml_xmin.childNodes[0].data = xmin-roi_xmin
xml_xmax.childNodes[0].data = x_width
#7.
elif ymin < roi_ymin and ymax > roi_ymax \
and xmin < roi_xmin and xmax > roi_xmax:
# print("******************************7-1")
xml_ymin.childNodes[0].data = 0
xml_ymax.childNodes[0].data = y_height
xml_xmin.childNodes[0].data = 0
xml_xmax.childNodes[0].data = x_width
else:
rootNode.removeChild(xml_object)
with open(os.path.join(save_xml,name+'.xml'),'w') as f:
doc.writexml(f,encoding='utf-8')
if __name__=="__main__":
list_image = os.listdir(image_file)
for image_path in list_image:
dir_image = os.path.join(image_file, image_path)
name, _ = os.path.splitext(image_path)
dir_xml = os.path.join(xml_file, name+'.xml')
img = cv2.imread(dir_image)
height, width,_ = img.shape
#手动设定roi中心点
x_center, y_center = width/2, height/2 + height * 1/4
#设定roi区域宽度和高度 width:height = 2:1
x_width = int(width/2)
y_height = int(x_width/2)
#得到roi左上角和右下角坐标
roi_ymin = int(y_center-y_height/2)
roi_ymax = int(y_center+y_height/2)
roi_xmin = int(x_center-x_width/2)
roi_xmax = int(x_center+x_width/2)
xml_solve(dir_xml, [roi_ymin, roi_ymax, roi_xmin, roi_xmax],x_width,y_height)
img = img[roi_ymin:roi_ymax,roi_xmin:roi_xmax]
cv2.imwrite(os.path.join(save_file,image_path), img)