mirror of
https://gitlab.ub.uni-bielefeld.de/sfb1288inf/ocr.git
synced 2025-01-13 06:00:34 +00:00
298 lines
12 KiB
Python
Executable File
298 lines
12 KiB
Python
Executable File
#!/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 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, pdftk, pdftoppm, poppler-utils, \
|
|
pyflow, python2.7, tesseract"
|
|
)
|
|
|
|
parser.add_argument("-i",
|
|
dest="inputDir",
|
|
help="Input directory.",
|
|
required=True)
|
|
parser.add_argument("-l",
|
|
dest='lang',
|
|
help="Language for OCR",
|
|
required=True)
|
|
parser.add_argument("-o",
|
|
dest="outputDir",
|
|
help="Output directory.",
|
|
required=True)
|
|
parser.add_argument("--skip-binarization",
|
|
action='store_true',
|
|
default=False,
|
|
dest="skipBinarization",
|
|
help="Skip binarization.",
|
|
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, jobs, skipBinarization, keepIntermediates, lang, nCores):
|
|
self.jobs = jobs
|
|
self.skipBinarization = skipBinarization
|
|
self.keepIntermediates = keepIntermediates
|
|
self.lang = lang
|
|
self.nCores = nCores
|
|
self.defaultNCores = min(nCores, max(1, int(nCores / len(jobs))))
|
|
|
|
|
|
def workflow(self):
|
|
###
|
|
# Task "create_output_directories_job": create output directories
|
|
# Dependencies: None
|
|
###
|
|
create_output_directories_jobs = []
|
|
create_output_directories_job_number = 0
|
|
for job in self.jobs:
|
|
create_output_directories_job_number += 1
|
|
cmd = 'mkdir -p "%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.skipBinarization:
|
|
cmd += ' "%s" "%s"' % (
|
|
os.path.join(job["output_dir"], "tmp", "binarized_png"),
|
|
os.path.join(job["output_dir"], "tmp", "normalized_png"),
|
|
)
|
|
create_output_directories_jobs.append(self.addTask(label="create_output_directories_job_-_%i" % (create_output_directories_job_number), command=cmd, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "split_job": split input file into one tiff file per page
|
|
# Dependencies: create_output_directories_jobs
|
|
###
|
|
split_jobs = []
|
|
split_job_number = 0
|
|
for job in self.jobs:
|
|
split_job_number += 1
|
|
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(label="split_job_-_%i" % (split_job_number), command=cmd, dependencies=create_output_directories_jobs, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "ocropus_nlbin_job": binarize tiff files from previous split
|
|
# Dependencies: split_jobs
|
|
###
|
|
binarization_jobs = []
|
|
binarization_job_number = 0
|
|
'''
|
|
' We run ocropus-nlbin with either four or, if there are less then four
|
|
' cores available for this workflow, the available core number.
|
|
'''
|
|
binarization_job_nCores = min(4, self.nCores)
|
|
if not self.skipBinarization:
|
|
for job in self.jobs:
|
|
binarization_job_number += 1
|
|
cmd = 'ocropus-nlbin --output "%s" --parallel "%i" $(ls "%s"/*.tif | sort -V)' % (
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
binarization_job_nCores,
|
|
os.path.join(job["output_dir"], "tmp")
|
|
)
|
|
binarization_jobs.append(self.addTask(label="binarization_job_-_%i" % (binarization_job_number), command=cmd, dependencies=split_jobs, nCores=binarization_job_nCores))
|
|
|
|
###
|
|
# Task "post_binarization_job": Normalize file names from binarization
|
|
# Dependencies: binarization_jobs
|
|
###
|
|
self.waitForTasks()
|
|
post_binarization_jobs = []
|
|
post_binarization_job_number = 0
|
|
if not self.skipBinarization:
|
|
for job in self.jobs:
|
|
for file in filter(lambda x: x.endswith((".bin.png", ".nrm.png")), os.listdir(os.path.join(job["output_dir"], "tmp"))):
|
|
post_binarization_job_number += 1
|
|
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_binarization_jobs.append(self.addTask(label="post_binarization_job_-_%i" % (post_binarization_job_number), command=cmd, dependencies=binarization_jobs, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "ocr_job": perform OCR
|
|
# Dependencies: waitForTasks
|
|
###
|
|
self.waitForTasks()
|
|
ocr_jobs = []
|
|
ocr_job_number = 0
|
|
'''
|
|
' 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 job in self.jobs:
|
|
for file in filter(lambda x: x.endswith(".tif") if self.skipBinarization else x.endswith(".bin.png"), os.listdir(os.path.join(job["output_dir"], "tmp"))):
|
|
ocr_job_number += 1
|
|
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.skipBinarization else 2)[0]),
|
|
self.lang
|
|
)
|
|
ocr_jobs.append(self.addTask(label="ocr_job_-_%i" % (ocr_job_number), command=cmd, dependencies=post_binarization_jobs, nCores=ocr_job_nCores))
|
|
|
|
###
|
|
# Task "hocr_to_tei_job": create TEI P5 file from hocr files
|
|
# Dependencies: ocr_jobs
|
|
###
|
|
hocr_to_tei_jobs = []
|
|
hocr_to_tei_job_number = 0
|
|
for job in self.jobs:
|
|
hocr_to_tei_job_number += 1
|
|
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(label="hocr_to_tei_job_-_%i" % (hocr_to_tei_job_number), command=cmd, dependencies=ocr_jobs, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "pdf_merge_job": Merge PDF files
|
|
# Dependencies: ocr_jobs
|
|
###
|
|
pdf_merge_jobs = []
|
|
pdf_merge_job_number = 0
|
|
for job in self.jobs:
|
|
pdf_merge_job_number += 1
|
|
cmd = 'pdftk $(ls "%s"/*.pdf | sort -V) cat output "%s"' % (
|
|
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(label="pdf_merge_job_-_%i" % (pdf_merge_job_number), command=cmd, dependencies=ocr_jobs, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "txt_merge_job": Merge .txt files
|
|
# Dependencies: ocr_jobs
|
|
###
|
|
txt_merge_jobs = []
|
|
txt_merge_job_number = 0
|
|
for job in self.jobs:
|
|
txt_merge_job_number += 1
|
|
cmd = 'cat $(ls "%s"/*.txt | sort -V) > "%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(label="txt_merge_job_-_%i" % (txt_merge_job_number), command=cmd, dependencies=ocr_jobs, nCores=self.defaultNCores))
|
|
|
|
###
|
|
# Task "cleanup_job": remove temporary files
|
|
# Dependencies: hocr_to_tei_jobs + pdf_merge_jobs + txt_merge_jobs
|
|
###
|
|
cleanup_jobs = []
|
|
cleanup_job_counter = 0
|
|
if self.keepIntermediates:
|
|
for job in self.jobs:
|
|
cleanup_job_counter += 1
|
|
cmd = 'mv "%s"/*.hocr "%s" && mv "%s"/*.pdf "%s" && mv "%s"/*.tif "%s" && mv "%s"/*.txt "%s"' % (
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "hocr"),
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "pdf"),
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "tiff"),
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "txt")
|
|
)
|
|
if not self.skipBinarization:
|
|
cmd += ' && mv "%s"/*.bin.png "%s" && mv "%s"/*.nrm.png "%s"' % (
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "binarized_png"),
|
|
os.path.join(job["output_dir"], "tmp"),
|
|
os.path.join(job["output_dir"], "tmp", "normalized_png"),
|
|
)
|
|
cleanup_jobs.append(self.addTask(label="cleanup_job_-_%i" % (cleanup_job_counter), command=cmd, dependencies=hocr_to_tei_jobs + pdf_merge_jobs + txt_merge_jobs, nCores=self.defaultNCores))
|
|
else:
|
|
for job in self.jobs:
|
|
cleanup_job_counter += 1
|
|
cmd = 'rm -r "%s"' % (
|
|
os.path.join(job["output_dir"], "tmp")
|
|
)
|
|
cleanup_jobs.append(self.addTask(label="cleanup_job_-_%i" % (cleanup_job_counter), command=cmd, dependencies=hocr_to_tei_jobs + pdf_merge_jobs + txt_merge_jobs), nCores=self.defaultNCores)
|
|
|
|
|
|
def analyze_jobs(inputDir, outputDir):
|
|
jobs = []
|
|
|
|
for file in os.listdir(inputDir):
|
|
if os.path.isdir(os.path.join(inputDir, file)):
|
|
jobs += analyze_jobs(
|
|
os.path.join(inputDir, file),
|
|
os.path.join(outputDir, file)
|
|
)
|
|
elif file.endswith((".pdf", ".tif", ".tiff")):
|
|
jobs.append({"filename": file, "output_dir": os.path.join(outputDir, file), "path": os.path.join(inputDir, file)})
|
|
|
|
return jobs
|
|
|
|
|
|
def main():
|
|
args = parse_arguments()
|
|
|
|
wflow = OCRWorkflow(
|
|
analyze_jobs(args.inputDir, args.outputDir),
|
|
args.skipBinarization,
|
|
args.keepIntermediates,
|
|
args.lang,
|
|
args.nCores
|
|
)
|
|
|
|
retval = wflow.run(nCores=args.nCores)
|
|
sys.exit(retval)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|