ocr/ocr

#!/usr/bin/env python2.7
# coding=utf-8


"""
ocr

Usage:  For usage instructions run with option --help
Author: Patrick Jentsch <p.jentsch@uni-bielefeld.de>
"""


from argparse import ArgumentParser
from pyflow import WorkflowRunner
import multiprocessing
import os
import re
import sys


TESSERACT_MODELS = ['deu', 'eng', 'enm', 'fra', 'frk', 'frm', 'ita', 'por',
                    'spa']


def parse_args():
    parser = ArgumentParser(description='OCR Pipeline utilizing tesseract.')
    parser.add_argument('i')
    parser.add_argument('o')
    parser.add_argument('-l', '--language', choices=TESSERACT_MODELS,
                        required=True)
    parser.add_argument('--binarize', action='store_true',
                        help='use ocropy binarisation')
    parser.add_argument('--keep-intermediates', action='store_true',
                        help='keep intermediate files')
    parser.add_argument('--n-cores',
                        default=min(4, multiprocessing.cpu_count()),
                        help='total number of cores available', type=int)
    parser.add_argument('--log-dir')
    parser.add_argument('--zip')
    return parser.parse_args()


class OCRPipelineJob:
    def __init__(self, file, output_dir):
        self.file = file
        self.name = os.path.basename(file).rsplit('.', 1)[0]
        self.output_dir = output_dir


