Das deutsche Python-Forum

Code: Alles auswählen

import os
import random
import xml.etree.ElementTree as etree
from string import digits

path = '/home/richard/MII/datasets/pascal/VOCdevkit/VOC2008/Annotations'
directoy = os.listdir(path)

list_xmls = []
set_instruments = set([])

for xml_file in directoy:

    base = os.path.splitext(xml_file)[0]
    list_xmls.append(base)
    file_path = os.path.join(path, xml_file)
    tree = etree.parse(file_path)
    root = tree.getroot()

    for object in root.findall('object'):
        name = object.find('name')
        set_instruments.add(name.text)
        
random.shuffle(list_xmls)

for instrument in set_instruments:
    txt_path = os.path.join('/home/richard/MII/datasets/pascal/VOCdevkit/VOC2008/ImageSets/Main/', instrument + '_train.txt')
    file = open(txt_path, 'a+')
    
    for xml in list_xmls:
        remove_digits = str.maketrans('', '', digits)
        res = xml.translate(remove_digits)
        if (res == instrument):
            file.write(xml + ' 1 \n')
        else:
            file.write(xml + ' -1 \n')
    file.close()

Code: Alles auswählen

import random
from pathlib import Path
import xml.etree.ElementTree as etree
from string import digits

BASEPATH = Path('/home/richard/MII/datasets/pascal/VOCdevkit/VOC2008')

def read_xml_files(path):
    filenames = []
    instruments = {}
    for xml_filename in path.iterdir():
        filenames.append(xml_filename.stem)
        tree = etree.parse(xml_filename)
        root = tree.getroot()
        for object in root.findall('object'):
            instruments.add(object.findtext('name'))
    return filenames, instruments

def main():
    filenames, instruments = read_xml_files(BASEPATH / 'Annotations'):
    random.shuffle(filenames)
    for instrument in instruments:
        text_path = BASEPATH / 'ImageSets' / 'Main' / instrument + '_train.txt'
        with text_path.open('a') as output:
            for xml_filename in filenames:
                remove_digits = str.maketrans('', '', digits)
                result = xml_filename.translate(remove_digits)
                output.write(f"{xml_filename} {1 if result == instrument else -1} \n")

if __name__ == '__main__':
    main()