random-stuff/python_scripts/image_pipeline.py

import cv2, numpy, glob


def read_img(file):
    img = cv2.imread(file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img

def order_corners(corners):
    xSorted = corners[numpy.argsort(corners[:, 0]), :]
    left = xSorted[:2, :]
    tl, bl = left[numpy.argsort(left[:, 1]), :]

    right = xSorted[2:, :]
    tr, br = right[numpy.argsort(right[:, 1]), :]

    return numpy.array([tl, tr, bl, br], 'float32')


def autocrop(img):
    # save copy to apply transform
    orig_img = img.copy()

    # smoothing to remove text and other small sharp edges
    kernel = numpy.ones((5,5), numpy.uint8)
    img = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel, iterations=5)

    # edge detection + best contours
    median_pixel_inten = numpy.median(img)
    img = cv2.Canny(img, int(max(0, 0.7*median_pixel_inten)), int(min(255, 1.3*median_pixel_inten)), 3)
    img = cv2.dilate(img, None, iterations=1)
    img = cv2.erode(img, None, iterations=1)
    conts, _ = cv2.findContours(img, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    conts = sorted(conts, key=cv2.contourArea, reverse=True)[:10]

    # compute estimated bounding corners
    for cont in conts:
        corners = cv2.approxPolyDP(cont, cv2.arcLength(cont, True) * 0.02, True)

        # if 4 corners found and cover an area > 1k pixels
        if len(corners) == 4 and cv2.contourArea(corners) >= 1000:
            corners = corners.reshape((4,2))
            break

        # default incase none found
        corners = numpy.array([[0,0],[img.shape[1], 0], [0, img.shape[0]], [img.shape[1], img.shape[0]]])

    # proper ordering for transformation calc
    corners = order_corners(corners)

    # compute new dimensions
    w1 = numpy.sqrt((corners[3][0] - corners[2][0]) ** 2 + (corners[3][1] - corners[3][1]) ** 2)
    w2 = numpy.sqrt((corners[1][0] - corners[0][0]) ** 2 + (corners[1][1] - corners[0][1]) ** 2)
    h1 = numpy.sqrt((corners[3][0] - corners[1][0]) ** 2 + (corners[3][1] - corners[1][1]) ** 2)
    h2 = numpy.sqrt((corners[2][0] - corners[0][0]) ** 2 + (corners[2][1] - corners[0][1]) ** 2)
    max_w = int(max(w1, w2))
    max_h = int(max(h1, h2))
    dest_corners = numpy.array([[0, 0], [max_w - 1, 0], [0, max_h - 1] ,[max_w - 1, max_h - 1]], 'float32')

    # compute and apply transformation
    transf_mat = cv2.getPerspectiveTransform(corners, dest_corners)
    img = cv2.warpPerspective(orig_img, transf_mat, (max_w-1, max_h-1), flags = cv2.INTER_LINEAR)
    # rescale for next steps
    img = cv2.resize(img, (int(img.shape[1] * 1200 / img.shape[0]), 1200))
    return img

def is_low_contrast(img, thresh):
    # reduce noise
    pixel_range = numpy.percentile(img, [5,95])

    # calc spread mesure
    val = (pixel_range[1] - pixel_range[0]) / 255
    return val < thresh

def is_blurry(img, thresh):
    # compute laplacian
    lap = cv2.Laplacian(img, cv2.CV_64F)

    # remove noise and find largest value
    blur = numpy.percentile(lap, [0, 99])[1]
    return blur < thresh

def remove_shadows(img):
    # spread out pixels
    dil = cv2.dilate(img, (7,7))

    # blur to infer background
    blur = cv2.medianBlur(dil, 25)

    # remove smoothed background (shadows)
    img = 225 - cv2.absdiff(img, blur)

    # threshold for b&w
    _, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    return img

def write_img(img, file):
    cv2.imwrite('outs/' + file, img)

def process_img(img):
    if is_low_contrast(img, 0.6):
        return -1, img
    img = autocrop(img)
    if is_blurry(img, 70):
        return -1, img
    img = remove_shadows(img)
    return 0, img

def main():
    files = glob.glob('*.jpg')
    for file in files:
        img = read_img(file)
        ret, img = process_img(img)
        if ret != -1:
            write_img(img, file)

if __name__ == '__main__':
    main()