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
This commit is contained in:
Esa Kataja
2024-12-03 21:38:55 +02:00
parent 6e1d90de2f
commit 682a1f13ee
2 changed files with 121 additions and 25 deletions
+98 -18
View File
@@ -12,6 +12,16 @@ import img2pdf
import numpy as np import numpy as np
from pdf2image import convert_from_path 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(): def ensure_temp_dir():
"""Ensure temporary directory exists and return its path.""" """Ensure temporary directory exists and return its path."""
temp_dir = "temp_processed_images" temp_dir = "temp_processed_images"
@@ -125,16 +135,30 @@ def deskew():
def trim_whitespace(image_path: str) -> bool: def trim_whitespace(image_path: str) -> bool:
"""Remove white space from around the image. """Remove white space from around the image.
Handles both RGB and RGBA images, treating transparent pixels as white.
Returns: Returns:
bool: True if successful, False otherwise bool: True if successful, False otherwise
""" """
try: 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: if img is None:
print(f"Warning: Could not read image {image_path}") print(f"Warning: Could not read image {image_path}")
return False 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 # Convert to grayscale
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
@@ -150,8 +174,8 @@ def trim_whitespace(image_path: str) -> bool:
# Get bounding rectangle # Get bounding rectangle
x, y, w, h = cv2.boundingRect(coords) 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] height, width = img.shape[:2]
x = max(0, x - padding) x = max(0, x - padding)
y = max(0, y - padding) y = max(0, y - padding)
@@ -159,15 +183,44 @@ def trim_whitespace(image_path: str) -> bool:
h = min(height - y, h + 2 * padding) h = min(height - y, h + 2 * padding)
# Crop the image # Crop the image
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] cropped = img[y:y+h, x:x+w]
# Save the cropped image # Save the cropped image with transparency preserved
cv2.imwrite(image_path, cropped) cv2.imwrite(image_path, cropped)
return True return True
except Exception as e: except Exception as e:
print(f"Error processing {image_path}: {str(e)}") print(f"Error processing {image_path}: {str(e)}")
return False 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: def check_optipng_installed() -> bool:
"""Check if optipng is installed.""" """Check if optipng is installed."""
try: try:
@@ -184,7 +237,7 @@ def run_optipng(file_path: str) -> bool:
""" """
try: try:
result = subprocess.run( result = subprocess.run(
['optipng', '-o7', file_path], ['optipng', f'-o{settings.OPTIPNG_OPTIMIZATION_LEVEL}', file_path],
capture_output=True, capture_output=True,
text=True, text=True,
check=True check=True
@@ -198,37 +251,62 @@ def run_optipng(file_path: str) -> bool:
return False return False
@cli.command() @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_dir = ensure_temp_dir()
temp_files = get_temp_files() temp_files = get_temp_files()
if not temp_files: if not temp_files:
print("No pages found in temporary directory. Run 'extract' first.") print("No pages found in temporary directory. Run 'extract' first.")
return return
opt_level = OptimizationLevel(level)
total_files = len(temp_files) 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...") print("\nTrimming whitespace from images...")
for i, file_path in enumerate(temp_files, 1): for i, file_path in enumerate(temp_files, 1):
print(f"[{i}/{total_files}] Processing {os.path.basename(file_path)}...", end='', flush=True) print(f"[{i}/{total_files}] Processing {os.path.basename(file_path)}...", end='', flush=True)
if trim_whitespace(file_path): if trim_whitespace(file_path):
successful_trims += 1 successful['trim'] += 1
print(" ") print(" ")
else: else:
print(" ") print(" ")
# Optimize with optipng if available # 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(): if check_optipng_installed():
print("\nOptimizing PNG files with optipng...") print("\nOptimizing PNG files with optipng...")
for i, file_path in enumerate(temp_files, 1): 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}] Optimizing {os.path.basename(file_path)}...", end='', flush=True)
if run_optipng(file_path): if run_optipng(file_path):
successful_opts += 1 successful['optipng'] += 1
print(" ") print(" ")
else: else:
print(" ") print(" ")
@@ -237,9 +315,11 @@ def optimize():
# Print summary # Print summary
print("\nOptimization complete!") 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() @cli.command()
@click.argument('output_pdf', type=click.Path()) @click.argument('output_pdf', type=click.Path())
+16
View File
@@ -2,6 +2,19 @@
import os import os
import multiprocessing 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 # Number of parallel processes to use for operations that support parallelization
# Defaults to number of CPU cores - 1, but never less than 1 # 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 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 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 # Temporary directory settings
TEMP_DIR_PREFIX = 'notes_cleaner_' TEMP_DIR_PREFIX = 'notes_cleaner_'