import cv2
import numpy as np
import os

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

def remove_white_bg(file_path):
    print(f"Processing: {file_path}")
    if not os.path.exists(file_path):
        print(f"File not found: {file_path}")
        return

    # Read image
    img = cv2.imread(file_path, cv2.IMREAD_COLOR)
    if img is None:
        print("Failed to load image.")
        return

    # Convert to Gray
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # Create mask: White (240-255) becomes transparent (0), everything else opaque (255)
    # THRESH_BINARY_INV: Values > 240 become 0, others 255.
    _, mask = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)

    # refine mask to remove jagged edges if possible (optional)
    # kernel = np.ones((2,2), np.uint8)
    # mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) 

    # Add Alpha Channel
    b, g, r = cv2.split(img)
    rgba = cv2.merge([b, g, r, mask])

    # Save back to same path
    cv2.imwrite(file_path, rgba)
    print("Saved with transparency.")

if __name__ == "__main__":
    for f in FILES_TO_PROCESS:
        remove_white_bg(f)