from rembg import remove
import cv2
from glob import glob
import os
import numpy as np

def remove_bg(image_path, output_path, extension):

    img_glob = glob(os.path.join(image_path, f"*.{extension}"))
    os.makedirs(output_path, exist_ok=True)
    for input_path in img_glob:
        input = cv2.imread(input_path)

        basename = os.path.basename(input_path)

        # img_empty = np.zeros((2048, 2592, 4))

        # input = input[700:1650, 0:2520]
        output = remove(input)
        # convert to BGR or GRAY to remove transparent
        output = cv2.cvtColor(output, cv2.COLOR_RGBA2BGR)
        # img_empty[700:1650, 0:2520] = output

        cv2.imwrite(os.path.join(output_path,basename), output)

if __name__ == "__main__":
    image_path = "../../../../../data/offlineData2/"
    output_path = "./output"
    extension = "bmp"
    remove_bg(image_path, output_path, extension)