目標檢測數據擴增——翻轉和旋轉

原文鏈接:https://blog.csdn.net/Mr_health/article/details/82667538

標題HBB格式的數據擴增

該格式下目標被標註爲xmin,ymin,xmax,ymax

import os
import cv2
import xml.dom.minidom
from xml.dom.minidom import Document
import math
 
#獲取路徑下所有文件的完整路徑,用於讀取文件用
def GetFileFromThisRootDir(dir,ext = None):
  allfiles = []
  needExtFilter = (ext != None)
  for root,dirs,files in os.walk(dir):
    for filespath in files:
      filepath = os.path.join(root, filespath)
      extension = os.path.splitext(filepath)[1][1:]
      if needExtFilter and extension in ext:
        allfiles.append(filepath)
      elif not needExtFilter:
        allfiles.append(filepath)
  return allfiles
 
def limit_value(a,b):
    if a<1:
        a = 1
    if a>b:
        a = b-1
    return a
    
#讀取xml文件,xmlfile參數表示xml的路徑
def readXml(xmlfile):
    DomTree = xml.dom.minidom.parse(xmlfile)  
    annotation = DomTree.documentElement  
    sizelist = annotation.getElementsByTagName('size') #[<DOM Element: filename at 0x381f788>]  
    heights = sizelist[0].getElementsByTagName('height')
    height = int(heights[0].childNodes[0].data)
    widths =sizelist[0].getElementsByTagName('width')
    width = int(widths[0].childNodes[0].data)
    depths = sizelist[0].getElementsByTagName('depth')
    depth = int(depths[0].childNodes[0].data)
    objectlist = annotation.getElementsByTagName('object')        
    bboxes = []
    for objects in objectlist:  
        namelist = objects.getElementsByTagName('name')  
        class_label = namelist[0].childNodes[0].data  
        bndbox = objects.getElementsByTagName('bndbox')[0]     
        x1_list = bndbox.getElementsByTagName('xmin')  
        x1 = int(float(x1_list[0].childNodes[0].data))  
        y1_list = bndbox.getElementsByTagName('ymin') 
        y1 = int(float(y1_list[0].childNodes[0].data)) 
        x2_list = bndbox.getElementsByTagName('xmax')  
        x2 = int(float(x2_list[0].childNodes[0].data))  
        y2_list = bndbox.getElementsByTagName('ymax')  
        y2 = int(float(y2_list[0].childNodes[0].data))
        #這裏我box的格式【xmin,ymin,xmax,ymax,classname】
        bbox = [x1,y1,x2,y2,class_label]
        bboxes.append(bbox)
    return bboxes,width,height,depth
 
 
#圖像旋轉用,裏面的angle是角度制的
def im_rotate(im,angle,center = None,scale = 1.0):
    h,w = im.shape[:2]
    if center is None:
        center = (w/2,h/2)
    M = cv2.getRotationMatrix2D(center,angle,scale)
    im_rot = cv2.warpAffine(im,M,(w,h))
    return im_rot
 
 
def im_flip(im, method = 'H'):   #翻轉圖像    
    if method == 'H':  # Flipped Horizontally 水平翻轉
        im_flip = cv2.flip(im, 1)
    elif method == 'V':  # Flipped Vertically 垂直翻轉
        im_flip = cv2.flip(im, 0)
    #elif method == 'HV':# Flipped Horizontally & Vertically 水平垂直翻轉
    #    im_flip = cv2.flip(im, -1)    
    return im_flip
 
