extract_image.py

import cv2 as cv
import numpy as np

import os

import sys

IMG_CUT_SIZE = 800
IMG_CROP_MARGIN = 20

show_count = 0


def show_image(image, window_name=None):
    global show_count
    show_count += 1
    window_name = "img - " + str(show_count) if window_name is None else window_name
    cv.imshow(window_name, image)


def load_image(file_path):
    return cv.imread(file_path)


def save_image(image, file_path):
    cv.imwrite(file_path, image)


def convert_to_gray_scale(image):
    return cv.cvtColor(image, cv.COLOR_BGR2GRAY)


def blur_image(image):
    return cv.GaussianBlur(image, (5, 5), 0)


def convert_to_binary_image(binary_image, threshold=127, max_value=255):
    return cv.threshold(binary_image, threshold, max_value, cv.THRESH_BINARY)[1]


def canny_edge(image, threshold1=30, threshold2=200):
    return cv.Canny(image, threshold1, threshold2)


def average_contours_area(contours):
    return sum([cv.contourArea(contour) for contour in contours]) / len(contours)


def normalize_contours(contours, average_ratio=0.5):
    contours = contours[1:]
    contours = [contour for contour in contours if cv.contourArea(contour) > 50]
    average_area = average_contours_area(contours)
    return [contour for contour in sorted(contours, key=cv.contourArea, reverse=True)
            if cv.contourArea(contour) > average_area * average_ratio]


def approx_rectangle(contour):
    epsilon = 0.02 * cv.arcLength(contour, True)
    return cv.approxPolyDP(contour, epsilon, True)


def perspective_transform(original_image, contour):
    approx = approx_rectangle(contour)

    from_points = np.float32([approx[0], approx[3], approx[1], approx[2]])
    to_points = np.float32([[-IMG_CROP_MARGIN, -IMG_CROP_MARGIN],
                            [IMG_CUT_SIZE + IMG_CROP_MARGIN, -IMG_CROP_MARGIN],
                            [-IMG_CROP_MARGIN, IMG_CROP_MARGIN + IMG_CUT_SIZE],
                            [IMG_CROP_MARGIN + IMG_CUT_SIZE, IMG_CROP_MARGIN + IMG_CUT_SIZE]])

    transform_matrix = cv.getPerspectiveTransform(from_points, to_points)
    return cv.warpPerspective(original_image, transform_matrix, (IMG_CUT_SIZE, IMG_CUT_SIZE))


def extract_image(file_name):
    original_image = load_image(file_name)
    gray_image = convert_to_gray_scale(original_image)
    blurred_image = blur_image(gray_image)

    flatten_pixels = np.array(blurred_image).flatten()
    max_white_level = max(flatten_pixels)
    avg_level = sum(flatten_pixels) / len(flatten_pixels)

    print max_white_level, avg_level

    binary_image = convert_to_binary_image(blurred_image, avg_level+20)

    edged_image = canny_edge(binary_image, 50, 180)

    show_image(original_image)
    show_image(blurred_image)
    show_image(binary_image)
    show_image(edged_image)

    im2, contours, hierarchy = cv.findContours(edged_image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)
    contours = normalize_contours(contours, average_ratio=0.5)

    paper = original_image.copy()

    cnt = 1
    for contour in contours:
        try:
            perspective_image = perspective_transform(original_image, contour)

            show_image(perspective_image)
            save_image(perspective_image, os.path.join("output", str(cnt) + ".jpg"))
            cnt += 1
        except Exception:
            # have not 4 points
            print("Unexpected error:", sys.exc_info()[0])

        x, y, width, height = cv.boundingRect(contour)
        cv.rectangle(paper, (x, y), (x + width, y + height), (0, 100, 255), 4)

    cv.drawContours(paper, contours, -1, (0, 255, 0), 2)
    show_image(paper, "paper")


extract_image("sample2.jpg")

cv.waitKey()