TRIPS/fixBrokenText.py at master · HLR/TRIPS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import os
import argparse
from shutil import copyfile
import shutil
import xml.etree.ElementTree as ET


def run(logFile,fileDir):
    inDir = os.path.join(fileDir,"input")
    outDir = os.path.join(fileDir,"output")
    badDir = os.path.join(fileDir,"bad")
    with open(logFile,'r') as file:
        for line in file.read().split('\n'):
            try:
                if (len(line)<1):
                    continue
                if (line[-3:]!='xml'):
                    continue
                sourceTree = ET.parse(os.path.join(inDir,line))
                genTree = ET.parse(os.path.join(outDir,line))
                sourceText=sourceTree.getroot().attrib['input']
                while sourceText[-1]==' ':
                    sourceText=sourceText[:-1]
                if len(genTree.getroot().findall('SCENE'))>0:
                    genTree.getroot().find('SCENE').find('SENTENCE').find('TEXT').text=sourceText
                else:
                    genTree.getroot().find('SENTENCE').find('TEXT').text=sourceText
                genTree.write(os.path.join(outDir,line))
            except Exception as e:
                print(e)
    return
def runWithoutLogFile(fileDir):
    inDir = os.path.join(fileDir,"input")
    outDir = os.path.join(fileDir,"output")
    badDir = os.path.join(fileDir,"bad")
    for fileName in os.listdir(outDir):
        try:
            if (len(fileName)<1):
                continue
            if (fileName[-3:]!='xml'):
                continue
            sourceTree = ET.parse(os.path.join(inDir,fileName))
            genTree = ET.parse(os.path.join(outDir,fileName))
            sourceText=sourceTree.getroot().attrib['input']
            if len(genTree.getroot().findall('SCENE'))>0:
                genTree.getroot().find('SCENE').find('SENTENCE').find('TEXT').text=sourceText
            else:
                genTree.getroot().find('SENTENCE').find('TEXT').text=sourceText
            genTree.write(os.path.join(outDir,fileName))
        except Exception as e:
            print(e)
    return

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='This is a script to find SpRL file data that did not find its way into a core file.')
    parser.add_argument(
        "--path",
        dest="path",
        required=True,
        help='Path to the input files')
    parser.add_argument(
        "--file",
        dest="file",
        required=False,
        help='Core SpRL log file')
    args=parser.parse_args()
    if args.file==None:
        runWithoutLogFile(args.path)
    else:
        run(args.file,args.path)