87 lines
2.3 KiB
Python
87 lines
2.3 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Background Removal for DolinaSmrti Assets
|
|
Uses rembg to remove white backgrounds and create transparent PNGs
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
from PIL import Image
|
|
|
|
try:
|
|
from rembg import remove
|
|
except ImportError:
|
|
print("❌ rembg not installed. Installing now...")
|
|
os.system("pip3 install rembg")
|
|
from rembg import remove
|
|
|
|
|
|
def remove_background(input_path: Path, output_path: Path = None):
|
|
"""Remove background from image and save as transparent PNG"""
|
|
|
|
if output_path is None:
|
|
output_path = input_path
|
|
|
|
print(f" 🎨 Processing: {input_path.name}")
|
|
|
|
# Load image
|
|
with open(input_path, 'rb') as input_file:
|
|
input_data = input_file.read()
|
|
|
|
# Remove background
|
|
output_data = remove(input_data)
|
|
|
|
# Save as PNG
|
|
with open(output_path, 'wb') as output_file:
|
|
output_file.write(output_data)
|
|
|
|
print(f" ✅ Saved: {output_path.name}")
|
|
|
|
|
|
def process_directory(directory: Path, recursive: bool = False):
|
|
"""Process all PNG files in directory"""
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"📁 Processing directory: {directory}")
|
|
print(f"{'='*60}")
|
|
|
|
pattern = "**/*.png" if recursive else "*.png"
|
|
png_files = list(directory.glob(pattern))
|
|
|
|
total = len(png_files)
|
|
print(f"\nFound {total} PNG files")
|
|
|
|
for i, png_file in enumerate(png_files, 1):
|
|
print(f"\n[{i}/{total}]")
|
|
remove_background(png_file)
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"✅ Completed! Processed {total} images")
|
|
print(f"{'='*60}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) < 2:
|
|
print("Usage: python3 remove_background.py <image_or_directory>")
|
|
print("\nExamples:")
|
|
print(" python3 remove_background.py image.png")
|
|
print(" python3 remove_background.py assets/images/buildings/")
|
|
print(" python3 remove_background.py assets/images/ --recursive")
|
|
sys.exit(1)
|
|
|
|
path = Path(sys.argv[1])
|
|
recursive = "--recursive" in sys.argv or "-r" in sys.argv
|
|
|
|
if not path.exists():
|
|
print(f"❌ Path not found: {path}")
|
|
sys.exit(1)
|
|
|
|
if path.is_file():
|
|
remove_background(path)
|
|
elif path.is_dir():
|
|
process_directory(path, recursive=recursive)
|
|
else:
|
|
print(f"❌ Invalid path: {path}")
|
|
sys.exit(1)
|