DocumentAIKit/Dewarp/demo.py

''' Document Localization using Recursive CNN
 Maintainer : Khurram Javed
 Email : kjaved@ualberta.ca '''
import math

import cv2
import numpy as np

import evaluation


def args_processor():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--imagePath", default="../058.jpg", help="Path to the document image")
    parser.add_argument("-o", "--outputPath", default="../output.jpg", help="Path to store the result")
    parser.add_argument("-rf", "--retainFactor", help="Floating point in range (0,1) specifying retain factor",
                        default="0.85")
    parser.add_argument("-cm", "--cornerModel", help="Model for corner point refinement",
                        default="../cornerModelWell")
    parser.add_argument("-dm", "--documentModel", help="Model for document corners detection",
                        default="../documentModelWell")
    return parser.parse_args()


if __name__ == "__main__":
    args = args_processor()

    corners_extractor = evaluation.corner_extractor.GetCorners(args.documentModel)
    corner_refiner = evaluation.corner_refiner.corner_finder(args.cornerModel)

    img = cv2.imread(args.imagePath)

    oImg = img

    extracted_corners = corners_extractor.get(oImg)
    corner_address = []
    # Refine the detected corners using corner refiner
    image_name = 0
    for corner in extracted_corners:
        image_name += 1
        corner_img = corner[0]
        refined_corner = np.array(corner_refiner.get_location(corner_img, 0.85))

        # Converting from local co-ordinate to global co-ordinates of the image
        refined_corner[0] += corner[1]
        refined_corner[1] += corner[2]

        # Final results
        corner_address.append(refined_corner)
    print(corner_address)
    width = int(math.sqrt(math.pow(corner_address[1][0]-corner_address[0][0], 2) + math.pow(corner_address[1][1]-corner_address[0][1], 2)))
    height = int(math.sqrt(math.pow(corner_address[2][0]-corner_address[1][0], 2) + math.pow(corner_address[2][1]-corner_address[1][1], 2)))
    print(width, height)
    # if width > height:
    #     t = width
    #     width = height
    #     height = t
    # print(width, height)
    pts1 = np.float32([[corner_address[0][0], corner_address[0][1]], [corner_address[1][0], corner_address[1][1]], [corner_address[3][0], corner_address[3][1]], [corner_address[2][0], corner_address[2][1]]])
    pts2 = np.float32([[0, 0], [width, 0], [0, height], [width, height]])
    M = cv2.getPerspectiveTransform(pts1, pts2)
    warped = cv2.warpPerspective(img, M, (width, height))
    # cv2.imshow('ww', warped)
    # cv2.waitKey()
    # cv2.destroyAllWindows()
    cv2.imwrite(str(args.outputPath).split('.')[0] + '_warped.png', warped)

    for a in range(0, len(extracted_corners)):
        cv2.line(oImg, tuple(corner_address[a % 4]), tuple(corner_address[(a + 1) % 4]), (255, 0, 0), 4)

    cv2.imwrite(args.outputPath, oImg)