def flip_image(imgs_path, method, anno_new_dir, img_new_dir):
    for img_path in imgs_path :
        #讀取原圖像
        im = cv2.imread(img_path)     
        flip_img = im_flip(im, method) #翻轉
        (H,W,D) = flip_img.shape           #得到翻轉後的圖像的高、寬、深度,用於書寫xml
        file_name = os.path.basename(os.path.splitext(img_path)[0])     #得到原圖的名稱 
        ext = os.path.splitext(img_path)[-1] #得到原圖的後綴
        #保存翻轉後圖像 
        new_img_name = '%s_%s%s'%(method, file_name,ext)
        cv2.imwrite(os.path.join(img_new_dir, new_img_name), flip_img) #新的命名方式爲H/V+原圖名稱
        #讀取anno標籤數據,返回相應的信息
        anno = os.path.join(anno_path,'%s.xml'%file_name)
        [gts,w,h,d] =readXml(anno) 
        gt_new = []
        for gt in gts:
            x1 = gt[0] #xmin
            y1 = gt[1]  #ymin
            x2 = gt[2]   #xmax
            y2 = gt[3]   #ymax
            classname = str(gt[4])
            if method == 'H':
                x1 = w - 1 - x1 #xmax
                x2 = w - 1 - x2  #xmin            
                x1 = limit_value(x1,w)
                x2 = limit_value(x2,w)               
                gt_new.append([x2,y1,x1,y2,classname])
            elif method == 'V':  
                y1 = h - 1 - y1 #ymax 
                y2 = h - 1 - y2  #ymin               
                y1 = limit_value(y1,h)
                y2 = limit_value(y2,h)  
                gt_new.append([x1,y2,x2,y1,classname])
    writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)
                
def rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir):
    j=0 # 計數用
    angle_num = len(angles)
    for img_path in imgs_path :
        #讀取原圖像
        im = cv2.imread(img_path)     
        for i in range(angle_num):
            gt_new = []
            im_rot = im_rotate(im,angles[i]) #旋轉
            (H,W,D) = im_rot.shape           #得到旋轉後的圖像的高、寬、深度,用於書寫xml
            file_name = os.path.basename(os.path.splitext(img_path)[0])     #得到原圖的名稱
            ext = os.path.splitext(img_path)[-1] #得到原圖的後綴
            #保存旋轉後圖像
            new_img_name = 'P%s_%s%s'%(angles[i],file_name,ext)
            cv2.imwrite(os.path.join(img_new_dir, new_img_name),im_rot) #新的命名方式爲P+角度+原圖名稱
            #讀取anno標籤數據,返回相應的信息
            anno = os.path.join(anno_path,'%s.xml'%file_name)
            [gts,w,h,d] =readXml(anno) 
            #計算旋轉後gt框四點的座標變換
            [xc,yc] = [float(w)/2,float(h)/2]
            for gt in gts:
                #計算左上角點的變換
                x1 = (gt[0]-xc)*math.cos(angle_rad[i]) - (yc-gt[1])*math.sin(angle_rad[i]) + xc
                if int(x1)<=0 : x1=1.0
                if int(x1)>w-1 : x1=w-1                        
                y1 = yc - (gt[0]-xc)*math.sin(angle_rad[i]) - (yc-gt[1])*math.cos(angle_rad[i]) 
                if int(y1)<=0 : y1=1.0       
                if int(y1)>h-1 : y1=h-1 
                #計算右上角點的變換   
                x2 = (gt[2]-xc)*math.cos(angle_rad[i]) - (yc-gt[1])*math.sin(angle_rad[i]) + xc
                if int(x2)<=0 : x2=1.0
                if int(x2)>w-1 : x2=w-1                        
                y2 = yc - (gt[2]-xc)*math.sin(angle_rad[i]) - (yc-gt[1])*math.cos(angle_rad[i])
                if int(y2)<=0 : y2=1.0         
                if int(y2)>h-1 : y2=h-1 
                #計算左下角點的變換
                x3 = (gt[0]-xc)*math.cos(angle_rad[i]) - (yc-gt[3])*math.sin(angle_rad[i]) + xc
                if int(x3)<=0 : x3=1.0
                if int(x3)>w-1 : x3=w-1                        
                y3 = yc - (gt[0]-xc)*math.sin(angle_rad[i]) - (yc-gt[3])*math.cos(angle_rad[i])
                if int(y3)<=0 : y3=1.0         
                if int(y3)>h-1 : y3=h-1
                #計算右下角點的變換
                x4 = (gt[2]-xc)*math.cos(angle_rad[i]) - (yc-gt[3])*math.sin(angle_rad[i]) + xc
                if int(x4)<=0 : x4=1.0          
                if int(x4)>w-1 : x4=w-1
                y4 = yc - (gt[2]-xc)*math.sin(angle_rad[i]) - (yc-gt[3])*math.cos(angle_rad[i])
                if int(y4)<=0 : y4=1.0     
                if int(y4)>h-1 : y4=h-1
                xmin = min(x1,x2,x3,x4)
                xmax = max(x1,x2,x3,x4)
                ymin = min(y1,y2,y3,y4)
                ymax = max(y1,y2,y3,y4)
                #把因爲旋轉導致的特別小的 長線型的去掉
                #w_new = xmax-xmin+1
                #h_new = ymax-ymin+1
                #ratio1 = float(w_new)/h_new
                #ratio2 = float(h_new)/w_new
                #if(1.0/6.0<ratio1<6 and 1.0/6.0<ratio2<6 and w_new>9 and h_new>9):
                classname = str(gt[4])
                gt_new.append([xmin,ymin,xmax,ymax,classname])
                #寫出新的xml
                writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)
            j = j+1
            if j%100==0 : print ('----%s----'%j)
            
            
