diff --git a/scripts/deskew_demo.py b/scripts/deskew_demo.py index 7a05806..ddbc22c 100644 --- a/scripts/deskew_demo.py +++ b/scripts/deskew_demo.py @@ -1,4 +1,4 @@ -#sample usage: python3 scripts/deskew_demo.py /path/to/crooked.pdf 0 +# sample usage: python3 scripts/deskew_demo.py /path/to/crooked.pdf 0 import argparse import numpy as np import pdf2image @@ -11,7 +11,7 @@ def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("pdf_path") parser.add_argument("page_index", type=int) - + args = parser.parse_args() return args diff --git a/vidocp/utils/deskew.py b/vidocp/utils/deskew.py index 70c36ff..0022f1e 100644 --- a/vidocp/utils/deskew.py +++ b/vidocp/utils/deskew.py @@ -2,40 +2,42 @@ import numpy as np import cv2 -def detect_angle(im: np.array, max_skew_deg=10, min_nlines=5) -> int: +def detect_angle(im: np.array, max_skew_deg=10, min_skew_deg=0.1, min_nlines=5) -> int: max_skew_rad = np.deg2rad(max_skew_deg) + min_skew_rad = np.deg2rad(min_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 - ) + 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] + raw_angle = np.arctan2(y2 - y1, x2 - x1) + angles.append(min(raw_angle, np.pi / 2 - raw_angle)) + angles = [angle for angle in angles if (abs(angle) < max_skew_rad)] + nonzero = list(filter(lambda x: x != 0, angles)) - if len(angles) < min_nlines: - return 0 + # empirically found this ad hoc approach to work + robust_avg = (np.mean(angles) + np.mean(nonzero) + np.median(nonzero)) / 3 + # slightly lower alternative: + # robust_avg = (np.mean(angles) + np.mean(nonzero) + np.median(angles) + np.median(nonzero)) / 4 - return np.rad2deg(np.median(angles)) + if robust_avg < min_skew_rad or min(len(angles), len(nonzero)) < min_nlines: + return 0.0 + return np.rad2deg(robust_avg) 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 - ) - + + rotated = cv2.warpAffine(im, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE) return rotated @@ -43,6 +45,7 @@ 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 + if skew_angle: + deskewed = rotate_straight(image, skew_angle) + return deskewed + return image