FasterRCNN畫框小程序——VOC2007格式(python)

用Faster RCNN訓練數據,手動標註好辛苦,好在我的數據是二值的,找到對應的像素值爲255的(x_min,y_min,x_max,y_max)然後替換xml中的對應值就好了

首先要有一個VOC2007格式的xml文件,在這個基礎上進行修改

學python的日子加起來不超過24小時,編這個小程序花了一天的時間,希望有個好結果。加油

# get the gt's x_min,y_min,x_max,y_max and replace the xml's values
# first traverse the image to find the pixel==255's position
# second find the x_min,y_min,x_max,y_max
# third read the xml file and replace the values
# by LYS 6/28/2017 
from PIL import Image,ImageDraw
import xml.etree.cElementTree as ET
import os

# according to the image's path to locate the position
def get_positions(image_path):
    im = Image.open(image_path).convert('L')#open the image and convert to gray image
    draw = ImageDraw.Draw(im)
    width = im.size[0]
    height = im.size[1]
    x = []
    y = []
    for w in range(0, width):
        for h in range(0, height):
	    pixel = im.getpixel((w, h))#get the pixel at position (w,h)
	    if pixel>=200:
                #dynamic allocate the list's memory
                x.append(w) 
                y.append(h)

    if(len(x)!=0 and len(y)!=0):
    #remove false elements
        for i in x[::-1]:
            if i < 10 or i >300:
                x.remove(i)
                continue
    
        for j in y[::-1]:
            if j < 10 or j >300:
                y.remove(j)
                continue 
    #find the max,min in x,y
        x_min = min(x) 
        y_min = min(y) 
        x_max = max(x) 
        y_max = max(y) 
        box = [x_min, y_min, x_max, y_max]
        draw.rectangle(box,outline=255)
        im.save(image_path)
    #call the modify xml file's function
        modify_xmlfile(box, xml_path)


#modify the xml file
def modify_xmlfile(box):
    #the file is an already exit xml file which is corresponding to the VOC2007 format
    tree = ET.ElementTree(file='003042.xml') 
    for elem in tree.iter(tag = 'xmin'):
        elem.text = `box[0]` 
    for elem in tree.iter(tag = 'ymin'):
        elem.text = `box[1]` 
    for elem in tree.iter(tag = 'xmax'):
        elem.text = `box[2]` 
    for elem in tree.iter(tag = 'ymax'):
        elem.text = `box[3]` 
    for elem in tree.iter(tag = 'filename'):
   # just save the image name without the postfix .jpg
        elem.text = image_name.split('.')[0] 
    for elem in tree.iter(tag = 'width'):
        elem.text = `image.size[0]`
    for elem in tree.iter(tag = 'height'):
        elem.text = `image.size[1]`
    #write the xml
    tree.write(xml_path)

if __name__ == '__main__':
# the original image path
    image_read_path = '/home/lys/lys/MRI_data_train/gt1_9/'
# the xml file's target save path
    xml_save_path = '/home/lys/py-faster-rcnn/data/VOCdevkit2007/VOC2007/Annotations/'
    image_names = os.listdir(image_read_path) 
    for image_name in image_names:
        xml_name = image_name.split('.')[0] + '.xml'
        xml_path = os.path.join(xml_save_path, xml_name)
        image_path = os.path.join(image_read_path, image_name)
        get_positions(image_path)

      xml文件格式主要是文件名要和圖片名匹配,對應座標值要正確

發佈了32 篇原創文章 · 獲贊 16 · 訪問量 7萬+
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章