File size: 1,265 Bytes
5e15524 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 |
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)
|