fix: improve content detection in trim_whitespace

- Add thresholding to better separate content from background
- Invert image so content becomes white for getbbox detection
- Fix issue where white and transparent areas weren't being trimmed
This commit is contained in:
Esa Kataja
2024-12-03 22:32:19 +02:00
parent a979712c8c
commit 0f105880e6
+20 -11
View File
@@ -97,18 +97,26 @@ def trim_whitespace(image_path: str, is_final: bool = False) -> bool:
# Open image with PIL
image = Image.open(image_path)
# Create flattened version with white background for finding bounds
flattened = Image.new('RGB', image.size, 'white')
# If image has transparency, flatten it first
if image.mode == 'RGBA':
flattened.paste(image, mask=image.split()[3])
else:
flattened.paste(image)
# Create a white background
background = Image.new('RGB', image.size, 'white')
# Paste using alpha channel as mask
background.paste(image, mask=image.split()[3])
image = background
# Convert to grayscale for finding bounds
flattened = flattened.convert('L')
# Convert to grayscale
image = image.convert('L')
# Get the bounding box of non-white pixels
bbox = flattened.getbbox()
# Threshold to make all light pixels white and everything else black
# This helps with finding content bounds
image = image.point(lambda x: 255 if x > 250 else 0)
# Invert so content is white on black background
image = Image.eval(image, lambda x: 255 - x)
# Get the bounding box of content (now white pixels)
bbox = image.getbbox()
if not bbox:
print(f"Warning: No content found in {image_path}")
return False
@@ -122,8 +130,9 @@ def trim_whitespace(image_path: str, is_final: bool = False) -> bool:
x2 = min(width, x2 + padding)
y2 = min(height, y2 + padding)
# Crop the original image (preserving transparency)
cropped = image.crop((x1, y1, x2, y2))
# Open original image again and crop it using the bounds
original = Image.open(image_path)
cropped = original.crop((x1, y1, x2, y2))
# Save the cropped image, optimizing only if this is the final operation
cropped.save(image_path, "PNG", optimize=is_final)