#寫xml文件,參數中tmp表示路徑,imgname是文件名(沒有尾綴)ps有尾綴也無所謂
def writeXml(tmp, imgname, w, h, d, bboxes):  
    doc = Document()  
    #owner  
    annotation = doc.createElement('annotation')  
    doc.appendChild(annotation)  
    #owner  
    folder = doc.createElement('folder')  
    annotation.appendChild(folder)  
    folder_txt = doc.createTextNode("VOC2007")  
    folder.appendChild(folder_txt)  
  
    filename = doc.createElement('filename')  
    annotation.appendChild(filename)  
    filename_txt = doc.createTextNode(imgname)  
    filename.appendChild(filename_txt)  
    #ones#  
    source = doc.createElement('source')  
    annotation.appendChild(source)  
  
    database = doc.createElement('database')  
    source.appendChild(database)  
    database_txt = doc.createTextNode("My Database")  
    database.appendChild(database_txt)  
  
    annotation_new = doc.createElement('annotation')  
    source.appendChild(annotation_new)  
    annotation_new_txt = doc.createTextNode("VOC2007")  
    annotation_new.appendChild(annotation_new_txt)  
  
    image = doc.createElement('image')  
    source.appendChild(image)  
    image_txt = doc.createTextNode("flickr")  
    image.appendChild(image_txt) 
    #owner
    owner = doc.createElement('owner')  
    annotation.appendChild(owner)  
  
    flickrid = doc.createElement('flickrid')  
    owner.appendChild(flickrid)  
    flickrid_txt = doc.createTextNode("NULL")  
    flickrid.appendChild(flickrid_txt) 
    
    ow_name = doc.createElement('name')  
    owner.appendChild(ow_name)  
    ow_name_txt = doc.createTextNode("idannel")  
    ow_name.appendChild(ow_name_txt)
    #onee#  
    #twos#  
    size = doc.createElement('size')  
    annotation.appendChild(size)  
  
    width = doc.createElement('width')  
    size.appendChild(width)  
    width_txt = doc.createTextNode(str(w))  
    width.appendChild(width_txt)  
  
    height = doc.createElement('height')  
    size.appendChild(height)  
    height_txt = doc.createTextNode(str(h))  
    height.appendChild(height_txt)  
  
    depth = doc.createElement('depth')  
    size.appendChild(depth)  
    depth_txt = doc.createTextNode(str(d))  
    depth.appendChild(depth_txt)  
    #twoe#  
    segmented = doc.createElement('segmented')  
    annotation.appendChild(segmented)  
    segmented_txt = doc.createTextNode("0")  
    segmented.appendChild(segmented_txt)  
  
    for bbox in bboxes:
        #threes#  
        object_new = doc.createElement("object")  
        annotation.appendChild(object_new)  
        
        name = doc.createElement('name')  
        object_new.appendChild(name)  
        name_txt = doc.createTextNode(str(bbox[4]))  
        name.appendChild(name_txt)  
  
        pose = doc.createElement('pose')  
        object_new.appendChild(pose)  
        pose_txt = doc.createTextNode("Unspecified")  
        pose.appendChild(pose_txt)  
  
        truncated = doc.createElement('truncated')  
        object_new.appendChild(truncated)  
        truncated_txt = doc.createTextNode("0")  
        truncated.appendChild(truncated_txt)  
  
        difficult = doc.createElement('difficult')  
        object_new.appendChild(difficult)  
        difficult_txt = doc.createTextNode("0")  
        difficult.appendChild(difficult_txt)  
        #threes-1#  
        bndbox = doc.createElement('bndbox')  
        object_new.appendChild(bndbox)  
  
        xmin = doc.createElement('xmin')  
        bndbox.appendChild(xmin)  
        xmin_txt = doc.createTextNode(str(bbox[0]))  
        xmin.appendChild(xmin_txt)  
  
        ymin = doc.createElement('ymin')  
        bndbox.appendChild(ymin)  
        ymin_txt = doc.createTextNode(str(bbox[1]))
        ymin.appendChild(ymin_txt)    
  
        xmax = doc.createElement('xmax')  
        bndbox.appendChild(xmax)  
        xmax_txt = doc.createTextNode(str(bbox[2]))
        xmax.appendChild(xmax_txt)  
        
        ymax = doc.createElement('ymax')  
        bndbox.appendChild(ymax)  
        ymax_txt = doc.createTextNode(str(bbox[3]))
        ymax.appendChild(ymax_txt)  
        
        print(bbox[0],bbox[1],bbox[2],bbox[3])
 
    xmlname = os.path.splitext(imgname)[0]
    tempfile = tmp +"/%s.xml"%xmlname
    with open(tempfile, 'wb') as f:
        f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))
    return  
 
 
