diff --git a/pdf_cleaner.py b/pdf_cleaner.py old mode 100644 new mode 100755 index 1c4e60a..57a2a74 --- a/pdf_cleaner.py +++ b/pdf_cleaner.py @@ -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() diff --git a/pyproject.toml b/pyproject.toml index 4585c71..716c825 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,5 +5,6 @@ description = "Add your description here" readme = "README.md" requires-python = ">=3.12" dependencies = [ + "click>=8.1.7", "pydantic>=2.10.3", ] diff --git a/uv.lock b/uv.lock index 0c4effc..d9bf036 100644 --- a/uv.lock +++ b/uv.lock @@ -10,16 +10,41 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643 }, ] +[[package]] +name = "click" +version = "8.1.7" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "platform_system == 'Windows'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/96/d3/f04c7bfcf5c1862a2a5b845c6b2b360488cf47af55dfa79c98f6a6bf98b5/click-8.1.7.tar.gz", hash = "sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de", size = 336121 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/00/2e/d53fa4befbf2cfa713304affc7ca780ce4fc1fd8710527771b58311a3229/click-8.1.7-py3-none-any.whl", hash = "sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28", size = 97941 }, +] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 }, +] + [[package]] name = "notes-cleaner" version = "0.1.0" source = { virtual = "." } dependencies = [ + { name = "click" }, { name = "pydantic" }, ] [package.metadata] -requires-dist = [{ name = "pydantic", specifier = ">=2.10.3" }] +requires-dist = [ + { name = "click", specifier = ">=8.1.7" }, + { name = "pydantic", specifier = ">=2.10.3" }, +] [[package]] name = "pydantic"