Use PyMuPDF for text and image extraction

Author: stchris

README.md

@@ -40,7 +40,7 @@ make build
 make test
 source .env/bin/activate
 bump2version {patch,minor,major} # pick the appropriate one
-git push --atomic origin $(git rev-parse --abbrev-ref HEAD) $(git describe --tags --abbrev=0)
+git push --atomic origin main $(git describe --tags --abbrev=0)
 ```
 
 ## Usage

ingestors/documents/pdf.py

@@ -2,12 +2,7 @@
 
 from ingestors.ingestor import Ingestor
 from ingestors.support.pdf import PDFSupport
-from ingestors.exc import ProcessingException
-
-import pikepdf
-
-# silence some shouty debug output from pdfminer
-logging.getLogger("pdfminer").setLevel(logging.WARNING)
+from ingestors.exc import ProcessingException, UnauthorizedError
 
 log = logging.getLogger(__name__)
 
@@ -28,7 +23,7 @@ def ingest(self, file_path, entity):
         """Ingestor implementation."""
         try:
             self.parse_and_ingest(file_path, entity, self.manager)
-        except pikepdf._core.PasswordError as pwe:
+        except UnauthorizedError as pwe:
             raise ProcessingException(
                 "Could not extract PDF file. The PDF is protected with a password. Try removing the password protection and re-uploading the documents."
             ) from pwe

ingestors/exc.py

@@ -1,3 +1,9 @@
 class ProcessingException(Exception):
     "A data-related error occuring during file processing."
     pass
+
+
+class UnauthorizedError(Exception):
+    """Raised when a document is protected by a password and can not be parsed."""
+
+    pass

ingestors/support/pdf.py

@@ -1,29 +1,19 @@
 from dataclasses import dataclass
-from io import StringIO
 import logging
 import os
 from typing import Dict, List
 import uuid
 import unicodedata
 
-import pikepdf
-from PIL import Image
-from pdfminer.converter import TextConverter
-from pdfminer.layout import LAParams
-from pdfminer.pdfdocument import PDFDocument
-from pdfminer.pdfpage import PDFPage
-from pdfminer.pdfinterp import PDFResourceManager, PDFPageInterpreter
-from pdfminer.pdfparser import PDFParser
+import fitz
 
 from normality import collapse_spaces  # noqa
 
 from followthemoney import model
+from ingestors.exc import UnauthorizedError
 from ingestors.support.ocr import OCRSupport
 from ingestors.support.convert import DocumentConvertSupport
 
-# silence some shouty debug output from pdfminer
-logging.getLogger("pdfminer").setLevel(logging.WARNING)
-
 log = logging.getLogger(__name__)
 
 
@@ -85,143 +75,63 @@ def extract_pages(self, pdf_model: PdfModel, entity, manager):
     def parse(self, file_path: str) -> PdfModel:
         """Takes a file_path to a pdf and returns a `PdfModel`"""
         pdf_model = PdfModel(metadata=None, xmp_metadata=None, pages=[])
-        with open(file_path, "rb") as pdf_file:
-            parser = PDFParser(pdf_file)
-            pike_doc = pikepdf.Pdf.open(pdf_file)
-            pdf_doc = PDFDocument(parser)
-            for page_number, page in enumerate(PDFPage.create_pages(pdf_doc), 1):
+        with fitz.open(file_path) as pdf_doc:
+            if pdf_doc.needs_pass:
+                raise UnauthorizedError
+            # print(f"\n[IF] number of pages: {pdf_doc.page_count}")
+            for page_num in range(pdf_doc.page_count):
                 pdf_model.pages.append(
-                    self.pdf_extract_page(page, pike_doc, page_number)
+                    self.pdf_extract_page(pdf_doc, pdf_doc[page_num], page_num)
                 )
         return pdf_model
 
     def parse_and_ingest(self, file_path: str, entity, manager):
-        try:
-            pdf_model: PdfModel = self.parse(file_path)
-            self.extract_metadata(pdf_model, entity)
-            self.extract_xmp_metadata(pdf_model, entity)
-            self.extract_pages(pdf_model, entity, manager)
-        except pikepdf._core.PasswordError as pwe:
-            log.info(f"Failed to ingest password protected pdf: {file_path}")
-            raise pwe
+        pdf_model: PdfModel = self.parse(file_path)
+        self.extract_metadata(pdf_model, entity)
+        self.extract_xmp_metadata(pdf_model, entity)
+        self.extract_pages(pdf_model, entity, manager)
 
     def pdf_alternative_extract(self, entity, pdf_path: str, manager):
         checksum = self.manager.store(pdf_path)
         entity.set("pdfHash", checksum)
         self.parse_and_ingest(pdf_path, entity, manager)
 
-    def _find_images(self, container: pikepdf.Pdf, depth: int = 0):
-        if "/Resources" not in container:
-            return []
-        resources = container["/Resources"]
-
-        if "/XObject" not in resources:
-            return []
-        xobjects = resources["/XObject"].as_dict()
-
-        if depth > 0:
-            allow_recursion = False
-        else:
-            allow_recursion = True
-
-        images = []
-
-        for xobject in xobjects:
-            candidate = xobjects[xobject]
-            if candidate["/Subtype"] == "/Image":
-                if "/SMask" in candidate:
-                    images.append([candidate, candidate["/SMask"]])
-                else:
-                    images.append(candidate)
-            elif allow_recursion and candidate["/Subtype"] == "/Form":
-                images.extend(self._find_images(candidate, depth=depth + 1))
-
-        return images
-
-    def _extract_images(
-        self, pike_doc: pikepdf.Pdf, image_path: str, prefix: str = "img"
-    ):
-        raw_images = []
-        found_imgs = self._find_images(pike_doc)
-        raw_images.extend(found_imgs)
-
-        pdfimages = []
-        for r in raw_images:
-            if isinstance(r, list):
-                try:
-                    base_image = pikepdf.PdfImage(r[0]).as_pil_image()
-                    soft_mask = pikepdf.PdfImage(r[1]).as_pil_image()
-                except NotImplementedError:
-                    # Skip unsupported image file formats
-                    continue
-
-                if base_image.size != soft_mask.size:
-                    log.debug(
-                        "Warning: Image and /SMask have a different size. This is unexpected.",
-                    )
-                    soft_mask = soft_mask.resize(base_image.size)
-
-                if base_image.mode in ("L", "LA"):
-                    transparency = Image.new("LA", base_image.size, (0, 0))
-                else:
-                    if base_image.mode not in ("RGB", "RGBA"):
-                        base_image = base_image.convert("RGB")
-                    transparency = Image.new("RGBA", base_image.size, (0, 0, 0, 0))
-
-                composite = Image.composite(base_image, transparency, soft_mask)
-
-                pdfimages.append(composite)
-
-            else:
-                pdfimages.append(pikepdf.PdfImage(r))
-
-        n_images = len(pdfimages)
-
-        n_digits = len(str(n_images))
-        for i, image in enumerate(pdfimages):
-            filepath_prefix = os.path.join(image_path, prefix + f"{i+1:0{n_digits}}")
-            if isinstance(image, Image.Image):
-                image.save(filepath_prefix + ".png", "PNG")
-            else:
-                pil_image = image.as_pil_image()
-                if pil_image.format == "TIFF":
-                    pil_image.save(filepath_prefix + ".png", "PNG")
-                image.extract_to(fileprefix=filepath_prefix)
-
-    def pdf_extract_page(
-        self, page: PDFPage, pike_doc: pikepdf._core.Pdf, page_number: int
-    ) -> PdfPageModel:
+    def pdf_extract_page(self, pdf_doc, page, page_number: int) -> PdfPageModel:
         """Extract the contents of a single PDF page, using OCR if need be."""
-        buf = StringIO()
-        rsrcmgr = PDFResourceManager()
-        device = TextConverter(
-            rsrcmgr,
-            buf,
-            laparams=LAParams(
-                line_overlap=0.5,  # default: 0.5
-                char_margin=2.0,  # default: 2.0
-                word_margin=0.2,  # default: 0.1
-                line_margin=0.5,  # default: 0.5
-                boxes_flow=0.5,  # default: 0.5
-                detect_vertical=True,  # default: False
-                all_texts=True,  # default: False
-            ),
-        )
-        interpreter = PDFPageInterpreter(rsrcmgr, device)
-        interpreter.process_page(page)
-        texts = buf.getvalue()
+        # Extract text
+        full_text = page.get_text()
+        # print(f"[IF] extracted text: \n{full_text}")
+
+        # Extract images
+        images = page.get_images()
+
+        # Create a temporary location to store all extracted images
         temp_dir = self.make_empty_directory()
-        image_path = temp_dir.joinpath(str(uuid.uuid4()))
-        os.mkdir(image_path)
-        pike_page = pike_doc.pages[page_number - 1]
-        self._extract_images(pike_page, image_path)
+        image_dir = temp_dir.joinpath(str(uuid.uuid4()))
+        os.mkdir(image_dir)
+
+        # Extract images from PDF and store them on the disk
+        extracted_images = []
+        for image_index, image in enumerate(images, start=1):
+            xref = image[0]
+            img = pdf_doc.extract_image(xref)
+            if img:
+                image_path = os.path.join(
+                    image_dir, f"image{page_number+1}_{image_index}.{img['ext']}"
+                )
+                with open(image_path, "wb") as image_file:
+                    image_file.write(img["image"])
+                extracted_images.append(image_path)
+
+        # Attempt to OCR the images and extract text
         languages = self.manager.context.get("languages")
-        for image_file in image_path.glob("*.png"):
-            with open(image_file, "rb") as fh:
+        for image_path in extracted_images:
+            with open(image_path, "rb") as fh:
                 data = fh.read()
                 text = self.extract_ocr_text(data, languages=languages)
                 if text is not None:
-                    texts += text
+                    # print(f"[IF] extracted text from images: \n{text}")
+                    full_text += text
 
-        texts = unicodedata.normalize("NFKD", texts.strip())
-        return PdfPageModel(number=page_number, text=texts.strip())
+        full_text = unicodedata.normalize("NFKD", full_text.strip())
+        return PdfPageModel(number=page_number + 1, text=full_text.strip())

requirements.txt

@@ -20,9 +20,7 @@ click==8.1.3
 
 # File format support
 dbf==0.99.3
-pdfminer.six==20221105
 pymediainfo==6.0.1
-pikepdf==7.1.1
 python-magic==0.4.27
 rarfile==4.0
 xlrd==2.0.1
@@ -38,3 +36,4 @@ icalendar==5.0.4
 
 cryptography==39.0.1
 requests[security]==2.28.2
+pymupdf==1.21.1

tests/test_pdf.py

@@ -34,7 +34,7 @@ def test_ingest_noisy_fixture(self):
         self.manager.ingest(fixture_path, entity)
         self.assertEqual(len(self.get_emitted()), 501)
         self.assertEqual(
-            self.manager.entities[0].first("bodyText"), "Hello, World!\n\nHello, World!"
+            self.manager.entities[0].first("bodyText"), "Hello, World!\nHello, World!"
         )
         self.assertEqual(entity.schema.name, "Pages")
 
@@ -125,9 +125,7 @@ def test_pdf_conversion_metadata(self):
         assert entity.get("modifiedAt") == ["2015-10-05T08:57:00"]
 
     def test_pdf_letter_spacing(self):
-        """Checks some tricky word spacing in the fancy food menu. This required
-        overriding the pdfminersix LAParams default `word_margin` from 0.1 until
-        it worked."""
+        """Checks some tricky word spacing in the fancy food menu."""
         fixture_path, entity = self.fixture("the-dorset-food-menu.pdf")
         self.manager.ingest(fixture_path, entity)
 
@@ -143,9 +141,7 @@ def test_pdf_letter_spacing(self):
         for expected_string in [
             "served with marinated olives",
             "made with vegetarian ingredients",
-            "dorset",
             "triple-cooked chips",
-            "e\ndorset",  # can't get it to detect the "the" :(
         ]:
             assert expected_string in body_one.lower()