microsoft/TATR 학습용 xml파일 만들기

이전 포스팅의 TATR 학습 방법에 대해 추가 설명을 하겠습니다.

TATR 학습방법

• 디렉토리 구조

└─det_train_data
    │
    ├─images
    ├─train
    │─val
    │─ train_filelist.txt
    └─ val_filelist.txt

- images : Train, Valid Image Set
- train : Train에 사용될 xml파일
- val : valid에 사용될 xml파일
- train_filelist.txt : train폴더에 있는 xml파일 리스트
- val_filelist.txt : val폴더에 있는 xml파일 리스트

• Detection model

python main.py --data_type detection --config_file detection_config.json --data_root_dir ./det_train_data

• Recognition model

python main.py --data_type structure --config_file structure_config.json --data_root_dir ./rec_train_data

train, valid xml파일 만드는 방법

OCR관련 포스팅 내용 중 라벨링 프로그램에 대해 소개한 적이 있습니다. PPOCRLabel이라는 프로그램인데 이 프로그램으로 라벨링을 진행하면 Label.txt파일이 생성되고 구조는 아래와 같습니다.

train/001.jpg	[{"transcription": "TEMPORARY", "points": [[85, 1007], [817, 1016], [810, 2023], [80, 2016]], "difficult": false}, {"transcription": "TEMPORARY", "points": [[849, 1015], [1580, 1023], [1572, 2030], [837, 2016]], "difficult": false}]

이제 이 구조를 TATR 학습 데이터 구조에 맞는 xml파일을 생성 해야하는데 Label.txt파일을 통해 xml파일로 변환시키는 코드를 공유 하겠습니다.

import json
import xml.etree.ElementTree as ET
from xml.dom import minidom
import os
from PIL import Image
import argparse

def read_input_file(file_path):
    input_data = {}
    with open(file_path, 'r', encoding='utf-8') as file:
        lines = file.readlines()
        
        for line in lines:
            line = line.strip()
            if line:
                key, value = line.split('\t', 1)
                input_data[key] = json.loads(value)
    
    return input_data

def create_xml_annotation(filepath, objects, image_dir):
    filename = os.path.basename(filepath)
    annotation = ET.Element('annotation')
    
    folder = ET.SubElement(annotation, 'folder')
    filename_elem = ET.SubElement(annotation, 'filename')
    filename_elem.text = filename
    
    path = ET.SubElement(annotation, 'path')
    path.text = filename
    
    source = ET.SubElement(annotation, 'source')
    database = ET.SubElement(source, 'database')
    database.text = 'PubTables1M-Detection'
    
    # Read image size
    image_path = os.path.join(image_dir, filename)
    with Image.open(image_path) as img:
        width, height = img.size
    
    size = ET.SubElement(annotation, 'size')
    width_elem = ET.SubElement(size, 'width')
    width_elem.text = str(width)
    height_elem = ET.SubElement(size, 'height')
    height_elem.text = str(height)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'
    
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'
    
    for obj in objects:
        object_elem = ET.SubElement(annotation, 'object')
        
        name = ET.SubElement(object_elem, 'name')
        name.text = 'table'
        
        pose = ET.SubElement(object_elem, 'pose')
        pose.text = 'Frontal'
        
        truncated = ET.SubElement(object_elem, 'truncated')
        truncated.text = '0'
        
        difficult = ET.SubElement(object_elem, 'difficult')
        difficult.text = '0'
        
        occluded = ET.SubElement(object_elem, 'occluded')
        occluded.text = '0'
        
        bndbox = ET.SubElement(object_elem, 'bndbox')
        
        xmin = ET.SubElement(bndbox, 'xmin')
        ymin = ET.SubElement(bndbox, 'ymin')
        xmax = ET.SubElement(bndbox, 'xmax')
        ymax = ET.SubElement(bndbox, 'ymax')
        
        points = obj['points']
        xmin.text = str(min(point[0] for point in points))
        ymin.text = str(min(point[1] for point in points))
        xmax.text = str(max(point[0] for point in points))
        ymax.text = str(max(point[1] for point in points))
    
    return annotation

def prettify(elem):
    """Return a pretty-printed XML string for the Element."""
    rough_string = ET.tostring(elem, 'utf-8')
    reparsed = minidom.parseString(rough_string)
    return reparsed.toprettyxml(indent="   ")

