feat: add triage step
remove tqdm demo
This commit is contained in:
@@ -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
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user