想付费请作者修改增加功能

[qq14909]

增加识别水印位置生成蒙版
或者增加调用https://github.com/Sanster/IOPaint?tab=readme-ov-file这个自动去除水印
作者有空吗

[qq14909]

另外想请教下，怎么才能返回识别到的所有水印位置呢

[Kamino666]

watermark-tracer/trace.py

Line 209 in d8cabef

results = model(imgs)

result里面就有yolov5模型检测到的所有信息了。

抱歉，平时也比较忙，这个库基本就不更新了hhhh

[qq14909]

我用ai改了下，支持识别多个水印位置生成遮罩了，但是yolov5模型识别不了我的水印，我用yolov8生成了一个模型，但是不知道你的代码如何改成支持yolov8，如果大佬愿意修改，小弟愿意付费的，`import logging
import random
import tkinter as tk
import tkinter.filedialog
from pathlib import Path
import argparse
import sys
import os
import cv2
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import skimage.draw as draw
import torch
from PIL import Image, ImageTk
from mmcv import VideoReader

from baidu import BaiduAPI

配置logger

logger = logging.getLogger('watermark-tracer')
plt.rcParams['font.family'] = 'SimHei'

def detect_watermark_from_video_result(frames, res, threshold=0.1):
res: pd.DataFrame = res.sort_values(by='confidence', ascending=False)
frames_np = [np.array(i) for i in frames]
# 提取最高置信度
res = res[res['confidence'] > threshold]
print("检测结果：\n", res)
w1, h1, w2, h2 = [int(i) for i in res.loc[0].to_list()[:4]]
wms = [i[h1:h2, w1:w2] for i in frames_np] # watermarks
# 增强水印
wm = estimate_watermark_from_images(wms)
return wm, [(w1, h1, w2, h2), ]

def detect_watermark_from_img_result(img, res, err_ratio=0.05, threshold=0.1):
res: pd.DataFrame = res.sort_values(by='confidence', ascending=False)
img_np = np.array(img)
boxes = [list(map(int, i[1:5])) for i in res.itertuples()]
wms = [] # watermarks
for w1, h1, w2, h2 in boxes:
i = img_np[h1:h2, w1:w2]
wms.append(np.array(i))
print("检测结果：\n", res) # 移动到这里
return wms, boxes
boxes = [list(map(int, i[1:5])) for i in res_less.itertuples()]
# 假如少于等于5个，直接返回，否则根据多幅图像提取水印
if len(res) <= 5:
print("未使用增强")
# w1, h1, w2, h2 = boxes[0]
w1, h1, w2, h2 = random.choice(boxes)
return img_np[h1:h2, w1:w2], boxes
else:
print("增强")
# 把所有子图都resize到相同大小
wms = [] # watermarks
for w1, h1, w2, h2 in boxes:
i = img_np[h1:h2, w1:w2]
i = Image.fromarray(i).resize((int(width), int(height)))
wms.append(np.array(i))
# 增强水印
wm = estimate_watermark_from_images(wms)
return wm, [list(map(int, i[1:5])) for i in res.itertuples()]

def estimate_watermark_from_images(imgs: list, enhance: int = 50):
# 估计水印
grad_x = list(map(lambda x: cv2.Sobel(x, cv2.CV_64F, 1, 0, ksize=3), imgs))
grad_y = list(map(lambda x: cv2.Sobel(x, cv2.CV_64F, 0, 1, ksize=3), imgs))
Wm_x = np.median(np.array(grad_x), axis=0)
Wm_y = np.median(np.array(grad_y), axis=0)

# plt.subplot(311)  # DEBUG
# plt.imshow(np.abs(Wm_x ** 2 + Wm_y ** 2) / np.max(Wm_x ** 2 + Wm_y ** 2))  # DEBUG
# ax = plt.gca()
# ax.axes.xaxis.set_visible(False)
# ax.axes.yaxis.set_visible(False)

est = poisson_reconstruct(Wm_x, Wm_y)
# 转换成255的
est: np.ndarray = 255 * (est - np.min(est)) / (np.max(est) - np.min(est))
est = est.astype(np.uint8)
# DEBUG
# plt.subplot(312)  # DEBUG
# plt.imshow(est)  # DEBUG
# ax = plt.gca()
# ax.axes.xaxis.set_visible(False)
# ax.axes.yaxis.set_visible(False)

# 寻找增强区域的模版
channels = []
for i in range(est.shape[-1]):
    # 二值化
    blur = cv2.GaussianBlur(est[:, :, i], (5, 5), 0)
    ret, th = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    channels.append(th)
mask = np.zeros_like(channels[0]).astype(bool)
for c in channels:
    mask = mask | c.astype(bool)
mask = mask[:, :, np.newaxis].repeat(3, axis=2)
# print(mask.shape, est.shape)
# print(mask.dtype, est.dtype)
# plt.figure(2)
# plt.subplot(211)
# plt.imshow(mask.astype(int)*255)
# plt.subplot(212)
# plt.imshow(est)
# plt.show()
# 增强
est = est + enhance * mask
est: np.ndarray = 255 * (est - np.min(est)) / (np.max(est) - np.min(est))
est = est.astype(np.uint8)
# DEBUG
# plt.subplot(313)
# plt.imshow(est)
# ax = plt.gca()
# ax.axes.xaxis.set_visible(False)
# ax.axes.yaxis.set_visible(False)
# plt.show()
return est

def poisson_reconstruct(gradx, grady, kernel_size=3, num_iters=100, h=0.1,
boundary_image=None, boundary_zero=True):
"""
Iterative algorithm for Poisson reconstruction.
Given the gradx and grady values, find laplacian, and solve for images
Also return the squared difference of every step.
h = convergence rate
"""
fxx = cv2.Sobel(gradx, cv2.CV_64F, 1, 0, ksize=kernel_size)
fyy = cv2.Sobel(grady, cv2.CV_64F, 0, 1, ksize=kernel_size)
laplacian = fxx + fyy
m, n, p = laplacian.shape

if boundary_zero is True:
    est = np.zeros(laplacian.shape)
else:
    assert (boundary_image is not None)
    assert (boundary_image.shape == laplacian.shape)
    est = boundary_image.copy()

est[1:-1, 1:-1, :] = np.random.random((m - 2, n - 2, p))
loss = []

for i in range(num_iters):
    old_est = est.copy()
    est[1:-1, 1:-1, :] = 0.25 * (
            est[0:-2, 1:-1, :] + est[1:-1, 0:-2, :] +
            est[2:, 1:-1, :] + est[1:-1, 2:, :] -
            h * h * laplacian[1:-1, 1:-1, :]
    )
    error = np.sum(np.square(est - old_est))
    loss.append(error)
return est

def get_file_path(args):
if args.image is not None:
return Path(args.image)
else:
root = tk.Tk()
file_path = tkinter.filedialog.askopenfilename(
multiple=False,
filetypes=[('图片', '.jpg'), ('图片', '.png'), ('视频', '.mp4')]
)
root.destroy()
return Path(file_path)

def get_images(file_path):
if file_path.suffix in ['.jpg', '.png', '.jpeg', '.webp']: # 图片
return 'image', [Image.open(file_path).convert('RGB')]
else: # 视频（平均抽取10帧）
video = VideoReader(str(file_path))
indices = np.linspace(2, video.frame_cnt - 2, 10).astype(int)
return 'video', [Image.fromarray(video.get_frame(i)[:, :, ::-1]) for i in indices]

def load_model():
logger.info("开始加载YoloV5模型")
try:
model = torch.hub.load('yolov5', 'custom', path='yolov5/best.pt', source='local')
#从本地加载一个名为custom的预训练模型，该模型在'yolov5'这个仓库中。
return model.cpu()
except Exception as e:
logger.error("YoloV5加载失败: " + str(e))
sys.exit()

def get_watermark(file_type, imgs, results):
results = results.pandas().xyxy
if file_type == 'image':
if len(results[0]) == 0:
logger.error("Yolo检测失败")
sys.exit()
return detect_watermark_from_img_result(imgs[0], results[0])
elif file_type == 'video':
idx = -1
for i, result_item in enumerate(results):
if len(result_item) != 0:
idx = i
break
if idx == -1:
logger.error("Yolo检测失败")
print(results)
sys.exit()
return detect_watermark_from_video_result(imgs, results[idx])

if name == "main":
parser = argparse.ArgumentParser()
parser.add_argument('-m', '--image', required=False, type=str, help="输入图片路径，假如没有则调用对话框进行选择")
args = parser.parse_args()

file_path = get_file_path(args)
if file_path.is_file() is False:
    logger.error("未选择资源")
    sys.exit()

file_type, imgs = get_images(file_path)

model = load_model()

matplotlib.use('Qt5Agg')
logger.info("YoloV5加载成功")

logger.info("检测中")
results = model(imgs)

test_wms, boxes = get_watermark(file_type, imgs, results)

mask = np.zeros_like(imgs[0])
for box in boxes:
    cv2.rectangle(mask, (box[0], box[1]), (box[2], box[3]), (255, 255, 255), -1)

mask_path = os.path.splitext(file_path)[0] + '_mask.png'
cv2.imwrite(mask_path, mask)
logger.info(f"遮罩蒙版已保存到 {mask_path}")`

