Watermark-Tracer

一个基于可视水印检测识别的数字媒体溯源应用系统，是我的大作业项目，包含这个系统以及一个开源的大规模常见水印图像数据集（Large-scale Common Watermark Dataset, LCWD）。

输入一个带有可视水印的图片或视频，系统会检测定位到水印所在的区域，然后将其提取出来，然后借助百度AI开放平台的OCR和logo识别以及Bing搜索引擎，溯源到这个图片或视频的源头。

效果

快速开始

环境准备

先安装好Python3.8和PyTorch1.7+。

# 要Clone到子模块！（可能会比较大，应为包含好几个yolo的模型，避免之后网络不好又要去下载）
git clone https://github.com/Kamino666/watermark-tracer.git --recurse-submodules
# 安装相应的库
pip install -r requirements.txt

百度开放AI平台API申请*

⚠️不申请API也可以使用除了溯源以外的功能

接着为了使用百度开放AI平台的接口，要去申请两个API（都有免费额度）：

• 网络图片识别 图片文字识别 图片转文字-百度AI开放平台 (baidu.com)
• 品牌logo识别_拍照识别2万多类商品logo-百度AI开放平台 (baidu.com)

然后在项目目录新建一个baidu_cfg.json的文件，将key填入

{
  "ak": "Your API key",
  "sk": "Your Secret key",
}

运行

python trace.py
# 一些可选参数
--no_api          # 不使用百度API，不使用溯源功能
--no_tk           # 不使用tkinter，使用matplotlib进行显示
-m <图片视频路径>   # 在命令行指定路径，不指定也可以

训练新的模型

数据集准备

本项目训练的数据集是一个新构建的大规模常见水印图像数据集（Large-scale Common Watermark Dataset, LCWD），可以从Large-scale Common Watermark Dataset | Kaggle下载。

数据集的介绍可见我的论文。

训练

根据情况修改run_train.sh和run_detect.sh中的路径

cd yolov5
./run_train.sh  # 训练
./run_detect.sh  # 测试

更多训练参数可见ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)

若要使用新的模型，可在trace.py中搜索torch.hub.load，更改path参数。

生成新的数据集

假如对公布的数据集不满意，也可尝试自己调整参数，生成新的数据集，数据集的详细构建方法可见我的论文，使用以下命令生成：

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

若要调整水印生成的格式，请参考代码WatermarkedImageGenerator类。

此外，在tools目录下有连个jupyter notebook可以参考用来划分和可视化数据集。

引用

@misc{Liu2022watermark,
    title={Watermark-Tracer},
    url={https://github.com/Kamino666/watermark-tracer},
    journal={GitHub},
    author={Zihao, Liu},
    year={2022},
    month={6}
} 

Credits

中国传媒大学 信息与通信工程学院 田佳音老师、杨成老师和和于瀛老师（无顺序）

ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite (github.com)

ImageNet 1000 (mini) | Kaggle

Douban Movie Short Comments Dataset | Kaggle

iconfont-阿里巴巴矢量图标库

百度AI开放平台-全球领先的人工智能服务平台 (baidu.com)

zhu733756/searchengine: 元搜索引擎 searchengine 元数据 元搜索 (github.com)

X-CCS/remove-watermarks: implementation of "On the Effectiveness of Visible Watermarks" (github.com)