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 # Open image with PIL
image = Image.open(image_path) image = Image.open(image_path)
# Create flattened version with white background for finding bounds # If image has transparency, flatten it first
flattened = Image.new('RGB', image.size, 'white')
if image.mode == 'RGBA': if image.mode == 'RGBA':
flattened.paste(image, mask=image.split()[3]) # Create a white background
else: background = Image.new('RGB', image.size, 'white')
flattened.paste(image) # Paste using alpha channel as mask
background.paste(image, mask=image.split()[3])
image = background
# Convert to grayscale for finding bounds # Convert to grayscale
flattened = flattened.convert('L') image = image.convert('L')
# Get the bounding box of non-white pixels # Threshold to make all light pixels white and everything else black
bbox = flattened.getbbox() # 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: if not bbox:
print(f"Warning: No content found in {image_path}") print(f"Warning: No content found in {image_path}")
return False return False
@@ -122,8 +130,9 @@ def trim_whitespace(image_path: str, is_final: bool = False) -> bool:
x2 = min(width, x2 + padding) x2 = min(width, x2 + padding)
y2 = min(height, y2 + padding) y2 = min(height, y2 + padding)
# Crop the original image (preserving transparency) # Open original image again and crop it using the bounds
cropped = image.crop((x1, y1, x2, y2)) original = Image.open(image_path)
cropped = original.crop((x1, y1, x2, y2))
# Save the cropped image, optimizing only if this is the final operation # Save the cropped image, optimizing only if this is the final operation
cropped.save(image_path, "PNG", optimize=is_final) cropped.save(image_path, "PNG", optimize=is_final)