[qq14909]

我主要是没看懂yolov5是怎么训练模型的，以及目录中该怎么存放结构，我有下载WatermarkDataset压缩包查看，还是没看懂
不知道是用这个
cd yolov5
./run_train.sh # 训练
./run_detect.sh # 测试

还是这个

cd data
python generator.py \
--images_dir <背景图片目录> \
--logo_dir <水印logo目录> \ # 可在data/logos查看目录结构
--output_dir <输出目录> \
--num_workers <使用的进程数> \
--seed <随机种子> \
-n <生成的数量> \
--test # 测试用，详情请看代码

[qq14909]

不然用yolov5训练也行

[qq14909]

我用工具生成了训练集和验证集
txt里面存放格式2 0.726316 0.665730 0.433684 0.101124

E:\YOLOTRAIN\DATASETS\LOGO
├─images
│ ├─train
│ │ 10F6GA68016.jpg
│ │ 10S3VA91051.jpg
│ │ 11DKWH68867.jpg
│ │ 12SBNF67282.jpg
│ │ 13NVRM48064.jpg
│ │ 154Q4R27522.jpg
│ │
│ └─val
│ 11DKWH68867.jpg
│ 154Q4R27522.jpg
│ 161GVG63540.jpg
│ 346M684758.jpg
│ 3SAIA10701.jpg
│ 4TKP415628.jpg
│
└─labels
├─train
│ 10F6GA68016.txt
│ 10S3VA91051.txt
│ 11DKWH68867.txt
│ 12SBNF67282.txt
│ 13NVRM48064.txt
│ 154Q4R27522.txt
│ classes.txt
│
└─var
11DKWH68867.txt
154Q4R27522.txt
161GVG63540.txt
346M684758.txt
3SAIA10701.txt
4TKP415628.txt
classes.txt

