Initial commit

script/process.py

@@ -0,0 +1,127 @@
+# 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()