90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import os
|
|
import sys
|
|
import subprocess
|
|
|
|
# Define the source directory on the Desktop
|
|
SOURCE_DIR = "/Users/davidkotnik/Desktop/referencne slike 2"
|
|
|
|
def install_dependencies():
|
|
"""Attempt to install required packages if they are missing."""
|
|
packages = ["rembg", "Pillow", "onnxruntime"]
|
|
for package in packages:
|
|
try:
|
|
__import__(package)
|
|
except ImportError:
|
|
if package == "Pillow":
|
|
# Pillow is imported as PIL
|
|
try:
|
|
__import__("PIL")
|
|
continue
|
|
except ImportError:
|
|
pass
|
|
|
|
print(f"Installing {package}...")
|
|
try:
|
|
subprocess.check_call([sys.executable, "-m", "pip", "install", package])
|
|
except subprocess.CalledProcessError as e:
|
|
print(f"Failed to install {package}: {e}")
|
|
print("Please try installing it manually: pip install rembg Pillow onnxruntime")
|
|
sys.exit(1)
|
|
|
|
def process_images():
|
|
# Verify directory
|
|
if not os.path.exists(SOURCE_DIR):
|
|
print(f"Error: Directory {SOURCE_DIR} not found.")
|
|
return
|
|
|
|
# Import libraries AFTER installation
|
|
try:
|
|
from rembg import remove
|
|
from PIL import Image
|
|
except ImportError:
|
|
print("Error: Library import failed even after installation attempt.")
|
|
return
|
|
|
|
# Get PNG files
|
|
files = [f for f in os.listdir(SOURCE_DIR) if f.lower().endswith('.png')]
|
|
total = len(files)
|
|
|
|
if total == 0:
|
|
print("No PNG files found to process.")
|
|
return
|
|
|
|
print(f"Found {total} PNG images. Starting AI background removal...")
|
|
print("WARNING: This makes permanent changes to the PNG files (JPG backups are safe).")
|
|
print("This might take a while depending on your computer speed.")
|
|
|
|
success = 0
|
|
errors = 0
|
|
|
|
for i, filename in enumerate(files):
|
|
file_path = os.path.join(SOURCE_DIR, filename)
|
|
|
|
try:
|
|
print(f"[{i+1}/{total}] Processing: {filename}")
|
|
|
|
# Open the image
|
|
with open(file_path, "rb") as input_file:
|
|
input_data = input_file.read()
|
|
|
|
# Remove background
|
|
output_data = remove(input_data)
|
|
|
|
# Save the result back to the same file (overwrite)
|
|
with open(file_path, "wb") as output_file:
|
|
output_file.write(output_data)
|
|
|
|
success += 1
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error processing {filename}: {e}")
|
|
errors += 1
|
|
|
|
print(f"\n--- DONE ---")
|
|
print(f"Successful: {success}")
|
|
print(f"Errors: {errors}")
|
|
|
|
if __name__ == "__main__":
|
|
install_dependencies()
|
|
process_images()
|