Try to get an homogeneous bg

cf jlsutherland#13

Author: rcatajar

doc2text/page.py

@@ -3,11 +3,13 @@
 import os
 import traceback
 import sys
+import tempfile
 
 import cv2
 import numpy as np
 from PIL import Image
 from scipy.ndimage.filters import rank_filter
+from scipy.optimize import minimize
 import pytesseract
 
 
@@ -42,7 +44,7 @@ def deskew(self):
             self.healthy = False
 
     def extract_text(self):
-        temp_path = 'text_temp.png'
+        _, temp_path = tempfile.mkstemp(suffix="png")
         cv2.imwrite(temp_path, self.image)
         self.text = pytesseract.image_to_string(Image.open(temp_path), lang=self.lang)
         os.remove(temp_path)
@@ -77,7 +79,7 @@ def downscale_image(im, max_dim=2048):
         return 1.0, im
 
     scale = 1.0 * max_dim / max(a, b)
-    new_im = cv2.resize(im, (int(b * scale), int(a * scale)), cv2.INTER_AREA)
+    new_im = cv2.resize(im, (int(b * scale), int(a * scale)), interpolation=cv2.INTER_AREA)
     return scale, new_im
 
 
@@ -146,6 +148,8 @@ def rect_area(crop):
 
 
 def crop_image(im, rect, scale):
+    if rect is None:
+        rect = [0, 0, im.shape[0], im.shape[1]]
     xmin, ymin, xmax, ymax = rect
     crop = [xmin, ymin, xmax, ymax]
     xmin, ymin, xmax, ymax = [int(x / scale) for x in crop]
@@ -283,6 +287,52 @@ def compute_skew(theta):
 
 def process_skewed_crop(image):
     theta = compute_skew(estimate_skew(image))
+    _, _, new_im = optimize_light(image)
     ret, thresh = cv2.threshold(image.copy(), 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
     rotated = rotate(thresh, theta)
     return (rotated, theta)
+
+
+def get_op_im(pars, dwn_im):
+    """
+    Reshape the parameters and compute the corresponding corrected image
+    """
+    offset = np.reshape(pars, dwn_im.shape)
+    new_im = offset + dwn_im
+    return offset, new_im
+
+
+def cost(pars, dwn_im, orig_lap, k1, k2):
+    """"
+    Cost function getting optimized.
+    Two parameters are considered:
+        - Getting the image close to white (bg_term)
+        - Preserving the laplacian of the image
+    """
+    offset, new_im = get_op_im(pars, dwn_im)
+    new_lap = cv2.Laplacian(new_im, cv2.CV_64F)
+    bg_term = k1 * np.sum(np.square(new_im - 255))
+    contrast_term = k2 * np.sum(np.square(new_lap - orig_lap))
+    return bg_term + contrast_term
+
+
+def optimize_light(image, k1=1, k2=50):
+    """"
+    Try to get an homogeneous white background while preserving good contrast.
+    As we down sample the image a lot, it will only work for slowly varying
+    illuminations
+    """
+    _, im = downscale_image(image, 50)
+    orig_lap = cv2.Laplacian(im, cv2.CV_64F)
+    nb_px = np.prod(im.shape)
+
+    offset = np.zeros(nb_px) + (255 - np.mean(image))
+    res = minimize(cost, offset, (im, orig_lap, k1, k2), method='CG',
+                   options={'maxiter': 5})
+
+    op_pars = res.x
+    offset, new_im = get_op_im(op_pars, im)
+    s_off = cv2.resize(offset, image.shape[::-1], interpolation=cv2.INTER_CUBIC)
+    new_im = cv2.convertScaleAbs(image + s_off)
+
+    return res, s_off, new_im