最近在研究海陆分割,下载了两套样本数据集,但是两套数据标注信息的色调不一致,因此需要处理为统一的标注形式。
原始的标注为rgb彩色的mask, 如下左图所示,我希望将其统一为0和1的二值图,如下右图所示。
下面的代码可以实现将rgb彩色的三波段mask转为单波段的0-1mask:
import os
from PIL import Image
from osgeo import gdal
def rgb2mask(rgb_dir, mask_dir):
"""
:param rgb_dir:原始的rgb mask所在的文件夹路径
:param mask_dir:存储单波段mask文件的路径
"""
img_list = os.listdir(rgb_dir)
for img_name in img_list:
img = Image.open(rgb_dir + img_name)
thresh = 223 ## 通过阈值判断为0还是为1
fn = lambda x: 255 if x == thresh else 0 # 这里x与阈值的关系可以是多种,如><=等
r = img.convert('L').point(fn, mode='1')
# convert的参数L表示用8个bit表示,0为黑,255为白,逐像素处理为0和1
r.save(mask_dir + img_name)
if __name__ == '__main__':
img_path = 'labels-color/'
save_path = 'labels/'
img_list = os.listdir(img_path)
rgb2mask(img_path, save_path)