Esa Kataja
75dc1c3e2c
Simplify transparency handling by flattening image with white background for bounds detection while preserving original transparency in the final crop. This fixes the issue where transparent areas were being treated as black during trimming.
380 lines
13 KiB
Python
Executable File
380 lines
13 KiB
Python
Executable File
#!/usr/bin/env -S uv run
|
|
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from pathlib import Path
|
|
from typing import List, Tuple
|
|
|
|
import click
|
|
import cv2
|
|
import img2pdf
|
|
import numpy as np
|
|
from pdf2image import convert_from_path
|
|
|
|
import settings
|
|
from settings import OptimizationLevel
|
|
from PIL import Image
|
|
|
|
class Settings:
|
|
TRIM_PADDING_PIXELS = 20
|
|
MONOCHROME_THRESHOLD = 127
|
|
OPTIPNG_OPTIMIZATION_LEVEL = 7
|
|
|
|
settings = Settings()
|
|
|
|
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) -> 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
|
|
|
|
Returns:
|
|
bool: True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Open image with PIL
|
|
image = Image.open(image_path)
|
|
|
|
# Create flattened version with white background for finding bounds
|
|
flattened = Image.new('RGB', image.size, 'white')
|
|
if image.mode == 'RGBA':
|
|
flattened.paste(image, mask=image.split()[3])
|
|
else:
|
|
flattened.paste(image)
|
|
|
|
# Convert to grayscale for finding bounds
|
|
flattened = flattened.convert('L')
|
|
|
|
# Get the bounding box of non-white pixels
|
|
bbox = flattened.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)
|
|
|
|
# Crop the original image (preserving transparency)
|
|
cropped = image.crop((x1, y1, x2, y2))
|
|
|
|
# Save the cropped image
|
|
cropped.save(image_path, "PNG", optimize=False)
|
|
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) -> bool:
|
|
"""Convert image to 1-bit monochrome.
|
|
|
|
Handles RGBA images by converting transparent pixels to white before thresholding.
|
|
|
|
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
|
|
image.save(image_path, "PNG", optimize=False)
|
|
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:
|
|
result = 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_dir = ensure_temp_dir()
|
|
temp_files = get_temp_files()
|
|
if not temp_files:
|
|
print("No pages found in temporary directory. Run 'extract' first.")
|
|
return
|
|
|
|
opt_level = 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: Always trim whitespace
|
|
print("\nTrimming whitespace from images...")
|
|
for i, file_path in enumerate(temp_files, 1):
|
|
print(f"[{i}/{total_files}] Processing {os.path.basename(file_path)}...", end='', flush=True)
|
|
if trim_whitespace(file_path):
|
|
successful['trim'] += 1
|
|
print(" ")
|
|
else:
|
|
print(" ")
|
|
|
|
# Step 2: Convert to monochrome if level >= 2
|
|
if int(opt_level) >= int(OptimizationLevel.MONOCHROME):
|
|
print("\nConverting to monochrome...")
|
|
for i, file_path in enumerate(temp_files, 1):
|
|
print(f"[{i}/{total_files}] Converting {os.path.basename(file_path)}...", end='', flush=True)
|
|
if convert_to_monochrome(file_path):
|
|
successful['monochrome'] += 1
|
|
print(" ")
|
|
else:
|
|
print(" ")
|
|
|
|
# Step 3: Run optipng if level = 3
|
|
if opt_level == OptimizationLevel.FULL:
|
|
if check_optipng_installed():
|
|
print("\nOptimizing PNG files with optipng...")
|
|
for i, file_path in enumerate(temp_files, 1):
|
|
print(f"[{i}/{total_files}] Optimizing {os.path.basename(file_path)}...", end='', flush=True)
|
|
if run_optipng(file_path):
|
|
successful['optipng'] += 1
|
|
print(" ")
|
|
else:
|
|
print(" ")
|
|
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(OptimizationLevel.MONOCHROME):
|
|
print(f"Successfully converted to monochrome: {successful['monochrome']}/{total_files} images")
|
|
if opt_level == 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}...")
|
|
|
|
# Convert to PDF
|
|
with open(output_pdf, "wb") as f:
|
|
f.write(img2pdf.convert(temp_files))
|
|
|
|
# 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()
|