import os
import cv2
import numpy as np


def rgb_to_binary_mask(img_path, output_path):
    '''
    Convert an RGB image to a binary mask where non-black pixels are set to 1, black pixels are 0.
    '''

    img_bgr = cv2.imread(img_path)

    img_binary_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8)

    non_black_pixels = np.any(img_bgr != [0, 0, 0], axis=-1)
    img_binary_mask[non_black_pixels] = 1

    cv2.imwrite(output_path, img_binary_mask)


def convert_rgb_to_binary(Dataset_Path):
    '''
    Convert all RGB images in the dataset to binary mask images.
    '''

    img_dir = os.path.join(Dataset_Path, 'annotation')

    ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')
        if not os.path.exists(ann_target_dir):
            os.mkdir(ann_target_dir)


    for img_name in os.listdir(img_dir):
        try:
            img_path = os.path.join(img_dir, img_name)
            mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')

            rgb_to_binary_mask(img_path, mask_path)

        except Exception as e:
            print(f"Failed to convert {img_name}: {e}")

    print("Conversion completed.")


if __name__ == '__main__':
    Dataset_Path = 'CVRPDataset'
    convert_rgb_to_binary(Dataset_Path)