python 基于openCV去水印封装

• 解决去水印封装问题
• 去半透明水印（颜色中和）
• 去不规则水印（sift）
• pdf图片去水印

# 由于版权原因使用sift算法，需要安装opencv3.4以下版本
pip install opencv-python

import cv2
import numpy as np
import os
import glob
from matplotlib import pyplot as plt
class WaterMark:
    def __init__(self):
        pass

    def remove(self, src_url, mask_url=None):
        result_url = src_url.replace('.jpg', '_out.jpg')
        src_url_arr = str(src_url).split('/')
        mask_url_arr = []
        if mask_url:
            mask_url_arr.append(mask_url)
        else:
            src_platform = src_url_arr[-1].split('_')[0]
            mask_dir = f"{'/'.join(src_url_arr[:-1])}/mask/"
            for file in glob.glob(f'{mask_dir}/{src_platform}*_mask*'):
                mask_url_arr.append(file)
        print('src_url', src_url)
        src_img = cv2.imread(src_url)
        for _, v in enumerate(mask_url_arr):
            print('mask_url', v)
            mask_img = cv2.imread(v)
            rect = self.find_water_sift(src_img=src_img, mask_img=mask_img)
            if len(rect[0]) > 0:
                result = self.inpaint(src_img=src_img, rect=rect, enlarge=4)
            else:
                rect = self.find_water_template(src_img=src_img, mask_img=mask_img)
                if len(rect[0]) > 0:
                    result = self.inpaint(src_img=src_img, rect=rect, enlarge=10)
            # result = self.remove_base(src_img=src_img, mask_img=mask_img)
            if len(result) > 0:
                cv2.imwrite(result_url, result)
                cv2.imshow('result.jpg', result)
            else:
                print('没有找到水印')

        cv2.waitKey(0)

    #  反色中和去水印
    def remove_base(self, src_img, mask_img):
        result = np.zeros(src_img.shape, np.uint8)
        result[mask_img != 0] = np.clip(255 - (255 - src_img[mask_img != 0]) / mask_img[mask_img != 0] * 256, 0, 255)
        return result

    #  腐蚀方式修复水印区域
    def inpaint(self, src_img, rect, enlarge=10):
        rect_min, rect_max = rect[0], rect[1]
        thresh_img = np.zeros(src_img.shape, np.uint8)
        rect_min -= enlarge
        rect_max += enlarge
        thresh_img[rect_min[1]:rect_max[1], rect_min[0]:rect_max[0]] = \
            src_img[rect_min[1]:rect_max[1], rect_min[0]:rect_max[0]]
        # 创建形状和尺寸的结构元素
        kernel = np.ones((3, 3), np.uint8)
        thresh = cv2.inRange(thresh_img, np.array([1, 1, 1]), np.array([255, 255, 255]))
        # 扩张待修复区域
        hi_mask = cv2.dilate(thresh, kernel, iterations=1)
        result = cv2.inpaint(src_img, hi_mask, 5, flags=cv2.INPAINT_TELEA)
        return result

    def remove_threshold(self, src_img):
        hight, width, depth = src_img.shape[0:3]
        # 图片二值化处理，把[240, 240, 240]~[255, 255, 255]以外的颜色变成0
        thresh = cv2.inRange(src_img, np.array([180, 180, 180]), np.array([255, 255, 255]))
        # 创建形状和尺寸的结构元素
        kernel = np.ones((3, 3), np.uint8)
        # 扩张待修复区域
        hi_mask = cv2.dilate(thresh, kernel, iterations=1)
        specular = cv2.inpaint(src_img, hi_mask, 5, flags=cv2.INPAINT_TELEA)
        cv2.namedWindow("Image", 0)
        cv2.resizeWindow("Image", int(width / 2), int(hight / 2))
        cv2.imshow("Image", src_img)

        cv2.namedWindow("newImage", 0)
        cv2.resizeWindow("newImage", int(width / 2), int(hight / 2))
        cv2.imshow("newImage", specular)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    #  去PDF水印
    def remove_pdf(self, src_img):
        new = np.clip(2.0 * src_img - 160, 0, 255).astype(np.uint8)
        cv2.imshow('cleaned.jpg', new)
        cv2.waitKey(0)

    #  检测图片是否存在单水印
    def find_water_template(self, src_img, mask_img):
        # 获得模板图片的高宽尺寸
        theight, twidth = mask_img.shape[:2]
        # 执行模板匹配，采用的匹配方式cv2.TM_SQDIFF_NORMED
        result = cv2.matchTemplate(src_img, mask_img, cv2.TM_SQDIFF_NORMED)
        # 归一化处理
        cv2.normalize(result, result, 0, 1, cv2.NORM_MINMAX, -1)
        # 寻找矩阵（一维数组当做向量，用Mat定义）中的最大值和最小值的匹配结果及其位置
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        # 匹配值转换为字符串
        # 对于cv2.TM_SQDIFF及cv2.TM_SQDIFF_NORMED方法min_val越趋近与0匹配度越好，匹配位置取min_loc
        # 对于其他方法max_val越趋近于1匹配度越好，匹配位置取max_loc
        strmin_val = str(min_val)
        # 绘制矩形边框，将匹配区域标注出来
        # min_loc：矩形定点
        # (min_loc[0]+twidth,min_loc[1]+theight)：矩形的宽高
        # (0,0,225)：矩形的边框颜色；2：矩形边框宽度
        # cv2.rectangle(src_img, min_loc, (min_loc[0] + twidth, min_loc[1] + theight), (0, 0, 225), 2)
        return [np.array(min_loc), np.array([min_loc[0] + twidth, min_loc[1] + theight])]
        # 显示结果,并将匹配值显示在标题栏上
        # cv2.imshow("MatchResult----MatchingValue=" + strmin_val, src_img)
        # cv2.waitKey()

    #  检测图片是否存在多水印
    def find_waters_template(self, src_img, mask_img):
        # 获得模板图片的高宽尺寸
        theight, twidth = mask_img.shape[:2]
        # 执行模板匹配，采用的匹配方式cv2.TM_SQDIFF_NORMED
        result = cv2.matchTemplate(src_img, mask_img, cv2.TM_SQDIFF_NORMED)
        # 归一化处理
        # cv2.normalize( result, result, 0, 1, cv2.NORM_MINMAX, -1 )
        # 寻找矩阵（一维数组当做向量，用Mat定义）中的最大值和最小值的匹配结果及其位置
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
        # 绘制矩形边框，将匹配区域标注出来
        # min_loc：矩形定点
        # (min_loc[0]+twidth,min_loc[1]+theight)：矩形的宽高
        # (0,0,225)：矩形的边框颜色；2：矩形边框宽度
        cv2.rectangle(src_img, min_loc, (min_loc[0] + twidth, min_loc[1] + theight), (0, 0, 225), 2)
        # 匹配值转换为字符串
        # 对于cv2.TM_SQDIFF及cv2.TM_SQDIFF_NORMED方法min_val越趋近与0匹配度越好，匹配位置取min_loc
        # 对于其他方法max_val越趋近于1匹配度越好，匹配位置取max_loc
        strmin_val = str(min_val)
        # 初始化位置参数
        temp_loc = min_loc
        other_loc = min_loc
        numOfloc = 1
        # 第一次筛选----规定匹配阈值，将满足阈值的从result中提取出来
        # 对于cv2.TM_SQDIFF及cv2.TM_SQDIFF_NORMED方法设置匹配阈值为0.01
        threshold = 0.01
        loc = np.where(result < threshold)
        # 遍历提取出来的位置
        for other_loc in zip(*loc[::-1]):
            # 第二次筛选----将位置偏移小于5个像素的结果舍去
            if (temp_loc[0] + 5 < other_loc[0]) or (temp_loc[1] + 5 < other_loc[1]):
                numOfloc = numOfloc + 1
                temp_loc = other_loc
                cv2.rectangle(src_img, other_loc, (other_loc[0] + twidth, other_loc[1] + theight), (0, 0, 225), 2)
        str_numOfloc = str(numOfloc)
        # 显示结果,并将匹配值显示在标题栏上
        strText = "MatchResult----MatchingValue=" + strmin_val + "----NumberOfPosition=" + str_numOfloc
        cv2.imshow(strText, src_img)
        cv2.waitKey()

    #  sift特征算法检测图片是否存在不规则水印
    def find_water_sift(self, src_img, mask_img):
        MIN_MATCH_COUNT = 6  # 设置最低特征点匹配数量为10
        template = cv2.cvtColor(mask_img, cv2.COLOR_RGB2GRAY)  # queryImage
        target = cv2.cvtColor(src_img, cv2.COLOR_RGB2GRAY)  # trainImage
        # Initiate SIFT detector创建sift检测器
        sift = cv2.xfeatures2d.SIFT_create()
        # find the keypoints and descriptors with SIFT
        kp1, des1 = sift.detectAndCompute(template, None)
        kp2, des2 = sift.detectAndCompute(target, None)
        # 创建设置FLANN匹配
        FLANN_INDEX_KDTREE = 0
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        flann = cv2.FlannBasedMatcher(index_params, search_params)
        matches = flann.knnMatch(des1, des2, k=2)
        # store all the good matches as per Lowe's ratio test.
        good = []
        # 舍弃大于0.7的匹配
        for m, n in matches:
            if m.distance < 0.8 * n.distance:
                good.append(m)
        if len(good) >= MIN_MATCH_COUNT:
            # 获取关键点的坐标
            src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
            dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
            # 计算变换矩阵和MASK
            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
            matchesMask = mask.ravel().tolist()
            h, w = template.shape
            # 使用得到的变换矩阵对原图像的四个角进行变换，获得在目标图像上对应的坐标
            pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
            dst = cv2.perspectiveTransform(pts, M)
            rect_min = dst.astype(np.int)[0, 0]
            rect_max = dst.astype(np.int)[2, 0]
            return [rect_min, rect_max]
            cv2.polylines(target, [np.int32(dst)], True, 0, 2, cv2.LINE_AA)
        else:
            print("Not enough matches are found - %d/%d" % (len(good), MIN_MATCH_COUNT))
            matchesMask = None
            return [], []
        draw_params = dict(matchColor=(0, 255, 0),
                           singlePointColor=None,
                           matchesMask=matchesMask,
                           flags=2)
        result = cv2.drawMatches(template, kp1, target, kp2, good, None, **draw_params)
        plt.imshow(result, 'gray')
        plt.show()

if __name__ == '__main__':
    src_url = f'../../data/water/QFNAG_.jpg'
    mask_url = f'../../data/water/mask/LYJ_mask_L.jpg'
    # src_url = f'../../data/water/QFANG_1.jpg'
    # mask_url = f'../../data/water/mask/QFANG_mask.jpg'
    water = WaterMark()
    water.remove(src_url=src_url, mask_url=mask_url)

• 去水印效果
  
  
• 半透明中和去水印
  
  
• pdf去水印