多张图片进行拼成一张图片

此案例将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()

  • 0
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值