# pip install opencv-contrib-python
# pip install numpy
import argparse
from pathlib import Path

from cv2 import cv2
import numpy as np


def mathc_img(src_path, target_path, logo_path, output_path=None):
    """
    图片去水印，加水印处理
    :param src_path: 原图
    :param target_path: 水印模板图
    :param logo_path: 新水印图
    :param output_path: 输出路径。为空时直接显示
    :return:
    """

    if not Path(src_path).is_file():
        ValueError("原文件不存在")
        return
    if not Path(logo_path).is_file():
        ValueError("Logo文件不存在")
        return
    if not Path(target_path).is_file():
        ValueError("模板文件不存在")
        return

    # 原图带颜色
    img_rgb = cv2.imread(src_path)

    # 复制一份原图 用来显示边框
    img_rgb_rect = cv2.imread(src_path)

    # 创建一个原始图像的灰度版本，所有操作在灰度版本中处理，然后在RGB图像中使用相同坐标还原
    img_gray = cv2.cvtColor(img_rgb, cv2.COLOR_BGR2GRAY)

    # 原图的大小
    img_w, img_h = img_gray.shape[::-1]

    if img_w <= 150 or img_h <= 150:
        print("图片尺寸过小")
        return

    # 加载将要搜索的图像模板
    template = cv2.imread(target_path, 0)

    # 模板图的坐标
    # w, h = template.shape[::-1]
    w, h = 120, 50

    # 读取logo并缩放
    img_logo = cv2.imread(logo_path)
    img_logo = cv2.resize(img_logo, (w, h), interpolation=cv2.INTER_NEAREST)

    # cv2.TM_CCOEFF_NORMED方法查询模板
    res = cv2.matchTemplate(img_gray, template, cv2.TM_CCOEFF_NORMED)

    # 制作一个和原图一样大小的黑底模板图
    mask = np.zeros([img_h, img_w], dtype=np.uint8)

    # 设定阈值
    threshold = 0.7
    loc = np.where(res >= threshold)
    logos = []
    for pt in zip(*loc[::-1]):
        x, y = pt[0] - 60, pt[1] - 10
        # 在原图上面进行标记出水印的位置
        cv2.rectangle(img_rgb_rect, (x, y), (x + w, y + h), (255, 255, 255), 1)

        # 模板图上面水印处填充白色
        mask[y:y + h, x:x + w] = 255

        # 水印坐标限定，防止超边界
        y = y if y + h < img_h else img_h - h
        w = w if x + w < img_w else img_w - x

        logos.append((x, y, w, h))

    # 使用黑白模板图对原图进行修复
    img_rm_watermark = cv2.inpaint(img_rgb, mask, 20, cv2.INPAINT_TELEA)

    # 复制一份去水印图 用来加新水印
    img_add_logo = img_rm_watermark.copy()

    # 遍历所有水印坐标进行加新水印
    for logo in logos:
        x, y, w, h = logo
        app_logo = cv2.addWeighted(img_rm_watermark[y:y + h, x:x + w], 1, img_logo, 0.5, 0)
        img_add_logo[y:y + h, x:x + w] = app_logo[:, :]

    if output_path:
        cv2.imwrite(output_path, img_add_logo)
        return
    cv2.imshow('src', img_rgb)  # 原图
    cv2.imshow('img_rgb_rect', img_rgb_rect)  # 在原图上标记出水印的位置
    cv2.imshow('mask', mask)  # 根据水印位置制作出来的蒙版图
    cv2.imshow('img_rm_watermark', img_rm_watermark)  # 去水印之后的图
    cv2.imshow('img_add_logo', img_add_logo)  # 添加新水印后的图

    cv2.waitKey(0)
    cv2.destroyAllWindows()


def test():
    src_img = "src/9.jpg"
    mask_img = "mask.jpg"
    logo_img = "img/logo.png"
    mathc_img(src_img, mask_img, logo_img)


def main():
    # 读取输入参数
    ap = argparse.ArgumentParser()
    ap.add_argument("-s", "--src", required=True, help="原图片文件路径")
    ap.add_argument("-l", "--logo", required=True, help="要覆盖的Logo图片文件路径")
    ap.add_argument("-t", "--template", required=True, help="要寻找的模板图片文件路径")
    ap.add_argument("-o", "--out", required=True, help="输出的图片文件路径")
    args = ap.parse_args()

    mathc_img(args.src, args.template, args.logo, args.out)


# 调试使用
if __name__ == '__main__':
    main()