robust fit 2D polynomial to an image. Given a photo of a paper, makes it perfectly white.

normalize_background.py

from sys import argv
import cv2
import numpy as np
import scipy

iterations = 2

def grid(shape):
  y = np.linspace(0, 1, shape[0])
  x = np.linspace(0, 1, shape[1])
  xx, yy = np.meshgrid(x, y)
  return np.dstack([xx**0, xx, yy, xx**2, xx*yy, yy**2])

def solve(mat, b):
  x, res, rank, sing = np.linalg.lstsq(mat, b, rcond=None)
  return x

def paint(params, shape):
  mat = grid(shape)
  return (mat.reshape((-1, 6)) @ params.reshape((6, 1))).reshape(shape)

def fit(img, mask):
  vmask = mask
  mat = grid(mask.shape)[vmask]
  b = img[vmask]
  return solve(mat, b)

def background(color_img, quantile=0.9):
  result = list()
  for c in range(color_img.shape[2]):
    img = color_img[..., c]
    bg = cv2.dilate(img, None)
    for i in range(iterations):
      diff = np.abs(img - bg)
      mask = diff <= np.quantile(diff, quantile)
      params = fit(img, mask)
      bg = paint(params, img.shape[:2])
    result.append(paint(params, img.shape[:2]))
  return np.dstack(result)

def augment(img, quantile=0.9):
  mask = np.random.uniform(0, 1, img.shape) > quantile
  img[mask] = 0

def normalize(img, bg, black=0.05, white=0.5):
  result = img - bg
  low = np.quantile(result, black)
  high = np.quantile(result, white)
  return (result - low) / (high - low)

is_visual = False

if __name__ == "__main__":
  if len(argv) > 1:
    for filename in argv[1:]:
      img = cv2.imread(filename) / 255.0
      cv2.imshow("img", img)
      bg = background(img, 0.5)
      if is_visual:
        cv2.imshow("background", bg)
        cv2.waitKey()
      else:
        result = normalize(img, bg)
        cv2.imwrite(f"{filename[:-4]}.norm.jpg", result * 255.0)
  else:
    params = np.array([1, 0, 0, -1, 2, -1])
    img = paint(params, (300, 200))
    augment(img, 0.5)
    if is_visual:
      cv2.imshow("img", img)
      cv2.waitKey()
    else:
      cv2.imwrite("normalize_background.png", img * 255.0)
      print(params)