From 2e6879ee515722930ad03eca65fbd0fde67e8ee3 Mon Sep 17 00:00:00 2001
From: "James R. Barlow" <jim@purplerock.ca>
Date: Mon, 8 Feb 2016 14:52:01 -0800
Subject: [PATCH] Introduce Leptonica class for Pix

---
 ocrmypdf/leptonica.py | 64 +++++++++++++++++++++++++++++++++++++++----
 tests/test_main.py    | 10 +++----
 2 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/ocrmypdf/leptonica.py b/ocrmypdf/leptonica.py
index ff905f57..617be4cb 100644
--- a/ocrmypdf/leptonica.py
+++ b/ocrmypdf/leptonica.py
@@ -82,6 +82,65 @@ class LeptonicaIOError(LeptonicaError):
     pass
 
 
+class Pix:
+    def __init__(self, cpix):
+        self.cpix = ffi.gc(cpix, _pix_destroy)
+
+    @classmethod
+    def read(cls, filename):
+        with LeptonicaErrorTrap():
+            return cls(lept.pixRead(
+                filename.encode(sys.getfilesystemencoding())))
+
+    def write_implied_format(
+            self, filename, jpeg_quality=0, jpeg_progressive=0):
+        fileroot, extension = os.path.splitext(filename)
+        fix_pnm = False
+        if extension.lower() in ('.pbm', '.pgm', '.ppm'):
+            # Leptonica does not process handle these extensions correctly, but
+            # does handle .pnm correctly.  Add another .pnm suffix.
+            filename += '.pnm'
+            fix_pnm = True
+
+        with LeptonicaErrorTrap():
+            lept.pixWriteImpliedFormat(
+                filename.encode(sys.getfilesystemencoding()),
+                self.cpix, jpeg_quality, jpeg_progressive)
+
+        if fix_pnm:
+            from shutil import move
+            move(filename, filename[:-4])   # Remove .pnm suffix
+
+    def deskew(self, reduction_factor=0):
+        return Pix(lept.pixDeskew(self.cpix, reduction_factor))
+
+    def scale(self, scalex, scaley):
+        with LeptonicaErrorTrap():
+            return Pix(lept.pixScale(self.cpix, scalex, scaley))
+
+    def rotate180(self):
+        with LeptonicaErrorTrap():
+            return Pix(lept.pixRotate180(ffi.NULL, self.cpix))
+
+    def find_skew(self):
+        with LeptonicaErrorTrap():
+            angle = ffi.new('float *', 0.0)
+            confidence = ffi.new('float *', 0.0)
+            result = lept.pixFindSkew(self.cpix, angle, confidence)
+            if result == 0:
+                return (angle[0], confidence[0])
+            else:
+                return (None, None)
+
+    @staticmethod
+    def correlation_binary(pix1, pix2):
+        correlation = ffi.new('float *', 0.0)
+        result = lept.pixCorrelationBinary(pix1.cpix, pix2.cpix, correlation)
+        if result != 0:
+            raise LeptonicaError("Correlation failed")
+        return correlation[0]
+
+
 def _pix_destroy(pix):
     ptr_to_pix = ffi.new('PIX **', pix)
     lept.pixDestroy(ptr_to_pix)
@@ -201,11 +260,6 @@ def correlation_binary(pix1, pix2):
     return correlation[0]
 
 
-def rotate180(pix):
-    pix_rotated = lept.pixRotate180(ffi.NULL, pix)
-    return ffi.gc(pix_rotated, _pix_destroy)
-
-
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(
         description="Python wrapper to access Leptonica")
diff --git a/tests/test_main.py b/tests/test_main.py
index c4e900f0..0d1b1347 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -260,14 +260,14 @@ def test_autorotate(spoof_tesseract_cache, renderer):
             _infile('cardinal.pdf'), _outfile('reference.png'),
             xres=100, yres=100, raster_device='pngmono', log=gslog, pageno=1)
 
-    pix_ref = leptonica.pixRead(_outfile('reference.png'))
-    pix_ref_180 = leptonica.rotate180(pix_ref)
-    correlation = leptonica.correlation_binary(pix_ref, pix_ref_180)
+    pix_ref = leptonica.Pix.read(_outfile('reference.png'))
+    pix_ref_180 = pix_ref.rotate180()
+    correlation = leptonica.correlation_binary(pix_ref.cpix, pix_ref_180.cpix)
     assert correlation < 0.10
 
     for n in range(1, 4+1):
-        pix_other = leptonica.pixRead(_outfile('cardinal-%i.png' % n))
-        correlation = leptonica.correlation_binary(pix_ref, pix_other)
+        pix_other = leptonica.Pix.read(_outfile('cardinal-%i.png' % n))
+        correlation = leptonica.Pix.correlation_binary(pix_ref, pix_other)
         assert correlation > 0.80
 
 @pytest.mark.parametrize('renderer', [