pipeline: Move weave* to its own file

This commit is contained in:
James R. Barlow
2018-05-16 22:08:31 -07:00
parent 2b5f23a2d1
commit 0843b5939c
3 changed files with 342 additions and 271 deletions
+332
View File
@@ -0,0 +1,332 @@
# © 2018 James R. Barlow: github.com/jbarlow83
#
# This file is part of OCRmyPDF.
#
# OCRmyPDF is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OCRmyPDF is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with OCRmyPDF. If not, see <http://www.gnu.org/licenses/>.
from pathlib import Path
from itertools import groupby
import pikepdf
from .helpers import flatten_groups, page_number
def _update_page_resources(*, page, font, font_key, procset):
"Update this page's fonts with a reference to the Glyphless font"
if '/Resources' not in page:
page['/Resources'] = pikepdf.Dictionary({})
resources = page['/Resources']
try:
fonts = resources['/Font']
except KeyError:
fonts = pikepdf.Dictionary({})
if font_key not in fonts:
fonts[font_key] = font
resources['/Font'] = fonts
# Reassign /ProcSet to one that just lists everything - ProcSet is
# obsolete and doesn't matter but recommended for old viewer support
resources['/ProcSet'] = procset
def _weave_layers_graft(
*, pdf_base, page_num, text, font, font_key, procset, rotation, log):
"""Insert the text layer from text page 0 on to pdf_base at page_num"""
log.debug("Grafting")
if Path(text).stat().st_size == 0:
return
# This is a pointer indicating a specific page in the base file
pdf_text = pikepdf.open(text)
pdf_text_contents = pdf_text.pages[0].Contents.read_bytes()
base_page = pdf_base.pages.p(page_num)
# The text page always will be oriented up by this stage but the original
# content may have a rotation applied. Wrap the text stream with a rotation
# so it will be oriented the same way as the rest of the page content.
# (Previous versions OCRmyPDF rotated the content layer to match the text.)
mediabox = [float(pdf_text.pages[0].MediaBox[v].decode())
for v in range(4)]
wt, ht = mediabox[2] - mediabox[0], mediabox[3] - mediabox[1]
mediabox = [float(base_page.MediaBox[v].decode())
for v in range(4)]
wp, hp = mediabox[2] - mediabox[0], mediabox[3] - mediabox[1]
translate = pikepdf.PdfMatrix((1, 0, 0, 1, -wt / 2, -ht / 2))
untranslate = pikepdf.PdfMatrix((1, 0, 0, 1, wp / 2, hp / 2))
# -rotation because the input is a clockwise angle and this formula
# uses CCW
rotation = -rotation % 360
if rotation == 0:
c, s = 1, 0
elif rotation == 90:
c, s = 0, 1
elif rotation == 180:
c, s = -1, 0
elif rotation == 270:
c, s = 0, -1
else:
raise NotImplementedError("rotation to arbitrary angle")
rotate = pikepdf.PdfMatrix((c, s, -s, c, 0, 0))
# Because of rounding of DPI, we might get a text layer that is not
# identically sized to the target page. Scale to adjust. Normally this
# is within 0.998.
scale_x = wp / wt
scale_y = hp / ht
if rotation % 90 == 0:
scale_x, scale_y = scale_y, scale_x
log.debug('%r', (scale_x, scale_y))
scale = pikepdf.PdfMatrix((scale_x, 0, 0, scale_y, 0, 0))
# Translate the text so it is centered at (0, 0), rotate it there, adjust
# for a size different between initial and text PDF, then untranslate
ctm = translate @ rotate @ scale @ untranslate
pdf_text_contents = (
b'q %s cm\n' % ctm.encode() +
pdf_text_contents +
b'\nQ\n'
)
new_text_layer = pikepdf.Stream(pdf_base, pdf_text_contents)
base_page.page_contents_add(new_text_layer, prepend=True)
_update_page_resources(
page=base_page, font=font, font_key=font_key, procset=procset
)
def _find_font(text, pdf_base):
"Copy a font from the filename text into pdf_base"
font, font_key = None, None
possible_font_names = ('/f-0-0', '/F1')
try:
pdf_text = pikepdf.open(text)
pdf_text_fonts = pdf_text.pages[0].Resources.get('/Font', {})
except Exception:
return None, None
for f in possible_font_names:
pdf_text_font = pdf_text_fonts.get(f, None)
if pdf_text_font is not None:
font_key = f
break
if pdf_text_font:
font = pdf_base.copy_foreign(pdf_text_font)
return font, font_key
def _fix_toc(pdf_base, pageref_remap, log):
"""Repair the table of contents
Whenever we replace a page wholesale, it gets assigned a new objgen number
and other references to it within the PDF become invalid, most notably in
the table of contents (/Outlines in PDF-speak). In weave_layers we collect
pageref_remap, a mapping that describes the new objgen number given an old
one. (objgen is a tuple, and the gen is almost always zero.)
The /Outlines data structure is a messy data structure, but rather than
navigating hierarchically we just track unique nodes. Enqueue nodes when
we find them, and never visit them again. set() is awesome. We look for
the two types of object in the table of contents that can be page bookmarks
and update the page entry.
It may ultimately be better to find a way to rebuild a page in place.
"""
visited = set()
queue = set()
link_keys = ('/Parent', '/First', '/Last', '/Prev', '/Next')
if not '/Outlines' in pdf_base.root:
return
if not pageref_remap:
return
def remap_dest(dest_node):
if not isinstance(dest_node, pikepdf.Array):
return
pageref = dest_node[0]
if pageref['/Type'] == '/Page' and \
pageref._objgen in pageref_remap:
new_objgen = pageref_remap[pageref._objgen]
dest_node[0] = pdf_base._get_object_id(*new_objgen)
queue.add(pdf_base.root.Outlines._objgen)
while queue:
objgen = queue.pop()
visited.add(objgen)
node = pdf_base._get_object_id(*objgen)
log.debug('fix toc: visiting %r', objgen)
# Enumerate other nodes we could visit from here
for key in link_keys:
if key not in node:
continue
item = node[key]
if not item.is_indirect:
continue
objgen = item._objgen
if objgen not in visited:
queue.add(objgen)
if '/Dest' in node:
remap_dest(node['/Dest'])
elif '/A' in node:
if '/S' in node['/A'] and node['/A']['/S'] == '/GoTo':
remap_dest(node['/A']['/D'])
def weave_layers(
infiles,
output_file,
log,
context):
"""Apply text layer and/or image layer changes to baseline file
This is where the magic happens. infiles will be the main PDF to modify,
and optional .text.pdf and .image-layer.pdf files, organized however ruffus
organizes them.
From .text.pdf, we copy the content stream (which contains the Tesseract
OCR results), and rotate it into place. The first time we do this, we also
copy the GlyphlessFont, and then reference that font again.
For .image-layer.pdf, we check if this is a "pointer" to the original file,
or a new file. If a new file, we replace the page and remember that we
replaced this page.
Every 100 open files, we save intermediate results, to avoid any resource
limits, since pikepdf/qpdf need to keep a lot of open file handles in the
background. When objects are copied from one file to another qpdf, qpdf
doesn't actually copy the data until asked to write, so all the resources
it may need to remain available.
For completeness, we set up a /ProcSet on every page, although it's
unlikely any PDF viewer cares about this anymore.
"""
def input_sorter(key):
try:
return page_number(key)
except ValueError:
return -1
flat_inputs = sorted(flatten_groups(infiles), key=input_sorter)
groups = groupby(flat_inputs, key=input_sorter)
# Extract first item
_, basegroup = next(groups)
base = list(basegroup)[0]
path_base = Path(base).resolve()
pdf_base = pikepdf.open(path_base)
keep_open = []
font, font_key, procset = None, None, None
pdfinfo = context.get_pdfinfo()
pagerefs = {}
procset = pdf_base.make_indirect(
pikepdf.Object.parse(b'[ /PDF /Text /ImageB /ImageC /ImageI ]'))
# Iterate rest
for page_num, layers in groups:
layers = list(layers)
log.debug(page_num)
log.debug(layers)
text = next(
(ii for ii in layers if ii.endswith('.text.pdf')), None
)
image = next(
(ii for ii in layers if ii.endswith('.image-layer.pdf')), None
)
if text and not font:
font, font_key = _find_font(text, pdf_base)
replacing = False
content_rotation = pdfinfo[page_num - 1].rotation
path_image = Path(image).resolve() if image else None
if path_image is not None and path_image != path_base:
# We are replacing the old page with a rasterized PDF of the new
# page
log.debug("Replace")
old_objgen = pdf_base.pages[page_num - 1]._objgen
pdf_image = pikepdf.open(image)
keep_open.append(pdf_image)
image_page = pdf_image.pages[0]
pdf_base.pages[page_num - 1] = image_page
# We're adding a new page, which will get a new objgen number pair,
# so we need to update any references to it. qpdf did not like
# my attempt to update the old object in place, but that is an
# option to consider
pagerefs[old_objgen] = pdf_base.pages[page_num - 1]._objgen
replacing = True
autorotate_correction = context.get_rotation(page_num - 1)
if replacing:
content_rotation = autorotate_correction
text_rotation = autorotate_correction
text_misaligned = (text_rotation - content_rotation) % 360
log.debug('%r', [
text_rotation, autorotate_correction, text_misaligned,
content_rotation]
)
if text and font:
# Graft the text layer onto this page, whether new or old
_weave_layers_graft(
pdf_base=pdf_base, page_num=page_num, text=text, font=font,
font_key=font_key, rotation=text_misaligned, procset=procset,
log=log
)
# Correct the rotation if applicable
pdf_base.pages[page_num - 1].Rotate = \
(content_rotation - autorotate_correction) % 360
if len(keep_open) > 100:
# qpdf limitations require us to keep files open when we intend
# to copy content from them before saving. However, we want to keep
# a lid on file handles and memory usage, so for big files we're
# going to stop and save periodically. Attach the font to page 1
# even if page 1 doesn't use it, so we have a way to get it back.
page0 = pdf_base.pages[0]
_update_page_resources(
page=page0, font=font, font_key=font_key, procset=procset)
interim = output_file + '_working{}.pdf'.format(page_num)
pdf_base.save(interim)
del pdf_base
keep_open = []
pdf_base = pikepdf.open(interim)
procset = pdf_base.pages[0].Resources.ProcSet
font = pdf_base.pages[0].Resources.Font.get(font_key)
_fix_toc(pdf_base, pagerefs, log)
pdf_base.save(output_file)
+8
View File
@@ -168,3 +168,11 @@ if sys.version_info[0:2] <= (3, 5):
else:
universal_open = open
fspath = os.fspath
def flatten_groups(groups):
for obj in groups:
if is_iterable_notstr(obj):
yield from obj
else:
yield obj
+2 -271
View File
@@ -19,7 +19,6 @@ from contextlib import suppress
from shutil import copyfileobj
from pathlib import Path
from datetime import datetime, timezone
from itertools import groupby
import sys
import os
@@ -36,7 +35,7 @@ from ruffus import formatter, regex, Pipeline, suffix
from .hocrtransform import HocrTransform
from .pdfinfo import PdfInfo, Encoding, Colorspace
from .pdfa import generate_pdfa_ps, encode_pdf_date
from .helpers import re_symlink, is_iterable_notstr, page_number
from .helpers import re_symlink, is_iterable_notstr, page_number, flatten_groups
from .exec import ghostscript, tesseract, qpdf
from .lib import fitz
from .exceptions import PdfMergeFailedError, UnsupportedImageFormatError, \
@@ -44,6 +43,7 @@ from .exceptions import PdfMergeFailedError, UnsupportedImageFormatError, \
from . import leptonica
from . import PROGRAM_NAME, VERSION
from ._optimize import optimize
from ._weave import weave_layers
VECTOR_PAGE_DPI = 400
@@ -679,275 +679,6 @@ def render_hocr_page(
interwordSpaces=True)
def flatten_groups(groups):
for obj in groups:
if is_iterable_notstr(obj):
yield from obj
else:
yield obj
def _update_page_resources(*, page, font, font_key, procset):
# Update page fonts with reference to Glyphless
if '/Resources' not in page:
page['/Resources'] = pikepdf.Dictionary({})
resources = page['/Resources']
try:
fonts = resources['/Font']
except KeyError:
fonts = pikepdf.Dictionary({})
if font_key not in fonts:
fonts[font_key] = font
resources['/Font'] = fonts
# Reassign /ProcSet to one that just lists everything - ProcSet is
# obsolete and doesn't matter but recommended for old viewer support
resources['/ProcSet'] = procset
def _weave_layers_graft(
*, pdf_base, page_num, text, font, font_key, procset, rotation, log):
log.debug("Grafting")
if Path(text).stat().st_size == 0:
return
# This is a pointer indicating a specific page in the base file
pdf_text = pikepdf.open(text)
pdf_text_contents = pdf_text.pages[0].Contents.read_bytes()
base_page = pdf_base.pages.p(page_num)
# The text page always will be oriented up by this stage but the original
# content may have a rotation applied. Wrap the text stream with a rotation
# so it will be oriented the same way as the rest of the page content.
# (Previous versions OCRmyPDF rotated the content layer to match the text.)
mediabox = [float(pdf_text.pages[0].MediaBox[v].decode())
for v in range(4)]
wt, ht = mediabox[2] - mediabox[0], mediabox[3] - mediabox[1]
mediabox = [float(base_page.MediaBox[v].decode())
for v in range(4)]
wp, hp = mediabox[2] - mediabox[0], mediabox[3] - mediabox[1]
translate = pikepdf.PdfMatrix((1, 0, 0, 1, -wt / 2, -ht / 2))
untranslate = pikepdf.PdfMatrix((1, 0, 0, 1, wp / 2, hp / 2))
# -rotation because the input is a clockwise angle and this formula
# uses CCW
rotation = -rotation % 360
if rotation == 0:
c, s = 1, 0
elif rotation == 90:
c, s = 0, 1
elif rotation == 180:
c, s = -1, 0
elif rotation == 270:
c, s = 0, -1
else:
raise NotImplementedError("rotation to arbitrary angle")
rotate = pikepdf.PdfMatrix((c, s, -s, c, 0, 0))
# Because of rounding of DPI, we might get a text layer that is not
# identically sized to the target page. Scale to adjust. Normally this
# is within 0.998.
scale_x = wp / wt
scale_y = hp / ht
if rotation % 90 == 0:
scale_x, scale_y = scale_y, scale_x
log.debug('%r', (scale_x, scale_y))
scale = pikepdf.PdfMatrix((scale_x, 0, 0, scale_y, 0, 0))
# Translate the text so it is centered at (0, 0), rotate it there, adjust
# for a size different between initial and text PDF, then untranslate
ctm = translate @ rotate @ scale @ untranslate
pdf_text_contents = (
b'q %s cm\n' % ctm.encode() +
pdf_text_contents +
b'\nQ\n'
)
new_text_layer = pikepdf.Stream(pdf_base, pdf_text_contents)
base_page.page_contents_add(new_text_layer, prepend=True)
_update_page_resources(
page=base_page, font=font, font_key=font_key, procset=procset
)
def _find_font(text, pdf_base):
"Copy a font from the filename text into pdf_base"
font, font_key = None, None
possible_font_names = ('/f-0-0', '/F1')
try:
pdf_text = pikepdf.open(text)
pdf_text_fonts = pdf_text.pages[0].Resources.get('/Font', {})
except Exception:
return None, None
for f in possible_font_names:
pdf_text_font = pdf_text_fonts.get(f, None)
if pdf_text_font is not None:
font_key = f
break
if pdf_text_font:
font = pdf_base.copy_foreign(pdf_text_font)
return font, font_key
def _fix_toc(pdf_base, pageref_remap, log):
visited = set()
queue = set()
link_keys = ('/Parent', '/First', '/Last', '/Prev', '/Next')
if not '/Outlines' in pdf_base.root:
return
if not pageref_remap:
return
def remap_dest(dest_node):
if not isinstance(dest_node, pikepdf.Array):
return
pageref = dest_node[0]
if pageref['/Type'] == '/Page' and \
pageref._objgen in pageref_remap:
new_objgen = pageref_remap[pageref._objgen]
dest_node[0] = pdf_base._get_object_id(*new_objgen)
queue.add(pdf_base.root.Outlines._objgen)
while queue:
objgen = queue.pop()
visited.add(objgen)
node = pdf_base._get_object_id(*objgen)
log.debug('fix toc: visiting %r', objgen)
# Enumerate other nodes we could visit from here
for key in link_keys:
if key not in node:
continue
item = node[key]
if not item.is_indirect:
continue
objgen = item._objgen
if objgen not in visited:
queue.add(objgen)
log.debug(repr(node))
if '/Dest' in node:
remap_dest(node['/Dest'])
elif '/A' in node:
if '/S' in node['/A'] and node['/A']['/S'] == '/GoTo':
remap_dest(node['/A']['/D'])
def weave_layers(
infiles,
output_file,
log,
context):
"Apply text layer and/or image layer changes to baseline file"
def input_sorter(key):
try:
return page_number(key)
except ValueError:
return -1
flat_inputs = sorted(flatten_groups(infiles), key=input_sorter)
groups = groupby(flat_inputs, key=input_sorter)
# Extract first item
_, basegroup = next(groups)
base = list(basegroup)[0]
path_base = Path(base).resolve()
pdf_base = pikepdf.open(path_base)
keep_open = []
font, font_key, procset = None, None, None
pdfinfo = context.get_pdfinfo()
pagerefs = {}
procset = pdf_base.make_indirect(
pikepdf.Object.parse(b'[ /PDF /Text /ImageB /ImageC /ImageI ]'))
# Iterate rest
for page_num, layers in groups:
layers = list(layers)
log.debug(page_num)
log.debug(layers)
text = next(
(ii for ii in layers if ii.endswith('.text.pdf')), None
)
image = next(
(ii for ii in layers if ii.endswith('.image-layer.pdf')), None
)
if text and not font:
font, font_key = _find_font(text, pdf_base)
replacing = False
content_rotation = pdfinfo[page_num - 1].rotation
path_image = Path(image).resolve() if image else None
if path_image is not None and path_image != path_base:
# We are replacing the old page
log.debug("Replace")
old_objgen = pdf_base.pages[page_num - 1]._objgen
pdf_image = pikepdf.open(image)
keep_open.append(pdf_image)
image_page = pdf_image.pages[0]
pdf_base.pages[page_num - 1] = image_page
pagerefs[old_objgen] = pdf_base.pages[page_num - 1]._objgen
replacing = True
autorotate_correction = context.get_rotation(page_num - 1)
if replacing:
content_rotation = autorotate_correction
text_rotation = autorotate_correction
text_misaligned = (text_rotation - content_rotation) % 360
log.debug('%r', [
text_rotation, autorotate_correction, text_misaligned,
content_rotation]
)
if text and font:
# Graft the text layer onto this page, whether new or old
_weave_layers_graft(
pdf_base=pdf_base, page_num=page_num, text=text, font=font,
font_key=font_key, rotation=text_misaligned, procset=procset,
log=log
)
# Correct the rotation if applicable
pdf_base.pages[page_num - 1].Rotate = \
(content_rotation - autorotate_correction) % 360
if len(keep_open) > 100:
# qpdf limitations require us to keep files open when we intend
# to copy content from them before saving. However, we want to keep
# a lid on file handles and memory usage, so for big files we're
# going to stop and save periodically. Attach the font to page 1
# even if page 1 doesn't use it, so we have a way to get it back.
page0 = pdf_base.pages[0]
_update_page_resources(
page=page0, font=font, font_key=font_key, procset=procset)
interim = output_file + '_working{}.pdf'.format(page_num)
pdf_base.save(interim)
del pdf_base
keep_open = []
pdf_base = pikepdf.open(interim)
procset = pdf_base.pages[0].Resources.ProcSet
font = pdf_base.pages[0].Resources.Font.get(font_key)
_fix_toc(pdf_base, pagerefs, log)
pdf_base.save(output_file)
def ocr_tesseract_textonly_pdf(
infiles,
outfiles,