380 lines
13 KiB
Python
380 lines
13 KiB
Python
|
# SPDX-FileCopyrightText: 2024 geisserml <geisserml@gmail.com>
|
||
|
|
# SPDX-License-Identifier: Apache-2.0 OR BSD-3-Clause
|
||
|
|
|
||
|
|
import os
|
||
|
|
import math
|
||
|
|
import logging
|
||
|
|
import functools
|
||
|
|
from pathlib import Path
|
||
|
|
import multiprocessing as mp
|
||
|
|
import concurrent.futures as ft
|
||
|
|
|
||
|
|
try:
|
||
|
|
import cv2
|
||
|
|
except ImportError:
|
||
|
|
cv2 = None
|
||
|
|
|
||
|
|
import pypdfium2._helpers as pdfium
|
||
|
|
import pypdfium2.internal as pdfium_i
|
||
|
|
import pypdfium2.raw as pdfium_r
|
||
|
|
# TODO? consider dotted access
|
||
|
|
from pypdfium2._cli._parsers import (
|
||
|
|
add_input, get_input,
|
||
|
|
setup_logging,
|
||
|
|
BooleanOptionalAction,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def _bitmap_wrapper_foreign_simple(width, height, format, *args, **kwargs):
|
||
|
|
if format == pdfium_r.FPDFBitmap_BGRx:
|
||
|
|
use_alpha = False
|
||
|
|
elif format == pdfium_r.FPDFBitmap_BGRA:
|
||
|
|
use_alpha = True
|
||
|
|
else:
|
||
|
|
raise RuntimeError(f"Cannot create foreign_simple bitmap with bitmap type {pdfium_i.BitmapTypeToStr[format]}.")
|
||
|
|
return pdfium.PdfBitmap.new_foreign_simple(width, height, use_alpha, *args, **kwargs)
|
||
|
|
|
||
|
|
BitmapMakers = dict(
|
||
|
|
native = pdfium.PdfBitmap.new_native,
|
||
|
|
foreign = pdfium.PdfBitmap.new_foreign,
|
||
|
|
foreign_packed = functools.partial(pdfium.PdfBitmap.new_foreign, force_packed=True),
|
||
|
|
foreign_simple = _bitmap_wrapper_foreign_simple,
|
||
|
|
)
|
||
|
|
|
||
|
|
CsFields = ("path_fill", "path_stroke", "text_fill", "text_stroke")
|
||
|
|
ColorOpts = dict(metavar="C", nargs=4, type=int)
|
||
|
|
SampleTheme = dict(
|
||
|
|
# TODO improve colors - currently it's just some random ones to distinguish the different drawings
|
||
|
|
path_fill = (170, 100, 0, 255), # dark orange
|
||
|
|
path_stroke = (0, 150, 255, 255), # sky blue
|
||
|
|
text_fill = (255, 255, 255, 255), # white
|
||
|
|
text_stroke = (150, 255, 0, 255), # green
|
||
|
|
)
|
||
|
|
|
||
|
|
def attach(parser):
|
||
|
|
add_input(parser, pages=True)
|
||
|
|
parser.add_argument(
|
||
|
|
"--output", "-o",
|
||
|
|
type = Path,
|
||
|
|
required = True,
|
||
|
|
help = "Output directory where the serially numbered images shall be placed.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--prefix",
|
||
|
|
help = "Custom prefix for the images. Defaults to the input filename's stem.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--format", "-f",
|
||
|
|
default = "jpg",
|
||
|
|
type = str.lower,
|
||
|
|
help = "The image format to use.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--engine",
|
||
|
|
dest = "engine_cls",
|
||
|
|
type = lambda k: {"pil": PILEngine, "numpy+cv2": NumpyCV2Engine}[k.lower()],
|
||
|
|
help = "The saver engine to use (pil, numpy+cv2)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--scale",
|
||
|
|
default = 1,
|
||
|
|
type = float,
|
||
|
|
help = "Define the resolution of the output images. By default, one PDF point (1/72in) is rendered to 1x1 pixel. This factor scales the number of pixels that represent one point.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--rotation",
|
||
|
|
default = 0,
|
||
|
|
type = int,
|
||
|
|
choices = (0, 90, 180, 270),
|
||
|
|
help = "Rotate pages by 90, 180 or 270 degrees.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--fill-color",
|
||
|
|
help = "Color the bitmap will be filled with before rendering. It shall be given in RGBA format as a sequence of integers ranging from 0 to 255. Defaults to white.",
|
||
|
|
**ColorOpts,
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--optimize-mode",
|
||
|
|
choices = ("lcd", "print"),
|
||
|
|
help = "The rendering optimisation mode. None if not given.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--crop",
|
||
|
|
nargs = 4,
|
||
|
|
type = float,
|
||
|
|
default = (0, 0, 0, 0),
|
||
|
|
help = "Amount to crop from (left, bottom, right, top).",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--draw-annots",
|
||
|
|
action = BooleanOptionalAction,
|
||
|
|
default = True,
|
||
|
|
help = "Whether annotations may be shown (default: true).",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--draw-forms",
|
||
|
|
action = BooleanOptionalAction,
|
||
|
|
default = True,
|
||
|
|
help = "Whether forms may be shown (default: true).",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--no-antialias",
|
||
|
|
nargs = "+",
|
||
|
|
default = [],
|
||
|
|
choices = ("text", "image", "path"),
|
||
|
|
type = str.lower,
|
||
|
|
help = "Item types that shall not be smoothed.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--force-halftone",
|
||
|
|
action = "store_true",
|
||
|
|
help = "Always use halftone for image stretching.",
|
||
|
|
)
|
||
|
|
|
||
|
|
bitmap = parser.add_argument_group(
|
||
|
|
title = "Bitmap options",
|
||
|
|
description = "Bitmap config, including pixel format.",
|
||
|
|
)
|
||
|
|
bitmap.add_argument(
|
||
|
|
"--bitmap-maker",
|
||
|
|
choices = BitmapMakers.keys(),
|
||
|
|
default = "native",
|
||
|
|
help = "The bitmap maker to use.",
|
||
|
|
type = str.lower,
|
||
|
|
)
|
||
|
|
bitmap.add_argument(
|
||
|
|
"--grayscale",
|
||
|
|
action = "store_true",
|
||
|
|
help = "Whether to render in grayscale mode (no colors).",
|
||
|
|
)
|
||
|
|
bitmap.add_argument(
|
||
|
|
"--byteorder",
|
||
|
|
dest = "rev_byteorder",
|
||
|
|
type = lambda v: {"bgr": False, "rgb": True}[v.lower()],
|
||
|
|
help = "Whether to use BGR or RGB byteorder (default: conditional).",
|
||
|
|
)
|
||
|
|
bitmap.add_argument(
|
||
|
|
"--x-channel",
|
||
|
|
dest = "prefer_bgrx",
|
||
|
|
action = BooleanOptionalAction,
|
||
|
|
help = "Whether to prefer BGRx/RGBx over BGR/RGB (default: conditional).",
|
||
|
|
)
|
||
|
|
|
||
|
|
parallel = parser.add_argument_group(
|
||
|
|
title = "Parallelization",
|
||
|
|
description = "Options for rendering with multiple processes.",
|
||
|
|
)
|
||
|
|
parallel.add_argument(
|
||
|
|
"--linear",
|
||
|
|
nargs = "?",
|
||
|
|
type = int,
|
||
|
|
const = math.inf,
|
||
|
|
help = "Render non-parallel if page count is less or equal to the specified value (default is conditional). If this flag is given without a value, then render linear regardless of document length.",
|
||
|
|
)
|
||
|
|
parallel.add_argument(
|
||
|
|
"--processes",
|
||
|
|
default = os.cpu_count(),
|
||
|
|
type = int,
|
||
|
|
help = "The maximum number of parallel rendering processes. Defaults to the number of CPU cores.",
|
||
|
|
)
|
||
|
|
parallel.add_argument(
|
||
|
|
"--parallel-strategy",
|
||
|
|
choices = ("spawn", "forkserver", "fork"),
|
||
|
|
default = "spawn",
|
||
|
|
type = str.lower,
|
||
|
|
help = "The process start method to use. ('fork' is discouraged due to stability issues.)",
|
||
|
|
)
|
||
|
|
parallel.add_argument(
|
||
|
|
"--parallel-lib",
|
||
|
|
choices = ("mp", "ft"),
|
||
|
|
default = "mp",
|
||
|
|
type = str.lower,
|
||
|
|
help = "The parallelization module to use (mp = multiprocessing, ft = concurrent.futures).",
|
||
|
|
)
|
||
|
|
parallel.add_argument(
|
||
|
|
"--parallel-map",
|
||
|
|
type = str.lower,
|
||
|
|
help = "The map function to use (backend specific, the default is an iterative map)."
|
||
|
|
)
|
||
|
|
|
||
|
|
color_scheme = parser.add_argument_group(
|
||
|
|
title = "Forced color scheme",
|
||
|
|
description = "Options for using pdfium's forced color scheme renderer. Deprecated, considered not useful.",
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--sample-theme",
|
||
|
|
action = "store_true",
|
||
|
|
help = "Use a dark background sample theme as base. Explicit color params override selectively."
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--path-fill",
|
||
|
|
**ColorOpts
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--path-stroke",
|
||
|
|
**ColorOpts
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--text-fill",
|
||
|
|
**ColorOpts
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--text-stroke",
|
||
|
|
**ColorOpts
|
||
|
|
)
|
||
|
|
color_scheme.add_argument(
|
||
|
|
"--fill-to-stroke",
|
||
|
|
action = "store_true",
|
||
|
|
help = "Only draw borders around fill areas using the `path_stroke` color, instead of filling with the `path_fill` color.",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class SavingEngine:
|
||
|
|
|
||
|
|
def __init__(self, path_parts):
|
||
|
|
self._path_parts = path_parts
|
||
|
|
|
||
|
|
def _get_path(self, i):
|
||
|
|
output_dir, prefix, n_digits, format = self._path_parts
|
||
|
|
return output_dir / f"{prefix}{i+1:0{n_digits}d}.{format}"
|
||
|
|
|
||
|
|
def __call__(self, bitmap, i):
|
||
|
|
out_path = self._get_path(i)
|
||
|
|
self._saving_hook(out_path, bitmap)
|
||
|
|
logger.info(f"Wrote page {i+1} as {out_path.name}")
|
||
|
|
|
||
|
|
|
||
|
|
class PILEngine (SavingEngine):
|
||
|
|
def _saving_hook(self, out_path, bitmap):
|
||
|
|
bitmap.to_pil().save(out_path)
|
||
|
|
|
||
|
|
class NumpyCV2Engine (SavingEngine):
|
||
|
|
def _saving_hook(self, out_path, bitmap):
|
||
|
|
cv2.imwrite(str(out_path), bitmap.to_numpy())
|
||
|
|
|
||
|
|
|
||
|
|
def _render_parallel_init(extra_init, input, password, may_init_forms, kwargs, engine):
|
||
|
|
|
||
|
|
if extra_init:
|
||
|
|
extra_init()
|
||
|
|
|
||
|
|
logger.info(f"Initializing data for process {os.getpid()}")
|
||
|
|
|
||
|
|
pdf = pdfium.PdfDocument(input, password=password, autoclose=True)
|
||
|
|
if may_init_forms:
|
||
|
|
pdf.init_forms()
|
||
|
|
|
||
|
|
global ProcObjs
|
||
|
|
ProcObjs = (pdf, kwargs, engine)
|
||
|
|
|
||
|
|
|
||
|
|
def _render_job(i, pdf, kwargs, engine):
|
||
|
|
# logger.info(f"Started page {i+1} ...")
|
||
|
|
page = pdf[i]
|
||
|
|
bitmap = page.render(**kwargs)
|
||
|
|
engine(bitmap, i)
|
||
|
|
|
||
|
|
def _render_parallel_job(i):
|
||
|
|
global ProcObjs; _render_job(i, *ProcObjs)
|
||
|
|
|
||
|
|
|
||
|
|
def main(args):
|
||
|
|
|
||
|
|
# TODO turn into a python-usable API yielding output paths as they are written
|
||
|
|
|
||
|
|
pdf = get_input(args, init_forms=args.draw_forms)
|
||
|
|
|
||
|
|
# TODO move to parsers?
|
||
|
|
pdf_len = len(pdf)
|
||
|
|
if not all(0 <= i < pdf_len for i in args.pages):
|
||
|
|
raise ValueError("Out-of-bounds page indices are prohibited.")
|
||
|
|
if len(args.pages) != len(set(args.pages)):
|
||
|
|
raise ValueError("Duplicate page indices are prohibited.")
|
||
|
|
|
||
|
|
if args.prefix is None:
|
||
|
|
args.prefix = f"{args.input.stem}_"
|
||
|
|
if args.fill_color is None:
|
||
|
|
args.fill_color = (0, 0, 0, 255) if args.sample_theme else (255, 255, 255, 255)
|
||
|
|
if args.linear is None:
|
||
|
|
args.linear = 6 if args.format == "jpg" else 3
|
||
|
|
|
||
|
|
# numpy+cv2 is much faster for PNG, and PIL faster for JPG, but this might simply be due to different encoding defaults
|
||
|
|
if args.engine_cls is None:
|
||
|
|
if cv2 != None and args.format == "png":
|
||
|
|
args.engine_cls = NumpyCV2Engine
|
||
|
|
else:
|
||
|
|
args.engine_cls = PILEngine
|
||
|
|
|
||
|
|
# PIL is faster with rev_byteorder and prefer_bgrx = True, as this achieves a natively supported pixel format. For numpy+cv2 there doesn't seem to be a difference.
|
||
|
|
if args.rev_byteorder is None:
|
||
|
|
args.rev_byteorder = args.engine_cls is PILEngine
|
||
|
|
if args.prefer_bgrx is None:
|
||
|
|
# PIL can't save BGRX as PNG
|
||
|
|
args.prefer_bgrx = args.engine_cls is PILEngine and args.format != "png"
|
||
|
|
|
||
|
|
cs_kwargs = dict()
|
||
|
|
if args.sample_theme:
|
||
|
|
cs_kwargs.update(**SampleTheme)
|
||
|
|
cs_kwargs.update(**{f: getattr(args, f) for f in CsFields if getattr(args, f)})
|
||
|
|
cs = pdfium.PdfColorScheme(**cs_kwargs) if len(cs_kwargs) > 0 else None
|
||
|
|
|
||
|
|
kwargs = dict(
|
||
|
|
scale = args.scale,
|
||
|
|
rotation = args.rotation,
|
||
|
|
crop = args.crop,
|
||
|
|
grayscale = args.grayscale,
|
||
|
|
fill_color = args.fill_color,
|
||
|
|
color_scheme = cs,
|
||
|
|
fill_to_stroke = args.fill_to_stroke,
|
||
|
|
optimize_mode = args.optimize_mode,
|
||
|
|
draw_annots = args.draw_annots,
|
||
|
|
may_draw_forms = args.draw_forms,
|
||
|
|
force_halftone = args.force_halftone,
|
||
|
|
rev_byteorder = args.rev_byteorder,
|
||
|
|
prefer_bgrx = args.prefer_bgrx,
|
||
|
|
bitmap_maker = BitmapMakers[args.bitmap_maker],
|
||
|
|
)
|
||
|
|
for type in args.no_antialias:
|
||
|
|
kwargs[f"no_smooth{type}"] = True
|
||
|
|
|
||
|
|
# TODO dump all args except password?
|
||
|
|
logger.info(f"{args.engine_cls.__name__}, Format: {args.format}, rev_byteorder: {args.rev_byteorder}, prefer_bgrx {args.prefer_bgrx}")
|
||
|
|
|
||
|
|
n_digits = len(str(pdf_len))
|
||
|
|
path_parts = (args.output, args.prefix, n_digits, args.format)
|
||
|
|
engine = args.engine_cls(path_parts)
|
||
|
|
|
||
|
|
if len(args.pages) <= args.linear:
|
||
|
|
|
||
|
|
logger.info("Linear rendering ...")
|
||
|
|
for i in args.pages:
|
||
|
|
_render_job(i, pdf, kwargs, engine)
|
||
|
|
|
||
|
|
else:
|
||
|
|
|
||
|
|
logger.info("Parallel rendering ...")
|
||
|
|
|
||
|
|
ctx = mp.get_context(args.parallel_strategy)
|
||
|
|
# TODO unify using mp.pool.Pool(context=...) ?
|
||
|
|
pool_backends = dict(
|
||
|
|
mp = (ctx.Pool, "imap"),
|
||
|
|
ft = (functools.partial(ft.ProcessPoolExecutor, mp_context=ctx), "map"),
|
||
|
|
)
|
||
|
|
pool_ctor, map_attr = pool_backends[args.parallel_lib]
|
||
|
|
if args.parallel_map:
|
||
|
|
map_attr = args.parallel_map
|
||
|
|
|
||
|
|
extra_init = (setup_logging if args.parallel_strategy in ("spawn", "forkserver") else None)
|
||
|
|
pool_kwargs = dict(
|
||
|
|
initializer = _render_parallel_init,
|
||
|
|
initargs = (extra_init, pdf._input, args.password, args.draw_forms, kwargs, engine),
|
||
|
|
)
|
||
|
|
|
||
|
|
n_procs = min(args.processes, len(args.pages))
|
||
|
|
with pool_ctor(n_procs, **pool_kwargs) as pool:
|
||
|
|
map_func = getattr(pool, map_attr)
|
||
|
|
for _ in map_func(_render_parallel_job, args.pages):
|
||
|
|
pass
|