import cv2
import numpy as np
import os
import glob

def fix_pink_edges(directory):
    print(f"Scanning {directory} for pink edges...")
    
    # Target specific problematic file if needed, or all
    files = glob.glob(os.path.join(directory, "*.png"))
    
    for path in files:
        img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        if img is None:
            continue
            
        if img.shape[2] < 4:
            continue
            
        b, g, r, a = cv2.split(img)
        
        # 1. Wider Pink/Magenta Range
        # Catch anything with high Red/Blue and relatively lower Green
        pink_mask = (r > 150) & (g < 100) & (b > 150)
        
        # Turn pink pixels transparent
        a[pink_mask] = 0
        
        # 2. Aggressive Alpha Erosion
        # Increase iterations to eat 2 pixels into the edge
        # This removes the "halo" left by anti-aliasing against pink
        kernel = np.ones((3,3), np.uint8) 
        a_eroded = cv2.erode(a, kernel, iterations=1) # 1 iteration with 3x3 kernel = 1-2px border
        
        # Merge back
        img_fixed = cv2.merge((b, g, r, a_eroded))
        
        cv2.imwrite(path, img_fixed)
        print(f"Aggressively Fixed: {os.path.basename(path)}")

if __name__ == "__main__":
    target_dir = "assets/DEMO_FAZA1/Vegetation"
    fix_pink_edges(target_dir)