目标检测数据增强--roi区域xml数据增强

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)


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

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值