ImageProcess/deploy/python/doc_scan/doc_scan_infer.py

import cv2
import numpy as np
import glob
import evaluation
import os
import shutil
import time


def args_processor():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--images", default="example_imgs", help="Document image folder")
    parser.add_argument('--model-type', default="resnet",
                    help='model type to be used. Example : resnet32, resnet20, densenet, test')
    parser.add_argument("-o", "--output", default="example_imgs/output", help="The folder to store results")
    parser.add_argument("-rf", "--retainFactor", help="Floating point in range (0,1) specifying retain factor",
                        default="0.85", type=float)
    parser.add_argument("-cm", "--cornerModel", help="Model for corner point refinement",
                        default="../cornerModelWell")
    parser.add_argument("-cp", "--cornerParams", help="Model for corner point refinement",
                        default="../cornerModelWell")
    parser.add_argument("-dm", "--documentModel", help="Model for document corners detection",
                        default="../documentModelWell")
    parser.add_argument("-dp", "--documentParams", 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, args.documentParams)
    corner_refiner = evaluation.corner_refiner.corner_finder(args.cornerModel, args.cornerParams)
    outputPath = f"{args.output}/res_scan.jpg"
    img = cv2.imread(args.images)
    oImg = img
    extracted_corners = corners_extractor.get(oImg)
    corner_address = []
    # Refine the detected corners using corner refiner
    for corner in extracted_corners:
        corner_img = corner[0]
        refined_corner = np.array(corner_refiner.get_location(corner_img, args.retainFactor))

        # 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)

    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(outputPath, oImg)