From 03669183d79d6e51a69916fb271f360398db2241 Mon Sep 17 00:00:00 2001 From: "James R. Barlow" Date: Mon, 20 Nov 2023 15:31:44 -0800 Subject: [PATCH] Rationalize canvas interface --- src/ocrmypdf/hocrtransform/_canvas.py | 107 +++++++++++---- src/ocrmypdf/hocrtransform/_hocr.py | 181 +++++++++++++------------- tests/test_hocrtransform.py | 3 +- 3 files changed, 171 insertions(+), 120 deletions(-) diff --git a/src/ocrmypdf/hocrtransform/_canvas.py b/src/ocrmypdf/hocrtransform/_canvas.py index deb3fca2..945292ba 100644 --- a/src/ocrmypdf/hocrtransform/_canvas.py +++ b/src/ocrmypdf/hocrtransform/_canvas.py @@ -1,8 +1,12 @@ +# SPDX-FileCopyrightText: 2023 James R. Barlow +# SPDX-License-Identifier: MPL-2.0 + from __future__ import annotations import logging import unicodedata from contextlib import contextmanager +from dataclasses import dataclass from importlib.resources import files as package_files from pathlib import Path @@ -15,6 +19,7 @@ from pikepdf import ( Pdf, unparse_content_stream, ) +from PIL import Image log = logging.getLogger(__name__) @@ -243,32 +248,38 @@ class ContentStreamBuilder: self._instructions.append(inst) return self + def draw_form_xobject(self, name: Name): + inst = ContentStreamInstruction([name], Operator("Do")) + self._instructions.append(inst) + return self + def build(self): return self._instructions -class PikepdfCanvas: - def __init__(self, path, *, page_size): - self.path = path - self.page_size = page_size - self._pdf = Pdf.new() - self._page = self._pdf.add_blank_page(page_size=page_size) - self._cs = ContentStreamBuilder() - self._stack_depth = 0 - self.push() - self._font_name = Name("/f-0-0") +@dataclass +class LoadedImage: + name: Name + image: Image.Image - def set_stroke_color(self, color): + +class PikepdfCanvasAccessor: + def __init__(self, cs: ContentStreamBuilder, images=None): + self._cs = cs + self._images = images if images is not None else [] + self._stack_depth = 0 + + def stroke_color(self, color): r, g, b = color.red, color.green, color.blue self._cs.set_stroke_color(r, g, b) return self - def set_fill_color(self, color): + def fill_color(self, color): r, g, b = color.red, color.green, color.blue self._cs.set_fill_color(r, g, b) return self - def set_line_width(self, width): + def line_width(self, width): self._cs.set_line_width(width) return self @@ -285,8 +296,20 @@ class PikepdfCanvas: self._cs.stroke_and_close() return self - def begin_text(self, x=0, y=0, direction=None): - return PikepdfText(x, y, direction) + def draw_image(self, image: Path | str | Image.Image, x, y, width, height): + with self.enter_context(): + self.cm(Matrix(width, 0, 0, height, x, y)) + if isinstance(image, (Path, str)): + image = Image.open(image) + image.load() + if image.mode == "P": + image = image.convert("RGB") + if image.mode not in ("1", "L", "RGB"): + raise ValueError(f"Unsupported image mode: {image.mode}") + name = Name.random(prefix="Im") + li = LoadedImage(name, image) + self._images.append(li) + self._cs.draw_form_xobject(name) def draw_text(self, text: PikepdfText): self._cs._instructions.extend(text._cs.build()) @@ -295,14 +318,7 @@ class PikepdfCanvas: def _end_text(self): self._cs.end_text() - def draw_image(self, image: Path, x, y, width, height): - raise NotImplementedError() - - def string_width(self, text, fontname, fontsize): - # NFKC: split ligatures, combine diacritics - return len(unicodedata.normalize("NFKC", text)) * (fontsize / CHAR_ASPECT) - - def set_dashes(self, *args): + def dashes(self, *args): self._cs.set_dashes(*args) return self @@ -327,8 +343,42 @@ class PikepdfCanvas: self._cs.cm(matrix) return self - def save(self): - self._cs.pop() + +class PikepdfCanvas: + def __init__(self, *, page_size: tuple[int, int]): + self.page_size = page_size + self._pdf = Pdf.new() + self._page = self._pdf.add_blank_page(page_size=page_size) + self._cs = ContentStreamBuilder() + self._images: list[LoadedImage] = [] + self._accessor = PikepdfCanvasAccessor(self._cs, self._images) + self._stack_depth = 0 + self._font_name = Name("/f-0-0") + self.do.push() + + @property + def do(self) -> PikepdfCanvasAccessor: + return self._accessor + + def string_width(self, text, fontname, fontsize): + # NFKC: split ligatures, combine diacritics + return len(unicodedata.normalize("NFKC", text)) * (fontsize / CHAR_ASPECT) + + def _save_image(self, li: LoadedImage): + return self._pdf.make_stream( + li.image.tobytes(), + Width=li.image.width, + Height=li.image.height, + ColorSpace=Name.DeviceGray + if li.image.mode in ("1", "L") + else Name.DeviceRGB, + Type=Name.XObject, + Subtype=Name.Image, + BitsPerComponent=1 if li.image.mode == '1' else 8, + ) + + def save(self, output_file: Path): + self.do.pop() if self._stack_depth != 0: log.warning( "Graphics state stack is not empty when page saved - " @@ -337,9 +387,12 @@ class PikepdfCanvas: self._page.Contents = self._pdf.make_stream( unparse_content_stream(self._cs.build()) ) - self._page.Resources = Dictionary(Font=Dictionary()) + self._page.MediaBox = [0, 0, *self.page_size] + self._page.Resources = Dictionary(Font=Dictionary(), XObject=Dictionary()) self._page.Resources.Font[self._font_name] = register_glyphlessfont(self._pdf) - self._pdf.save(self.path) + for li in self._images: + self._page.Resources.XObject[li.name] = self._save_image(li) + self._pdf.save(output_file) class PikepdfText: diff --git a/src/ocrmypdf/hocrtransform/_hocr.py b/src/ocrmypdf/hocrtransform/_hocr.py index f8b250de..7113ed6e 100644 --- a/src/ocrmypdf/hocrtransform/_hocr.py +++ b/src/ocrmypdf/hocrtransform/_hocr.py @@ -173,57 +173,54 @@ class HocrTransform: """ # create the PDF file # page size in points (1/72 in.) - canvas = Canvas( - out_filename, - page_size=(self.width, self.height), - ) - canvas.push() - page_matrix = ( - Matrix() - .translated(0, self.height) - .scaled(1, -1) - .scaled(INCH / self.dpi, INCH / self.dpi) - ) - canvas.cm(page_matrix) - log.debug(page_matrix) - - self._debug_draw_paragraph_boxes(canvas) - - 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( - canvas, - line, - "ocrx_word", - fontname, - invisible_text, + canvas = Canvas(page_size=(self.width, self.height)) + with canvas.do.enter_context(): + page_matrix = ( + Matrix() + .translated(0, self.height) + .scaled(1, -1) + .scaled(INCH / self.dpi, INCH / self.dpi) ) + canvas.do.cm(page_matrix) + log.debug(page_matrix) - 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( - canvas, - root, - "ocrx_word", - fontname, - invisible_text, - ) - canvas.pop() + self._debug_draw_paragraph_boxes(canvas) + + 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( + canvas, + line, + "ocrx_word", + fontname, + invisible_text, + ) + + 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( + canvas, + root, + "ocrx_word", + fontname, + invisible_text, + ) # put the image on the page, scaled to fill the page if image_filename is not None: - canvas.draw_image( + canvas.do.draw_image( image_filename, 0, 0, width=self.width, height=self.height ) # finish up the page and save it - canvas.save() + canvas.save(out_filename) @classmethod def polyval(cls, poly, x): # pragma: no cover @@ -261,42 +258,43 @@ class HocrTransform: # Setup a new coordinate system on the line box's intercept and rotated by # its slope. - canvas.push() - line_matrix = ( - Matrix() - .translated(*bottom_left_corner) - .translated(0, intercept) - .rotated(angle / pi * 180) - ) - canvas.cm(line_matrix) - log.debug(line_matrix) - text = canvas.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. - line_box_height = abs(line_box.height) / cos(angle) - fontsize = line_box_height + intercept - text.set_font(fontname, fontsize) - if invisible_text or True: - text.set_render_mode(3) # Invisible (indicates OCR text) - - self._debug_draw_baseline(canvas, line_matrix.inverse().transform(line_box), 0) - - canvas.set_fill_color(BLACK) # text in black - elements = line.findall(self._child_xpath('span', elemclass)) - for elem, next_elem in pairwise(elements + [None]): - self._do_line_word( - canvas, - fontname, - line_matrix, - text, - fontsize, - elem, - next_elem, + with canvas.do.enter_context(): + line_matrix = ( + Matrix() + .translated(*bottom_left_corner) + .translated(0, intercept) + .rotated(angle / pi * 180) ) - canvas.draw_text(text) - canvas.pop() + canvas.do.cm(line_matrix) + log.debug(line_matrix) + text = PikepdfText() + + # 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. + line_box_height = abs(line_box.height) / cos(angle) + fontsize = line_box_height + intercept + text.set_font(fontname, fontsize) + if invisible_text or True: + text.set_render_mode(3) # Invisible (indicates OCR text) + + self._debug_draw_baseline( + canvas, line_matrix.inverse().transform(line_box), 0 + ) + + canvas.do.fill_color(BLACK) # text in black + elements = line.findall(self._child_xpath('span', elemclass)) + for elem, next_elem in pairwise(elements + [None]): + self._do_line_word( + canvas, + fontname, + line_matrix, + text, + fontsize, + elem, + next_elem, + ) + canvas.do.draw_text(text) def _do_line_word( self, @@ -351,16 +349,15 @@ class HocrTransform: """Draw boxes around paragraphs in the document.""" if not self.render_options.render_paragraph_bbox: # pragma: no cover return - with canvas.enter_context(): + with canvas.do.enter_context(): # draw box around paragraph - canvas.set_stroke_color(color) - canvas.set_line_width(0.1) # no line for bounding box + canvas.do.stroke_color(color).line_width(0.1) for elem in self.hocr.iterfind(self._child_xpath('p', 'ocr_par')): elemtxt = self._get_element_text(elem).strip() if len(elemtxt) == 0: continue ocr_par = self.element_coordinates(elem) - canvas.rect( + canvas.do.rect( ocr_par.llx, ocr_par.lly, ocr_par.width, ocr_par.height, fill=0 ) @@ -368,8 +365,8 @@ class HocrTransform: """Render the bounding box of a text line.""" if not self.render_options.render_line_bbox: # pragma: no cover return - with canvas.enter_context(): - canvas.set_stroke_color(color).set_line_width(0.15).rect( + with canvas.do.enter_context(): + canvas.do.stroke_color(color).line_width(0.15).rect( line_box.llx, line_box.lly, line_box.width, line_box.height, fill=0 ) @@ -379,8 +376,8 @@ class HocrTransform: """Render a triangle that conveys word height and drawing direction.""" if not self.render_options.render_triangle: # pragma: no cover return - with canvas.enter_context(): - canvas.set_stroke_color(color).set_line_width(line_width).line( + with canvas.do.enter_context(): + canvas.do.stroke_color(color).line_width(line_width).line( box.llx, box.lly, box.urx, box.lly ).line(box.urx, box.lly, box.llx, box.ury).line( box.llx, box.lly, box.llx, box.ury @@ -392,8 +389,8 @@ class HocrTransform: """Render a box depicting the word.""" if not self.render_options.render_word_bbox: # pragma: no cover return - with canvas.enter_context(): - canvas.set_stroke_color(color).set_line_width(line_width).rect( + with canvas.do.enter_context(): + canvas.do.stroke_color(color).line_width(line_width).rect( box.llx, box.lly, box.width, box.height, fill=0 ) @@ -403,8 +400,8 @@ class HocrTransform: """Render a box depicting the space between two words.""" if not self.render_options.render_space_bbox: # pragma: no cover return - with canvas.enter_context(): - canvas.set_fill_color(color).set_line_width(line_width).rect( + with canvas.do.enter_context(): + canvas.do.fill_color(color).line_width(line_width).rect( box.llx, box.lly, box.width, box.height, fill=1 ) @@ -419,8 +416,8 @@ class HocrTransform: """Render the text baseline.""" if not self.render_options.render_baseline: return - with canvas.enter_context(): - canvas.set_stroke_color(color).set_line_width(line_width).line( + with canvas.do.enter_context(): + canvas.do.stroke_color(color).line_width(line_width).line( line_box.llx, baseline_lly, line_box.urx, diff --git a/tests/test_hocrtransform.py b/tests/test_hocrtransform.py index 1bfb0f95..354cd04d 100644 --- a/tests/test_hocrtransform.py +++ b/tests/test_hocrtransform.py @@ -4,6 +4,7 @@ from __future__ import annotations import re +import shutil from io import StringIO import pytest @@ -68,7 +69,7 @@ def test_mono_image(blank_hocr, outdir): hocr.to_pdf( out_filename=str(outdir / 'mono.pdf'), image_filename=str(outdir / 'mono.tif') ) - + # shutil.copy(outdir / 'mono.pdf', 'mono.pdf') check_pdf(str(outdir / 'mono.pdf'))