From 682a1f13eeb3c96e127c332a84ca8ecf782ac474 Mon Sep 17 00:00:00 2001 From: Esa Kataja Date: Tue, 3 Dec 2024 21:38:55 +0200 Subject: [PATCH] fix: Improve transparency handling and optimization levels - Handle transparency in whitespace trimming - Use IntEnum for optimization levels - Add proper level comparisons - Treat transparent pixels as white for boundary detection - Preserve transparency in output - Fix black artifacts in transparent areas --- pdf_cleaner.py | 130 +++++++++++++++++++++++++++++++++++++++---------- settings.py | 16 ++++++ 2 files changed, 121 insertions(+), 25 deletions(-) diff --git a/pdf_cleaner.py b/pdf_cleaner.py index 168094e..2ca1494 100755 --- a/pdf_cleaner.py +++ b/pdf_cleaner.py @@ -12,6 +12,16 @@ import img2pdf import numpy as np from pdf2image import convert_from_path +import settings +from settings import OptimizationLevel + +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" @@ -125,16 +135,30 @@ def deskew(): 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. + Returns: bool: True if successful, False otherwise """ try: - # Read the image - img = cv2.imread(image_path) + # Read the image with alpha channel + img = cv2.imread(image_path, cv2.IMREAD_UNCHANGED) if img is None: print(f"Warning: Could not read image {image_path}") return False + # Convert to grayscale, handling transparency + if img.shape[-1] == 4: # RGBA + # Create a white background + white_background = np.ones_like(img, dtype=np.uint8) * 255 + + # Extract alpha channel and create mask + alpha = img[:, :, 3] + alpha_mask = alpha[:, :, np.newaxis] / 255.0 + + # Blend image with white background based on alpha + img = (img[:, :, :3] * alpha_mask + white_background[:, :, :3] * (1 - alpha_mask)).astype(np.uint8) + # Convert to grayscale gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) @@ -150,8 +174,8 @@ def trim_whitespace(image_path: str) -> bool: # Get bounding rectangle x, y, w, h = cv2.boundingRect(coords) - # Add small padding (20 pixels) - padding = 20 + # Add padding + padding = settings.TRIM_PADDING_PIXELS height, width = img.shape[:2] x = max(0, x - padding) y = max(0, y - padding) @@ -159,15 +183,44 @@ def trim_whitespace(image_path: str) -> bool: h = min(height - y, h + 2 * padding) # Crop the image - cropped = img[y:y+h, x:x+w] + if img.shape[-1] == 4: # If original was RGBA + cropped = img[:, :, :4][y:y+h, x:x+w] # Keep alpha channel + else: + cropped = img[y:y+h, x:x+w] - # Save the cropped image + # Save the cropped image with transparency preserved cv2.imwrite(image_path, cropped) return True except Exception as e: print(f"Error processing {image_path}: {str(e)}") return False +def convert_to_monochrome(image_path: str) -> bool: + """Convert image to 1-bit monochrome. + + 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) + + # Apply threshold + _, mono = cv2.threshold(gray, settings.MONOCHROME_THRESHOLD, 255, cv2.THRESH_BINARY) + + # Save the monochrome image + cv2.imwrite(image_path, mono) + 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: @@ -184,7 +237,7 @@ def run_optipng(file_path: str) -> bool: """ try: result = subprocess.run( - ['optipng', '-o7', file_path], + ['optipng', f'-o{settings.OPTIPNG_OPTIMIZATION_LEVEL}', file_path], capture_output=True, text=True, check=True @@ -198,48 +251,75 @@ def run_optipng(file_path: str) -> bool: return False @cli.command() -def optimize(): - """Optimize images by trimming whitespace and optionally running optipng.""" +@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_trims = 0 - successful_opts = 0 + successful = { + 'trim': 0, + 'monochrome': 0, + 'optipng': 0 + } - print(f"Processing {total_files} images...") + print(f"Processing {total_files} images at optimization level {level}...") - # Trim whitespace + # 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_trims += 1 + successful['trim'] += 1 print(" ") else: print(" ") - # Optimize with optipng if available - if check_optipng_installed(): - print("\nOptimizing PNG files with optipng...") + # 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}] Optimizing {os.path.basename(file_path)}...", end='', flush=True) - if run_optipng(file_path): - successful_opts += 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(" ") - else: - print("\nNote: optipng not found. Skipping PNG optimization.") + + # 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_trims}/{total_files} images") - if check_optipng_installed(): - print(f"Successfully optimized: {successful_opts}/{total_files} images") + 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()) diff --git a/settings.py b/settings.py index 044e084..a7785d4 100644 --- a/settings.py +++ b/settings.py @@ -2,6 +2,19 @@ import os import multiprocessing +from enum import IntEnum + +class OptimizationLevel(IntEnum): + """Optimization levels for image processing. + + Levels: + 1: Only trim whitespace + 2: Trim + convert to 1-bit monochrome + 3: All optimizations + optipng + """ + TRIM = 1 # Only trim whitespace + MONOCHROME = 2 # Trim + convert to 1-bit monochrome + FULL = 3 # All optimizations + optipng # Number of parallel processes to use for operations that support parallelization # Defaults to number of CPU cores - 1, but never less than 1 @@ -14,5 +27,8 @@ PARALLEL_PROCESSES = int(os.getenv('NOTES_CLEANER_PARALLEL_PROCESSES', DEFAULT_P OPTIPNG_OPTIMIZATION_LEVEL = int(os.getenv('NOTES_CLEANER_OPTIPNG_LEVEL', '7')) # 0-7, higher = better compression but slower TRIM_PADDING_PIXELS = int(os.getenv('NOTES_CLEANER_TRIM_PADDING', '20')) # Padding around content after trimming +# Image processing settings +MONOCHROME_THRESHOLD = int(os.getenv('NOTES_CLEANER_MONO_THRESHOLD', '200')) # 0-255, higher = more white + # Temporary directory settings TEMP_DIR_PREFIX = 'notes_cleaner_'