diff --git a/scripts/deskew_demo.py b/scripts/deskew_demo.py new file mode 100644 index 0000000..7a05806 --- /dev/null +++ b/scripts/deskew_demo.py @@ -0,0 +1,27 @@ +#sample usage: python3 scripts/deskew_demo.py /path/to/crooked.pdf 0 +import argparse +import numpy as np +import pdf2image + +from vidocp.utils.display import show_mpl +from vidocp.utils.deskew import deskew_image + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("pdf_path") + parser.add_argument("page_index", type=int) + + args = parser.parse_args() + + return args + + +if __name__ == "__main__": + args = parse_args() + page = pdf2image.convert_from_path(args.pdf_path, first_page=args.page_index + 1, last_page=args.page_index + 1)[0] + page = np.array(page) + + show_mpl(page) + page = deskew_image(page, verbose=True) + show_mpl(page) diff --git a/vidocp/utils/deskew.py b/vidocp/utils/deskew.py new file mode 100644 index 0000000..70c36ff --- /dev/null +++ b/vidocp/utils/deskew.py @@ -0,0 +1,48 @@ +import numpy as np +import cv2 + + +def detect_angle(im: np.array, max_skew_deg=10, min_nlines=5) -> int: + max_skew_rad = np.deg2rad(max_skew_deg) + width = im.shape[1] + + im_gs = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) + im_gs = cv2.fastNlMeansDenoising(im_gs, h=3) + im_bw = cv2.threshold(im_gs, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1] + + lines = cv2.HoughLinesP( + im_bw, 1, np.pi / 180, 200, minLineLength=width / 12, maxLineGap=width / 150 + ) + + angles = [] + for line in lines: + x1, y1, x2, y2 = line[0] + angles.append(np.arctan2(y2 - y1, x2 - x1)) + angles = [angle for angle in angles if abs(angle) < max_skew_rad] + + if len(angles) < min_nlines: + return 0 + + return np.rad2deg(np.median(angles)) + + +def rotate_straight(im: np.array, skew_angle: int) -> np.array: + h, w = im.shape[:2] + center = (w // 2, h // 2) + + M = cv2.getRotationMatrix2D(center, skew_angle, 1.0) + + rotated = cv2.warpAffine( + im, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE + ) + + return rotated + + +def deskew_image(image: np.array, verbose=False) -> np.array: + skew_angle = detect_angle(image) + if verbose: + print(f"Skew angle: {skew_angle}") + deskewed = rotate_straight(image, skew_angle) + return deskewed + \ No newline at end of file