if __name__ == '__main__': 
    
    #數據路徑
    root = '/home/yantianwang/clone/haha'
    img_dir = root + '/images' 
    anno_path = root + '/xml'
    imgs_path=GetFileFromThisRootDir(img_dir)#返回每一張原圖的路徑
    AUG = 'Flip'  #數據擴增的方式 Rotate代表旋轉,Flip表示翻轉
    
    #存儲新的anno位置
    anno_new_dir = os.path.join(root, AUG , 'xml') 
    if not os.path.isdir(anno_new_dir):
        os.makedirs(anno_new_dir)
    #擴增後圖片保存的位置
    img_new_dir = os.path.join(root, AUG , 'images') 
    if not os.path.isdir(img_new_dir):
        os.makedirs(img_new_dir)
    
    if AUG == 'Rotate':
        #旋轉角的大小,正數表示逆時針旋轉
        angles = [5,90,180,270,355]#角度im_rotate用到的是角度制
        angle_rad = [angle*math.pi/180.0 for angle in angles] #cos三角函數裏要用到弧度制的      
        #開始旋轉
        rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir)
    elif AUG == 'Flip':
        method = 'H'
        flip_image(imgs_path, method, anno_new_dir, img_new_dir)

標題OBB格式的數據擴增

該格式下目標被標註爲x0,y0,x1,y1,x2,y2,x3,y3

import os
import cv2
import xml.dom.minidom
from xml.dom.minidom import Document
import math
import pdb
#獲取路徑下所有文件的完整路徑,用於讀取文件用
def GetFileFromThisRootDir(dir,ext = None):
  allfiles = []
  needExtFilter = (ext != None)
  for root,dirs,files in os.walk(dir):
    for filespath in files:
      filepath = os.path.join(root, filespath)
      extension = os.path.splitext(filepath)[1][1:]
      if needExtFilter and extension in ext:
        allfiles.append(filepath)
      elif not needExtFilter:
        allfiles.append(filepath)
  return allfiles
 
def limit_value(a,b):
    if a<1:
        a = 1
    if a>b:
        a = b-1
    return a
 
