570 lines
19 KiB
Python
Executable File
570 lines
19 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# SPDX-FileCopyrightText: 2010 Jonathan Brinley
|
||
# SPDX-FileCopyrightText: 2013-2014 Julien Pfefferkorn
|
||
# SPDX-FileCopyrightText: 2022 James R. Barlow
|
||
# SPDX-License-Identifier: MIT
|
||
|
||
"""Transform .hocr and page image to text PDF."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import re
|
||
from dataclasses import dataclass
|
||
from math import atan, cos, sin
|
||
from pathlib import Path
|
||
from typing import Any, NamedTuple
|
||
from xml.etree import ElementTree
|
||
|
||
from pikepdf import PdfMatrix
|
||
|
||
from ocrmypdf.hocrtransform.backends import (
|
||
Canvas,
|
||
)
|
||
from ocrmypdf.hocrtransform.backends.pikepdf import PikepdfCanvas
|
||
from ocrmypdf.hocrtransform.backends.reportlab import (
|
||
ReportlabCanvas,
|
||
black,
|
||
cyan,
|
||
green,
|
||
inch,
|
||
magenta,
|
||
red,
|
||
)
|
||
|
||
# According to Wikipedia these languages are supported in the ISO-8859-1 character
|
||
# set, meaning reportlab can generate them and they are compatible with hocr,
|
||
# assuming Tesseract has the necessary languages installed. Note that there may
|
||
# not be language packs for them.
|
||
HOCR_OK_LANGS = frozenset(
|
||
[
|
||
# Languages fully covered by Latin-1:
|
||
'afr', # Afrikaans
|
||
'alb', # Albanian
|
||
'ast', # Leonese
|
||
'baq', # Basque
|
||
'bre', # Breton
|
||
'cos', # Corsican
|
||
'eng', # English
|
||
'eus', # Basque
|
||
'fao', # Faoese
|
||
'gla', # Scottish Gaelic
|
||
'glg', # Galician
|
||
'glv', # Manx
|
||
'ice', # Icelandic
|
||
'ind', # Indonesian
|
||
'isl', # Icelandic
|
||
'ita', # Italian
|
||
'ltz', # Luxembourgish
|
||
'mal', # Malay Rumi
|
||
'mga', # Irish
|
||
'nor', # Norwegian
|
||
'oci', # Occitan
|
||
'por', # Portugeuse
|
||
'roh', # Romansh
|
||
'sco', # Scots
|
||
'sma', # Sami
|
||
'spa', # Spanish
|
||
'sqi', # Albanian
|
||
'swa', # Swahili
|
||
'swe', # Swedish
|
||
'tgl', # Tagalog
|
||
'wln', # Walloon
|
||
# Languages supported by Latin-1 except for a few rare characters that OCR
|
||
# is probably not trained to recognize anyway:
|
||
'cat', # Catalan
|
||
'cym', # Welsh
|
||
'dan', # Danish
|
||
'deu', # German
|
||
'dut', # Dutch
|
||
'est', # Estonian
|
||
'fin', # Finnish
|
||
'fra', # French
|
||
'hun', # Hungarian
|
||
'kur', # Kurdish
|
||
'nld', # Dutch
|
||
'wel', # Welsh
|
||
]
|
||
)
|
||
|
||
|
||
Element = ElementTree.Element
|
||
|
||
|
||
class Rect(NamedTuple):
|
||
"""A rectangle for managing PDF coordinates."""
|
||
|
||
x1: Any
|
||
y1: Any
|
||
x2: Any
|
||
y2: Any
|
||
|
||
def transform(self, matrix, inverse=True):
|
||
"""Transform the rectangle by the given matrix."""
|
||
if inverse:
|
||
matrix = matrix.inverse()
|
||
return Rect._make(
|
||
[
|
||
matrix.a * self.x1 + matrix.c * self.y1 + matrix.e,
|
||
matrix.b * self.x1 + matrix.d * self.y1 + matrix.f,
|
||
matrix.a * self.x2 + matrix.c * self.y2 + matrix.e,
|
||
matrix.b * self.x2 + matrix.d * self.y2 + matrix.f,
|
||
]
|
||
)
|
||
|
||
|
||
@dataclass
|
||
class DebugRenderOptions:
|
||
"""A class for managing rendering options."""
|
||
|
||
render_paragraph_bbox: bool
|
||
render_baseline: bool
|
||
render_triangle: bool
|
||
render_word_bbox: bool
|
||
|
||
|
||
class HocrTransformError(Exception):
|
||
"""Error while applying hOCR transform."""
|
||
|
||
|
||
class HocrTransform:
|
||
"""A class for converting documents from the hOCR format.
|
||
|
||
For details of the hOCR format, see:
|
||
http://kba.cloud/hocr-spec/.
|
||
"""
|
||
|
||
box_pattern = re.compile(r'bbox((\s+\d+){4})')
|
||
baseline_pattern = re.compile(
|
||
r'''
|
||
baseline \s+
|
||
([\-\+]?\d*\.?\d*) \s+ # +/- decimal float
|
||
([\-\+]?\d+) # +/- int''',
|
||
re.VERBOSE,
|
||
)
|
||
ligatures = str.maketrans(
|
||
{'ff': 'ff', 'ffi': 'ffi', 'ffl': 'ffl', 'fi': 'fi', 'fl': 'fl'}
|
||
)
|
||
|
||
def __init__(self, *, hocr_filename: str | Path, dpi: float):
|
||
"""Initialize the HocrTransform object."""
|
||
self.dpi = dpi
|
||
self.hocr = ElementTree.parse(os.fspath(hocr_filename))
|
||
|
||
# if the hOCR file has a namespace, ElementTree requires its use to
|
||
# find elements
|
||
matches = re.match(r'({.*})html', self.hocr.getroot().tag)
|
||
self.xmlns = ''
|
||
if matches:
|
||
self.xmlns = matches.group(1)
|
||
|
||
# get dimension in pt (not pixel!!!!) of the OCRed image
|
||
self.width, self.height = None, None
|
||
for div in self.hocr.findall(self._child_xpath('div', 'ocr_page')):
|
||
coords = self.element_coordinates(div)
|
||
pt_coords = self.pt_from_pixel(coords)
|
||
self.width = pt_coords.x2 - pt_coords.x1
|
||
self.height = pt_coords.y2 - pt_coords.y1
|
||
# there shouldn't be more than one, and if there is, we don't want
|
||
# it
|
||
break
|
||
if self.width is None or self.height is None:
|
||
raise HocrTransformError("hocr file is missing page dimensions")
|
||
self.render_options = DebugRenderOptions(
|
||
render_baseline=True,
|
||
render_triangle=True,
|
||
render_word_bbox=False,
|
||
render_paragraph_bbox=False,
|
||
)
|
||
|
||
def __str__(self): # pragma: no cover
|
||
"""Return the textual content of the HTML body."""
|
||
if self.hocr is None:
|
||
return ''
|
||
body = self.hocr.find(self._child_xpath('body'))
|
||
if body:
|
||
return self._get_element_text(body)
|
||
else:
|
||
return ''
|
||
|
||
def _get_element_text(self, element: Element):
|
||
"""Return the textual content of the element and its children."""
|
||
text = ''
|
||
if element.text is not None:
|
||
text += element.text
|
||
for child in element:
|
||
text += self._get_element_text(child)
|
||
if element.tail is not None:
|
||
text += element.tail
|
||
return text
|
||
|
||
@classmethod
|
||
def element_coordinates(cls, element: Element) -> Rect:
|
||
"""Get coordinates of the bounding box around an element."""
|
||
out = Rect._make(0 for _ in range(4))
|
||
if 'title' in element.attrib:
|
||
matches = cls.box_pattern.search(element.attrib['title'])
|
||
if matches:
|
||
coords = matches.group(1).split()
|
||
out = Rect._make(int(coords[n]) for n in range(4))
|
||
return out
|
||
|
||
@classmethod
|
||
def baseline(cls, element: Element) -> tuple[float, float]:
|
||
"""Get baseline's slope and intercept."""
|
||
if 'title' in element.attrib:
|
||
matches = cls.baseline_pattern.search(element.attrib['title'])
|
||
if matches:
|
||
return float(matches.group(1)), int(matches.group(2))
|
||
return (0.0, 0.0)
|
||
|
||
def pt_from_pixel(self, pxl: Rect, bottomup=False) -> Rect:
|
||
"""Returns the quantity in PDF units (pt) given quantity in pixels."""
|
||
if bottomup:
|
||
return Rect._make(
|
||
[
|
||
(pxl.x1 / self.dpi * inch),
|
||
self.height - (pxl.y2 / self.dpi * inch), # swap y1/y2
|
||
(pxl.x2 / self.dpi * inch),
|
||
self.height - (pxl.y1 / self.dpi * inch),
|
||
]
|
||
)
|
||
else:
|
||
return Rect._make(
|
||
(c / self.dpi * inch) for c in (pxl.x1, pxl.y1, pxl.x2, pxl.y2)
|
||
)
|
||
|
||
def _child_xpath(self, html_tag: str, html_class: str | None = None) -> str:
|
||
xpath = f".//{self.xmlns}{html_tag}"
|
||
if html_class:
|
||
xpath += f"[@class='{html_class}']"
|
||
return xpath
|
||
|
||
@classmethod
|
||
def replace_unsupported_chars(cls, s: str) -> str:
|
||
"""Replaces characters with those available in the Helvetica typeface."""
|
||
return s.translate(cls.ligatures)
|
||
|
||
def to_pdf(
|
||
self,
|
||
*,
|
||
out_filename: Path,
|
||
image_filename: Path | None = None,
|
||
show_bounding_boxes: bool = False,
|
||
fontname: str = "Helvetica",
|
||
invisible_text: bool = False,
|
||
interword_spaces: bool = False,
|
||
) -> None:
|
||
"""Creates a PDF file with an image superimposed on top of the text.
|
||
|
||
Text is positioned according to the bounding box of the lines in
|
||
the hOCR file.
|
||
The image need not be identical to the image used to create the hOCR
|
||
file.
|
||
It can have a lower resolution, different color mode, etc.
|
||
|
||
Arguments:
|
||
out_filename: Path of PDF to write.
|
||
image_filename: Image to use for this file. If omitted, the OCR text
|
||
is shown.
|
||
show_bounding_boxes: Show bounding boxes around various text regions,
|
||
for debugging.
|
||
fontname: Name of font to use.
|
||
invisible_text: If True, text is rendered invisible so that is
|
||
selectable but never drawn. If False, text is visible and may
|
||
be seen if the image is skipped or deleted in Acrobat.
|
||
interword_spaces: If True, insert spaces between words rather than
|
||
drawing each word without spaces. Generally this improves text
|
||
extraction.
|
||
"""
|
||
# create the PDF file
|
||
# page size in points (1/72 in.)
|
||
pdf = PikepdfCanvas( # ReportlabCanvas(
|
||
out_filename,
|
||
page_size=(self.width, self.height),
|
||
)
|
||
|
||
for elem in self.hocr.iterfind(self._child_xpath('p', 'ocr_par')):
|
||
elemtxt = self._get_element_text(elem).rstrip()
|
||
if len(elemtxt) == 0:
|
||
continue
|
||
|
||
pxl_coords = self.element_coordinates(elem)
|
||
pt = self.pt_from_pixel(pxl_coords, bottomup=True)
|
||
# draw cyan box around paragraph
|
||
if self.render_options.render_paragraph_bbox:
|
||
# pragma: no cover
|
||
pdf.set_stroke_color(cyan)
|
||
pdf.set_line_width(0.1) # no line for bounding box
|
||
pdf.rect(pt.x1, pt.y2, pt.x2 - pt.x1, pt.y2 - pt.y1, fill=0)
|
||
|
||
found_lines = False
|
||
for line in (
|
||
element
|
||
for element in self.hocr.iterfind(self._child_xpath('span'))
|
||
if 'class' in element.attrib
|
||
and element.attrib['class'] in {'ocr_header', 'ocr_line', 'ocr_textfloat'}
|
||
):
|
||
found_lines = True
|
||
self._do_line(
|
||
pdf,
|
||
line,
|
||
"ocrx_word",
|
||
fontname,
|
||
invisible_text,
|
||
interword_spaces,
|
||
)
|
||
|
||
if not found_lines:
|
||
# Tesseract did not report any lines (just words)
|
||
root = self.hocr.find(self._child_xpath('div', 'ocr_page'))
|
||
self._do_line(
|
||
pdf,
|
||
root,
|
||
"ocrx_word",
|
||
fontname,
|
||
invisible_text,
|
||
interword_spaces,
|
||
)
|
||
# put the image on the page, scaled to fill the page
|
||
if image_filename is not None:
|
||
pdf.draw_image(image_filename, 0, 0, width=self.width, height=self.height)
|
||
|
||
# finish up the page and save it
|
||
pdf.save()
|
||
|
||
@classmethod
|
||
def polyval(cls, poly, x): # pragma: no cover
|
||
"""Calculate the value of a polynomial at a point."""
|
||
return x * poly[0] + poly[1]
|
||
|
||
def _do_line(
|
||
self,
|
||
pdf: Canvas,
|
||
line: Element | None,
|
||
elemclass: str,
|
||
fontname: str,
|
||
invisible_text: bool,
|
||
interword_spaces: bool,
|
||
):
|
||
if line is None:
|
||
return
|
||
pxl_line_coords = self.element_coordinates(line)
|
||
line_box = self.pt_from_pixel(pxl_line_coords, bottomup=True)
|
||
line_height = line_box.y2 - line_box.y1
|
||
assert line_box.y2 > line_box.y1
|
||
|
||
slope, pxl_intercept = self.baseline(line)
|
||
if abs(slope) < 0.005:
|
||
slope = 0.0
|
||
angle = atan(slope)
|
||
cos_a, sin_a = cos(angle), sin(angle)
|
||
intercept = pxl_intercept / self.dpi * inch
|
||
|
||
# Enter a new coordinate system with the linebox at the origin
|
||
pdf._cs.push()
|
||
line_matrix = PdfMatrix().translated(line_box.x1, line_box.y1).rotated(-angle)
|
||
pdf._cs.cm(*line_matrix.shorthand)
|
||
cos_a, sin_a = 1, 0
|
||
slope = 0
|
||
intercept = 0
|
||
|
||
cm_line_box = line_box.transform(line_matrix, inverse=True)
|
||
|
||
text = pdf.begin_text()
|
||
|
||
# Don't allow the font to break out of the bounding box. Division by
|
||
# cos_a accounts for extra clearance between the glyph's vertical axis
|
||
# on a sloped baseline and the edge of the bounding box.
|
||
fontsize = (line_height - abs(intercept)) / cos_a
|
||
text.set_font(fontname, fontsize)
|
||
if invisible_text or True:
|
||
text.set_render_mode(3) # Invisible (indicates OCR text)
|
||
|
||
# Intercept is normally negative. Subtracting it will raise the baseline
|
||
# above the bottom of the bounding box (y1). We're in page coordinates,
|
||
# origin bottom left, y2 > y1.
|
||
baseline_y1 = -intercept
|
||
|
||
self._do_debug_baseline(pdf, slope, cm_line_box, baseline_y1)
|
||
text.set_text_transform(1, 0, 0, 1, line_box.x1, baseline_y1)
|
||
pdf.set_fill_color(black) # text in black
|
||
|
||
elements = line.findall(self._child_xpath('span', elemclass))
|
||
for elem in elements:
|
||
self._do_line_word(
|
||
pdf,
|
||
fontname,
|
||
interword_spaces,
|
||
line_height,
|
||
line_matrix,
|
||
cm_line_box,
|
||
text,
|
||
fontsize,
|
||
baseline_y1,
|
||
elem,
|
||
)
|
||
pdf.draw_text(text)
|
||
pdf._cs.pop()
|
||
|
||
def _do_line_word(
|
||
self,
|
||
pdf,
|
||
fontname,
|
||
interword_spaces,
|
||
line_height,
|
||
line_matrix,
|
||
cm_line_box,
|
||
text,
|
||
fontsize,
|
||
baseline_y1,
|
||
elem,
|
||
):
|
||
elemtxt = self._get_element_text(elem).strip()
|
||
elemtxt = self.replace_unsupported_chars(elemtxt)
|
||
if elemtxt == '':
|
||
return
|
||
|
||
pxl_coords = self.element_coordinates(elem)
|
||
box = self.pt_from_pixel(pxl_coords, bottomup=True)
|
||
cm_box = box.transform(line_matrix, inverse=True)
|
||
if interword_spaces:
|
||
# if `--interword-spaces` is true, append a space
|
||
# to the end of each text element to allow simpler PDF viewers
|
||
# such as PDF.js to better recognize words in search and copy
|
||
# and paste. Do not remove space from last word in line, even
|
||
# though it would look better, because it will interfere with
|
||
# naive text extraction. \n does not work either.
|
||
elemtxt += ' '
|
||
cm_box = Rect._make(
|
||
(
|
||
cm_box.x1,
|
||
cm_line_box.y1,
|
||
cm_box.x2, # + pdf.string_width(' ', fontname, line_height),
|
||
cm_line_box.y2,
|
||
)
|
||
)
|
||
box_width = cm_box.x2 - cm_box.x1
|
||
font_width = pdf.string_width(elemtxt, fontname, fontsize)
|
||
|
||
# draw the bbox border
|
||
self._do_debug_word_triangle(pdf, cm_box)
|
||
self._do_debug_bbox(pdf, line_height, cm_line_box, cm_box, box_width)
|
||
|
||
# Adjust relative position of cursor
|
||
# This is equivalent to:
|
||
# text.setTextOrigin(pt.x1, self.height - line_box.y2)
|
||
# but the former generates a full text reposition matrix (Tm) in the
|
||
# content stream while this issues a "offset" (Td) command.
|
||
# .move_cursor() is relative to start of the text line, where the
|
||
# "text line" means whatever reportlab defines it as. Do not use
|
||
# use .getCursor(), since move_cursor() rather unintuitively plans
|
||
# its moves relative to .getStartOfLine().
|
||
# For skewed lines, in the text transform we set up a rotated
|
||
# coordinate system, so we don't have to account for the
|
||
# incremental offset. Surprisingly most PDF viewers can handle this.
|
||
if 0:
|
||
cursor = text.get_start_of_line()
|
||
dx = box.x1 - cursor[0]
|
||
dy = baseline_y1 - cursor[1]
|
||
text.move_cursor(dx, dy)
|
||
text.set_text_transform(
|
||
1,
|
||
0,
|
||
0,
|
||
1,
|
||
cm_box.x1,
|
||
cm_line_box.y1,
|
||
)
|
||
|
||
# If reportlab tells us this word is 0 units wide, our best seems
|
||
# to be to suppress this text
|
||
if font_width > 0:
|
||
text.set_horiz_scale(100 * box_width / font_width)
|
||
text.show(elemtxt)
|
||
|
||
def _do_debug_word_triangle(
|
||
self,
|
||
pdf,
|
||
cm_box,
|
||
):
|
||
if not self.render_options.render_triangle: # pragma: no cover
|
||
return
|
||
pdf._cs.push()
|
||
pdf.set_dashes()
|
||
pdf.set_stroke_color(red)
|
||
pdf.set_line_width(0.1)
|
||
# Draw a triangle that conveys word height and drawing direction
|
||
pdf.line(cm_box.x1, cm_box.y1, cm_box.x2, cm_box.y1) # across bottom
|
||
pdf.line(cm_box.x2, cm_box.y1, cm_box.x1, cm_box.y2) # diagonal
|
||
pdf.line(cm_box.x1, cm_box.y1, cm_box.x1, cm_box.y2) # rise
|
||
pdf._cs.pop()
|
||
|
||
def _do_debug_bbox(self, pdf, line_height, cm_line_box, cm_box, box_width):
|
||
if not self.render_options.render_word_bbox: # pragma: no cover
|
||
return
|
||
pdf._cs.push()
|
||
pdf.set_dashes()
|
||
pdf.set_stroke_color(green)
|
||
pdf.set_line_width(0.1)
|
||
pdf.rect(cm_box.x1, cm_line_box.y1, box_width, line_height, fill=0)
|
||
pdf._cs.pop()
|
||
|
||
def _do_debug_baseline(self, pdf, slope, line_box, baseline_y1):
|
||
if not self.render_options.render_baseline:
|
||
return
|
||
# draw the baseline in magenta, dashed
|
||
pdf.set_dashes()
|
||
pdf.set_stroke_color(magenta)
|
||
pdf.set_line_width(0.25)
|
||
# negate slope because it is defined as a rise/run in pixel
|
||
# coordinates and page coordinates have the y axis flipped
|
||
pdf.line(
|
||
line_box.x1,
|
||
baseline_y1,
|
||
line_box.x2,
|
||
self.polyval((-slope, baseline_y1), line_box.x2 - line_box.x1),
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(description='Convert hocr file to PDF')
|
||
parser.add_argument(
|
||
'-b',
|
||
'--boundingboxes',
|
||
action="store_true",
|
||
default=False,
|
||
help='Show bounding boxes borders',
|
||
)
|
||
parser.add_argument(
|
||
'-r',
|
||
'--resolution',
|
||
type=int,
|
||
default=300,
|
||
help='Resolution of the image that was OCRed',
|
||
)
|
||
parser.add_argument(
|
||
'-i',
|
||
'--image',
|
||
default=None,
|
||
help='Path to the image to be placed above the text',
|
||
)
|
||
parser.add_argument(
|
||
'--interword-spaces',
|
||
action='store_true',
|
||
default=False,
|
||
help='Add spaces between words',
|
||
)
|
||
parser.add_argument('hocrfile', help='Path to the hocr file to be parsed')
|
||
parser.add_argument('outputfile', help='Path to the PDF file to be generated')
|
||
args = parser.parse_args()
|
||
|
||
hocr = HocrTransform(hocr_filename=args.hocrfile, dpi=args.resolution)
|
||
hocr.to_pdf(
|
||
out_filename=args.outputfile,
|
||
image_filename=args.image,
|
||
show_bounding_boxes=args.boundingboxes,
|
||
interword_spaces=args.interword_spaces,
|
||
)
|