98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
|
|
import os
|
|
import sys
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
def remove_background(image_path, tolerance=30):
|
|
try:
|
|
img = Image.open(image_path).convert("RGBA")
|
|
datas = img.getdata()
|
|
|
|
# Sample the corner pixel to guess the "background" theme
|
|
# But since it's a checkerboard, precise color matching fails.
|
|
# We'll use a seed-fill (flood fill) approach from the corners.
|
|
|
|
# Convert to numpy array for faster processing
|
|
arr = np.array(img)
|
|
|
|
# Create a mask for visited pixels (background)
|
|
h, w = arr.shape[:2]
|
|
mask = np.zeros((h, w), dtype=bool)
|
|
|
|
# Stack for flood fill: (r, c)
|
|
stack = [(0, 0), (0, w-1), (h-1, 0), (h-1, w-1)]
|
|
|
|
# We need to be careful not to delete the object.
|
|
# Assumption: The checkerboard is grey-scale or near grey-scale.
|
|
# And the object (tree) is colorful (green/brown).
|
|
|
|
# Helper to check if a pixel is "grey-ish" (background candidate)
|
|
def is_background_candidate(pixel):
|
|
r, g, b, a = pixel
|
|
# Check for low saturation (grey/white/black)
|
|
# max(r,g,b) - min(r,g,b) should be small for greys
|
|
saturation = max(r, g, b) - min(r, g, b)
|
|
return saturation < tolerance
|
|
|
|
# New approach: smart flood fill with color tolerance is hard on noisy checkerboard.
|
|
# Let's try a simpler heuristic first:
|
|
# Iterate all pixels. If a pixel is GREY-SCALE (within tolerance), make it transparent.
|
|
# This risks deleting grey parts of the object (e.g. stones, bark).
|
|
|
|
newData = []
|
|
for item in datas:
|
|
# item is (r, g, b, a)
|
|
r, g, b, a = item
|
|
|
|
# Check if it's a grey/white/black pixel
|
|
# Checkerboard usually consists of light grey/white and dark grey blocks.
|
|
# We enforce that R, G, and B are close to each other.
|
|
if abs(r - g) < tolerance and abs(g - b) < tolerance and abs(r - b) < tolerance:
|
|
# Also check brightness to avoid deleting dark black outlines if they are pure black
|
|
# Let's say we only delete "light" greys/whites?
|
|
# The "dark" squares in the checkerboard from the log were around (77, 77, 77).
|
|
# This is quite dark.
|
|
|
|
# If we just delete all greys, we might lose outlines (usually (0,0,0)).
|
|
# So we check if it's NOT pure black.
|
|
if r > 20 and g > 20 and b > 20:
|
|
newData.append((0, 0, 0, 0)) # Transparent
|
|
else:
|
|
newData.append(item)
|
|
else:
|
|
newData.append(item)
|
|
|
|
img.putdata(newData)
|
|
|
|
# Save back
|
|
img.save(image_path, "PNG")
|
|
print(f"Processed: {image_path}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"Error processing {image_path}: {e}")
|
|
return False
|
|
|
|
# List of specific files to fix first (Trees)
|
|
targets = [
|
|
"assets/references/trees/apple/apple_tree.png"
|
|
]
|
|
|
|
# Or scan directory
|
|
target_dir = "assets/references"
|
|
|
|
if __name__ == "__main__":
|
|
# If arguments provided, use them
|
|
if len(sys.argv) > 1:
|
|
files = sys.argv[1:]
|
|
for f in files:
|
|
remove_background(f)
|
|
else:
|
|
# Recursive scan for testing
|
|
for root, dirs, files in os.walk(target_dir):
|
|
for file in files:
|
|
if file.endswith("apple_tree.png"): # Safety: only target the known bad one first
|
|
full_path = os.path.join(root, file)
|
|
remove_background(full_path)
|