Skip to content

Model for watermark classification implemented with PyTorch

Notifications You must be signed in to change notification settings

boomb0om/watermark-detection

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

19 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

watermark-detection

Installation

git clone https://github.com/boomb0om/watermark-detection
cd watermark-detection
pip install -r requirements.txt

Model weights are automatically downloaded when you create model with get_watermarks_detection_model. Also you can find them on huggingface

Basic usage

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

Test dataset

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%

Training and train dataset

Training dataset

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.

Training a model

You can find traning code in this notebook. Good luck!

About

Model for watermark classification implemented with PyTorch

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published