415 lines
15 KiB
Python
Executable File
415 lines
15 KiB
Python
Executable File
#!/usr/bin/env -S uv run
|
|
|
|
import os
|
|
import subprocess
|
|
import click
|
|
import cv2
|
|
import img2pdf
|
|
import numpy as np
|
|
from pdf2image import convert_from_path
|
|
|
|
import settings
|
|
from PIL import Image
|
|
import concurrent.futures
|
|
|
|
class Settings:
|
|
TRIM_PADDING_PIXELS = 20
|
|
MONOCHROME_THRESHOLD = 127
|
|
OPTIPNG_OPTIMIZATION_LEVEL = 7
|
|
PDF_BORDER_SIZE = 50
|
|
|
|
def ensure_temp_dir():
|
|
"""Ensure temporary directory exists and return its path."""
|
|
temp_dir = "temp_processed_images"
|
|
os.makedirs(temp_dir, exist_ok=True)
|
|
return temp_dir
|
|
|
|
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()
|
|
|
|
# Convert PDF to images
|
|
print(f"Extracting pages from {input_pdf}...")
|
|
pages = convert_from_path(input_pdf, dpi=400)
|
|
|
|
# Save each page
|
|
for i, page in enumerate(pages):
|
|
# Convert to grayscale
|
|
page = page.convert('L')
|
|
|
|
output_path = os.path.join(temp_dir, f"page_{i:03d}.png")
|
|
|
|
# Save initial version
|
|
page.save(output_path, "PNG", optimize=False)
|
|
|
|
# Trim whitespace
|
|
trim_whitespace(output_path)
|
|
|
|
# Reload the trimmed image
|
|
page = Image.open(output_path)
|
|
|
|
# Calculate new height maintaining aspect ratio
|
|
width = 2048
|
|
ratio = width / page.width
|
|
height = int(page.height * ratio)
|
|
|
|
# Resize using Lanczos
|
|
page = page.resize((width, height), Image.Resampling.LANCZOS)
|
|
|
|
# Save final version
|
|
page.save(output_path, "PNG", optimize=False)
|
|
print(f"Saved page {i+1}/{len(pages)}")
|
|
|
|
print(f"Extracted {len(pages)} pages to {temp_dir}/")
|
|
|
|
def trim_whitespace(image_path: str, is_final: bool = False) -> bool:
|
|
"""Remove white space from around the image.
|
|
|
|
Handles both RGB and RGBA images, treating transparent pixels as white.
|
|
|
|
Args:
|
|
image_path: Path to the image file
|
|
is_final: Whether this is the final operation on the image
|
|
|
|
Returns:
|
|
bool: True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Open image with PIL
|
|
image = Image.open(image_path)
|
|
|
|
# If image has transparency, flatten it first
|
|
if image.mode == 'RGBA':
|
|
# Create a white background
|
|
background = Image.new('RGB', image.size, 'white')
|
|
# Paste using alpha channel as mask
|
|
background.paste(image, mask=image.split()[3])
|
|
image = background
|
|
|
|
# Convert to grayscale
|
|
image = image.convert('L')
|
|
|
|
# Threshold to make all light pixels white and everything else black
|
|
# This helps with finding content bounds
|
|
image = image.point(lambda x: 255 if x > 250 else 0)
|
|
|
|
# Invert so content is white on black background
|
|
image = Image.eval(image, lambda x: 255 - x)
|
|
|
|
# Get the bounding box of content (now white pixels)
|
|
bbox = image.getbbox()
|
|
if not bbox:
|
|
print(f"Warning: No content found in {image_path}")
|
|
return False
|
|
|
|
# Add padding
|
|
padding = settings.TRIM_PADDING_PIXELS
|
|
width, height = image.size
|
|
x1, y1, x2, y2 = bbox
|
|
x1 = max(0, x1 - padding)
|
|
y1 = max(0, y1 - padding)
|
|
x2 = min(width, x2 + padding)
|
|
y2 = min(height, y2 + padding)
|
|
|
|
# Open original image again and crop it using the bounds
|
|
original = Image.open(image_path)
|
|
cropped = original.crop((x1, y1, x2, y2))
|
|
|
|
# Save the cropped image, optimizing only if this is the final operation
|
|
cropped.save(image_path, "PNG", optimize=is_final)
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"Error processing {image_path}: {str(e)}")
|
|
return False
|
|
|
|
@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
|
|
|
|
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")
|
|
|
|
def convert_to_monochrome(image_path: str, is_final: bool = False) -> bool:
|
|
"""Convert image to 1-bit monochrome.
|
|
|
|
Handles RGBA images by converting transparent pixels to white before thresholding.
|
|
|
|
Args:
|
|
image_path: Path to the image file
|
|
is_final: Whether this is the final operation on the image
|
|
|
|
Returns:
|
|
bool: True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Open image with PIL
|
|
image = Image.open(image_path)
|
|
|
|
# Convert to RGBA if not already
|
|
if image.mode != 'RGBA':
|
|
image = image.convert('RGBA')
|
|
|
|
# Get the image data as a list of pixels
|
|
data = image.getdata()
|
|
|
|
# Create new image data, replacing transparent pixels with white
|
|
new_data = []
|
|
for item in data:
|
|
# If pixel is transparent (alpha < 128), make it white
|
|
if item[3] < 128:
|
|
new_data.append((255, 255, 255, 255))
|
|
else:
|
|
new_data.append(item)
|
|
|
|
# Create new image with modified data
|
|
image.putdata(new_data)
|
|
|
|
# Convert to grayscale
|
|
image = image.convert('L')
|
|
|
|
# Convert to 1-bit using threshold
|
|
image = image.point(lambda x: 255 if x > settings.MONOCHROME_THRESHOLD else 0, '1')
|
|
|
|
# Save the monochrome image, optimizing only if this is the final operation
|
|
image.save(image_path, "PNG", optimize=is_final)
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"Error converting to monochrome {image_path}: {str(e)}")
|
|
return False
|
|
|
|
def check_optipng_installed() -> bool:
|
|
"""Check if optipng is installed."""
|
|
try:
|
|
result = subprocess.run(['optipng', '-v'], capture_output=True, text=True)
|
|
return result.returncode == 0
|
|
except FileNotFoundError:
|
|
return False
|
|
|
|
def run_optipng(file_path: str) -> bool:
|
|
"""Run optipng on a file with error handling.
|
|
|
|
Returns:
|
|
bool: True if successful, False otherwise
|
|
"""
|
|
try:
|
|
subprocess.run(
|
|
['optipng', f'-o{settings.OPTIPNG_OPTIMIZATION_LEVEL}', file_path],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True
|
|
)
|
|
return True
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Error optimizing {file_path}: {e.stderr}")
|
|
return False
|
|
except Exception as e:
|
|
print(f"Unexpected error optimizing {file_path}: {str(e)}")
|
|
return False
|
|
|
|
@cli.command()
|
|
@click.option('--level', type=click.IntRange(1, 3), default=1,
|
|
help='Optimization level: 1=trim, 2=monochrome, 3=full with optipng')
|
|
def optimize(level):
|
|
"""Optimize images with specified level of processing.
|
|
|
|
Optimization Levels:
|
|
1: Only trim whitespace
|
|
2: Level 1 + convert to 1-bit monochrome
|
|
3: Level 2 + optipng optimization
|
|
"""
|
|
temp_files = get_temp_files()
|
|
if not temp_files:
|
|
print("No pages found in temporary directory. Run 'extract' first.")
|
|
return
|
|
|
|
opt_level = settings.OptimizationLevel(level)
|
|
|
|
total_files = len(temp_files)
|
|
successful = {
|
|
'trim': 0,
|
|
'monochrome': 0,
|
|
'optipng': 0
|
|
}
|
|
|
|
print(f"Processing {total_files} images at optimization level {level}...")
|
|
|
|
# Step 1: Convert to monochrome if level >= 2
|
|
if int(opt_level) >= int(settings.OptimizationLevel.MONOCHROME):
|
|
print("\nConverting to monochrome...")
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
futures = {executor.submit(convert_to_monochrome, file_path, (opt_level == settings.OptimizationLevel.MONOCHROME)): file_path for file_path in temp_files}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
file_path = futures[future]
|
|
try:
|
|
if future.result():
|
|
successful['monochrome'] += 1
|
|
print(f"Converted {os.path.basename(file_path)}")
|
|
else:
|
|
print(f"Failed to convert {os.path.basename(file_path)}")
|
|
except Exception as exc:
|
|
print(f"{os.path.basename(file_path)} generated an exception: {exc}")
|
|
|
|
# Step 2: Always trim whitespace
|
|
print("\nTrimming whitespace from images...")
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
futures = {executor.submit(trim_whitespace, file_path, (opt_level == settings.OptimizationLevel.TRIM)): file_path for file_path in temp_files}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
file_path = futures[future]
|
|
try:
|
|
if future.result():
|
|
successful['trim'] += 1
|
|
print(f"Trimmed {os.path.basename(file_path)}")
|
|
else:
|
|
print(f"Failed to trim {os.path.basename(file_path)}")
|
|
except Exception as exc:
|
|
print(f"{os.path.basename(file_path)} generated an exception: {exc}")
|
|
|
|
# Step 3: Run optipng if level = 3
|
|
if opt_level == settings.OptimizationLevel.FULL:
|
|
if check_optipng_installed():
|
|
print("\nOptimizing PNG files with optipng...")
|
|
with concurrent.futures.ThreadPoolExecutor() as executor:
|
|
futures = {executor.submit(run_optipng, file_path): file_path for file_path in temp_files}
|
|
for future in concurrent.futures.as_completed(futures):
|
|
file_path = futures[future]
|
|
try:
|
|
if future.result():
|
|
successful['optipng'] += 1
|
|
print(f"Optimized {os.path.basename(file_path)}")
|
|
else:
|
|
print(f"Failed to optimize {os.path.basename(file_path)}")
|
|
except Exception as exc:
|
|
print(f"{os.path.basename(file_path)} generated an exception: {exc}")
|
|
else:
|
|
print("\nNote: optipng not found. Skipping PNG optimization.")
|
|
|
|
# Print summary
|
|
print("\nOptimization complete!")
|
|
print(f"Successfully trimmed: {successful['trim']}/{total_files} images")
|
|
if int(opt_level) >= int(settings.OptimizationLevel.MONOCHROME):
|
|
print(f"Successfully converted to monochrome: {successful['monochrome']}/{total_files} images")
|
|
if opt_level == settings.OptimizationLevel.FULL and check_optipng_installed():
|
|
print(f"Successfully optimized with optipng: {successful['optipng']}/{total_files} images")
|
|
|
|
@cli.command()
|
|
@click.argument('output_pdf', type=click.Path())
|
|
def finalize(output_pdf):
|
|
"""Combine processed pages into final PDF and clean up."""
|
|
temp_dir = ensure_temp_dir()
|
|
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}...")
|
|
|
|
# A4 size in millimeters
|
|
A4_WIDTH_MM = 210
|
|
A4_HEIGHT_MM = 297
|
|
|
|
# Convert to PDF with border
|
|
with open(output_pdf, "wb") as f:
|
|
f.write(img2pdf.convert(
|
|
temp_files,
|
|
with_pdfrw=True,
|
|
layout_fun=img2pdf.get_layout_fun(
|
|
pagesize=(img2pdf.mm_to_pt(A4_WIDTH_MM), img2pdf.mm_to_pt(A4_HEIGHT_MM)),
|
|
border=(settings.PDF_BORDER_SIZE,) * 4, # Same border size for all sides (top, right, bottom, left)
|
|
fit=img2pdf.FitMode.into
|
|
)
|
|
))
|
|
|
|
# Clean up temporary files
|
|
print("Cleaning up temporary files...")
|
|
for file_path in temp_files:
|
|
os.remove(file_path)
|
|
os.rmdir(temp_dir)
|
|
|
|
print("PDF created successfully and temporary files removed!")
|
|
|
|
if __name__ == '__main__':
|
|
cli()
|