python port of https://github.com/daisygao/ScannerLite

interactive.py

import cv2
import numpy as np

canny = rho = threshold = minLen = maxGap = None


def draw():
    lines = cv2.HoughLinesP(canny, rho, np.pi / 180,
                            threshold, None, minLen, maxGap)
    dst = cv2.cvtColor(canny, cv2.COLOR_GRAY2BGR)
    for l in lines[0]:
        x1, y1, x2, y2 = l
        cv2.line(dst, (x1, y1), (x2, y2), (0, 0, 255), 1)
    cv2.imshow('demo', dst)


def onRho(v):
    global rho
    rho = v
    draw()


def onThreshold(v):
    global threshold
    threshold = v
    draw()


def onMinLen(v):
    global minLen
    minLen = v
    draw()


def onMaxGap(v):
    global maxGap
    maxGap = v
    draw()


def main():
    global canny, rho, threshold, minLen, maxGap

    im = cv2.imread('doc1.jpg')
    h, w, _ = im.shape
    min_w = 200
    scale = w * 1. / min_w
    h_proc = int(h * 1. / scale)
    w_proc = int(w * 1. / scale)

    im_dis = cv2.resize(im, (w_proc, h_proc))
    gray = cv2.cvtColor(im_dis, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (3, 3), 0)

    high_thres = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[0]
    low_thres = high_thres * 0.5
    canny = cv2.Canny(gray, low_thres, high_thres)

    rho, threshold, minLen, maxGap = 1, w_proc / 3, w_proc / 3, 20

    cv2.namedWindow('demo', cv2.WINDOW_NORMAL)
    cv2.createTrackbar('rho', 'demo', rho, 5, onRho)
    cv2.createTrackbar('threshold', 'demo', threshold, w_proc, onThreshold)
    cv2.createTrackbar('minLen', 'demo', minLen, w_proc, onMinLen)
    cv2.createTrackbar('maxGap', 'demo', maxGap, 50, onMaxGap)

    draw()
    cv2.waitKey(0)
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()

scannerLite.py

import sys
import cv2
import numpy as np
from matplotlib import pyplot as plt


class Line:

    def __init__(self, l):
        self.point = l
        x1, y1, x2, y2 = l
        self.c_x = (x1 + x2) / 2
        self.c_y = (y1 + y2) / 2


def show(im):
    msg = 'press any key to continue'
    cv2.namedWindow(msg, cv2.WINDOW_NORMAL)
    cv2.imshow(msg, im)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def intersection(l1, l2):
    x1, y1, x2, y2 = l1.point
    x3, y3, x4, y4 = l2.point
    a1, b1 = y2 - y1, x1 - x2
    c1 = a1 * x1 + b1 * y1
    a2, b2 = y4 - y3, x3 - x4
    c2 = a2 * x3 + b2 * y3
    det = a1 * b2 - a2 * b1
    assert det, "lines are parallel"
    return (1. * (b2 * c1 - b1 * c2) / det, 1. * (a1 * c2 - a2 * c1) / det)


def scannerLite(im, debug=False):
    # resize
    h, w, _ = im.shape
    min_w = 200
    scale = min(10., w * 1. / min_w)
    h_proc = int(h * 1. / scale)
    w_proc = int(w * 1. / scale)
    im_dis = cv2.resize(im, (w_proc, h_proc))

    # gray
    gray = cv2.cvtColor(im_dis, cv2.COLOR_BGR2GRAY)

    # blur
    #gray = cv2.blur(gray, (3, 3))
    #gray = cv2.GaussianBlur(gray, (3,3), 0)

    # canny edges detection
    high_thres = cv2.threshold(
        gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)[0]
    low_thres = high_thres * 0.5
    canny = cv2.Canny(gray, low_thres, high_thres)
    if debug:
        show(canny)

    # lines detection
    lines = cv2.HoughLinesP(
        canny, 1, np.pi / 180, w_proc / 3, None, w_proc / 3, 20)

    if debug:
        t = cv2.cvtColor(canny, cv2.COLOR_GRAY2BGR)

    # classify lines
    hori, vert = [], []
    for l in lines[0]:
        x1, y1, x2, y2 = l
        if abs(x1 - x2) > abs(y1 - y2):
            hori.append(Line(l))
        else:
            vert.append(Line(l))
        if debug:
            cv2.line(t, (x1, y1), (x2, y2), (0, 0, 255), 1)
    if debug:
        show(t)

    # not enough
    if len(hori) < 2:
        if not hori or hori[0].c_y > h_proc / 2:
            hori.append(Line((0, 0, w_proc - 1, 0)))
        if not hori or hori[0].c_y <= h_proc / 2:
            hori.append(Line((0, h_proc - 1, w_proc - 1, h_proc - 1)))

    if len(vert) < 2:
        if not vert or vert[0].c_x > w_proc / 2:
            vert.append(Line((0, 0, 0, h_proc - 1)))
        if not vert or vert[0].c_x <= w_proc / 2:
            vert.append(Line((w_proc - 1, 0, w_proc - 1, h_proc - 1)))

    hori.sort(key=lambda l: l.c_y)
    vert.sort(key=lambda l: l.c_x)

    # corners
    if debug:
        for l in [hori[0], vert[0], hori[-1], vert[-1]]:
            x1, y1, x2, y2 = l.point
            cv2.line(t, (x1, y1), (x2, y2), (0, 255, 255), 1)

    img_pts = [intersection(hori[0], vert[0]), intersection(hori[0], vert[-1]),
               intersection(hori[-1], vert[0]), intersection(hori[-1], vert[-1])]

    for i, p in enumerate(img_pts):
        x, y = p
        img_pts[i] = (x * scale, y * scale)
        if debug:
            cv2.circle(t, (int(x), int(y)), 1, (255, 255, 0), 3)
    if debug:
        show(t)

    # perspective transform
    w_a4, h_a4 = 1654, 2339
    #w_a4, h_a4 = 600, 800
    dst_pts = np.array(
        ((0, 0), (w_a4 - 1, 0), (0, h_a4 - 1), (w_a4 - 1, h_a4 - 1)),
        np.float32)
    img_pts = np.array(img_pts, np.float32)
    transmtx = cv2.getPerspectiveTransform(img_pts, dst_pts)
    return cv2.warpPerspective(im, transmtx, (w_a4, h_a4))


if __name__ == '__main__':
    im = cv2.imread('doc1.jpg')
    show(im)
    dst = scannerLite(im, debug=True)
    show(dst)