#!/usr/bin/env python3
"""
Total Background Removal - Čiščenje vseh ozadij
Odstrani: zeleno, belo, checkerboard, sivo
"""

from PIL import Image
import numpy as np
import os

def clean_background(input_path, output_path):
    """
    Totalno čiščenje ozadja - zeleno, belo, checkerboard
    """
    print(f"Cleaning: {input_path}")
    
    # Odpri sliko
    img = Image.open(input_path)
    img = img.convert("RGBA")
    
    # Pretvori v numpy array
    data = np.array(img)
    
    # Ekstrahiraj kanale
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Maska 1: Zeleno ozadje (chroma key)
    green_mask = (g > 200) & (g > r + 30) & (g > b + 30)
    
    # Maska 2: Belo ozadje
    white_mask = (r > 240) & (g > 240) & (b > 240)
    
    # Maska 3: Checkerboard pattern (sivo-belo) - detektira oba tona
    # Svetlo sivo v checkerboard
    light_gray_mask = (r > 180) & (r < 210) & (g > 180) & (g < 210) & (b > 180) & (b < 210) & (np.abs(r.astype(int) - g.astype(int)) < 5) & (np.abs(g.astype(int) - b.astype(int)) < 5)
    
    # Temno sivo v checkerboard
    dark_gray_mask = (r > 140) & (r < 170) & (g > 140) & (g < 170) & (b > 140) & (b < 170) & (np.abs(r.astype(int) - g.astype(int)) < 5) & (np.abs(g.astype(int) - b.astype(int)) < 5)
    
    # Kombinirana maska
    remove_mask = green_mask | white_mask | light_gray_mask | dark_gray_mask
    
    # Nastavi alfa kanal na 0 za remove piksle
    data[remove_mask, 3] = 0
    
    # Pretvori nazaj v sliko
    result = Image.fromarray(data, 'RGBA')
    
    # Shrani
    result.save(output_path, "PNG")
    print(f"✓ Cleaned: {output_path}")

# Najdi vse PNG slike v NOVE_SLIKE
for root, dirs, files in os.walk("NOVE_SLIKE"):
    for file in files:
        if file.endswith(".png"):
            filepath = os.path.join(root, file)
            clean_background(filepath, filepath)

print("\n✅ VSE SLIKE V NOVE_SLIKE SČIŠČENE!")