import os
import cv2

stamp_root = '/data1/lxl/data/ocr/stamp/src'
mask_root = '/data1/lxl/data/ocr/stamp/mask'


def read_img_list(root):
    img_full_path_list = list()
    img_list = os.listdir(root)
    for img_name in img_list:
        img_full_path_list.append(os.path.join(root, img_name))

    return img_full_path_list


def get_mask_list(stamp_list):
    for stamp in stamp_list:
        img_name = stamp.split('/')[-1]
        img = cv2.imread(stamp, -1)
        mask = img[:, :, -1]
        cv2.imwrite(os.path.join(mask_root, img_name), mask)


if __name__ == '__main__':
    full_path_list = read_img_list(stamp_root)
    get_mask_list(full_path_list)