#讀取xml文件,xmlfile參數表示xml的路徑
def readXml(xmlfile):
    DomTree = xml.dom.minidom.parse(xmlfile)  
    annotation = DomTree.documentElement  
    sizelist = annotation.getElementsByTagName('size') #[<DOM Element: filename at 0x381f788>]  
    heights = sizelist[0].getElementsByTagName('height')
    height = int(heights[0].childNodes[0].data)
    widths =sizelist[0].getElementsByTagName('width')
    width = int(widths[0].childNodes[0].data)
    depths = sizelist[0].getElementsByTagName('depth')
    depth = int(depths[0].childNodes[0].data)
    objectlist = annotation.getElementsByTagName('object')        
    bboxes = []
    for objects in objectlist:  
        namelist = objects.getElementsByTagName('name')  
        class_label = namelist[0].childNodes[0].data  
        bndbox = objects.getElementsByTagName('bndbox')[0]  
        
        x0_list = bndbox.getElementsByTagName('x0')  
        x0 = int(float(x0_list[0].childNodes[0].data))  
        y0_list = bndbox.getElementsByTagName('y0') 
        y0 = int(float(y0_list[0].childNodes[0].data)) 
               
        x1_list = bndbox.getElementsByTagName('x1')  
        x1 = int(float(x1_list[0].childNodes[0].data))  
        y1_list = bndbox.getElementsByTagName('y1') 
        y1 = int(float(y1_list[0].childNodes[0].data)) 
 
        x2_list = bndbox.getElementsByTagName('x2')  
        x2 = int(float(x2_list[0].childNodes[0].data))  
        y2_list = bndbox.getElementsByTagName('y2')  
        y2 = int(float(y2_list[0].childNodes[0].data))
        
        x3_list = bndbox.getElementsByTagName('x3')  
        x3 = int(float(x3_list[0].childNodes[0].data))  
        y3_list = bndbox.getElementsByTagName('y3')  
        y3 = int(float(y3_list[0].childNodes[0].data))
        #這裏我box的格式【xmin,ymin,xmax,ymax,classname】
        bbox = [x0,y0,x1,y1,x2,y2,x3,y3,class_label]
        bboxes.append(bbox)
    return bboxes,width,height,depth
 
 
#圖像旋轉用,裏面的angle是角度制的
def im_rotate(im,angle,center = None,scale = 1.0):
    h,w = im.shape[:2]
    if center is None:
        center = (w/2,h/2)
    M = cv2.getRotationMatrix2D(center,angle,scale)
    im_rot = cv2.warpAffine(im,M,(w,h))
    return im_rot
 
 
def im_flip(im, method = 'H'):   #翻轉圖像    
    if method == 'H':  # Flipped Horizontally 水平翻轉
        im_flip = cv2.flip(im, 1)
    elif method == 'V':  # Flipped Vertically 垂直翻轉
        im_flip = cv2.flip(im, 0)
    #elif method == 'HV':# Flipped Horizontally & Vertically 水平垂直翻轉
    #    im_flip = cv2.flip(im, -1)    
    return im_flip
 
def flip_image(imgs_path, method, anno_new_dir, img_new_dir):
    for img_path in imgs_path :
        #讀取原圖像
        im = cv2.imread(img_path)     
        flip_img = im_flip(im, method) #翻轉
        (H,W,D) = flip_img.shape           #得到翻轉後的圖像的高、寬、深度,用於書寫xml
        file_name = os.path.basename(os.path.splitext(img_path)[0])     #得到原圖的名稱 
        ext = os.path.splitext(img_path)[-1] #得到原圖的後綴
        #保存翻轉後圖像 
        new_img_name = '%s_%s%s'%(method, file_name,ext)
        cv2.imwrite(os.path.join(img_new_dir, new_img_name), flip_img) #新的命名方式爲H/V+原圖名稱
        #讀取anno標籤數據,返回相應的信息
        anno = os.path.join(anno_path,'%s.xml'%file_name)
        [gts,w,h,d] =readXml(anno) 
        gt_new = []
        for gt in gts:
            x0 = gt[0]
            y0 = gt[1]  
            x1 = gt[2]   
            y1 = gt[3]  
            x2 = gt[4] 
            y2 = gt[5]  
            x3 = gt[6]   
            y3 = gt[7]  
            classname = gt[-1]
            if method == 'H':
                x0 = w - x0   
                x1 = w - x1
                x2 = w - x2  
                x3 = w - x3         
                x0 = limit_value(x0,w)
                x1 = limit_value(x1,w)
                x2 = limit_value(x2,w)
                x3 = limit_value(x3,w)
            elif method == 'V':  
                y0 = h - y0 
                y1 = h - y1 
                y2 = h - y2  
                y3 = h - y3         
                y0 = limit_value(y0,h)
                y1 = limit_value(y1,h)
                y2 = limit_value(y2,h)
                y3 = limit_value(y3,h)
            gt_new.append([x0,y0,x1,y1,x2,y2,x3,y3,classname])
        writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)
                
def rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir):
    j=0 # 計數用
    angle_num = len(angles)
    for img_path in imgs_path :
        #讀取原圖像
        im = cv2.imread(img_path)     
        for i in range(angle_num):
            gt_new = []
            im_rot = im_rotate(im,angles[i]) #旋轉
            (H,W,D) = im_rot.shape           #得到旋轉後的圖像的高、寬、深度,用於書寫xml
            file_name = os.path.basename(os.path.splitext(img_path)[0])     #得到原圖的名稱
            ext = os.path.splitext(img_path)[-1] #得到原圖的後綴
            #保存旋轉後圖像
            new_img_name = 'P%s_%s%s'%(angles[i],file_name,ext)
            cv2.imwrite(os.path.join(img_new_dir, new_img_name),im_rot) #新的命名方式爲P+角度+原圖名稱
            #讀取anno標籤數據,返回相應的信息
            anno = os.path.join(anno_path,'%s.xml'%file_name)
            [gts,w,h,d] =readXml(anno) 
            #計算旋轉後gt框四點的座標變換
            [xc,yc] = [float(w)/2,float(h)/2]
            for gt in gts:
                #計算左上角點的變換
                x1 = (gt[0]-xc)*math.cos(angle_rad[i]) - (yc-gt[1])*math.sin(angle_rad[i]) + xc
                if int(x1)<=0 : x1=1
                if int(x1)>w-1 : x1=w-1                        
                y1 = yc - (gt[0]-xc)*math.sin(angle_rad[i]) - (yc-gt[1])*math.cos(angle_rad[i]) 
                if int(y1)<=0 : y1=1.0       
                if int(y1)>h-1 : y1=h-1 
                #計算右上角點的變換   
                x2 = (gt[2]-xc)*math.cos(angle_rad[i]) - (yc-gt[3])*math.sin(angle_rad[i]) + xc
                if int(x2)<=0 : x2=1
                if int(x2)>w-1 : x2=w-1                        
                y2 = yc - (gt[2]-xc)*math.sin(angle_rad[i]) - (yc-gt[3])*math.cos(angle_rad[i])
                if int(y2)<=0 : y2=1         
                if int(y2)>h-1 : y2=h-1 
                #計算左下角點的變換
                x3 = (gt[4]-xc)*math.cos(angle_rad[i]) - (yc-gt[5])*math.sin(angle_rad[i]) + xc
                if int(x3)<=0 : x3=1
                if int(x3)>w-1 : x3=w-1                        
                y3 = yc - (gt[4]-xc)*math.sin(angle_rad[i]) - (yc-gt[5])*math.cos(angle_rad[i])
                if int(y3)<=0 : y3=1         
                if int(y3)>h-1 : y3=h-1
                #計算右下角點的變換
                x4 = (gt[6]-xc)*math.cos(angle_rad[i]) - (yc-gt[7])*math.sin(angle_rad[i]) + xc
                if int(x4)<=0 : x4=1          
                if int(x4)>w-1 : x4=w-1
                y4 = yc - (gt[6]-xc)*math.sin(angle_rad[i]) - (yc-gt[7])*math.cos(angle_rad[i])
                if int(y4)<=0 : y4=1     
                if int(y4)>h-1 : y4=h-1
                #xmin = min(x1,x2,x3,x4)
                #xmax = max(x1,x2,x3,x4)
                #ymin = min(y1,y2,y3,y4)
                #ymax = max(y1,y2,y3,y4)
                #把因爲旋轉導致的特別小的 長線型的去掉
                #w_new = xmax-xmin+1
                #h_new = ymax-ymin+1
                #ratio1 = float(w_new)/h_new
                #ratio2 = float(h_new)/w_new
                #if(1.0/6.0<ratio1<6 and 1.0/6.0<ratio2<6 and w_new>9 and h_new>9):
                classname = gt[-1]
                gt_new.append([int(x1),int(y1),int(x2),int(y2),int(x3),int(y3),int(x4),int(y4),classname])
                #寫出新的xml
            writeXml(anno_new_dir, new_img_name, W, H, D, gt_new)
            j = j+1
            if j%100==0 : print ('----%s----'%j)
            
            
