fix angle detection to make more sensitive to small angles; format with black

This commit is contained in:
Isaac Riley 2022-02-21 16:52:24 +01:00
parent 08cce36940
commit 59e082379c
2 changed files with 23 additions and 20 deletions

View File

@ -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

View File

@ -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
if skew_angle:
deskewed = rotate_straight(image, skew_angle)
return deskewed
return image