Implement new preprocessing feature, background removal

ocrmypdf/__main__.py

@@ -32,6 +32,7 @@ from .pdfa import generate_pdfa_def, file_claims_pdfa
 from . import ghostscript
 from . import tesseract
 from . import qpdf
+from . import leptonica
 from . import ExitCode, page_number, is_iterable_notstr
 from collections.abc import Sequence
 
@@ -204,6 +205,10 @@ preprocessing = parser.add_argument_group(
 preprocessing.add_argument(
     '-r', '--rotate-pages', action='store_true',
     help="automatically rotate pages based on detected text orientation")
+preprocessing.add_argument(
+    '--remove-background', action='store_true',
+    help="attempt to remove background from gray or color pages, setting it "
+         "to white ")
 preprocessing.add_argument(
     '-d', '--deskew', action='store_true',
     help="deskew each page before performing OCR")
@@ -350,7 +355,8 @@ if set(options.language) & {'chi_sim', 'chi_tra'} \
 
 lossless_reconstruction = False
 if options.pdf_renderer == 'hocr':
-    if not options.deskew and not options.clean_final and not options.force_ocr:
+    if not options.deskew and not options.clean_final and \
+            not options.force_ocr and not options.remove_background:
         lossless_reconstruction = True
 
 
@@ -820,6 +826,32 @@ def rasterize_with_ghostscript(
 @transform(
     input=rasterize_with_ghostscript,
     filter=suffix(".page.png"),
+    output=".pp-background.png",
+    extras=[_log, _pdfinfo, _pdfinfo_lock])
+def preprocess_remove_background(
+        input_file,
+        output_file,
+        log,
+        pdfinfo,
+        pdfinfo_lock):
+
+    if not options.remove_background:
+        re_symlink(input_file, output_file, log)
+        return
+
+    pageinfo = get_pageinfo(input_file, pdfinfo, pdfinfo_lock)
+
+    if any(image['bpc'] > 1 for image in pageinfo['images']):
+        leptonica.remove_background(input_file, output_file)
+    else:
+        log.info("{0:4d}: background removal skipped on mono page".format(
+            pageinfo['pageno']))
+        re_symlink(input_file, output_file, log)
+
+
+@transform(
+    input=preprocess_remove_background,
+    filter=suffix(".pp-background.png"),
     output=".pp-deskew.png",
     extras=[_log, _pdfinfo, _pdfinfo_lock])
 def preprocess_deskew(
@@ -836,7 +868,6 @@ def preprocess_deskew(
     pageinfo = get_pageinfo(input_file, pdfinfo, pdfinfo_lock)
     dpi = get_page_square_dpi(pageinfo)
 
-    from . import leptonica
     leptonica.deskew(input_file, output_file, dpi)
 
 
@@ -890,8 +921,9 @@ def ocr_tesseract_hocr(
 
 
 @collate(
-    input=[rasterize_with_ghostscript, preprocess_deskew, preprocess_clean],
-    filter=regex(r".*/(\d{6})(?:\.page|\.pp-deskew|\.pp-clean)\.png"),
+    input=[rasterize_with_ghostscript, preprocess_remove_background,
+           preprocess_deskew, preprocess_clean],
+    filter=regex(r".*/(\d{6})(?:\.page|\.pp-.*)\.png"),
     output=os.path.join(work_folder, r'\1.image'),
     extras=[_log, _pdfinfo, _pdfinfo_lock])
 @graphviz(shape='diamond')
@@ -905,6 +937,8 @@ def select_image_for_pdf(
         image_suffix = '.pp-clean.png'
     elif options.deskew:
         image_suffix = '.pp-deskew.png'
+    elif options.remove_background:
+        image_suffix = '.pp-background.png'
     else:
         image_suffix = '.page.png'
     image = next(ii for ii in infiles if ii.endswith(image_suffix))

ocrmypdf/leptonica.py

@@ -1,12 +1,9 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 #
-# © 2013-15: jbarlow83 from Github (https://github.com/jbarlow83)
+# © 2013-16: jbarlow83 from Github (https://github.com/jbarlow83)
 #
-#
-# Use Leptonica to detect find and remove page skew.  Leptonica uses the method
-# of differential square sums, which its author claim is faster and more robust
-# than the Hough transform used by ImageMagick.
+# Python FFI wrapper for Leptonica library
 
 from __future__ import print_function, absolute_import, division
 import argparse
@@ -505,6 +502,22 @@ def deskew(infile, outfile, dpi):
         raise LeptonicaIOError("Failed to open destination file: %s" % outfile)
 
 
+def remove_background(infile, outfile, tile_size=(40, 60), gamma=1.0,
+                      black_threshold=70, white_threshold=190):
+    try:
+        pix = Pix.read(infile)
+    except LeptonicaIOError:
+        raise LeptonicaIOError("Failed to open file: %s" % infile)
+
+    pix = pix.background_norm(tile_size=tile_size).gamma_trc(
+            gamma, black_threshold, white_threshold)
+
+    try:
+        pix.write_implied_format(outfile)
+    except LeptonicaIOError:
+        raise LeptonicaIOError("Failed to open destination file: %s" % outfile)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(
         description="Python wrapper to access Leptonica")

tests/test_main.py

@@ -162,6 +162,38 @@ def test_clean(spoof_tesseract_noop):
                    env=spoof_tesseract_noop)
 
 
+def test_remove_background(spoof_tesseract_noop):
+    from PIL import Image
+
+    # Ensure the input image does not contain pure white/black
+    im = Image.open(_infile('congress.jpg'))
+    assert im.getextrema() != ((0, 255), (0, 255), (0, 255))
+
+    output_pdf = check_ocrmypdf(
+        'congress.jpg', 'test_remove_bg.pdf', '--remove-background',
+        '--image-dpi', '150',
+        env=spoof_tesseract_noop)
+
+    from ocrmypdf.ghostscript import rasterize_pdf
+    import logging
+    log = logging.getLogger()
+
+    output_png = _outfile('remove_bg.png')
+
+    rasterize_pdf(
+        output_pdf,
+        output_png,
+        xres=100,
+        yres=100,
+        raster_device='png16m',
+        log=log)
+
+
+    # The output image should contain pure white and black
+    im = Image.open(output_png)
+    assert im.getextrema() == ((0, 255), (0, 255), (0, 255))
+
+
 # This will run 5 * 2 * 2 = 20 test cases
 @pytest.mark.parametrize(
     "pdf",
@@ -404,6 +436,7 @@ def test_maximum_options(spoof_tesseract_cache, renderer, output_type):
     check_ocrmypdf(
         'multipage.pdf', 'test_multipage%s.pdf' % renderer,
         '-d', '-c', '-i', '-g', '-f', '-k', '--oversample', '300',
+        '--remove-background',
         '--skip-big', '10', '--title', 'Too Many Weird Files',
         '--author', 'py.test', '--pdf-renderer', renderer,
         '--output-type', output_type,