import os
from transformers import pipeline

def remove_background_with_bria(input_folder, output_folder):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)

    for filename in os.listdir(input_folder):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
            input_path = os.path.join(input_folder, filename)
            output_path = os.path.join(output_folder, os.path.splitext(filename)[0] + '.png')
            
            # Use pipeline to remove the background
            pillow_image = pipe(input_path) # applies mask on input and returns a pillow image
            pillow_image.save(output_path, format="PNG")
            
            # Print the filename of the processed image
            print(f"Processed {filename}")

# Example usage
remove_background_with_bria("D:\Ardha\Kerja\AISensum\ROX\Task 1.2\Raw_PSD", "bria_output")