def generate_xml_files(input_file, image_dir, result_dir):
    input_data = read_input_file(input_file)
    
    for filepath, objects in input_data.items():
        annotation = create_xml_annotation(filepath, objects, image_dir)
        xml_str = prettify(annotation)
        
        xml_filename = os.path.basename(filepath).replace('.jpg', '.xml')
        with open(result_dir + xml_filename, 'w', encoding='utf-8') as f:
            f.write(xml_str)
    
    print("XML files created successfully.")

if __name__ == "__main__":

    # 스크립트 파일의 디렉토리를 기준으로 경로를 설정
    script_dir = os.path.dirname(os.path.abspath(__file__))

    parser = argparse.ArgumentParser(description="Generate XML annotations from input file and image directory")
    parser.add_argument('--imgdir', required=False, default=script_dir+'/Images/', help="Path to the image directory (default is ./Images/)")
    parser.add_argument('--inputdir', required=False, default=script_dir+'/Label.txt', help="Path to the input text file (default is ./Label.txt)")
    parser.add_argument('--resultdir', required=False, default=script_dir+'/xml_result/', help="Path to the result Folder (default is ./xml_result)")
    
    args = parser.parse_args()
    
    input_file = os.path.join(script_dir, args.inputdir)
    print(input_file)
    image_dir = os.path.join(script_dir, args.imgdir)
    result_dir = os.path.join(script_dir, args.resultdir)
    
    generate_xml_files(input_file, image_dir, result_dir)

• txt2xml 사용 예시

python txt2xml.py --imgdir ./Images/ --inputdir ./Label.txt --resultdir ./xml_result/

• --imgdir : 학습 이미지 경로
• --inputdir : PPOCRLabel을 통해 나온 Label.txt파일 경로
• --resultdir : xml파일을 저장할 폴더 경로

result 디렉토리인 ./xml_result 폴더에 이미지 별로 xml파일들이 생성 되어 있을텐데 학습을 위해서 해당 폴더의 xml파일 리스트를 나타내는 txt파일이 필요하다.
이 file list를 작성 해주는 코드도 공유 하겠습니다.

import os
import argparse

def list_xml_files(xml_dir, output_file, mode, root_dir):
    xml_files = [f for f in os.listdir(xml_dir) if f.endswith('.xml')]
    xml_files.sort()  # Sort the files alphabetically
    
    with open(output_file, 'w', encoding='utf-8') as f:
        for xml_file in xml_files:
            f.write(f"{root_dir}{mode}/{xml_file}\n")
    
    print(f"{output_file} created successfully with {len(xml_files)} entries.")

if __name__ == "__main__":
    # 스크립트 파일의 디렉토리를 기준으로 경로를 설정
    script_dir = os.path.dirname(os.path.abspath(__file__))

    parser = argparse.ArgumentParser(description="List XML files in a directory")
    parser.add_argument('--mode', required=False, default='train', help="train or val (default=train)")
    parser.add_argument('--data_root_dir', required=False, default='./det_train_data/', help="train or val (default=./det_train_data/)")
    parser.add_argument('--xmldir', required=False, default=script_dir+'/xml_result/', help="Path to the directory containing XML files (default=./xml_result/)")
    parser.add_argument('--outputfile', required=False, default=script_dir+'/output.txt', help="Path to the output file (default=./)")
    
    args = parser.parse_args()
    
    xml_dir = os.path.join(script_dir, args.xmldir)
    output_file = os.path.join(script_dir, args.outputfile)
    
    list_xml_files(xml_dir, output_file, args.mode, args.data_root_dir)

• make_xml_list 사용 예시

python make_xml_list.py --mode train --xmldir ./xml_result/ --outputfile ./train_filelist.txt --data_root_dir ./det_train_data/

• --mode : train 또는 val
• --xmldir : xml이 있는 폴더 경로
• --outputfile : 결과 txt파일경로 및 파일이름
• --data_root_dir : train데이터 root디렉토리 경로

outputfile에는 xml파일의 경로가 작성되어 있는데 TATR를 통해 학습 하려면 데이터 구조를 맞춰줘야 한다. --data_root_dir은 outputfile의 경로 작성을 위한 arguments입니다.