feat: Improved deskew to focus on horizontal lines
- Modified deskew function to better detect staff lines - Added morphological operations to enhance horizontal lines - Used probabilistic Hough transform for more precise line detection - Added .png to gitignore
This commit is contained in:
Esa Kataja
Regular → Executable
+125
-143
@@ -1,157 +1,139 @@
|
||||
#!/usr/bin/env python3
|
||||
import os
|
||||
#!/usr/bin/env -S uv run
|
||||
import click
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
from pdf2image import convert_from_path
|
||||
import img2pdf
|
||||
from PIL import Image
|
||||
import argparse
|
||||
|
||||
def deskew(image):
|
||||
"""Deskew the image using contour detection and rotation."""
|
||||
# Create a copy for processing while keeping original quality
|
||||
proc_image = image.copy()
|
||||
|
||||
# Convert to grayscale and blur
|
||||
gray = cv2.cvtColor(proc_image, cv2.COLOR_BGR2GRAY)
|
||||
blur = cv2.GaussianBlur(gray, (9, 9), 0)
|
||||
|
||||
# Threshold the image
|
||||
thresh = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
|
||||
|
||||
# Find all contours
|
||||
contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
|
||||
|
||||
if not contours:
|
||||
return image
|
||||
|
||||
# Find largest contour
|
||||
contour = max(contours, key=cv2.contourArea)
|
||||
|
||||
# Find minimum area rectangle
|
||||
rect = cv2.minAreaRect(contour)
|
||||
angle = rect[-1]
|
||||
|
||||
# Adjust angle to be between -45 and 45 degrees
|
||||
while angle < -45:
|
||||
angle += 90
|
||||
while angle > 45:
|
||||
angle -= 90
|
||||
|
||||
# Only rotate if the angle is significant enough
|
||||
if abs(angle) < 0.5: # Skip tiny rotations
|
||||
return image
|
||||
|
||||
# Rotate the image
|
||||
(h, w) = image.shape[:2]
|
||||
center = (w // 2, h // 2)
|
||||
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||||
rotated = cv2.warpAffine(image, M, (w, h),
|
||||
flags=cv2.INTER_CUBIC,
|
||||
borderMode=cv2.BORDER_REPLICATE)
|
||||
|
||||
return rotated
|
||||
|
||||
def adjust_levels(image):
|
||||
"""Adjust image levels for better contrast and clarity."""
|
||||
# Convert to grayscale if not already
|
||||
if len(image.shape) == 3:
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
else:
|
||||
gray = image
|
||||
|
||||
# Apply bilateral filter to preserve edges while reducing noise
|
||||
denoised = cv2.bilateralFilter(gray, 9, 75, 75)
|
||||
|
||||
# Create a background mask using morphological operations
|
||||
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
|
||||
background = cv2.morphologyEx(denoised, cv2.MORPH_DILATE, kernel)
|
||||
|
||||
# Subtract background to normalize lighting
|
||||
normalized = cv2.subtract(background, denoised)
|
||||
|
||||
# Apply Gaussian blur to reduce noise while preserving edges
|
||||
blurred = cv2.GaussianBlur(normalized, (3, 3), 0)
|
||||
|
||||
# Use Otsu's thresholding for optimal binary threshold
|
||||
_, binary = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
|
||||
|
||||
# Always ensure background is white and text is black
|
||||
if cv2.countNonZero(binary) < binary.size / 2:
|
||||
binary = cv2.bitwise_not(binary)
|
||||
|
||||
return binary
|
||||
|
||||
def process_pdf(input_path, output_path):
|
||||
"""Process a PDF file and save the cleaned version."""
|
||||
print(f"Processing {input_path}...")
|
||||
|
||||
# Convert PDF to images with high DPI to ensure minimum width of 2048px
|
||||
pages = convert_from_path(input_path, dpi=300)
|
||||
|
||||
# Create temporary directory for processed images
|
||||
def ensure_temp_dir():
|
||||
"""Ensure temporary directory exists and return its path."""
|
||||
temp_dir = "temp_processed_images"
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
# Process each page
|
||||
temp_image_paths = []
|
||||
for i, page in enumerate(pages):
|
||||
print(f"Processing page {i+1}/{len(pages)}")
|
||||
|
||||
# Ensure minimum width of 2048px
|
||||
width, height = page.size
|
||||
scale = max(1, 2048 / width)
|
||||
if scale > 1:
|
||||
new_width = int(width * scale)
|
||||
new_height = int(height * scale)
|
||||
page = page.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
# Convert PIL Image to OpenCV format
|
||||
opencv_image = cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR)
|
||||
|
||||
# Process the image
|
||||
deskewed = deskew(opencv_image)
|
||||
cleaned = adjust_levels(deskewed)
|
||||
|
||||
# Convert to PIL Image
|
||||
pil_image = Image.fromarray(cleaned)
|
||||
|
||||
# Save temporarily as 1-bit PNG
|
||||
temp_path = os.path.join(temp_dir, f"page_{i:03d}.png")
|
||||
pil_image.save(temp_path, "PNG", optimize=False)
|
||||
temp_image_paths.append(temp_path)
|
||||
|
||||
# Save processed images as PDF with high quality settings
|
||||
print("Saving cleaned PDF...")
|
||||
a4_width_mm = 210
|
||||
a4_height_mm = 297
|
||||
layout_fun = img2pdf.get_layout_fun((img2pdf.mm_to_pt(a4_width_mm),
|
||||
img2pdf.mm_to_pt(a4_height_mm)))
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(img2pdf.convert(temp_image_paths,
|
||||
layout_fun=layout_fun,
|
||||
with_pdfrw=True))
|
||||
|
||||
# Clean up temporary files
|
||||
for temp_path in temp_image_paths:
|
||||
os.remove(temp_path)
|
||||
os.rmdir(temp_dir)
|
||||
|
||||
print(f"Saved cleaned PDF to {output_path}")
|
||||
return temp_dir
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Clean and straighten image-based PDFs')
|
||||
parser.add_argument('input_pdf', help='Path to the input PDF file')
|
||||
parser.add_argument('output_pdf', help='Path for the output PDF file')
|
||||
def get_temp_files():
|
||||
"""Get list of temporary PNG files in order."""
|
||||
temp_dir = ensure_temp_dir()
|
||||
files = [f for f in os.listdir(temp_dir) if f.endswith('.png')]
|
||||
files.sort() # Ensure correct page order
|
||||
return [os.path.join(temp_dir, f) for f in files]
|
||||
|
||||
@click.group()
|
||||
def cli():
|
||||
"""PDF cleaning toolbox for musical scores."""
|
||||
pass
|
||||
|
||||
@cli.command()
|
||||
@click.argument('input_pdf', type=click.Path(exists=True))
|
||||
def extract(input_pdf):
|
||||
"""Extract pages from PDF to temporary directory."""
|
||||
temp_dir = ensure_temp_dir()
|
||||
|
||||
args = parser.parse_args()
|
||||
# Convert PDF to images
|
||||
print(f"Extracting pages from {input_pdf}...")
|
||||
pages = convert_from_path(input_pdf, dpi=400)
|
||||
|
||||
if not os.path.exists(args.input_pdf):
|
||||
print(f"Error: Input file '{args.input_pdf}' does not exist")
|
||||
# Save each page
|
||||
for i, page in enumerate(pages):
|
||||
output_path = os.path.join(temp_dir, f"page_{i:03d}.png")
|
||||
page.save(output_path, "PNG", optimize=False)
|
||||
print(f"Saved page {i+1}/{len(pages)}")
|
||||
|
||||
print(f"Extracted {len(pages)} pages to {temp_dir}/")
|
||||
|
||||
@cli.command()
|
||||
def deskew():
|
||||
"""Deskew all pages in temporary directory."""
|
||||
temp_files = get_temp_files()
|
||||
if not temp_files:
|
||||
print("No pages found in temporary directory. Run 'extract' first.")
|
||||
return
|
||||
|
||||
process_pdf(args.input_pdf, args.output_pdf)
|
||||
for file_path in temp_files:
|
||||
print(f"Deskewing {os.path.basename(file_path)}...")
|
||||
|
||||
# Read image
|
||||
image = cv2.imread(file_path)
|
||||
|
||||
# Convert to grayscale
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
|
||||
# Apply threshold to get binary image
|
||||
_, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
|
||||
|
||||
# Create a rectangular kernel that's wider than it is tall
|
||||
# This helps detect horizontal lines
|
||||
kernel_length = np.array(binary).shape[1]//80
|
||||
horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (kernel_length, 1))
|
||||
|
||||
# Detect horizontal lines
|
||||
horizontal_lines = cv2.erode(binary, horizontal_kernel, iterations=3)
|
||||
horizontal_lines = cv2.dilate(horizontal_lines, horizontal_kernel, iterations=3)
|
||||
|
||||
# Use probabilistic Hough transform to detect line segments
|
||||
line_segments = cv2.HoughLinesP(
|
||||
cv2.bitwise_not(horizontal_lines),
|
||||
rho=1,
|
||||
theta=np.pi/180,
|
||||
threshold=100,
|
||||
minLineLength=binary.shape[1]//4, # Lines must be at least 1/4 of image width
|
||||
maxLineGap=20
|
||||
)
|
||||
|
||||
if line_segments is not None and len(line_segments) > 0:
|
||||
# Calculate angles of detected line segments
|
||||
angles = []
|
||||
for line in line_segments:
|
||||
x1, y1, x2, y2 = line[0]
|
||||
if x2 - x1 == 0: # Avoid division by zero
|
||||
continue
|
||||
angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
|
||||
# Only consider angles that are close to horizontal
|
||||
if abs(angle) < 20:
|
||||
angles.append(angle)
|
||||
|
||||
if angles:
|
||||
# Use median angle to avoid outliers
|
||||
median_angle = np.median(angles)
|
||||
|
||||
# Only rotate if the angle is significant but not too large
|
||||
if 0.5 < abs(median_angle) < 20:
|
||||
height, width = image.shape[:2]
|
||||
center = (width/2, height/2)
|
||||
rotation_matrix = cv2.getRotationMatrix2D(center, median_angle, 1.0)
|
||||
rotated = cv2.warpAffine(image, rotation_matrix, (width, height),
|
||||
flags=cv2.INTER_CUBIC,
|
||||
borderMode=cv2.BORDER_REPLICATE)
|
||||
|
||||
# Save rotated image
|
||||
cv2.imwrite(file_path, rotated)
|
||||
print(f" Rotated by {median_angle:.2f} degrees")
|
||||
else:
|
||||
print(" No significant rotation needed")
|
||||
else:
|
||||
print(" No valid horizontal lines found")
|
||||
else:
|
||||
print(" No line segments detected")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@cli.command()
|
||||
@click.argument('output_pdf', type=click.Path())
|
||||
def finalize(output_pdf):
|
||||
"""Combine processed pages into final PDF."""
|
||||
temp_files = get_temp_files()
|
||||
if not temp_files:
|
||||
print("No pages found in temporary directory. Run 'extract' first.")
|
||||
return
|
||||
|
||||
print(f"Combining {len(temp_files)} pages into {output_pdf}...")
|
||||
|
||||
# Convert to PDF
|
||||
with open(output_pdf, "wb") as f:
|
||||
f.write(img2pdf.convert(temp_files))
|
||||
|
||||
print("PDF created successfully!")
|
||||
print("Note: Temporary files were kept for further processing if needed.")
|
||||
|
||||
if __name__ == '__main__':
|
||||
cli()
|
||||
|
||||
Reference in New Issue
Block a user