diff --git a/pdf_cleaner.py b/pdf_cleaner.py index 3bb3463..168094e 100755 --- a/pdf_cleaner.py +++ b/pdf_cleaner.py @@ -1,11 +1,16 @@ -import os #!/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 numpy as np -from PIL import Image -from pdf2image import convert_from_path import img2pdf +import numpy as np +from pdf2image import convert_from_path def ensure_temp_dir(): """Ensure temporary directory exists and return its path.""" @@ -117,6 +122,125 @@ def deskew(): else: print(" No line segments detected") +def trim_whitespace(image_path: str) -> bool: + """Remove white space from around the image. + + Returns: + bool: True if successful, False otherwise + """ + try: + # Read the image + img = cv2.imread(image_path) + if img is None: + print(f"Warning: Could not read image {image_path}") + return False + + # Convert to grayscale + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + + # Threshold the image + _, thresh = cv2.threshold(gray, 250, 255, cv2.THRESH_BINARY_INV) + + # Find non-zero points + coords = cv2.findNonZero(thresh) + if coords is None: + print(f"Warning: No content found in {image_path}") + return False + + # Get bounding rectangle + x, y, w, h = cv2.boundingRect(coords) + + # Add small padding (20 pixels) + padding = 20 + height, width = img.shape[:2] + x = max(0, x - padding) + y = max(0, y - padding) + w = min(width - x, w + 2 * padding) + h = min(height - y, h + 2 * padding) + + # Crop the image + cropped = img[y:y+h, x:x+w] + + # Save the cropped image + cv2.imwrite(image_path, cropped) + return True + except Exception as e: + print(f"Error processing {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', '-o7', 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() +def optimize(): + """Optimize images by trimming whitespace and optionally running optipng.""" + 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 + + total_files = len(temp_files) + successful_trims = 0 + successful_opts = 0 + + print(f"Processing {total_files} images...") + + # 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_trims += 1 + print(" ") + else: + print(" ") + + # Optimize with optipng if available + 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_opts += 1 + print(" ") + else: + print(" ") + else: + print("\nNote: optipng not found. Skipping PNG optimization.") + + # Print summary + print("\nOptimization complete!") + print(f"Successfully trimmed: {successful_trims}/{total_files} images") + if check_optipng_installed(): + print(f"Successfully optimized: {successful_opts}/{total_files} images") + @cli.command() @click.argument('output_pdf', type=click.Path()) def finalize(output_pdf):