使用Pytorch和深度图像先验去除水印

✅ 更新1

更新了模型架构，将参数数量从约300万减少到约50万，以加快推理速度。

✅ 更新2

添加了API端点以便于使用：


from api import remove_watermark

remove_watermark(
    image_path = 图像名称,
    mask_path = 掩码名称,
    max_dim = 最大维度,
    show_step = 显示步骤,
    reg_noise = 正则化噪声,
    input_depth = 输入深度,
    lr = 学习率,
    training_steps = 训练步骤,
    tqdm_length = 900
)

✅ 更新3

现在支持Mac用户的MPS！由@edwin-yan贡献。

这是Deep Image Prior论文的实现，所有功劳归于其作者。

我的贡献

我提供了一个超级简单但高效的解决方案，以处理原论文未涵盖的场景，即当水印在训练之前无法单独获得的情况。
一个易于使用的良好API封装 :)

简介

CNN在图像生成和恢复任务中非常常见。人们认为，它们的出色表现是因为它们能够通过对大型数据集的训练学习到真实的图像先验。这篇论文表明，生成器的结构本身就足以提供足够的低级图像统计信息，无需任何学习。因此，大多数图像恢复任务，例如去噪、超分辨率、伪影去除、水印去除等，都可以在不进行任何训练的情况下获得高度真实的结果。

在这个仓库中，我实现了水印去除任务，结果与作者声称的一样好。

当水印可用时

# 在这种情况下，要求是：

应用于带水印图像的水印是可用的。
水印的缩放、位置、旋转和其他空间变换与图像上应用的水印完全匹配。

任何带水印的图像都可以表示为原始图像和水印的哈达玛积。

因此，如果我们知道水印是什么，那么只需训练一个生成器，使其输出满足带水印图像等于生成的图像 * 水印。作者提议简单地使用L2损失来最小化它们之间的距离。生成器只需以带水印图像相同的高度和宽度作为输入随机噪声，我们可以对其进行正则化，然后生成输出。

这里是一个示例测试运行-

<table align="center"> <tr> <th>原始图像</th><th>水印</th><th>带水印图像</th> </tr> <tr> <td> <img src='https://yellow-cdn.veclightyear.com/ab5030c0/54ecefb4-9cb0-46ff-a284-b5f295d0ddf0.png' style="border-radius: 6px;"> </td> <td> <img src='https://yellow-cdn.veclightyear.com/ab5030c0/c43d7eb2-e046-49f4-9835-096daded0d0f.png' style="border-radius: 6px;"> </td> <td> <img src='https://yellow-cdn.veclightyear.com/ab5030c0/01f979fe-8f30-47ec-be75-469147709083.png' style="border-radius: 6px;"> </td> </tr> </table>

# 运行推理

运行 $ python inference.py 并使用以下参数-


去除水印

可选参数:
  -h, --help            显示此帮助信息并退出
  --image-path IMAGE_PATH
                        "带水印"图像的路径。
  --mask-path MASK_PATH
                        "水印"图像的路径。
  --input-depth INPUT_DEPTH
                        噪声输入的最大通道维度。根据您可用的GPU/设备内存设置。
  --lr LR               学习率。
  --training-steps TRAINING_STEPS
                        训练迭代次数。
  --show-step SHOW_STEP
                        可视化结果的间隔。
  --reg-noise REG_NOISE
                        正则化噪声输入的超参数。
  --max-dim MAX_DIM     最终输出图像的最大维度

# 输出

当水印不可用时

在这种情况下，我们只有带水印的图像。这是一个实际且高度现实的场景，原因显而易见。由于明显的原因，第一种情况比这种情况更容易处理。即使作者只提供了第一种情况的输出。

# 让我们看看这有多么困难：

只有带水印的图像可用。
我们绝对__无法__向生成器提供以下信息：
- 水印在哪里？
- 水印的具体部分是什么？
- 哪些部分不是水印？
而且我们想在没有任何训练的情况下做到这一点！！！为什么？好吧，这毫无意义，我的意思是我们知道生成器确实能够修复水印，只是我们无法为生成器提供上述问题的答案。

# 我的解决方案

再读一遍最后那句加粗的话，如果我们解决了这个问题，那么剩下的就只是按照第一种情况来处理，不是吗？我希望你能看出我的思路😉。如果你从更宏观的角度看水印去除，那么归根结底，它只是一个图像修复任务，对吧？所以，我们需要做的就是用任何绘图软件粗略地突出显示水印区域，然后你就可以开始了。

是的，就是这么简单。

没错，这不是一个全新的解决方案，我以前在图像修复中见过这种做法。但我从未见过有人将此应用于水印去除。

现在你可能在想，这不再是一个完全自动化的任务了，因为你必须手动坐下来...标记水印区域。不过我非常确信，在看到最终结果后，你肯定会享受这个过程的:)

此外，想想我们通过简单地做这件事解决了多少问题：

无需训练水印检测模型。这比典型的物体检测更难做到。
即使我们能检测到水印，也不会有太大帮助，因为水印可能画在整个图像上，而不仅仅是一个小区域。
无需在大型图像数据集上训练生成器以学习图像统计信息。
无需对生成器进行对抗性损失训练,这对于生成1024及以上更高分辨率的图像来说已经非常困难...
到目前为止,我所看到的所有其他试图自动检测和移除水印的解决方案,都会产生非常明显的人工痕迹。

好了,说得够多了。

第一步是创建一个包含笔触的覆盖层,用来隐藏水印。

我只是简单地使用MS Paint来做这个。每张图片大约只需要1分钟就能画好。

这里有一些覆盖层示例...

覆盖层示例

最后一步就是简单地遵循上述训练程序,用我们手绘的覆盖层替换假定已知的水印。

# 结果展示时间

生成器的进度示例:

实验0

实验结果

实验1

实验结果

实验2

实验结果

实验3

实验结果

实验4

实验结果

实验5

实验结果

实验6

实验结果

正如我之前所说,输出结果非常逼真,几乎看不出人工痕迹。

额外用途 :)

由于水印移除是图像修复的一种应用,所以使用这个笔记本你甚至可以根据需要编辑图像,也就是移除图像中任何不想要的部分。

以下是一些图像编辑示例的最终输出:

# 注意

看最后一张输出图像,如你所见,在这种情况下最终输出包含相当明显的人工痕迹。我们的模型并没有在任何数据集上训练以学习任何形状或图像统计信息,生成器的输出完全取决于单张图像乘以蒙版,所以如果蒙版覆盖了图像的太大部分,那么模型就很难对填充这些缺失像素做出最佳预测。因此,总的来说,蒙版越薄/越细致,输出效果就越好。