import cv2
import numpy as np

FILE_PATH = "/Users/davidkotnik/repos/novafarma/nova farma TRAE/assets/DEMO_FAZA1/Environment/water_clean_patch.png"

def remove_magenta_bg(file_path):
    print(f"Processing {file_path}")
    img = cv2.imread(file_path, cv2.IMREAD_UNCHANGED)
    if img is None:
        print("Failed to load image")
        return

    # If already RGBA, keep alpha. If RGB, add alpha.
    if img.shape[2] == 3:
        b, g, r = cv2.split(img)
        a = np.ones_like(b) * 255
        img = cv2.merge((b, g, r, a))
    
    # Convert to HSV
    hsv = cv2.cvtColor(img[:,:,0:3], cv2.COLOR_BGR2HSV)
    
    # Magenta #FF00FF is (300 deg, 100%, 100%). In OpenCV (0-179, 0-255, 0-255):
    # Hue ~ 150.
    # Define range for Magenta
    lower_magenta = np.array([145, 150, 150])
    upper_magenta = np.array([165, 255, 255])
    
    mask = cv2.inRange(hsv, lower_magenta, upper_magenta)
    
    # Invert mask (keep non-magenta)
    mask_inv = cv2.bitwise_not(mask)
    
    # Get current alpha
    b, g, r, a = cv2.split(img)
    
    # Combine existing alpha with new mask
    new_alpha = cv2.bitwise_and(a, mask_inv)
    
    # Soften edges (1px blur on alpha)
    # This helps remove jagged pink pixels
    new_alpha = cv2.GaussianBlur(new_alpha, (3, 3), 0)
    
    # Update image
    final_img = cv2.merge((b, g, r, new_alpha))
    
    cv2.imwrite(file_path, final_img)
    print("Background removed and saved.")

if __name__ == "__main__":
    remove_magenta_bg(FILE_PATH)