77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
import numpy as np
|
|
from scipy.ndimage import rotate
|
|
import cv2
|
|
|
|
|
|
def detect_angle_from_lines(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)
|
|
|
|
angles = []
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line[0]
|
|
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))
|
|
|
|
# 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
|
|
|
|
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)
|
|
return rotated
|
|
|
|
|
|
def deskew_linebased(image: np.array, verbose=False) -> np.array:
|
|
skew_angle = detect_angle_from_lines(image)
|
|
if verbose:
|
|
print(f"Skew angle from lines: {skew_angle}")
|
|
if skew_angle:
|
|
deskewed = rotate_straight(image, skew_angle)
|
|
return deskewed
|
|
return image
|
|
|
|
|
|
def deskew_histbased(page: np.array, max_abs_angle=1.5, delta=0.15, mode="nearest", verbose=False):
|
|
page = cv2.cvtColor(page, cv2.COLOR_BGR2GRAY)
|
|
page = cv2.fastNlMeansDenoising(page, h=3)
|
|
w, h = page.shape
|
|
|
|
def find_score(arr, angle):
|
|
data = rotate(arr, angle, reshape=False, order=0)
|
|
hist = np.sum(data, axis=1)
|
|
score = np.sum((hist[1:] - hist[:-1]) ** 2)
|
|
return score
|
|
|
|
angles = np.arange(-max_abs_angle, max_abs_angle + delta, delta)
|
|
scores = []
|
|
for angle in angles:
|
|
scores.append(find_score(page, angle))
|
|
|
|
best_angle = angles[scores.index(max(scores))]
|
|
if verbose:
|
|
print("Skew angle from pixel histogram: {}".format(best_angle))
|
|
|
|
rotated = rotate(page, best_angle, reshape=False, order=0, mode=mode)
|
|
return rotated
|