数据处理实用脚本二

目标检测（倾斜检测框），利用预训练模型标注得到TXT文件，将TXT文件转化为XML文件，以供调整标注框。

# -*- coding: utf-8 -*-
"""
Created on Wed Mar  7 16:21:25 2018

@author: XW
"""

from xml.dom.minidom import Document  
import os  
import os.path  
from PIL import Image  
import math

ann_path = "/home/xw/data/VOCdevkit/ship/inference_label/"  
img_path = "/home/xw/data/VOCdevkit/ship/JPEGImages_detected/" 
xml_path = "/home/xw/data/VOCdevkit/ship/Annotations/"  
  
if not os.path.exists(xml_path):  
    os.mkdir(xml_path)  
  
def writeXml(tmp, imgname, w, h, objbuds, wxml):  
    doc = Document()  
     
    annotation = doc.createElement('annotation')  
    doc.appendChild(annotation)  
      
    folder = doc.createElement('folder')  
    annotation.appendChild(folder)  
    folder_txt = doc.createTextNode("ship")  
    folder.appendChild(folder_txt)  
  
    filename = doc.createElement('filename')  
    annotation.appendChild(filename)  
    filename_txt = doc.createTextNode(imgname)  
    filename.appendChild(filename_txt)  
      
    source = doc.createElement('source')  
    annotation.appendChild(source)  
  
    database = doc.createElement('database')  
    source.appendChild(database)  
    database_txt = doc.createTextNode("The VOC2007 Database")  
    database.appendChild(database_txt)  
  
    annotation_new = doc.createElement('annotation')  
    source.appendChild(annotation_new)  
    annotation_new_txt = doc.createTextNode("PASCAL VOC2007")  
    annotation_new.appendChild(annotation_new_txt)  
  
    image = doc.createElement('image')  
    source.appendChild(image)  
    image_txt = doc.createTextNode("flickr")  
    image.appendChild(image_txt)  
      
      
    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("3")  
    depth.appendChild(depth_txt)  
      
    segmented = doc.createElement('segmented')  
    annotation.appendChild(segmented)  
    segmented_txt = doc.createTextNode("0")  
    segmented.appendChild(segmented_txt)  
  
    for i in range(len(objbuds)):
        if len(objbuds[i]) != 0:    
            objbud = objbuds[i]
            if (int(objbud[2]) >= int(objbud[3])):
                tmp = objbud[3]
                objbud[3] = objbud[2]
                objbud[2] = tmp
                objbud[4] = str(round(((90 + int(objbud[4])) * math.pi / 180) % math.pi , 6))
            else:
                objbud[4] = str(round(((360 + int(objbud[4])) * math.pi / 180) % math.pi , 6))
                
            object_new = doc.createElement("object")  
            annotation.appendChild(object_new)

            type_1 = doc.createElement('type')
            object_new.appendChild(type_1)
            type_1_txt = doc.createTextNode('robndbox')
            type_1.appendChild(type_1_txt)
  
            name = doc.createElement('name')  
            object_new.appendChild(name)  
            name_txt = doc.createTextNode('ship')  
            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)  
        
            bndbox = doc.createElement('robndbox')  
            object_new.appendChild(bndbox)  
  
            xmin = doc.createElement('cx')  
            bndbox.appendChild(xmin)  
            xmin_txt = doc.createTextNode(objbud[0])  
            xmin.appendChild(xmin_txt)  
  
            ymin = doc.createElement('cy')  
            bndbox.appendChild(ymin)  
            ymin_txt = doc.createTextNode(objbud[1])  
            ymin.appendChild(ymin_txt)  
  
            xmax = doc.createElement('w')  
            bndbox.appendChild(xmax)  
            xmax_txt = doc.createTextNode(objbud[2])  
            xmax.appendChild(xmax_txt)  
  
            ymax = doc.createElement('h')  
            bndbox.appendChild(ymax)  
            ymax_txt = doc.createTextNode(objbud[3])  
            ymax.appendChild(ymax_txt)

            angle = doc.createElement('angle')
            bndbox.appendChild(angle)
            angle_txt = doc.createTextNode(objbud[4])
            angle.appendChild(angle_txt)
      
    tempfile = tmp + "test.xml"  
    with open(tempfile, "w") as f:  
        f.write(doc.toprettyxml(indent = '\t'))  
  
    rewrite = open(tempfile, "r")  
    lines = rewrite.read().split('\n')  
    newlines = lines[1:len(lines)-1]  
      
    fw = open(wxml, "w")  
    for i in range(0, len(newlines)):  
        fw.write(newlines[i] + '\n')  
      
    fw.close()  
    rewrite.close()  
    os.remove(tempfile)  
    return  
  
for files in os.walk(ann_path):  
    temp = "/home/xw/data/VOCdevkit/ship/temp"  
    if not os.path.exists(temp):  
        os.mkdir(temp)  
    for file in files[2]:  
        objs = []
        print (file + "-->start!")  
        img_name = os.path.splitext(file)[0] + '.jpg'  
        fileimgpath = img_path + img_name  
        im=Image.open(fileimgpath)    
        width= int(im.size[0])  
        height= int(im.size[1])  
          
        filelabel = open(ann_path + file, "r")  
        lines = filelabel.readlines()  
        for line in lines:
            line = line.split()
            objs.append(line[:len(line)])
        filename = xml_path + os.path.splitext(file)[0] + '.xml'  
        writeXml(temp, img_name, width, height, objs, filename)  
    os.rmdir(temp)