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>
"""


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


''' TODO:
' Implement --end-page: Last page to ocr
' Implement --memMb: Total amount of memory (RAM) available for this workflow.
'                    Default: 2048 * nCores
' Implement --rotate: Rotate pages from input (90, 180, 270)
' Implement --split-pages: Split pages in half after possible rotation
' Implement --start-page: First page to ocr
'''


def parse_arguments():
    parser = argparse.ArgumentParser(
        'Performs OCR of (historical) documents utilizing OCRopus for \
        preprocessing and Tesseract OCR for OCR. Available outputs are HOCR, \
        PDF, shrinked PDF, and simple DTAbf (TEI P5 compliant). Software \
        requirements: imagemagick, ocropus, pdftoppm, pdfunite, \
        poppler-utils, pyflow, python2.7, python3.5, tesseract'
    )
    parser.add_argument(
        '-l',
        dest='lang',
        help='Language for OCR.',
        required=True
    )
    parser.add_argument(
        '--i',
        default=os.path.normpath('/files_for_ocr'),
        dest='inputDirectory',
        help='The input directory.',
        required=False
    )
    parser.add_argument(
        '--o',
        default=os.path.normpath('/files_from_ocr'),
        dest='outputDirectory',
        help='The output directory.',
        required=False
    )
    parser.add_argument(
        '--skip-binarisation',
        action='store_true',
        default=False,
        dest='skipBinarisation',
        help='Skip binarisation.',
        required=False
    )
    parser.add_argument(
        '--keep-intermediates',
        action='store_true',
        default=False,
        dest='keepIntermediates',
        help='Keep intermediate files.',
        required=False
    )
    parser.add_argument(
        '--nCores',
        default=min(4, multiprocessing.cpu_count()),
        dest='nCores',
        help='Total number of cores available.',
        required=False,
        type=int
    )
    return parser.parse_args()


class OCRWorkflow(WorkflowRunner):
    def __init__(self, args):
        self.jobs = analyze_jobs(args.inputDirectory, args.outputDirectory)
        self.skipBinarisation = args.skipBinarisation
        self.keepIntermediates = args.keepIntermediates
        self.lang = args.lang
        self.nCores = args.nCores

    def workflow(self):
        '''
        ' Creating output directories...
        '''
        create_output_directories_jobs = []
        for index, job in enumerate(self.jobs):
            cmd = 'mkdir -p "%s"' % (
                os.path.join(job['output_dir'], 'tmp')
            )
            if self.keepIntermediates:
                cmd += ' "%s" "%s" "%s" "%s"' % (
                    os.path.join(job['output_dir'], 'tmp', 'hocr'),
                    os.path.join(job['output_dir'], 'tmp', 'pdf'),
                    os.path.join(job['output_dir'], 'tmp', 'tiff'),
                    os.path.join(job['output_dir'], 'tmp', 'txt')
                )
            if not self.skipBinarisation:
                cmd += ' "%s"' % (
                    os.path.join(job['output_dir'], 'tmp', 'bin.png')
                )
            create_output_directories_jobs.append(
                self.addTask(
                    command=cmd,
                    label='create_output_directories_job_-_%i' % (index)
                )
            )

        '''
        ' Splitting...
        '''
        split_jobs = []
        split_job_nCores = min(
            self.nCores,
            max(1, int(self.nCores / len(self.jobs)))
        )
        for index, job in enumerate(self.jobs):
            if job['filename'].endswith(('.tif', '.tiff')):
                cmd = 'convert "%s" -compress LZW -density 300 -scene 1 "%s"/page-%%d.tif' % (
                    job['path'],
                    os.path.join(job['output_dir'], 'tmp')
                )
            else:
                cmd = 'pdftoppm -r 300 -tiff -tiffcompression lzw "%s" "%s"' % (
                    job['path'],
                    os.path.join(job['output_dir'], 'tmp', 'page')
                )
            split_jobs.append(
                self.addTask(
                    command=cmd,
                    dependencies='create_output_directories_job_-_%i' % (index),
                    label='split_job_-_%i' % (index),
                    nCores=split_job_nCores
                )
            )

        if not self.skipBinarisation:
            '''
            ' Binarising...
            '''
            binarisation_jobs = []
            '''
            ' We run ocropus-nlbin with either four or, if there are less then
            ' four cores available for this workflow, the available core
            ' number.
            '''
            binarisation_job_nCores = min(4, self.nCores)
            for index, job in enumerate(self.jobs):
                cmd = 'ls --quoting-style=shell-escape -v "%s"/*.tif | xargs ocropus-nlbin --output "%s" --parallel "%i"' % (
                    os.path.join(job['output_dir'], 'tmp'),
                    os.path.join(job['output_dir'], 'tmp'),
                    binarisation_job_nCores
                )
                binarisation_jobs.append(
                    self.addTask(
                        command=cmd,
                        dependencies='split_job_-_%i' % (index),
                        label='binarisation_job_-_%i' % (index),
                        nCores=binarisation_job_nCores
                    )
                )

            '''
            ' Normalising file names from binarisation...
            '''
            self.waitForTasks()
            post_binarisation_jobs = []
            for index, job in enumerate(self.jobs):
                number = 0
                files = os.listdir(os.path.join(job['output_dir'], 'tmp'))
                files = filter(lambda x: x.endswith('.bin.png'), files)
                files = sorted(
                    files,
                    key=lambda x: int(re.search(r'\d+', x).group(0))
                )
                for file in files:
                    cmd = 'mv "%s" "%s"' % (
                        os.path.join(job['output_dir'], 'tmp', file),
                        os.path.join(job['output_dir'], 'tmp', 'page-%i.%s' % (
                            int(file.split('.', 1)[0]),
                            file.split('.', 1)[1])
                        ),
                    )
                    post_binarisation_jobs.append(
                        self.addTask(
                            command=cmd,
                            dependencies='binarisation_job_-_%i' % (index),
                            label='post_binarisation_job_-_%i-%i' % (
                                index,
                                number
                            )
                        )
                    )
                    number += 1

        '''
        ' Performing OCR...
        '''
        self.waitForTasks()
        ocr_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.
        '''
        ocr_job_nCores = min(4, self.nCores)
        '''
        ' WORKAROUND: Tesseract only uses one core for the deu_frak language
        ' model, so the workflow will also only reserve one in this case.
        '''
        if self.lang == "deu_frak":
            ocr_job_nCores = 1
        for index, job in enumerate(self.jobs):
            number = 0
            files = os.listdir(os.path.join(job['output_dir'], 'tmp'))
            if self.skipBinarisation:
                files = filter(lambda x: x.endswith('.tif'), files)
            else:
                files = filter(lambda x: x.endswith('.bin.png'), files)
            files = sorted(
                files,
                key=lambda x: int(re.search(r'\d+', x).group(0))
            )
            for file in files:
                cmd = 'tesseract "%s" "%s" -l "%s" hocr pdf txt' % (
                    os.path.join(job['output_dir'], 'tmp', file),
                    os.path.join(
                        job['output_dir'],
                        'tmp',
                        file.rsplit('.', 1 if self.skipBinarisation else 2)[0]
                    ),
                    self.lang
                )
                if self.skipBinarisation:
                    ocr_job_dependencies = 'split_job_-_%i' % (index)
                else:
                    ocr_job_dependencies = filter(
                        lambda x: x == 'post_binarisation_job_-_%i-%i' % (
                            index,
                            number
                        ),
                        post_binarisation_jobs
                    )
                print(ocr_job_dependencies)
                ocr_jobs.append(
                    self.addTask(
                        command=cmd,
                        dependencies=ocr_job_dependencies,
                        label='ocr_job_-_%i-%i' % (index, number),
                        nCores=ocr_job_nCores
                    )
                )
                number += 1

        '''
        ' Creating TEI P5 files...
        '''
        hocr_to_tei_jobs = []
        for index, job in enumerate(self.jobs):
            cmd = 'hocrtotei "%s" "%s"' % (
                os.path.join(job['output_dir'], 'tmp'),
                os.path.join(
                    job['output_dir'],
                    job['filename'].rsplit('.', 1)[0] + '.xml'
                )
            )
            hocr_to_tei_jobs.append(
                self.addTask(
                    command=cmd,
                    dependencies=filter(
                        lambda x: x.startswith('ocr_job_-_%i' % (index)),
                        ocr_jobs
                    ),
                    label='hocr_to_tei_job_-_%i' % (index)
                )
            )

        '''
        ' Merging PDF files...
        '''
        pdf_merge_jobs = []
        for index, job in enumerate(self.jobs):
            cmd = '(ls --quoting-style=shell-escape -v "%s"/*.pdf && echo "\'%s\'") | xargs pdfunite' % (
                os.path.join(job['output_dir'], 'tmp'),
                os.path.join(
                    job['output_dir'],
                    job['filename'].rsplit('.', 1)[0] + '.pdf'
                )
            )
            pdf_merge_jobs.append(
                self.addTask(
                    command=cmd,
                    dependencies=filter(
                        lambda x: x.startswith('ocr_job_-_%i' % (index)),
                        ocr_jobs
                    ),
                    label='pdf_merge_job_-_%i' % (index)
                )
            )

        '''
        ' Merging text files...
        '''
        txt_merge_jobs = []
        for index, job in enumerate(self.jobs):
            cmd = 'ls --quoting-style=shell-escape -v "%s"/*.txt | xargs cat > "%s"' % (
                os.path.join(job['output_dir'], 'tmp'),
                os.path.join(
                    job['output_dir'],
                    job['filename'].rsplit('.', 1)[0] + '.txt'
                )
            )
            txt_merge_jobs.append(
                self.addTask(
                    command=cmd,
                    dependencies=filter(
                        lambda x: x.startswith('ocr_job_-_%i' % (index)),
                        ocr_jobs
                    ),
                    label='txt_merge_job_-_%i' % (index)
                )
            )

        '''
        ' Cleanup...
        '''
        cleanup_jobs = []
        if self.keepIntermediates:
            for index, job in enumerate(self.jobs):
                cleanup_job_dependencies = [
                    'hocr_to_tei_job_-_%i' % (index),
                    'pdf_merge_job_-_%i' % (index),
                    'txt_merge_job_-_%i' % (index)
                ]
                cmd = 'mv "%s"/*.hocr "%s"' % (
                    os.path.join(job['output_dir'], 'tmp'),
                    os.path.join(job['output_dir'], 'tmp', 'hocr'),
                )
                cmd += ' && mv "%s"/*.pdf "%s"' % (
                    os.path.join(job['output_dir'], 'tmp'),
                    os.path.join(job['output_dir'], 'tmp', 'pdf'),
                )
                cmd += ' && mv "%s"/*.tif "%s"' % (
                    os.path.join(job['output_dir'], 'tmp'),
                    os.path.join(job['output_dir'], 'tmp', 'tiff'),
                )
                cmd += ' && mv "%s"/*.txt "%s"' % (
                    os.path.join(job['output_dir'], 'tmp'),
                    os.path.join(job['output_dir'], 'tmp', 'txt'),
                )
                if not self.skipBinarisation:
                    cmd += ' && mv "%s"/*.bin.png "%s"' % (
                        os.path.join(job['output_dir'], 'tmp'),
                        os.path.join(job['output_dir'], 'tmp', 'bin.png'),
                    )
                    cmd += ' && rm "%s"/*.nrm.png' % (
                        os.path.join(job['output_dir'], 'tmp')
                    )
                cleanup_jobs.append(
                    self.addTask(
                        command=cmd,
                        dependencies=cleanup_job_dependencies,
                        label='cleanup_job_-_%i' % (index)
                    )
                )
        else:
            for index, job in enumerate(self.jobs):
                cleanup_job_dependencies = [
                    'hocr_to_tei_job_-_%i' % (index),
                    'pdf_merge_job_-_%i' % (index),
                    'txt_merge_job_-_%i' % (index)
                ]
                cmd = 'rm -r "%s"' % (
                    os.path.join(job['output_dir'], 'tmp')
                )
                cleanup_jobs.append(
                    self.addTask(
                        command=cmd,
                        dependencies=cleanup_job_dependencies,
                        label='cleanup_job_-_%i' % (index)
                    )
                )


def analyze_jobs(inputDirectory, outputDirectory):
    jobs = []

    for file in os.listdir(inputDirectory):
        if os.path.isdir(os.path.join(inputDirectory, file)):
            jobs += analyze_jobs(
                os.path.join(inputDirectory, file),
                os.path.join(outputDirectory, file)
            )
        elif file.endswith(('.pdf', '.tif', '.tiff')):
            jobs.append(
                {
                    'filename': file,
                    'output_dir': os.path.join(outputDirectory, file),
                    'path': os.path.join(inputDirectory, file)
                }
            )

    return jobs


def main():
    args = parse_arguments()

    wflow = OCRWorkflow(args)

    retval = wflow.run(dataDirRoot=args.outputDirectory, nCores=args.nCores)

    sys.exit(retval)


if __name__ == '__main__':
    main()