braindotai / Watermark-Removal-Pytorch

šŸ”„ CNN for Watermark Removal using Deep Image Prior with Pytorch šŸ”„.
MIT License
883 stars 146 forks source link
artefacts-removal deep-image-priors deep-learning paper pytorch restoration-tasks skip-connections watermark watermark-removal

Watermark Removal using Deep Image Priors with Pytorch

DOI License: MIT

āœ… Update 1

Updated the model architecture, and reduced the parameters count from ~3,000,000 to ~500,000 for faster inference.

āœ… Update 2

Added api endpoint for easy usage:


from api import remove_watermark

remove_watermark(
    image_path = IMAGE_NAME,
    mask_path = MASK_NAME,
    max_dim = MAX_DIM,
    show_step = SHOW_STEPS,
    reg_noise = REG_NOISE,
    input_depth = INPUT_DEPTH,
    lr = LR,
    training_steps = TRAINING_STEPS,
    tqdm_length = 900
)

āœ… Update 3

Now supports MPS for Mac Users! Contribution by - @edwin-yan

This is the implementation of paper Deep Image Prior, all credit goes its authors.

My Contribution

Introduction

CNNs are very common for image generation and restoration tasks. And it is believed that their great performance is because of their ability to learn realistic image priors from training on large datasets. This paper shows that the structure of a generator alone is sufficient to provide enough low-level image statistics without any learning. Thus most of the image restoration tasks, for example, denoising, super-resolution, artefacts removal, watermark removal etc can be done with highly realistic results without any training.

In this repo, I've implemented the watermark removal task, and the results are just as good as claimed by the authors.

When the watermark is available.

So in this scenario, the requirements are:

Any watermarked image can be represented as the Hadamard product of Original Image and Watermark.

So if we know what Watermark is, then its just a matter of training a generator that produces outputs, such that Watermarked Image is equal to Generated Image * Watermark. And the authors propose to simply use L2 loss to minimize the distance between them. The generator simply takes the random noise with same height and width as of watermarked image, which we can regularize, and produces the outputs.

Here's a sample test run-

Original ImageWatermarkWatermarked-Image

Running inference

Run $ python inference.py with following arguments-


Removing Watermark

optional arguments:
  -h, --help            show this help message and exit
  --image-path IMAGE_PATH
                        Path to the "watermarked" image.
  --mask-path MASK_PATH
                        Path to the "watermark" image.
  --input-depth INPUT_DEPTH
                        Max channel dimension of the noise input. Set it based
                        on gpu/device memory you have available.
  --lr LR               Learning rate.
  --training-steps TRAINING_STEPS
                        Number of training iterations.
  --show-step SHOW_STEP
                        Interval for visualizing results.
  --reg-noise REG_NOISE
                        Hyper-parameter for regularized noise input.
  --max-dim MAX_DIM     Max dimension of the final output image

Outputs

When the watermark is not available.

In this scenario, we'll have the watermarked image only. And this is the actual and highly realistic scenario because of obvious reasons. For very trivial causes, the first scenario was too easy to tackle than this one. Even the authors provided the outputs for the first scenario only.

Let us see how seriously difficult this is:

My solution

Read the last bold statement again, if we solve that issue, then it's just a matter of following the first scenario ain't it. I hope you can see where I'm going from this šŸ˜‰. If you look at the bigger picture of watermark removal, then, in a nutshell, its just an image inpainting task right? So, all we need to do is, roughly highlight the watermarked region from any paint software and you're good to go.

Yup, it's just that simple.

Yes, it's not a holy new solution, I've seen it done before in form of image inpainting. But I never saw anyone applying this for removing the watermark.

Now you might be thinking that it's not a fully automated task anymore, since you would have to manually sit..tahh and highlight the watermarked region. Though I'm pretttty sure that after seeing the final results, you surely would enjoy that sitting :)

Moreover, think about how many problems we are solving by just simply doing this:

Okay, enough talk.

The first step is to create an overlay containing strokes, that hides the watermark.

I'm simply using MS Paint for that. And per image, it hardly takes 1 minute to draw.

Here are some sample overlays...

Overlay example

Overlay example

Overlay example

And the final step is to simply follow the above training procedure, by replacing the supposedly known watermark with our hand drawn overlay.

Time for the results

A sample progress of the generator:

Experiment 0

Experiment result

Experiment 1

Experiment result

Experiment 2

Experiment result

Experiment 3

Experiment result

Experiment 4

Experiment result

Experiment 5

Experiment result

Experiment 6

Experiment result

As I said earlier, the outputs are highly realistic with almost unnoticeable artefacts.

A Bonus Usecase :)

Since watermark removal is one the application of image inpainting, so using this notebook you can even edit the images as per your need, that is removing any unwanted part of the image.

Below are the final outputs for some sample image editing examples:

Note

Look at the last output image, as you can see, the final outputs contains pretty visible artefacts in this case. Our model is not trained on any dataset to learn any shape or form of image statistics, the outputs of the generator solely depends on the single image, multiplied by the mask, so if the mask is covering too much portion of the image, then it becomes too hard for the model to make the best predictions for filling those missing pixels. So in conclusion, thinner/detailed the mask, better the outputs.

Read Full Article

Author - Rishik Mourya [Linkedin]

Star History

Star History Chart