yolo에서 라벨링으로 생성한 txt 파일을 tensorflow에 맞는 xml 파일로 변환하기

우선 윈도우와 리눅스에서 아나콘다로 실행하여 본 결과 모두 정상 작동을 한다.

경로 작성시 window에서

아래와 같은 오류가 표시 될 시 경로 앞에 r을 붙여주면 된다.
ex) IMG_PATH = r"D:\VisualStudio\WorkSpace\data\img-GBS-IMG"

txt 파일에서 xml 파일로 변환시에 xml파일 이름 앞에 경로가 추가되는 오류는 고쳤으나 저장 경로가 제대로 지정되지 않는 오류는 아직 수정하지 못 하였다.
그러니 따로 디렉토리를 생성하고 그 안에서 이 코드를 돌리는 것이 나을 것이다.

코드에서 IMG_PATH와 txt_folder는 각각 .jpg, .txt 파일만 넣어 두는 것이 나을 것이다.
한 폴더 안에 jpg, txt 파일이 같이 있으면 파일 정보를 읽어 오는데 오류가 발생할 것이다.(능력자 분은 고쳐서 사용 하시면 됩니다...)



from lxml import etree
from PIL import Image
import csv
import os
import io

IMG_PATH = r"D:\VisualStudio\WorkSpace\data\img-GBS-IMG"
fw = os.listdir(IMG_PATH)
#path of save xml file
save_path = r"D:\VisualStudio\WorkSpace\data\annotations"

# txt_folder is txt file root that using darknet rectbox
txt_folder = r"D:\VisualStudio\WorkSpace\data\img-GBS-TXT"

#edit your label set
labels = ['green','brown','sugar']

def csvread(fn):
    with open(fn, 'r') as csvfile:
      list_arr = []
      reader = csv.reader(csvfile, delimiter=' ')

      for row in reader:
        list_arr.append(row)
    return list_arr

def convert_label(txt_file):
    if((txt_file[0]) == str(0)):
      label = 'green'
    elif((txt_file[0]) == str(1)):
      label = 'brown'
    elif((txt_file[0]) == str(2)):
      label = 'sugar'

    return label

#core code = convert the yolo txt file to the x_min, x_max..

def extract_coor(txt_file, img_width, img_height):
    x_rect_mid = float(txt_file[1])
    y_rect_mid = float(txt_file[2])
    width_rect = float(txt_file[3])
    height_rect = float(txt_file[4])

    x_min_rect = ((2 * x_rect_mid * img_width) - (width_rect * img_width)) / 2
    x_max_rect = ((2 * x_rect_mid * img_width) + (width_rect * img_width)) / 2
    y_min_rect = ((2 * y_rect_mid * img_height) - (height_rect * img_height)) / 2
    y_max_rect = ((2 * y_rect_mid * img_height) + (height_rect * img_height)) / 2

    return x_min_rect, x_max_rect, y_min_rect, y_max_rect

for line in fw:
    root = etree.Element("annotation")

    #try debug to check your path
    img_style = IMG_PATH.split('/')[-1]
    img_name = line
    image_info = IMG_PATH + "/" + line
    img_txt_root = txt_folder + "/" + line[:-4]
    txt = ".txt"

    txt_path = img_txt_root + txt
    txt_file = csvread(txt_path)
    #####################################

    #read the image information
    img_size = Image.open(image_info).size

    img_width = img_size[0]
    img_height = img_size[1]
    img_depth = Image.open(image_info).layers
    #######################################

    folder = etree.Element("folder")
    folder.text = "%s" % (img_style)

    filename = etree.Element("filename")
    filename.text = "%s" % (img_name)

    path = etree.Element("path")
    path.text = "%s" % (IMG_PATH)

    source = etree.Element("source")
    ##################source - element##################
    source_database = etree.SubElement(source, "database")
    source_database.text = "Unknown"
    ####################################################

    size = etree.Element("size")
    ####################size - element##################
    image_width = etree.SubElement(size, "width")
    image_width.text = "%d" % (img_width)

    image_height = etree.SubElement(size, "height")
    image_height.text = "%d" % (img_height)

    image_depth = etree.SubElement(size, "depth")
    image_depth.text = "%d" % (img_depth)
    ####################################################

    segmented = etree.Element("segmented")
    segmented.text = "0"

    root.append(folder)
    root.append(filename)
    root.append(path)
    root.append(source)
    root.append(size)
    root.append(segmented)

    for ii in range(len(txt_file)):

        label = convert_label(txt_file[ii][0])
        x_min_rect, x_max_rect, y_min_rect, y_max_rect = extract_coor(
            txt_file[ii], img_width, img_height)

        object = etree.Element("object")
        ####################object - element##################
        name = etree.SubElement(object, "name")
        name.text = "%s" % (label)

        pose = etree.SubElement(object, "pose")
        pose.text = "Unspecified"

        truncated = etree.SubElement(object, "truncated")
        truncated.text = "0"

        difficult = etree.SubElement(object, "difficult")
        difficult.text = "0"

        bndbox = etree.SubElement(object, "bndbox")
        #####sub_sub########
        xmin = etree.SubElement(bndbox, "xmin")
        xmin.text = "%d" % (x_min_rect)
        ymin = etree.SubElement(bndbox, "ymin")
        ymin.text = "%d" % (y_min_rect)
        xmax = etree.SubElement(bndbox, "xmax")
        xmax.text = "%d" % (x_max_rect)
        ymax = etree.SubElement(bndbox, "ymax")
        ymax.text = "%d" % (y_max_rect)
        #####sub_sub########

        root.append(object)
        ####################################################

    file_output = etree.tostring(root, pretty_print=True, encoding='UTF-8')
    # print(file_output.decode('utf-8'))
    ff = io.open('%s.xml' % (img_name[:-4]), 'w', encoding="utf-8")
    ff.write(file_output.decode('utf-8'))



🔻 아 그리고 추가로 encoding 시에 오류가 발생하여 마지막에서 두번째 줄 수정

ff = open('%s.xml' % (img_name[:-4]), 'w', encoding="utf-8")
에서
ff = io.open('%s.xml' % (img_name[:-4]), 'w', encoding="utf-8")
로 수정

댓글

이 블로그의 인기 게시물

python 에서 resize 이벤트 만들기

PLC와 아두이노 그리고 온도 센서를 이용하여 램프 점등

C언어로 로또 프로그램 만들기