From 2974929b26036d0ac24956a48f0bf5996ceabc38 Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <jim@purplerock.ca>
Date: Mon, 2 Jul 2018 22:22:59 -0700
Subject: [PATCH] Make jpeg/png quality tunable args

---
 src/ocrmypdf/__main__.py | 24 ++++++++++++++++++++++--
 src/ocrmypdf/optimize.py | 33 +++++++++++++++++++++------------
 tests/test_main.py       |  1 +
 3 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/src/ocrmypdf/__main__.py b/src/ocrmypdf/__main__.py
index 21941401..b051d90d 100755
--- a/src/ocrmypdf/__main__.py
+++ b/src/ocrmypdf/__main__.py
@@ -80,7 +80,7 @@ if tesseract.version() < MINIMUM_TESS_VERSION:
 # Parser
 
 def numeric(basetype, min_=None, max_=None):
-    "Validator for numeric params"
+    """Validator for numeric params"""
     min_ = basetype(min_) if min_ is not None else None
     max_ = basetype(max_) if max_ is not None else None
     def _numeric(string):
@@ -276,12 +276,32 @@ optimizing.add_argument(
         "3 - do aggressive lossy optimizations"
     )
 )
+optimizing.add_argument(
+    '--jpeg-quality', type=numeric(int, 0, 100), default=0, metavar='Q',
+    help=("Adjust JPEG quality level for JPEG optimization. "
+          "100 is best quality and largest output size; "
+          "1 is lowest quality and smallest output"
+          "0 uses the default."
+    )
+)
+optimizing.add_argument(
+    '--jpg-quality', type=numeric(int, 0, 100), default=0, metavar='Q',
+    dest='jpeg_quality',
+    help=argparse.SUPPRESS  # Alias for --jpeg-quality
+)
+optimizing.add_argument(
+    '--png-quality', type=numeric(int, 0, 100), default=0, metavar='Q',
+    help=("Adjust PNG quality level to use when quantizing PNGs. "
+          "Values have same meaning as with --jpeg-quality"
+    )
+)
 
 advanced = parser.add_argument_group(
     "Advanced",
     "Advanced options to control Tesseract's OCR behavior")
 advanced.add_argument(
-    '--max-image-mpixels', action='store', type=numeric(float, 0), metavar='MPixels',
+    '--max-image-mpixels', action='store', type=numeric(float, 0),
+    metavar='MPixels',
     help="Set maximum number of pixels to unpack before treating an image as a "
          "decompression bomb",
     default=128.0)
diff --git a/src/ocrmypdf/optimize.py b/src/ocrmypdf/optimize.py
index 4af75f16..1584e2f1 100644
--- a/src/ocrmypdf/optimize.py
+++ b/src/ocrmypdf/optimize.py
@@ -31,8 +31,8 @@ from .helpers import re_symlink, fspath
 from .exec import pngquant, jbig2enc
 
 PAGE_GROUP_SIZE = 10
-JPEG_QUALITY = 75
-PNG_QUALITY = (65, 75)
+DEFAULT_JPEG_QUALITY = 75
+DEFAULT_PNG_QUALITY = (65, 75)
 
 
 def img_name(root, xref, ext):
@@ -231,7 +231,7 @@ def transcode_jpegs(pike, jpegs, root, log, options):
         with Image.open(fspath(in_jpg)) as im:
             im.save(fspath(opt_jpg),
                     optimize=True,
-                    quality=JPEG_QUALITY)
+                    quality=options.jpeg_quality)
         # pylint: disable=E1101
         if opt_jpg.stat().st_size > in_jpg.stat().st_size:
             log.debug("xref {}, jpeg, made larger - skip".format(xref))
@@ -247,13 +247,17 @@ def transcode_jpegs(pike, jpegs, root, log, options):
 
 def transcode_pngs(pike, pngs, root, log, options):
     if options.optimize >= 2:
+        png_quality = (
+            max(10, options.png_quality - 10),
+            min(100, options.png_quality + 10)
+        )
         with concurrent.futures.ThreadPoolExecutor(
                 max_workers=options.jobs) as executor:
             for xref in pngs:
                 executor.submit(
                     pngquant.quantize,
                     png_name(root, xref), png_name(root, xref),
-                    PNG_QUALITY[0], PNG_QUALITY[1])
+                    png_quality[0], png_quality[1])
 
     for xref in pngs:
         im_obj = pike.get_object(xref, 0)
@@ -314,12 +318,12 @@ def optimize(
         re_symlink(input_file, output_file, log)
         return
 
-    global PNG_QUALITY
-    global JPEG_QUALITY
-
-    if options.optimize == 3:
-        PNG_QUALITY = (20, 40)
-        JPEG_QUALITY = 40
+    if options.jpeg_quality == 0:
+        options.jpeg_quality = \
+                DEFAULT_JPEG_QUALITY if options.optimize < 3 else 40
+    if options.png_quality == 0:
+        options.png_quality = \
+                DEFAULT_PNG_QUALITY if options.optimize < 3 else 30
 
     pike = pikepdf.Pdf.open(input_file)
 
@@ -353,13 +357,18 @@ def optimize(
 def main(infile, outfile, level, jobs=1):
     from tempfile import TemporaryDirectory
     from shutil import copy
-    Options = namedtuple('Options', 'jobs optimize')
+    Options = namedtuple('Options', 'jobs optimize png_quality jpeg_quality')
 
     logging.basicConfig(level=logging.DEBUG)
     log = logging.getLogger()
 
     ctx = JobContext()
-    options = Options(jobs=jobs, optimize=int(level))
+    options = Options(
+        jobs=jobs,
+        optimize=int(level),
+        jpeg_quality=0,  # Use default
+        png_quality=0
+    )
     ctx.set_options(options)
 
     with TemporaryDirectory() as td:
diff --git a/tests/test_main.py b/tests/test_main.py
index 55a38493..48393a18 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -236,6 +236,7 @@ def test_maximum_options(spoof_tesseract_cache, renderer, output_type,
         '--remove-background',
         '--skip-big', '10', '--title', 'Too Many Weird Files',
         '--author', 'py.test', '--pdf-renderer', renderer,
+        '--optimize', '2', '--jpg-quality', '50', '--png-quality', '20',
         '--output-type', output_type,
         env=spoof_tesseract_cache)