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

@@ -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
+        # Read the image with alpha channel
-        img = cv2.imread(image_path)
+        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)
+        # Add padding
-        padding = 20
+        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():
+@click.option('--level', type=click.IntRange(1, 3), default=1,
-    """Optimize images by trimming whitespace and optionally running optipng."""
+              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 = {
-    successful_opts = 0
+        '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
+    # Step 2: Convert to monochrome if level >= 2
-    if check_optipng_installed():
+    if int(opt_level) >= int(OptimizationLevel.MONOCHROME):
-        print("\nOptimizing PNG files with optipng...")
+        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)
+            print(f"[{i}/{total_files}] Converting {os.path.basename(file_path)}...", end='', flush=True)
-            if run_optipng(file_path):
+            if convert_to_monochrome(file_path):
-                successful_opts += 1
+                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")
+    print(f"Successfully trimmed: {successful['trim']}/{total_files} images")
-    if check_optipng_installed():
+    if int(opt_level) >= int(OptimizationLevel.MONOCHROME):
-        print(f"Successfully optimized: {successful_opts}/{total_files} images")
+        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())

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_'