mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
87 lines
2.7 KiB
Python
87 lines
2.7 KiB
Python
"""Image preprocessing pipeline for OCR accuracy improvement.
|
|
|
|
Uses OpenCV for deskew, adaptive binarization, CLAHE contrast enhancement,
|
|
and denoising. All operations work on the image file in-place (overwrite).
|
|
"""
|
|
import cv2
|
|
import numpy as np
|
|
import sys
|
|
import json
|
|
|
|
|
|
def emit_progress(percent, stage):
|
|
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
|
|
|
|
|
|
def deskew(image):
|
|
"""Detect and correct rotation/skew using Hough line transform."""
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
|
|
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
|
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=100, maxLineGap=10)
|
|
|
|
if lines is None:
|
|
return image
|
|
|
|
angles = []
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line[0]
|
|
angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
|
|
if abs(angle) < 45:
|
|
angles.append(angle)
|
|
|
|
if not angles:
|
|
return image
|
|
|
|
median_angle = np.median(angles)
|
|
if abs(median_angle) < 0.5 or abs(median_angle) > 15:
|
|
return image
|
|
|
|
h, w = image.shape[:2]
|
|
center = (w // 2, h // 2)
|
|
matrix = cv2.getRotationMatrix2D(center, median_angle, 1.0)
|
|
rotated = cv2.warpAffine(image, matrix, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
|
|
return rotated
|
|
|
|
|
|
def binarize(image):
|
|
"""Adaptive thresholding for high-contrast black/white."""
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
|
|
binary = cv2.adaptiveThreshold(
|
|
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
|
|
)
|
|
return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
|
|
|
|
|
|
def enhance_contrast(image):
|
|
"""CLAHE contrast enhancement for uneven lighting."""
|
|
lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
|
l, a, b = cv2.split(lab)
|
|
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
|
l = clahe.apply(l)
|
|
enhanced = cv2.merge([l, a, b])
|
|
return cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
|
|
|
|
|
|
def denoise(image):
|
|
"""Bilateral filter to remove noise while preserving text edges."""
|
|
return cv2.bilateralFilter(image, 9, 75, 75)
|
|
|
|
|
|
def preprocess(input_path, output_path):
|
|
"""Run the full preprocessing pipeline and save result.
|
|
|
|
Steps: deskew -> enhance contrast -> denoise -> binarize
|
|
Order matters: binarize last because it strips color info needed by CLAHE.
|
|
"""
|
|
image = cv2.imread(input_path)
|
|
if image is None:
|
|
raise ValueError(f"Cannot read image: {input_path}")
|
|
|
|
image = deskew(image)
|
|
image = enhance_contrast(image)
|
|
image = denoise(image)
|
|
image = binarize(image)
|
|
|
|
cv2.imwrite(output_path, image)
|
|
return output_path
|