数据集xml标注截取数据增强

import numpy as np
import cv2
import xml.etree.ElementTree as ET
import os
import shutil
import random

def translate(image, x, y):
    M = np.float32([[1, 0, x], [0, 1, y]])
    shifted = cv2.warpAffine(image, M, (image.shape[1], image.shape[0]))
    return shifted

def modify_xml(xml_path,x,y):
    f = open(xml_path)
    tree = ET.parse(f)
    root = tree.getroot()

    for obj in root.iter('object'):
        name = obj.find('name').text
        print('正在修改目标{}'.format(name))

        xmin = obj.find('bndbox').find('xmin').text
        obj.find('bndbox').find('xmin').text = str(int(xmin) + x)
        xmax = obj.find('bndbox').find('xmin').text
        obj.find('bndbox').find('xmax').text = str(int(xmax) + x)
        ymin = obj.find('bndbox').find('ymin').text
        obj.find('bndbox').find('ymin').text = str(int(ymin) + y)
        ymax = obj.find('bndbox').find('ymax').text
        obj.find('bndbox').find('ymax').text = str(int(ymax) + y)

        tree.write(xml_path,encoding="utf-8",xml_declaration=True)
def modify_xml2(xml_path,x,y,save_xml_path):
    f = open(xml_path)
    tree = ET.parse(f)
    root = tree.getroot()
    width, height = int(root.find('size').find('width').text), int(root.find('size').find('height').text)
    for obj in root.iter('object'):
        name = obj.find('name').text
        #print('正在修改目标{}'.format(name))

        xmin = obj.find('bndbox').find('xmin').text
        xmax = obj.find('bndbox').find('xmax').text
        ymin = obj.find('bndbox').find('ymin').text
        ymax = obj.find('bndbox').find('ymax').text
        if (int(xmin) + x)>0 and (int(xmax) + x)<width and (int(ymin) + y)>0 and (int(ymax) + y)<height:
            obj.find('bndbox').find('xmin').text = str(int(xmin) + x)

            obj.find('bndbox').find('xmax').text = str(int(xmax) + x)

            obj.find('bndbox').find('ymin').text = str(int(ymin) + y)

            obj.find('bndbox').find('ymax').text = str(int(ymax) + y)
        else:
            print('超出范围，放弃该图片')
            return False
    tree.write(save_xml_path, encoding="utf-8", xml_declaration=True)
    print('成功创建新图片数据')
    return True

new_xml = r'E:\Files\pycharm\data-processing\data\new_xml'
new_img = r'E:\Files\pycharm\data-processing\data\new_img'
img_dir = r'E:\Files\pycharm\data-processing\data\Images'
xml_dir = r'E:\Files\pycharm\data-processing\data\Annotations'

D =[random.randint(30,60),random.randint(-60,-30)]
xml_list = os.listdir(xml_dir)
for xml_file in xml_list:
    x, y = D[random.randint(0,1)],D[random.randint(0,1)]#随机偏移量
    print('预设偏移量为x {}  y {}'.format(x,y))
    xml_path = os.path.join(xml_dir,xml_file)
    flag = modify_xml2(xml_path,x,y,os.path.join(new_xml, xml_file))
    #flag = modify_xml(xml_path, x, y)
    if flag:
        img_file = xml_file.split('.')[0] + '.jpg'
        img = cv2.imread(os.path.join(img_dir,img_file))
        img_move = translate(img,x,y)
        cv2.imwrite(os.path.join(new_img, img_file), img_move)

import numpy as np
import cv2
import xml.etree.ElementTree as ET
import os
import shutil
import random

def re_size(img,x,y):
    new_img = cv2.resize(img, None, fx=x, fy=y, interpolation=cv2.INTER_CUBIC)
    return new_img

def modify_xml2(xml_path,x,y,save_xml_path):
    f = open(xml_path)
    tree = ET.parse(f)
    root = tree.getroot()
    width, height = int(root.find('size').find('width').text), int(root.find('size').find('height').text)
    for obj in root.iter('object'):
        name = obj.find('name').text
        #print('正在修改目标{}'.format(name))

        xmin = obj.find('bndbox').find('xmin').text
        xmax = obj.find('bndbox').find('xmax').text
        ymin = obj.find('bndbox').find('ymin').text
        ymax = obj.find('bndbox').find('ymax').text
        if (int(xmax)-int(xmin))>int(width/6) and (int(ymax)-int(ymin))>int(height/6):
            obj.find('bndbox').find('xmin').text = str(int(int(xmin)*x))

            obj.find('bndbox').find('xmax').text = str(int(int(xmax)*x))

            obj.find('bndbox').find('ymin').text = str(int(int(ymin)*y))

            obj.find('bndbox').find('ymax').text = str(int(int(ymax)*y))
        else:
            print('存在过小目标，放弃该图片')
            return False
    tree.write(save_xml_path, encoding="utf-8", xml_declaration=True)
    print('成功创建新图片数据')
    return True

new_xml = r'E:\Files\pycharm\data-processing\data\new_xml'
new_img = r'E:\Files\pycharm\data-processing\data\new_img'
img_dir = r'E:\Files\pycharm\data-processing\data\Images'
xml_dir = r'E:\Files\pycharm\data-processing\data\Annotations'


xml_list = os.listdir(xml_dir)
for xml_file in xml_list:
    x, y = random.uniform(0.5,2.0),random.uniform(0.5,2.0)#随机生成缩放比例
    print('预设缩放比例为x： {}  y： {}'.format(x,y))
    xml_path = os.path.join(xml_dir,xml_file)
    flag = modify_xml2(xml_path,x,y,os.path.join(new_xml, xml_file))
    #flag = modify_xml(xml_path, x, y)
    if flag:
        img_file = xml_file.split('.')[0] + '.jpg'
        img = cv2.imread(os.path.join(img_dir,img_file))
        img_move = re_size(img,x,y)
        cv2.imwrite(os.path.join(new_img, img_file), img_move)