square-fixer.py

import cv2
import numpy as np
import sys
import math

"""

This finds quadrilateral objects, cuts them out, squares them, and fixes skew.
It works but has strict limitations:

- needs a good border (white page on black for example)
- can't deal with noise around paper
- resizing will probably be weird if angle is far from true

It's a simplified version of tools like Microsoft's OfficeLens or CamScanner.
Google Docs (at least on Android) also has a similar feature called "Scan".

Some references:

Automatic perspective correction for quadrilateral objects | OpenCV Code
https://web.archive.org/web/20151125040240/http://opencv-code.com/tutorials/automatic-perspective-correction-for-quadrilateral-objects/

ScannerLite/scannerLite.cpp at master · daisygao/ScannerLite
https://github.com/daisygao/ScannerLite/blob/master/scannerLite.cpp

Hough Line Transform — OpenCV-Python Tutorials 1 documentation
http://opencv-python-tutroals.readthedocs.io/en/latest/py_tutorials/py_imgproc/py_houghlines/py_houghlines.html?highlight=hough

"""

def ishorizontal(line):
    (x1,y1,x2,y2) = line
    return abs(x1 - x2) > abs(y1 - y2)

def center(line):
    (x1,y1,x2,y2) = line
    return [ (x1 + x2) / 2, (y1 + y2) / 2]

def intersection(line1, line2):
    x1,y1,x2,y2 = line1
    x3,y3,x4,y4 = line2

    d = ((x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4))
    if d:
        xx = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / d
        yy = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / d
        return [xx,yy]
    else: # no intersection
        return [-1, -1]

def get_corners(borders):
    corners = []
    corners.append(intersection(borders[0], borders[2]))
    corners.append(intersection(borders[0], borders[3]))
    corners.append(intersection(borders[1], borders[2]))
    corners.append(intersection(borders[1], borders[3]))
    return corners

def transformation_matrix(corners, size):
    corners = np.float32(corners)

    width, height = size
    dummy = np.float32([[0,0], [width, 0], [0, height], [width, height]])
    return cv2.getPerspectiveTransform(corners, dummy)

def dist(p1, p2):
    return math.sqrt( math.pow(p1[0] - p2[0], 2) + math.pow(p1[1] - p2[1], 2))
    
def newsize(corners):
    """Get the average size and assume that's a good one to correct to."""
    # alternate strategies:
    # - resize to known paper proportions
    # - optimistically use largest edge
    width = (dist(corners[0],corners[1]) + dist(corners[2],corners[3])) / 2
    height = (dist(corners[0],corners[2]) + dist(corners[1],corners[3])) / 2
    return (int(width), int(height))

img = cv2.imread(sys.argv[1])
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
edges = cv2.Canny(gray, 100, 150)
lines = cv2.HoughLinesP(edges, 1, np.pi/180, 50, 30, 10)

hori = []
vert = []

for line in lines:
    line = line[0]
    if ishorizontal(line):
        hori.append(line)
    else:
        vert.append(line)

if len(hori) < 2 or len(vert) < 2:
    print("Doesn't look like a quadrilateral")
    sys.exit()

# We want the lines closest to the edge of the image
# so sort by center points
hori.sort(key=lambda l: center(l)[1])
vert.sort(key=lambda l: center(l)[0])

top = hori[0]
bottom = hori[-1]
left = vert[0]
right = vert[-1]

borders = [top, bottom, left, right]
corners = get_corners(borders)
size = newsize(corners)

dst = cv2.warpPerspective(img, transformation_matrix(corners, size), size)

# for debugging
for line in borders:
    x1, y1, x2, y2 = line
    cv2.line(img, (x1, y1), (x2, y2), (0,255,0), 2)

cv2.imwrite('out.jpg',dst)