[代码记录]得到篡改图的mask

记录下。

import cv2
import numpy as np
import os
import re

# 腐蚀膨胀操作
def erode_dilate(img, num1, num2):
    erode = cv2.erode(img, None, iterations=num1)
    dilate = cv2.dilate(erode, None, iterations=num2)
    return dilate

# mask获取，主要运用cv2的absdiff函数
def get_mask(fake_img, real_img):
    # 转化为灰度图像
    img1 = cv2.cvtColor(fake_img, cv2.COLOR_RGB2GRAY)
    img2 = cv2.cvtColor(real_img, cv2.COLOR_RGB2GRAY)
    mask = cv2.absdiff(img1, img2)
    mask = ~mask # 取反 mask区域为黑
    # 阈值
    _, mask = cv2.threshold(mask, 248, 255, cv2.THRESH_BINARY)
    mask = erode_dilate(mask, 2, 2)
    return mask


if __name__ == "__main__":
    fake_path = "fake/"
    real_path = "real/"
    mask_path = "mask/"
    fake_list = os.listdir(fake_path)
    real_list = os.listdir(real_path)
    # 这里我是根据文件名特点来查找图片的
    for fake_name in fake_list:
        fake_img_path = fake_path + fake_name
        fake_img = cv2.imread(fake_img_path)
        match_fake_name = re.findall("\d+", fake_name)[0]

        for real_name in real_list:
            match_real_name = re.findall("\d+", real_name)[0]
            if match_fake_name == match_real_name:
                real_img_path = real_path + real_name
                real_img = cv2.imread(real_img_path)
                mask = get_mask(fake_img, real_img)
                mask_path = mask_path + fake_name.split(".")[0] + "_mask.jpg"
                cv2.imwrite(mask_path, mask)
            else:
                continue