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)