python数据集扩增之裁剪图片上下左右及中心并随之改变xml文件的边界框信息

import cv2
import xml.etree.ElementTree as ET
import os
import numpy as np
from tqdm import tqdm
from PIL import Image


# 数据扩增
# 读xml文档返回每个目标边界框信息
def read_xml_annotation(root, image_id):  # (源xml文件夹地址和文件夹内所有图像名称)
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []
    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)
        bndboxlist.append([xmin, ymin, xmax, ymax])
   

    bndbox = root.find('object').find('bndbox')
    return bndboxlist

# 更改xml信息将裁剪至每个小部分的对应目标边界框进行提取
def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    # 修改增强后的xml文件中的filename
    elem = tree.find('filename')
    elem.text = (str(id) + '.jpg')
    xmlroot = tree.getroot()
    # 修改增强后的xml文件中的path
    elem = tree.find('path')
    if elem != None:
        elem.text = (saveroot + str(id) + '.jpg')

    index = 0
    for object in xmlroot.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)
        if xmin.text=="0" and ymin.text=="0" and xmax.text=="0" and ymax.text=="0":
            xmlroot.remove(object)  # 由于在后续设置中将超出五个裁剪完成图像范围的目标边界框的四个参数全部设为0,此处将参数为0的object删除掉,剩余满足裁剪范围内的目标边界框
        index = index + 1

    tree.write(os.path.join(saveroot, str(id + '.xml')))


# 遍历指定目录,显示目录下的所有文件名
def CropImage4File(filepath, destpath):
    pathDir = os.listdir(filepath)  # 列出文件路径中的所有路径或文件
    for allDir in pathDir:
        child = os.path.join(filepath, allDir)
        dest = os.path.join(destpath, allDir)
        if os.path.isfile(child):
            image = cv2.imread(child)
        sp = image.shape  # 获取图像形状:返回【行数值,列数值】列表
        sz1 = sp[0]  # 图像的高度(行 范围)
        sz2 = sp[1]  # 图像的宽度(列 范围)
        # sz3 = sp[2]                #像素值由【RGB】三原色组成
        # 你想对文件的操作
        n, s = os.path.splitext(dest)

        cropImg = image
        cv2.imwrite(n + s, cropImg)  # 写入图像路径

        a1 = 0  # x start
        b1 = int(sz1 / 2)  # x end
        c1 = 0  # y start
        d1 = int(sz2 / 2)  # y end
        cropImg1 = image[a1:b1, c1:d1]  # 裁剪图像
        cv2.imwrite(n+'-1'+s, cropImg1)  # 写入图像路径

        a2 = 0  # x start
        b2 = int(sz1 / 2)  # x end
        c2 = int(sz2 / 2)  # y start
        d2 = int(sz2)  # y end
        cropImg2 = image[a2:b2, c2:d2]  # 裁剪图像
        cv2.imwrite(n+'-2'+s, cropImg2)  # 写入图像路径

        a3 = int(sz1 / 2)  # x start
        b3 = int(sz1)  # x end
        c3 = 0  # y start
        d3 = int(sz2 / 2)  # y end
        cropImg3 = image[a3:b3, c3:d3]  # 裁剪图像
        cv2.imwrite(n+'-3'+s, cropImg3)  # 写入图像路径

        a4 = int(sz1 / 2)  # x start
        b4 = int(sz1)  # x end
        c4 = int(sz2 / 2)  # y start
        d4 = int(sz2)  # y end
        cropImg4 = image[a4:b4, c4:d4]  # 裁剪图像
        cv2.imwrite(n+'-4'+s, cropImg4)  # 写入图像路径

        a5 = int(sz1 / 4)  # x start
        b5 = int(sz1*3/4)  # x end
        c5 = int(sz2 / 4)  # y start
        d5 = int(sz2*3/4)  # y end
        cropImg5 = image[a5:b5, c5:d5]  # 裁剪图像
        cv2.imwrite(n + '-5' + s, cropImg5)  # 写入图像路径


