将PNG标签中的非0像素值改成0并保存

from PIL import Image
import numpy as np
import os
from glob import glob
from tqdm import tqdm


def processPng(image_path):
    """ 后处理PNG标签，将其转为只包含[0,1]值的图片(在只有背景和1个类别下) """
    image = Image.open(image_path)
    img_array = np.array(image) # 转成numpy数组
    """ 将不是1的地方置位1 """
    for i in range(img_array.shape[0]):
        for j in range(img_array.shape[1]):
            if img_array[i,j]!=0:
                img_array[i,j] = 1

    print('1的像素个数为:',np.sum(img_array == 1)) # 计算图像中1的个数

    img_tr = Image.fromarray(img_array) # 转成PIL
    img_tr = img_tr.convert('P')
    base_name = os.path.basename(image_path)

    save_dir = 'OneZero/'
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    
    img_tr.save(save_dir+base_name)
    print(f'成功转换与保存{image_path}')


if __name__ == '__main__':
    work_dir = 'saveDir/'
    file_paths = glob(work_dir+'*.png')
    print(f'有{len(file_paths)}张图片')

    for filePath in tqdm(file_paths):
        processPng(filePath)