from PIL import Image
import os

# Lista slik za procesiranje
images_to_process = [
    'assets/tree_purple.png',
    'assets/tree_apple.png',
    'assets/tree_pear.png',
    'assets/tree_cherry.png',
    'assets/tree_green_final.png',
    'assets/tree_blue_final.png',
    'assets/tree_dead_final.png'
]

def ultra_remove_background(image_path):
    """ULTRA AGRESIVNO odstranjevanje ozadja"""
    try:
        # Naloži sliko
        img = Image.open(image_path).convert('RGBA')
        data = img.getdata()
        
        new_data = []
        for item in data:
            r, g, b, a = item
            
            # Če je že transparentno, ohrani
            if a == 0:
                new_data.append(item)
                continue
            
            # Izračunaj brightness
            brightness = (r + g + b) / 3
            
            # ULTRA AGRESIVNO: Odstrani VSE kar je svetlejše od 100
            # To bo odstranilo VSE šahovske kvadratke
            if brightness > 100:
                new_data.append((r, g, b, 0))  # Transparentno
                continue
            
            # Odstrani VSE sive barve (ne glede na brightness)
            is_grayscale = abs(r - g) < 25 and abs(g - b) < 25 and abs(r - b) < 25
            if is_grayscale:
                new_data.append((r, g, b, 0))  # Transparentno
                continue
            
            # Ohrani samo temne, nasičene barve (drevesa)
            max_rgb = max(r, g, b)
            min_rgb = min(r, g, b)
            saturation = 0 if max_rgb == 0 else (max_rgb - min_rgb) / max_rgb
            
            # Če je nizka saturacija (pastelno) -> odstrani
            if saturation < 0.2:
                new_data.append((r, g, b, 0))
                continue
            
            # Ohrani pixel
            new_data.append(item)
        
        # Posodobi sliko
        img.putdata(new_data)
        
        # Shrani
        img.save(image_path)
        print(f'✅ ULTRA Processed: {image_path}')
        
    except Exception as e:
        print(f'❌ Error processing {image_path}: {e}')

# Procesiraj vse slike
print('🔥 ULTRA AGGRESSIVE background removal...')
for img_path in images_to_process:
    if os.path.exists(img_path):
        ultra_remove_background(img_path)
    else:
        print(f'⚠️ File not found: {img_path}')

print('✅ Done! All backgrounds ULTRA removed.')