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
+105
-25
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user