class OCRPipeline(WorkflowRunner):
    def __init__(self, binarize, jobs, keep_intermediates, lang, n_cores,
                 output_dir, zip):
        self.binarize = binarize
        self.jobs = jobs
        self.keep_intermediates = keep_intermediates
        self.lang = lang
        self.n_cores = n_cores
        self.output_dir = output_dir
        self.zip = zip

    def workflow(self):
        if not self.jobs:
            return

        '''
        ' ##################################################
        ' # mkdir_jobs                                     #
        ' ##################################################
        '''
        mkdir_jobs = []
        for i, job in enumerate(self.jobs):
            output_dir = os.path.join(job.output_dir, 'tmp')
            cmd = 'mkdir'
            cmd += ' -p'
            cmd += ' "{}"'.format(output_dir)
            if self.keep_intermediates:
                cmd += ' "{}"'.format(os.path.join(output_dir, 'hocr'))
                cmd += ' "{}"'.format(os.path.join(output_dir, 'pdf'))
                cmd += ' "{}"'.format(os.path.join(output_dir, 'tiff'))
                cmd += ' "{}"'.format(os.path.join(output_dir, 'txt'))
                if self.binarize:
                    cmd += ' "{}"'.format(os.path.join(output_dir, 'bin.png'))
                    cmd += ' "{}"'.format(os.path.join(output_dir, 'nrm.png'))
            lbl = 'mkdir_job_-_{}'.format(i)
            mkdir_jobs.append(self.addTask(command=cmd, label=lbl))

        '''
        ' ##################################################
        ' # pdftoppm_jobs                                  #
        ' ##################################################
        '''
        pdftoppm_jobs = []
        n_cores = min(self.n_cores, max(1, int(self.n_cores / len(self.jobs))))
        for i, job in enumerate(self.jobs):
            output_dir = os.path.join(job.output_dir, 'tmp')
            output_file_base = os.path.join(output_dir, 'page')
            cmd = 'pdftoppm'
            cmd += ' -r 300'
            cmd += ' -tiff'
            cmd += ' -tiffcompression lzw'
            cmd += ' "{}" "{}"'.format(job.file, output_file_base)
            deps = 'mkdir_job_-_{}'.format(i)
            lbl = 'pdftoppm_job_-_{}'.format(i)
            pdftoppm_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                              label=lbl, nCores=n_cores))

        if self.binarize:
            '''
            ' The ocropus_nlbin_jobs list is created based on the output files
            ' of the pdftoppm_jobs. So wait until they are finished.
            '''
            self.waitForTasks()

            '''
            ' ##################################################
            ' # ocropus_nlbin_jobs                             #
            ' ##################################################
            '''
            ocropus_nlbin_jobs = []
            '''
            ' We run ocropus-nlbin with either four or, if there are less then
            ' four cores available for this workflow, the available core
            ' number.
            '''
            n_cores = min(4, self.n_cores)
            for i, job in enumerate(self.jobs):
                input_dir = os.path.join(job.output_dir, 'tmp')
                output_dir = input_dir
                files = filter(lambda x: x.endswith('.tif'),
                               os.listdir(input_dir))
                files.sort(key=lambda x: int(re.search(r'\d+', x).group(0)))
                files = map(lambda x: os.path.join(input_dir, x), files)
                cmd = 'ocropus-nlbin "{}"'.format('" "'.join(files))
                cmd += ' -o "{}"'.format(output_dir)
                cmd += ' -Q "{}"'.format(n_cores)
                deps = 'pdftoppm_job_-_{}'.format(i)
                lbl = 'ocropus_nlbin_job_-_{}'.format(i)
                ocropus_nlbin_jobs.append(
                    self.addTask(command=cmd, dependencies=deps, label=lbl,
                                 nCores=n_cores))

            '''
            ' The post_ocropus_nlbin_jobs are created based on the output files
            ' of the ocropus_nlbin_jobs. So wait until they are finished.
            '''
            self.waitForTasks()

            '''
            ' ##################################################
            ' # post_ocropus_nlbin_jobs                        #
            ' ##################################################
            '''
            post_ocropus_nlbin_jobs = []
            for i, job in enumerate(self.jobs):
                input_dir = os.path.join(job.output_dir, 'tmp')
                output_dir = input_dir
                number = 0
                files = filter(lambda x: x.endswith('.bin.png'),
                               os.listdir(input_dir))
                files.sort()
                for file in files:
                    # int conversion is done in order to trim leading zeros
                    output_file = os.path.join(output_dir, 'page-{}.bin.png'.format(int(file.split('.', 1)[0])))  # noqa
                    cmd = 'mv "{}" "{}"'.format(os.path.join(output_dir, file),
                                                output_file)
                    deps = 'ocropus_nlbin_job_-_{}'.format(i)
                    lbl = 'post_ocropus_nlbin_job_-_{}-{}'.format(i, number)
                    post_ocropus_nlbin_jobs.append(
                        self.addTask(command=cmd, dependencies=deps,
                                     label=lbl))
                    number += 1

        '''
        ' The tesseract_jobs are created based of the output files of either
        ' the pdftoppm_jobs or post_ocropus_nlbin_jobs. So wait until they are
        ' finished.
        '''
        self.waitForTasks()

        '''
        ' ##################################################
        ' # tesseract_jobs                                 #
        ' ##################################################
        '''
        tesseract_jobs = []
        '''
        ' Tesseract runs fastest with four cores. So we run it with either four
        ' or, if there are less then four cores available for this workflow,
        ' the available core number.
        '''
        n_cores = min(4, self.n_cores)
        for i, job in enumerate(self.jobs):
            input_dir = os.path.join(job.output_dir, 'tmp')
            output_dir = input_dir
            files = filter(lambda x: x.endswith('.bin.png' if self.binarize else '.tif'),  # noqa
                           os.listdir(input_dir))
            files.sort(key=lambda x: int(re.search(r'\d+', x).group(0)))
            files = map(lambda x: os.path.join(input_dir, x), files)
            number = 0
            for file in files:
                output_file_base = os.path.join(output_dir, file.rsplit('.', 2 if self.binarize else 1)[0])  # noqa
                cmd = 'tesseract "{}" "{}"'.format(file, output_file_base)
                cmd += ' -l "{}"'.format(self.lang)
                cmd += ' hocr pdf txt'
                if self.binarize:
                    deps = 'post_ocropus_nlbin_job_-_{}-{}'.format(i, number)
                else:
                    deps = 'pdftoppm_job_-_{}'.format(i)
                label = 'tesseract_jobs_-_{}-{}'.format(i, number)
                tesseract_jobs.append(
                    self.addTask(command=cmd, dependencies=deps, label=label,
                                 nCores=n_cores))
                number += 1

        '''
        ' The following jobs are created based of the output files of the
        ' ocr_jobs. So wait until they are finished.
        '''
        self.waitForTasks()

        '''
        ' ##################################################
        ' # hocrtotei_jobs                                 #
        ' ##################################################
        '''
        hocrtotei_jobs = []
        for i, job in enumerate(self.jobs):
            input_dir = os.path.join(job.output_dir, 'tmp')
            files = filter(lambda x: x.endswith('.hocr'),
                           os.listdir(input_dir))
            files.sort(key=lambda x: int(re.search(r'\d+', x).group(0)))
            files = map(lambda x: os.path.join(input_dir, x), files)
            output_file = os.path.join(job.output_dir,
                                       '{}.xml'.format(job.name))
            cmd = 'hocrtotei "{}" "{}"'.format('" "'.join(files), output_file)
            deps = filter(lambda x: x.startswith('ocr_job_-_{}'.format(i)),
                          tesseract_jobs)
            lbl = 'hocrtotei_job_-_{}'.format(i)
            hocrtotei_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                               label=lbl))

        '''
        ' ##################################################
        ' # pdfunite_jobs                                  #
        ' ##################################################
        '''
        pdfunite_jobs = []
        for i, job in enumerate(self.jobs):
            input_dir = os.path.join(job.output_dir, 'tmp')
            files = filter(lambda x: x.endswith('.pdf'), os.listdir(input_dir))
            files.sort(key=lambda x: int(re.search(r'\d+', x).group(0)))
            files = map(lambda x: os.path.join(input_dir, x), files)
            output_file = os.path.join(job.output_dir,
                                       '{}.pdf'.format(job.name))
            cmd = 'pdfunite "{}" "{}"'.format('" "'.join(files), output_file)
            deps = filter(lambda x: x.startswith('ocr_job_-_{}'.format(i)),
                          tesseract_jobs)
            lbl = 'pdfunite_job_-_{}'.format(i)
            pdfunite_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                              label=lbl))

        '''
        ' ##################################################
        ' # cat_jobs                                       #
        ' ##################################################
        '''
        cat_jobs = []
        for i, job in enumerate(self.jobs):
            input_dir = os.path.join(job.output_dir, 'tmp')
            files = filter(lambda x: x.endswith('.txt'), os.listdir(input_dir))
            files.sort(key=lambda x: int(re.search(r'\d+', x).group(0)))
            files = map(lambda x: os.path.join(input_dir, x), files)
            output_file = os.path.join(job.output_dir,
                                       '{}.txt'.format(job.name))
            cmd = 'cat "{}" > "{}"'.format('" "'.join(files), output_file)
            deps = filter(lambda x: x.startswith('ocr_job_-_{}'.format(i)),
                          tesseract_jobs)
            lbl = 'cat_job_-_{}'.format(i)
            cat_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                         label=lbl))

        '''
        ' ##################################################
        ' # zip_jobs                                       #
        ' ##################################################
        '''
        zip_jobs = []
        if self.zip is not None:
            # Remove .zip file extension if provided
            if self.zip.endswith('.zip'):
                self.zip = self.zip[:-4]
                self.zip = self.zip if self.zip else 'output'
            cmd = 'cd "{}"'.format(self.output_dir)
            cmd += ' && '
            cmd += 'zip'
            cmd += ' -r'
            cmd += ' "{}".zip .'.format(self.zip)
            cmd += ' -x "pyflow.data*" "*tmp*"'
            cmd += ' -i "*.pdf" "*.txt" "*.xml"'
            cmd += ' && '
            cmd += 'cd -'
            deps = hocrtotei_jobs + pdfunite_jobs + cat_jobs
            lbl = 'zip_job_-_all'
            zip_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                         label=lbl))

            cmd = 'cd "{}"'.format(self.output_dir)
            cmd += ' && '
            cmd += 'zip'
            cmd += ' -r'
            cmd += ' "{}".pdf.zip .'.format(self.zip)
            cmd += ' -x "pyflow.data*" "*tmp*"'
            cmd += ' -i "*.pdf"'
            cmd += ' && '
            cmd += 'cd -'
            deps = 'zip_job_-_all'
            lbl = 'zip_job_-_pdf'
            zip_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                         label=lbl))

            cmd = 'cd "{}"'.format(self.output_dir)
            cmd += ' && '
            cmd += 'zip'
            cmd += ' -r'
            cmd += ' "{}".txt.zip .'.format(self.zip)
            cmd += ' -x "pyflow.data*" "*tmp*"'
            cmd += ' -i "*.txt"'
            cmd += ' && '
            cmd += 'cd -'
            deps = 'zip_job_-_all'
            lbl = 'zip_job_-_txt'
            zip_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                         label=lbl))

            cmd = 'cd "{}"'.format(self.output_dir)
            cmd += ' && '
            cmd += 'zip'
            cmd += ' -r'
            cmd += ' "{}".xml.zip .'.format(self.zip)
            cmd += ' -x "pyflow.data*" "*tmp*"'
            cmd += ' -i "*.xml"'
            cmd += ' && '
            cmd += 'cd -'
            deps = 'zip_job_-_all'
            lbl = 'zip_job_-_xml'
            zip_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                         label=lbl))

        '''
        ' ##################################################
        ' # mv_jobs                                        #
        ' ##################################################
        '''
        mv_jobs = []
        if self.keep_intermediates:
            for i, job in enumerate(self.jobs):
                input_dir = os.path.join(job.output_dir, 'tmp')
                output_dir = input_dir
                cmd = 'mv "{}"/*.hocr "{}"'.format(
                    input_dir, os.path.join(output_dir, 'hocr'))
                cmd += ' && '
                cmd += 'mv "{}"/*.pdf "{}"'.format(input_dir, os.path.join(output_dir, 'pdf'))  # noqa
                cmd += ' && '
                cmd += 'mv "{}"/*.tif "{}"'.format(input_dir, os.path.join(output_dir, 'tiff'))  # noqa
                cmd += ' && '
                cmd += 'mv "{}"/*.txt "{}"'.format(input_dir, os.path.join(output_dir, 'txt'))  # noqa
                if self.binarize:
                    cmd += ' && '
                    cmd += 'mv "{}"/*.bin.png "{}"'.format(input_dir, os.path.join(output_dir, 'bin.png'))  # noqa
                    cmd += ' && '
                    cmd += 'mv "{}"/*.nrm.png "{}"'.format(input_dir, os.path.join(output_dir, 'nrm.png'))  # noqa
                deps = ['hocrtotei_job_-_{}'.format(i),
                        'pdfunite_job_-_{}'.format(i),
                        'cat_job_-_{}'.format(i)]
                lbl = 'mv_job_-_{}'.format(i)
                mv_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                            label=lbl))
        else:
            for i, job in enumerate(self.jobs):
                input_dir = os.path.join(job.output_dir, 'tmp')
                cmd = 'rm -r "{}"'.format(input_dir)
                deps = ['hocrtotei_job_-_{}'.format(i),
                        'pdfunite_job_-_{}'.format(i),
                        'cat_job_-_{}'.format(i)]
                lbl = 'mv_job_-_{}'.format(i)
                mv_jobs.append(self.addTask(command=cmd, dependencies=deps,
                                            label=lbl))


def collect_jobs(input_dir, output_dir):
    jobs = []
    for file in os.listdir(input_dir):
        if os.path.isdir(os.path.join(input_dir, file)):
            jobs += collect_jobs(os.path.join(input_dir, file),
                                 os.path.join(output_dir, file))
        elif file.endswith('.pdf'):
            jobs.append(OCRPipelineJob(os.path.join(input_dir, file),
                                       os.path.join(output_dir, file)))
    return jobs


def main():
    args = parse_args()
    jobs = collect_jobs(args.i, args.o)
    ocr_pipeline = OCRPipeline(args.binarize, jobs, args.keep_intermediates,
                               args.language, args.n_cores, args.o, args.zip)
    retval = ocr_pipeline.run(dataDirRoot=(args.log_dir or args.o),
                              nCores=args.n_cores)
    sys.exit(retval)


if __name__ == '__main__':
    main()