[qq14909]

我将yolov5的best.pt替换为yolov8的best.pt，就出错了，trace.py中搜索torch.hub.load这里已经修改为yolov8的best.pt

提示以下：

python watermark-trace.py -m E:\1\084102t147fe1774ftst45.jpg YOLOv5 2024-3-8 Python-3.8.18 torch-2.2.1+cpu CPU

YoloV5加载失败: No module named 'ultralytics'. Cache may be out of date, try force_reload=True or see https://github.com/ultralytics/yolov5/issues/36 for help.

安装pip install ultralytics也没用，提示

ultralytics 0.0.14 depends on psutil
ultralytics 0.0.13 depends on psutil

To fix this you could try to:

1. loosen the range of package versions you've specified
2. remove package versions to allow pip attempt to solve the dependency conflict

ERROR: ResolutionImpossible: for help visit https://pip.pypa.io/en/latest/topics/dependency-resolution/#dealing-with-dependency-conflicts

[qq14909]

pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/
设置阿里的镜像后可以安装pip install ultralytics了

再次执行提示>> python watermark-trace.py -m E:\1\084102t147fe1774ftst45.jpg
YOLOv5 2024-3-8 Python-3.8.18 torch-2.2.1+cpu CPU

Model summary (fused): 168 layers, 3006233 parameters, 0 gradients, 8.1 GFLOPs
Adding AutoShape...
YoloV5加载失败: 'Detect' object has no attribute 'grid'. Cache may be out of date, try force_reload=True or see ultralytics/yolov5#36 for help.

[qq14909]

算了，我都不知道自己在干啥，大佬能教教我怎么使用你的代码训练yolov5模型吗，可付费，如果没空的话就先不打扰了

[shiertier]

算了，我都不知道自己在干啥，大佬能教教我怎么使用你的代码训练yolov5模型吗，可付费，如果没空的话就先不打扰了

你直接用他的代码重新构建新的数据集并按照yolo8的方法训练就行了，按照他的逻辑来重构一个新的代码会比较好。

[Kamino666]

抱歉哦，精力不足，无法更新这个程序~

很久之前写的代码哈哈哈哈，代码不是很难，可以给大家的工作提供一些参考。