#寫xml文件,參數中tmp表示路徑,imgname是文件名(沒有尾綴)ps有尾綴也無所謂
def writeXml(tmp, imgname, w, h, d, bboxes):  
    doc = Document()  
    #owner  
    annotation = doc.createElement('annotation')  
    doc.appendChild(annotation)  
    #owner  
    folder = doc.createElement('folder')  
    annotation.appendChild(folder)  
    folder_txt = doc.createTextNode("VOC2007")  
    folder.appendChild(folder_txt)  
  
    filename = doc.createElement('filename')  
    annotation.appendChild(filename)  
    filename_txt = doc.createTextNode(imgname)  
    filename.appendChild(filename_txt)  
    #ones#  
    source = doc.createElement('source')  
    annotation.appendChild(source)  
  
    database = doc.createElement('database')  
    source.appendChild(database)  
    database_txt = doc.createTextNode("My Database")  
    database.appendChild(database_txt)  
  
    annotation_new = doc.createElement('annotation')  
    source.appendChild(annotation_new)  
    annotation_new_txt = doc.createTextNode("VOC2007")  
    annotation_new.appendChild(annotation_new_txt)  
  
    image = doc.createElement('image')  
    source.appendChild(image)  
    image_txt = doc.createTextNode("flickr")  
    image.appendChild(image_txt) 
    #owner
    owner = doc.createElement('owner')  
    annotation.appendChild(owner)  
  
    flickrid = doc.createElement('flickrid')  
    owner.appendChild(flickrid)  
    flickrid_txt = doc.createTextNode("NULL")  
    flickrid.appendChild(flickrid_txt) 
    
    ow_name = doc.createElement('name')  
    owner.appendChild(ow_name)  
    ow_name_txt = doc.createTextNode("idannel")  
    ow_name.appendChild(ow_name_txt)
    #onee#  
    #twos#  
    size = doc.createElement('size')  
    annotation.appendChild(size)  
  
    width = doc.createElement('width')  
    size.appendChild(width)  
    width_txt = doc.createTextNode(str(w))  
    width.appendChild(width_txt)  
  
    height = doc.createElement('height')  
    size.appendChild(height)  
    height_txt = doc.createTextNode(str(h))  
    height.appendChild(height_txt)  
  
    depth = doc.createElement('depth')  
    size.appendChild(depth)  
    depth_txt = doc.createTextNode(str(d))  
    depth.appendChild(depth_txt)  
    #twoe#  
    segmented = doc.createElement('segmented')  
    annotation.appendChild(segmented)  
    segmented_txt = doc.createTextNode("0")  
    segmented.appendChild(segmented_txt)  
  
    for bbox in bboxes:
        #threes#  
        object_new = doc.createElement("object")  
        annotation.appendChild(object_new)  
        
        name = doc.createElement('name')  
        object_new.appendChild(name)  
        name_txt = doc.createTextNode(bbox[-1])  
        name.appendChild(name_txt)  
  
        pose = doc.createElement('pose')  
        object_new.appendChild(pose)  
        pose_txt = doc.createTextNode("Unspecified")  
        pose.appendChild(pose_txt)  
  
        truncated = doc.createElement('truncated')  
        object_new.appendChild(truncated)  
        truncated_txt = doc.createTextNode("0")  
        truncated.appendChild(truncated_txt)  
  
        difficult = doc.createElement('difficult')  
        object_new.appendChild(difficult)  
        difficult_txt = doc.createTextNode("0")  
        difficult.appendChild(difficult_txt)  
        #threes-1#  
        bndbox = doc.createElement('bndbox')  
        object_new.appendChild(bndbox)  
  
        x0 = doc.createElement('x0')  
        bndbox.appendChild(x0)  
        x0_txt = doc.createTextNode(str(bbox[0]))
        x0.appendChild(x0_txt)  
  
        y0 = doc.createElement('y0')  
        bndbox.appendChild(y0)  
        y0_txt = doc.createTextNode(str(bbox[1]))
        y0.appendChild(y0_txt)    
  
        x1 = doc.createElement('x1')  
        bndbox.appendChild(x1)  
        x1_txt = doc.createTextNode(str(bbox[2]))
        x1.appendChild(x1_txt)  
  
        y1 = doc.createElement('y1')  
        bndbox.appendChild(y1)  
        y1_txt = doc.createTextNode(str(bbox[3]))
        y1.appendChild(y1_txt) 
        
        x2 = doc.createElement('x2')  
        bndbox.appendChild(x2)  
        x2_txt = doc.createTextNode(str(bbox[4]))
        x2.appendChild(x2_txt)  
  
        y2 = doc.createElement('y2')  
        bndbox.appendChild(y2)  
        y2_txt = doc.createTextNode(str(bbox[5]))
        y2.appendChild(y2_txt)
 
        x3 = doc.createElement('x3')  
        bndbox.appendChild(x3)  
        x3_txt = doc.createTextNode(str(bbox[6]))
        x3.appendChild(x3_txt)  
  
        y3 = doc.createElement('y3')  
        bndbox.appendChild(y3)  
        y3_txt = doc.createTextNode(str(bbox[7]))
        y3.appendChild(y3_txt)  
        
        #print(float(int(bbox[0])),float(int(bbox[1])),float(int(bbox[2])),float(int(bbox[3])))
 
    xmlname = os.path.splitext(imgname)[0]
    tempfile = tmp +"/%s.xml"%xmlname
    with open(tempfile, 'wb') as f:
        f.write(doc.toprettyxml(indent='\t', encoding='utf-8'))
    return  
 
 
