此案例将8张图片进行拼成一张图片,标注信息对应调整
import os
import cv2
import random
import numpy as np
from lxml.etree import Element,SubElement,tostring
import xml.dom.minidom as xmldom
from datetime import datetime
from tqdm import tqdm
from operator import itemgetter
"""
单路图像合并成8路,输入必须是标注好的数据,xml中必须包含标注信息,确保背景纯净
"""
def cake_img(img):
img = cv2.imread(img)
h = img.shape[0]
w = img.shape[1]
return h,w,img
def img_to_xml(img_name):
xml = str(img_name).replace(".jpg",".xml")
if os.path.isfile(xml) == False:
xml = str(img_name).replace(".jpg",".xml").replace("images","xml")
if os.path.isfile(xml) == False:
xml = str(img_name).replace(".jpg",".xml").replace("images","c_xml")
return xml
def cake_xml(xml_path):
xml_information = []
domobj = xmldom.parse(xml_path)
bndbox = [0, 0, 0, 0]
elementobj = domobj.documentElement
# object_sub = elementobj.getElementsByTagName('object')
# tree = ET.parse(xml_path)
# root = tree.getroot()
# size = root.find("size")
# width = float(size.find("width").text)
# height = float(size.find("height").text)
sub_element_obj_labes = elementobj.getElementsByTagName('name')
for i in range(0, len(sub_element_obj_labes)):
labes = sub_element_obj_labes[i]
# print(labes.firstChild.data)
str_labe = str(labes.firstChild.data)
sub_element_obj = elementobj.getElementsByTagName('bndbox')
try:
bndbox[0] = int(sub_element_obj[i].getElementsByTagName('xmin')[0].firstChild.data)
bndbox[1] = int(sub_element_obj[i].getElementsByTagName('ymin')[0].firstChild.data)
bndbox[2] = int(sub_element_obj[i].getElementsByTagName('xmax')[0].firstChild.data)
bndbox[3] = int(sub_element_obj[i].getElementsByTagName('ymax')[0].firstChild.data)
except IndexError as I:
print("\033[1;31;40mIndexError: list index out of range\033[0m")
continue
xmin = bndbox[0]
ymin = bndbox[1]
xmax = bndbox[2]
ymax = bndbox[3]
xml_information.append([str_labe,xmin,ymin,xmax,ymax])
return xml_information
def xml_adjust(xml_information,w1=0,max_w=0,w3=0,xml_information_2 = False,xml_information_3 = False,xml_information_4 = False):
information = []
if xml_information_2: # 右上区域
for i in xml_information:
information.append([i[0], int(i[1]) + w1, i[2], int(i[3]) + w1 ,i[4]]) # 960,960 w
return information
elif xml_information_3: #
for i in xml_information:
information.append([i[0], i[1], int(i[2]) + max_w, i[3], int(i[4])+max_w]) # 540,540 hmax
return information
elif xml_information_4: #
for i in xml_information:
information.append([i[0], int(i[1]) + w3, int(i[2]) + max_w, int(i[3]) + w3, int(i[4])+max_w])
return information
def xml_adjust_5_8(xml_information,max_h_1_4_y=0,add_w=0,max_y_add_to7_8=0,xml_information_5 = False,xml_information_6 = False,xml_information_7 = False,xml_information_8 = False):
information = []
if xml_information_5:
for i in xml_information:
information.append([i[0], i[1], i[2]+max_h_1_4_y, i[3] ,i[4]+max_h_1_4_y])
# label xmin ymin xmax ymax
return information
elif xml_information_6:
for i in xml_information:
information.append([i[0], i[1] + add_w, i[2] + max_h_1_4_y, i[3]+add_w, i[4]+max_h_1_4_y])
# label xmin ymin xmax ymax
return information
elif xml_information_7:
for i in xml_information:
information.append([i[0], i[1], int(i[2])+max_h_1_4_y+max_y_add_to7_8, int(i[3]), int(i[4])+max_h_1_4_y+max_y_add_to7_8])
# label xmin ymin xmax ymax
return information
elif xml_information_8: # 右下区域
for i in xml_information:
information.append([i[0], int(i[1]) + add_w , int(i[2])+max_h_1_4_y+max_y_add_to7_8, int(i[3]) + add_w, int(i[4])+max_h_1_4_y+max_y_add_to7_8])
# label xmin ymin xmax ymax
return information
# def img_h_w_cake(h1,w1,h2,w2,im_list,xml,img):
# if h1 != h2 or w1 != w2:
# while True:
# img_name = im_list[random.randint(0, num - 1)]
# xml_ = img_to_xml(img_name)
# h2, w2, img_ = cake_img(img_name)
# if h1 == h2 and w1 == w2:
# break
# return xml_,img_
# elif h1 == h2 and w1 == w2:
# return xml,img
if __name__ == '__main__':
"""
单路图像合并成8路,输入必须是标注好的数据,xml中最好都要包含标注信息,确保背景纯净
input_img_dir: xml和jpg共同放在同一个目录下,该目录作为输入的目录
output_dir: 最终输出想要保存的位置
"""
input_img_dir = r"D:\XRJC\DATA\PMT\Train_data_10\images" # 建议xml和jpg共同放在同一个目录下 images
output_dir = r"D:\XRJC\DATA\PMT\Train_data_10_make_8\images" # 图片存放位置 images 自动创建文件夹
cake_labes = ["ricc"] # 不参与训练的类别
if os.path.isdir(os.path.dirname(output_dir)) == False:
os.mkdir(os.path.dirname(output_dir))
if os.path.isdir(output_dir) == False:
os.mkdir(output_dir)
output_dir_xml = output_dir.replace("images","xml")
if os.path.isdir(output_dir_xml) == False:
os.mkdir(output_dir_xml)
im_list = []
for root,dir,data_list in os.walk(input_img_dir):
for im in [x for x in data_list if x.endswith(".jpg")]:
im_path = os.path.join(input_img_dir,im)
im_list.append(im_path)
# 加入负样本
for root, dir, data_list in os.walk(input_img_dir.replace("images", "fyb")):
for im in [x for x in data_list if x.endswith(".jpg")]:
im_path = os.path.join(input_img_dir.replace("images", "fyb"), im)
im_list.append(im_path)
num = len(im_list)
for index,img in enumerate(tqdm(im_list)):
# xml全局信息汇总
data_inf = []
# 获取各个视图的在列表中元素的下标对应的图片
img1_name = img
img2_name = im_list[random.randint(0,num-1)]
img3_name = im_list[random.randint(0,num-1)]
img4_name = im_list[random.randint(0,num-1)]
img5_name = im_list[random.randint(0, num - 1)]
img6_name = im_list[random.randint(0, num - 1)]
img7_name = im_list[random.randint(0, num - 1)]
img8_name = im_list[random.randint(0, num - 1)]
# 第一视图 左上区域
xml1 = img_to_xml(img1_name)
h1,w1,img1 = cake_img(img1_name)
if os.path.isfile(xml1) == False:
xml_information_1 = []
elif os.path.isfile(xml1) == True:
xml_information_1 = cake_xml(xml1)
# 第二视图 右上区域
xml2 = img_to_xml(img2_name)
h2,w2,img2 = cake_img(img2_name)
# xml2,img2 = img_h_w_cake(h1,w1,h2,w2,im_list,xml2,img2)
if os.path.isfile(xml2) == False:
xml_information_2 = []
elif os.path.isfile(xml2) == True:
xml_information_2 = cake_xml(xml2)
if len(xml_information_2) > 0:
xml_information_2 = xml_adjust(xml_information_2,w1,max_w=0,w3=0,xml_information_2=True)
elif len(xml_information_2) == 0:
xml_information_2 = []
max_w = max(img1.shape[0], img2.shape[0])
# 第三视图
xml3 = img_to_xml(img3_name)
h3,w3,img3 = cake_img(img3_name)
# xml3, img3 = img_h_w_cake(h1, w1, h3, w3, im_list, xml3, img3)
if os.path.isfile(xml3) ==False:
xml_information_3 = []
elif os.path.isfile(xml3) == True:
xml_information_3 = cake_xml(xml3)
if len(xml_information_3) > 0:
xml_information_3 = xml_adjust(xml_information_3,w1,max_w,w3,xml_information_3=True)
elif len(xml_information_3) == 0:
xml_information_3 = []
# 第四视图
xml4 = img_to_xml(img4_name)
h4,w4,img4 = cake_img(img4_name)
if os.path.isfile(xml4) == False:
xml_information_4=[]
elif os.path.isfile(xml4) == True:
xml_information_4 = cake_xml(xml4)
if len(xml_information_4) > 0:
xml_information_4 = xml_adjust(xml_information_4,w1,max_w,w3,xml_information_4=True)
elif len(xml_information_4) == 0:
xml_information_4 = []
max_h_1_4_y = max(h1+h3,h2+h4)
max_w_1_4_x = max(w1+w3,w2+w4)
# 第五视图
xml5 = img_to_xml(img5_name)
h5, w5, img5 = cake_img(img5_name)
if os.path.isfile(xml5) == False:
xml_information_5 = []
elif os.path.isfile(xml5) == True:
xml_information_5 = cake_xml(xml5)
if len(xml_information_5) > 0:
xml_information_5 = xml_adjust_5_8(xml_information_5, max_h_1_4_y, xml_information_5=True)
elif len(xml_information_5) == 0:
xml_information_5 = []
# 第六视图
xml6 = img_to_xml(img6_name)
h6, w6, img6 = cake_img(img6_name)
if os.path.isfile(xml6) == False:
xml_information_6 = []
elif os.path.isfile(xml6) == True:
xml_information_6 = cake_xml(xml6)
if len(xml_information_6) > 0:
xml_information_6 = xml_adjust_5_8(xml_information_6, max_h_1_4_y, w5, xml_information_6=True)
elif len(xml_information_6) == 0:
xml_information_6 = []
max_y_add_to7_8 = max(h5,h6)
# print(max_y_add_to7_8,max_h_1_4_y)
# quit()
# 第七视图
xml7 = img_to_xml(img7_name)
h7, w7, img7 = cake_img(img7_name)
if os.path.isfile(xml7) == False:
xml_information_7 = []
elif os.path.isfile(xml7) == True:
xml_information_7 = cake_xml(xml7)
if len(xml_information_7) > 0:
xml_information_7 = xml_adjust_5_8(xml_information_7, max_h_1_4_y,0,max_y_add_to7_8, xml_information_7=True)
elif len(xml_information_7) == 0:
xml_information_7 = []
# 第八视图
xml8 = img_to_xml(img8_name)
h8, w8, img8 = cake_img(img8_name)
if os.path.isfile(xml8) == False:
xml_information_8 = []
elif os.path.isfile(xml8) == True:
xml_information_8 = cake_xml(xml8)
if len(xml_information_8) > 0:
xml_information_8 = xml_adjust_5_8(xml_information_8, max_h_1_4_y,w7,max_y_add_to7_8, xml_information_8=True)
elif len(xml_information_8) == 0:
xml_information_8 = []
# 得到全部xml信息进行汇总
data_inf = xml_information_1 + xml_information_2 + xml_information_3 + xml_information_4 \
+xml_information_5+xml_information_6+xml_information_7+xml_information_8
# W_make = max([w1,w2]) + max([w3,w4])
# H_make = max([h1,h2]) + max([h3,h4])
W_make = max([w1+w2,w3+w4,w5+w6,w7+w8])
H_make = max([h1+h3+h5+h7,h2+h4+h6+h8])
# 建立空白图片
# print((H_make, W_make, 3))
num_l = random.randint(1, 256)
result = np.random.randint(num_l,size=(H_make, W_make, 3), dtype =np.uint8)
# h,w,c = img.shape
# 导入各个视图到空白图片中
try:
result[0:img1.shape[0] , 0:img1.shape[1]] = img1 # 0:540,0:960 左上角
result[0:img2.shape[0] , img1.shape[1]:img1.shape[1] + img2.shape[1]] = img2 # 0:540,960:1920 右上角
result[max_w:max_w + img3.shape[0],0:img3.shape[1]] = img3 # 540:1080,0:960 左下角
result[max_w:max_w + img4.shape[0],img3.shape[1]:img3.shape[1] + img4.shape[1]] = img4 # 540:1080,960:1920 右下角
result[max_h_1_4_y:img5.shape[0]+max_h_1_4_y, 0:img5.shape[1]] = img5 # 第五视图
result[max_h_1_4_y:img6.shape[0]+max_h_1_4_y, img5.shape[1]:img5.shape[1] + img6.shape[1]] = img6 # 第六视图
result[max_h_1_4_y+max_y_add_to7_8:img7.shape[0]+max_h_1_4_y+max_y_add_to7_8, 0:img7.shape[1]] = img7
result[max_h_1_4_y+max_y_add_to7_8:max_h_1_4_y+max_y_add_to7_8+img8.shape[0], img7.shape[1]:img7.shape[1] + img8.shape[1]] = img8
except:
continue
# 画框显示验证
# for i in data_inf:
# result = cv2.rectangle(result,(i[1],i[2]),(i[3],i[4]),(0,255,0), 2)
# cv2.namedWindow("img", cv2.WINDOW_NORMAL)
# cv2.resizeWindow("img", 1422, 800)
# cv2.imshow('img', result)
# cv2.waitKey(0)
# 保存图片
start = datetime.now() #获取当前系统时间
start_time1 = start.strftime('%Y-%m-%d %H:%M:%S').split(" ")[0]
result_name = start_time1 + "_" +str(index + 1).zfill(8) + ".jpg"
# save_path = os.path.join(output_dir,result_name)
# cv2.imwrite(save_path,result)
name_list = list(set(list(map(itemgetter(0), data_inf))))
cake_name_last = set(name_list).issubset(set(cake_labes)) # 判断是否全是不参与训练的类别
if len(data_inf) == 0 or cake_name_last == True: # 负样本不做输出xml操作
save_path = os.path.join(output_dir.replace("images", "fyb"), result_name)
if os.path.isdir(output_dir.replace("images", "fyb")) == False:
os.mkdir(output_dir.replace("images", "fyb"))
cv2.imwrite(save_path, result)
continue
else:
save_path = os.path.join(output_dir, result_name)
cv2.imwrite(save_path, result)
# xml生成
result_xml_name = result_name.replace(".jpg", ".xml")
imh, imw = result.shape[0:2]
node_root = Element('annotation')
node_folder = SubElement(node_root, 'folder')
node_folder.text = '1'
node_filename = SubElement(node_root, 'filename')
node_filename.text = result_name
node_size = SubElement(node_root, 'size')
node_width = SubElement(node_size, 'width')
node_width.text = str(imw)
node_height = SubElement(node_size, 'height')
node_height.text = str(imh)
node_depth = SubElement(node_size, 'depth')
node_depth.text = '3'
for i in range(len(data_inf)):
node_object = SubElement(node_root, 'object')
node_name = SubElement(node_object, 'name')
node_name.text = str(data_inf[i][0])
node_pose = SubElement(node_object, 'pose')
node_pose.text = "Unspecified"
node_truncated = SubElement(node_object, 'truncated')
node_truncated.text = "truncated"
node_difficult = SubElement(node_object, 'difficult')
node_difficult.text = '0'
node_bndbox = SubElement(node_object, 'bndbox')
node_xmin = SubElement(node_bndbox, 'xmin')
node_xmin.text = str(data_inf[i][1])
node_ymin = SubElement(node_bndbox, 'ymin')
node_ymin.text = str(data_inf[i][2])
node_xmax = SubElement(node_bndbox, 'xmax')
node_xmax.text = str(int(data_inf[i][3]))
node_ymax = SubElement(node_bndbox, 'ymax')
node_ymax.text = str(int(data_inf[i][4]))
xml = tostring(node_root, pretty_print=True) #
img_newxml = os.path.join(output_dir_xml, result_xml_name)
file_object = open(img_newxml, 'wb')
file_object.write(xml)
file_object.close()