53 lines
1.5 KiB
Python
53 lines
1.5 KiB
Python
"""Remove white background from an image by flood-filling from edges only."""
|
|
|
|
import sys
|
|
from collections import deque
|
|
from PIL import Image
|
|
|
|
def remove_white_bg(path: str, tolerance: int = 20, out: str | None = None):
|
|
img = Image.open(path).convert("RGBA")
|
|
pixels = img.load()
|
|
w, h = img.size
|
|
threshold = 255 - tolerance
|
|
|
|
visited = set()
|
|
queue = deque()
|
|
|
|
# Seed from all edge pixels
|
|
for x in range(w):
|
|
queue.append((x, 0))
|
|
queue.append((x, h - 1))
|
|
for y in range(h):
|
|
queue.append((0, y))
|
|
queue.append((w - 1, y))
|
|
|
|
while queue:
|
|
x, y = queue.popleft()
|
|
if (x, y) in visited or x < 0 or y < 0 or x >= w or y >= h:
|
|
continue
|
|
visited.add((x, y))
|
|
r, g, b, a = pixels[x, y]
|
|
if r >= threshold and g >= threshold and b >= threshold:
|
|
pixels[x, y] = (r, g, b, 0)
|
|
queue.append((x + 1, y))
|
|
queue.append((x - 1, y))
|
|
queue.append((x, y + 1))
|
|
queue.append((x, y - 1))
|
|
|
|
# Crop to bounding box of non-transparent pixels
|
|
bbox = img.getbbox()
|
|
if bbox:
|
|
img = img.crop(bbox)
|
|
|
|
dest = out or path
|
|
img.save(dest)
|
|
print(f"Saved to {dest} ({img.size[0]}x{img.size[1]})")
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python remove_bg.py <image> [tolerance] [output]")
|
|
sys.exit(1)
|
|
path = sys.argv[1]
|
|
tol = int(sys.argv[2]) if len(sys.argv) > 2 else 20
|
|
out = sys.argv[3] if len(sys.argv) > 3 else None
|
|
remove_white_bg(path, tol, out)
|