if __name__ == '__main__': 
    
    #數據路徑
    root = '/home/yantianwang/lala/20190225/oldship/examplesplit'
    img_dir = root + '/images' 
    anno_path = root + '/obbxml'
    imgs_path=GetFileFromThisRootDir(img_dir)#返回每一張原圖的路徑
    AUG = 'Flip'  #數據擴增的方式 Rotate代表旋轉,Flip表示翻轉
    
    #存儲新的anno位置
    anno_new_dir = os.path.join(root, AUG , 'NewAnnotations') 
    #anno_new_path = root + '/NewAnnotations' 
    if not os.path.isdir(anno_new_dir):
        os.makedirs(anno_new_dir)
    #擴增後圖片保存的位置
    img_new_dir = os.path.join(root, AUG , 'Newimage') 
    if not os.path.isdir(img_new_dir):
        os.makedirs(img_new_dir)
    
    if AUG == 'Rotate':
        #旋轉角的大小,正數表示逆時針旋轉
        angles = [5,90,180,270,355]#角度im_rotate用到的是角度制
        angle_rad = [angle*math.pi/180.0 for angle in angles] #cos三角函數裏要用到弧度制的      
        #開始旋轉
        rotate_image(angles, angle_rad, imgs_path, anno_new_dir, img_new_dir)
    elif AUG == 'Flip':
        method = 'H'
        flip_image(imgs_path, method, anno_new_dir, img_new_dir)
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章