if __name__ == '__main__':
    filepath = '../data/chap/chap2/datasets/03/'  # 源图像
    XML_DIR = '../data/chap/chap2/datasets/03-1/'  # 源xml
    destpath = '../data/chap/chap2/datasets/03-2/'  # 保存裁剪图片的文件夹地址
    AUG_XML_DIR = '../data/chap/chap2/datasets/03-3/'  # 保存xml地址
    new_bndbox_list = []
    for name in tqdm(os.listdir(XML_DIR), desc='Processing'):
        bndbox = read_xml_annotation(XML_DIR, name)  # 一张图片的边界框信息
        og_xml = open(os.path.join(XML_DIR, name))
        tree = ET.parse(og_xml)
        elem = tree.find('filename')
        elem.text = (name[:-4] + '.jpg')  # 每一次将xml文件中的filename节点信息进行更改
        tree.write(os.path.join(AUG_XML_DIR, name))
        img = Image.open(os.path.join(filepath, name[:-4] + '.jpg'))
        img = np.asarray(img)
        i_h, i_w, _= img.shape
        # 分别依据左上右上左下右下和中心位置对满足要求的目标边界框进行分别提取,并保存在各自对应的xml文件中,若目标边界框在裁剪图片边缘上,则舍弃
        for i in range(len(bndbox)):
            x1=bndbox[i][0]
            y1=bndbox[i][1]
            x2=bndbox[i][2]
            y2=bndbox[i][3]
            n_x1 = 0
            n_y1 = 0
            n_x2 = 0
            n_y2 = 0

            if x1 < i_w/2 and x2 < i_w/2 and y1 < i_h/2 and y2 < i_h/2:
                n_x1 = x1
                n_y1 = y1
                n_x2 = x2
                n_y2 = y2
            new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
        epoch=1
        # 存储变化后的XML
        change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
        # 将new_bndbox_list清零
        new_bndbox_list = []
        for i in range(len(bndbox)):
            x1=bndbox[i][0]
            y1=bndbox[i][1]
            x2=bndbox[i][2]
            y2=bndbox[i][3]
            n_x1 = 0
            n_y1 = 0
            n_x2 = 0
            n_y2 = 0

            if x1 > i_w/2 and x2 > i_w/2 and y1 < i_h/2 and y2 < i_h/2:
                n_x1 = x1
                n_y1 = y1
                n_x2 = x2
                n_y2 = y2
            new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
        epoch=2
        # 存储变化后的XML
        change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
        # 将new_bndbox_list清零
        new_bndbox_list = []
        for i in range(len(bndbox)):
            x1=bndbox[i][0]
            y1=bndbox[i][1]
            x2=bndbox[i][2]
            y2=bndbox[i][3]
            n_x1 = 0
            n_y1 = 0
            n_x2 = 0
            n_y2 = 0

            if x1 < i_w/2 and x2 < i_w/2 and y1 > i_h/2 and y2 > i_h/2:
                n_x1 = x1
                n_y1 = y1
                n_x2 = x2
                n_y2 = y2
            new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
        epoch=3
        # 存储变化后的XML
        change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
        # 将new_bndbox_list清零
        new_bndbox_list = []
        for i in range(len(bndbox)):
            x1=bndbox[i][0]
            y1=bndbox[i][1]
            x2=bndbox[i][2]
            y2=bndbox[i][3]
            n_x1 = 0
            n_y1 = 0
            n_x2 = 0
            n_y2 = 0

            if x1 > i_w/2 and x2 > i_w/2 and y1 > i_h/2 and y2 > i_h/2:
                n_x1 = x1
                n_y1 = y1
                n_x2 = x2
                n_y2 = y2
            new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
        epoch=4
        # 存储变化后的XML
        change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
        # 将new_bndbox_list清零
        new_bndbox_list = []
        for i in range(len(bndbox)):
            x1=bndbox[i][0]
            y1=bndbox[i][1]
            x2=bndbox[i][2]
            y2=bndbox[i][3]
            n_x1 = 0
            n_y1 = 0
            n_x2 = 0
            n_y2 = 0

            if x1 > i_w/4 and x2 > i_w/4 and y1 > i_h/4 and y2 > i_h/4 and x1 < i_w*3/4 and x2 < i_w*3/4 and y1 < i_h*3/4 and y2 < i_h*3/4:
                n_x1 = x1
                n_y1 = y1
                n_x2 = x2
                n_y2 = y2
            new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
        epoch=5
        # 存储变化后的XML
        change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR, str(name[:-4]) + '_' + str(epoch))
    CropImage4File(filepath, destpath)  # 裁剪图片保存

由于在csdn中并没有找寻到相关裁剪五块图像并使xml文档目标边界框信息跟随裁剪变化的文章,所以研究了一下相关数据扩增的内容,本代码可以完成对目标图片的比例裁剪,以及相应的xml裁剪

效果如下:

 

通过验证每个部分的边界框数量可以知道本代码是可行的

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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值