git clone https://github.com/boomb0om/watermark-detection
cd watermark-detection
pip install -r requirements.txtModel weights are automatically downloaded when you create model with get_watermarks_detection_model.
Also you can find them on huggingface
Checkout this notebook for basic usage.
from PIL import Image
from wmdetection.models import get_watermarks_detection_model
from wmdetection.pipelines.predictor import WatermarksPredictor
# checkpoint is automatically downloaded
model, transforms = get_watermarks_detection_model(
'convnext-tiny',
fp16=False,
cache_dir='/path/to/weights'
)
predictor = WatermarksPredictor(model, transforms, 'cuda:0')
result = predictor.predict_image(Image.open('images/watermark/1.jpg'))
print('watermarked' if result else 'clean') # prints "watermarked"Use predictor.run to run model on list of images:
results = predictor.run([
'images/watermark/1.jpg',
'images/watermark/2.jpg',
'images/watermark/3.jpg',
'images/watermark/4.jpg',
'images/clean/1.jpg',
'images/clean/2.jpg',
'images/clean/3.jpg',
'images/clean/4.jpg'
], num_workers=8, bs=8)
for result in results:
print('watermarked' if result else 'clean')Download test dataset from huggingface and put it into dataset/ folder. Then run evaluation code in this jupyter.
Test dataset is consists of 60 clean images and 62 watermarked images. Model accuracy for this dataset is shown below:
| Model name | Accuracy |
|---|---|
| convnext-tiny | 93,44% |
| resnext101_32x8d-large | 84,42% |
| ARKseal model | 77,86% |
| resnext50_32x4d-small | 76,22% |
Synthetic training data is generated using random watermark generator. Gather clean images for generator and put them into dataset/ folder. Then run notebook with synthetic dataset generation. Generator will randomly put watermark on every image and save it into another folder.
After synthetic images are generated, run notebook to create csv for all training images.
You can find traning code in this notebook. Good luck!