forked from all_user/stat_center
128 lines
3.9 KiB
Python
128 lines
3.9 KiB
Python
# 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()
|