Files
stat_center/script/process.py
T
2026-05-20 16:13:58 +08:00

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()