从RGBA图像中得到分割标签

在做动漫分割项目时，数据集非常稀缺；那就自己制作吧，这个网站removebg可以去除部分图片的背景，那怎么办得到标签呢？
可以这样理解： RGB图片—>removebg—>RGBA图片(作为我下面代码的输入)—>输入到下面代码运算—>RGB图片和P模式的分割标签

需求

输入RGBA图像，输出RGB图片和P模式的分割标签
四通道图片

代码如下

import numpy as np
from PIL import Image

def array2p_mode(alpha_channel):
    """alpha_channel是灰度图像，像素值在0~255"""
    # 1.预处理
    img_arr = np.array(alpha_channel)
    img_arr[img_arr >= 128] = 128
    img_arr[img_arr < 128] = 0
    # 2.生成rgb图，通道合并
    heigth, width = img_arr.shape
    img_arr_rgb = np.zeros((heigth, width, 3))
    img_arr_rgb[:, :, 0] = img_arr
    # 3.转换成"p"模式图片
    img = Image.fromarray(np.uint8(img_arr_rgb))
    img_p = img.convert("P")
    return img_p

png_dir = "./网图标签/"
img = Image.open(png_dir + "0116.png")
print("模式:", img.mode, "尺寸:", img.size, "通道数:", len(img.split())))
r,g,b,alpha_channel = img.split()
img_p = array2p_mode(alpha_channel)
img_p.save("img_p.png")

实现方法

1.读入RGBA图像，并取出alpha通道图像
2.由于alpha通道图像是灰度图像，进行处理，得到像素值只有128,0的图像
3.alpha通道图像作为纯黑色的rgb图像的第一通道，得到img_arr_rgb
4.将img_arr_rgb转换为p模式的图像，记得分割标签P图

效果展示

RGB原图

P模式的分割标签