feat: add triage step

remove tqdm demo
This commit is contained in:
mawi
2019-04-08 10:26:56 +02:00
parent 01bbf064e0
commit 6590875756
5 changed files with 40 additions and 88 deletions
+2 -4
View File
@@ -179,10 +179,8 @@ jobcontrol.add_argument(
jobcontrol.add_argument(
'-v',
'--verbose',
const="+",
default=[],
nargs='?',
action="append",
default=0,
action="count",
help="Print more verbose messages for each additional verbose level",
)
+13 -13
View File
@@ -29,13 +29,15 @@ DEBUG = 10
class PDFContext:
"""Holds our context for a particular run of the pipeline"""
def __init__(self, options, work_folder, origin, pdfinfo, tick=None):
def __init__(self, options, work_folder, origin, pdfinfo):
self.options = options
self.work_folder = work_folder
self.origin = origin
self.pdfinfo = pdfinfo
self.log = get_logger(options, '%s: ' % os.path.basename(origin))
self.tick_callback = tick
self.name = os.path.basename(options.input_file)
if self.name == '-':
self.name = 'stdin'
self.log = get_logger(options, '%s: ' % self.name)
def get_path(self, name):
return os.path.join(self.work_folder, name)
@@ -45,10 +47,6 @@ class PDFContext:
for n in range(npages):
yield PageContext(self, n)
def tick(self, times=1):
if self.tick_callback:
self.tick_callback(times)
class PageContext:
"""Holds our context for a page"""
@@ -58,14 +56,11 @@ class PageContext:
self.options = pdf_context.options
self.pageno = pageno
self.pageinfo = pdf_context.pdfinfo[pageno]
self.log = get_logger(pdf_context.options, '%s Page %d: ' % (os.path.basename(pdf_context.origin), pageno + 1))
self.log = get_logger(pdf_context.options, '%s Page %d: ' % (pdf_context.name, pageno + 1))
def get_path(self, name):
return os.path.join(self.pdf_context.work_folder, "%06d_%s" % (self.pageno + 1, name))
def tick(self, times=1):
self.pdf_context.tick(times)
def cleanup_working_files(work_folder, options):
if options.keep_temporary_files:
@@ -77,8 +72,13 @@ def cleanup_working_files(work_folder, options):
def get_logger(options=None, prefix=''):
level = ERROR # TODO: add option
if options is not None and options.output_file == '-' or options.sidecar == '-':
return NullLogger()
if options is not None:
if options.quiet or options.output_file == '-' or options.sidecar == '-':
return NullLogger()
if options.verbose > 0:
level = INFO
if options.verbose > 1:
level = DEBUG
return Logger(prefix, level)
+7 -16
View File
@@ -46,21 +46,12 @@ from .pdfinfo import Colorspace, PdfInfo
VECTOR_PAGE_DPI = 400
def triage_image_file(input_file, output_file, log, options):
def triage_image_file(input_file, output_file, options, log):
try:
log.info("Input file is not a PDF, checking if it is an image...")
im = Image.open(input_file)
except EnvironmentError as e:
msg = str(e)
# Recover the original filename
realpath = ''
if os.path.islink(input_file):
realpath = os.path.realpath(input_file)
elif os.path.isfile(input_file):
realpath = '<stdin>'
msg = msg.replace(input_file, realpath)
log.error(msg)
log.error(str(e))
raise UnsupportedImageFormatError() from e
else:
log.info("Input file is an image")
@@ -135,9 +126,7 @@ def _pdf_guess_version(input_file, search_window=1024):
return ''
def triage(input_file, output_file, log, context):
options = context.get_options()
def triage(input_file, output_file, options, log):
try:
if _pdf_guess_version(input_file):
if options.image_dpi:
@@ -145,13 +134,15 @@ def triage(input_file, output_file, log, context):
"Argument --image-dpi ignored because the "
"input file is a PDF, not an image."
)
# Origin file is a pdf create a symlink with pdf extension
re_symlink(input_file, output_file, log)
return
return output_file
except EnvironmentError as e:
log.error(e)
raise InputFileError() from e
triage_image_file(input_file, output_file, log, options)
triage_image_file(input_file, output_file, options, log)
return output_file
def get_pdfinfo(input_file, detailed_page_analysis=False):
+16 -47
View File
@@ -18,11 +18,11 @@
import os
import atexit
import concurrent.futures
from tqdm import tqdm
from tempfile import mkdtemp
from ._jobcontext import PDFContext, get_logger, cleanup_working_files
from ._weave import weave_layers
from ._pipeline import (
triage,
get_pdfinfo,
validate_pdfinfo_options,
is_ocr_required,
@@ -50,10 +50,10 @@ from .exceptions import (
ExitCode,
ExitCodeException,
)
from . import VERSION
from .helpers import available_cpu_count
from ._validation import (
check_closed_streams,
preamble,
check_options,
check_dependency_versions,
check_environ,
@@ -78,7 +78,7 @@ def exec_page_sync(page_context):
orientation_correction = get_orientation_correction(rasterize_preview_out, page_context)
rasterize_out = rasterize(page_context.pdf_context.origin, page_context, correction=orientation_correction)
page_context.tick()
preprocess_out = rasterize_out
if options.remove_background:
preprocess_out = preprocess_remove_background(preprocess_out, page_context)
@@ -90,7 +90,7 @@ def exec_page_sync(page_context):
preprocess_out = preprocess_clean(preprocess_out, page_context)
ocr_image_out = create_ocr_image(preprocess_out, page_context)
page_context.tick()
pdf_page_from_image_out = None
if not options.lossless_reconstruction:
visible_image_out = preprocess_out
@@ -104,9 +104,7 @@ def exec_page_sync(page_context):
if options.pdf_renderer == 'sandwich':
(ocr_out, text_out) = ocr_tesseract_textonly_pdf(ocr_image_out, page_context)
page_context.tick()
else:
page_context.tick(3)
return (page_context.pageno, pdf_page_from_image_out, ocr_out, text_out, orientation_correction)
@@ -120,43 +118,16 @@ def post_process(pdf_file, context):
return optimize_pdf(pdf_out, context)
def exec_sync(context):
"""Execute the pipeline single threaded"""
# TODO: triage
# Run exec_page_sync on every page context
layers = map(exec_page_sync, context.get_page_contexts())
# Output sidecar text
if context.options.sidecar:
sidecars = [layer[3] for layer in layers]
text = merge_sidecars(sidecars, context)
# Copy final text file to destination
copy_final(text, context.options.sidecar, context)
# Merge layers to one single pdf
pdf = weave_layers(layers, context)
# PDF/A and metadata
pdf = post_process(pdf, context)
# Copy final PDF file to destination
copy_final(pdf, context.options.output_file, context)
def exec_concurrent(context):
"""Execute the pipeline concurrent"""
# TODO: triage
context.tick()
# Run exec_page_sync on every page context
max_workers = min(len(context.pdfinfo), context.options.jobs)
context.log.info("Start processing %d pages concurrent" % max_workers)
if max_workers > 1:
context.log.info("Start processing %d pages concurrent" % max_workers)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
layers = executor.map(exec_page_sync, context.get_page_contexts())
context.tick()
# Output sidecar text
if context.options.sidecar:
sidecars = [layer[3] for layer in layers]
@@ -166,10 +137,10 @@ def exec_concurrent(context):
# Merge layers to one single pdf
pdf = weave_layers(layers, context)
context.tick()
# PDF/A and metadata
pdf = post_process(pdf, context)
context.tick()
# Copy PDF file to destination
copy_final(pdf, context.options.output_file, context)
@@ -178,8 +149,8 @@ def run_pipeline(options):
if not check_closed_streams(options):
return ExitCode.bad_args
log = get_logger(options, 'Pipeline')
preamble(log)
log = get_logger(options, 'Setup: ')
log.debug('ocrmypdf ' + VERSION)
check_code = check_options(options, log)
if check_code != ExitCode.ok:
return check_code
@@ -211,20 +182,18 @@ def run_pipeline(options):
os.nice(5)
try:
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(start_input_file)
steps = 5 + len(pdfinfo) * 3
t = tqdm(total=steps, bar_format='{l_bar}{bar}{n_fmt}/{total_fmt}')
# Triage image or pdf
origin_pdf = triage(start_input_file, os.path.join(work_folder, 'origin.pdf'), options, log)
context = PDFContext(options, work_folder, start_input_file, pdfinfo, tick=lambda n: t.update(n))
# Gather pdfinfo and create context
pdfinfo = get_pdfinfo(origin_pdf)
context = PDFContext(options, work_folder, origin_pdf, pdfinfo)
# Validate options are okey for this pdf
validate_pdfinfo_options(context)
# Execute the pipeline
exec_concurrent(context)
t.update()
t.close()
except ExitCodeException as e:
return e.exit_code
except Exception as e:
+2 -8
View File
@@ -26,8 +26,6 @@ from pathlib import Path
import PIL
from . import VERSION
from ._unicodefun import verify_python3_env
from .exec import (
@@ -361,10 +359,6 @@ def log_page_orientations(pdfinfo, _log):
_log.info('Page orientations detected: ' + ' '.join(orientations))
def preamble(_log):
_log.debug('ocrmypdf ' + VERSION)
def check_environ(options, _log):
old_envvars = (
'OCRMYPDF_TESSERACT',
@@ -387,14 +381,14 @@ def create_input_file(options, log, work_folder):
if options.input_file == '-':
# stdin
log.info('reading file from standard input')
target = os.path.join(work_folder, 'stdin.pdf')
target = os.path.join(work_folder, 'stdin')
with open(target, 'wb') as stream_buffer:
from shutil import copyfileobj
copyfileobj(sys.stdin.buffer, stream_buffer)
return target
else:
try:
target = os.path.join(work_folder, os.path.basename(options.input_file))
target = os.path.join(work_folder, 'origin')
re_symlink(options.input_file, target, log)
return target
except FileNotFoundError: