refactor: consolidate image trimming and standardize processing

- Replace OpenCV-based trim_whitespace with PIL implementation
- Remove redundant trim_whitespace_pil function
- Add initial trim step to extract command
- Standardize image processing to 2048px wide grayscale using Lanczos
This commit is contained in:
Esa Kataja
2024-12-03 21:47:03 +02:00
parent 682a1f13ee
commit 84e1a6a82e
+68 -63
View File
@@ -14,6 +14,7 @@ from pdf2image import convert_from_path
import settings
from settings import OptimizationLevel
from PIL import Image
class Settings:
TRIM_PADDING_PIXELS = 20
@@ -52,12 +53,79 @@ def extract(input_pdf):
# Save each page
for i, page in enumerate(pages):
# Convert to grayscale
page = page.convert('L')
output_path = os.path.join(temp_dir, f"page_{i:03d}.png")
# Save initial version
page.save(output_path, "PNG", optimize=False)
# Trim whitespace
trim_whitespace(output_path)
# Reload the trimmed image
page = Image.open(output_path)
# Calculate new height maintaining aspect ratio
width = 2048
ratio = width / page.width
height = int(page.height * ratio)
# Resize using Lanczos
page = page.resize((width, height), Image.Resampling.LANCZOS)
# Save final version
page.save(output_path, "PNG", optimize=False)
print(f"Saved page {i+1}/{len(pages)}")
print(f"Extracted {len(pages)} pages to {temp_dir}/")
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.
Args:
image_path: Path to the image file
Returns:
bool: True if successful, False otherwise
"""
try:
# Open image with PIL
image = Image.open(image_path)
# Convert to grayscale if not already
if image.mode != 'L':
image = image.convert('L')
# Get the bounding box of non-white pixels
bbox = image.getbbox()
if not bbox:
print(f"Warning: No content found in {image_path}")
return False
# Add padding
padding = settings.TRIM_PADDING_PIXELS
width, height = image.size
x1, y1, x2, y2 = bbox
x1 = max(0, x1 - padding)
y1 = max(0, y1 - padding)
x2 = min(width, x2 + padding)
y2 = min(height, y2 + padding)
# Crop the image
image = image.crop((x1, y1, x2, y2))
# Save the image
image.save(image_path, "PNG", optimize=False)
return True
except Exception as e:
print(f"Error processing {image_path}: {str(e)}")
return False
@cli.command()
def deskew():
"""Deskew all pages in temporary directory."""
@@ -132,69 +200,6 @@ def deskew():
else:
print(" No line segments detected")
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 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)
# 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 padding
padding = settings.TRIM_PADDING_PIXELS
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
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 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.