Coder Social home page Coder Social logo

braindotai / watermark-removal-pytorch Goto Github PK

View Code? Open in Web Editor NEW
739.0 14.0 137.0 44.21 MB

πŸ”₯ CNN for Watermark Removal using Deep Image Prior with Pytorch πŸ”₯.

License: MIT License

Python 46.59% Jupyter Notebook 53.41%
deep-image-priors watermark-removal pytorch watermark skip-connections artefacts-removal paper restoration-tasks deep-learning

watermark-removal-pytorch's Introduction

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

  • I provide a super simple yet highly effective solution to handle the scenario, which is not covered by the original paper, that is, when the watermark is not available separately to us prior to any training.
  • A nice api wrapper which is easy to work with :)

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:

  • The watermark that is applied to the watermarked image, is available to you.

  • The scale, position, rotation and other spatial transformations of the watermark, exactly matches the applied watermark of the image.

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:

  • Only the watermarked image is available.

  • We can provide absolutely no info to the generator regarding:

    • Where is the watermark?
    • What exactly is the part of watermark?
    • What exactly is not the part of watermark?
  • And we want to do this without any training!!! Why? Well there's no point, I mean we know for a fact that the generator is indeed capable of inpainting the watermark, it's just us who are not able to provide the answers to the generator for questions above.

# 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:

  • No need to train a watermark detection model. It's hard to do than typical object detections.

  • Even if we can detect the watermark, it still doesn't help that much, cuz the watermark can be drawn on the whole image, not on just a small region.

  • No need to train the generator on huge image datasets for learning image statistics.

  • No need to train the generator over an adversarial loss, which is already very difficult for producing higher resolution images like 1024 and more...

  • And, all other solutions I've seen to far, which try to automate the whole procedure of detecting and removing the watermark, produces very visible artefacts.

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.

Author - Rishik Mourya [Linkedin]

watermark-removal-pytorch's People

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

watermark-removal-pytorch's Issues

Broken can't install it even with build tools

gi\gimodule.c: fatal error C1083: Impossible d'ouvrir le fichier includeΓΏ: 'msvc_recommended_pragmas.h'ΓΏ: No such file or directory
      error: command 'C:\\Program Files (x86)\\Microsoft Visual Studio\\2022\\BuildTools\\VC\\Tools\\MSVC\\14.37.32822\\bin\\HostX86\\x86\\cl.exe' failed with exit code 2
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
  ERROR: Failed building wheel for PyGObject
Failed to build PyGObject
ERROR: Could not build wheels for PyGObject, which is required to install pyproject.toml-based projects 

retrain

is it required that every time retrain for each image?

[RuntimeError: CUDA out of memory] when max dimensions of final image are increased

I am getting the following error when I increase the --max-dim parameter above 700 on a 6GiB GPU.

RuntimeError: CUDA out of memory. Tried to allocate 208.00 MiB (GPU 0; 6.00 GiB total capacity; 5.29 GiB already allocated; 0 bytes free; 5.33 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I tried reducing the --input-depth to see if this would reduce the memory usage, but it didn't help.

Are there any other possible tweaks that would allow for processing of larger images?

Licence?

Hello,

really nice project!
Under which licence is it published?

Not an issue. wondering if the set up is correct.

after changing the image in the watermarked folder.
and masks in the masks folder.

then runpython3 inference.py

(base) x@xnoMacBook-Pro Watermark-Removal-Pytorch % python3 inference.py   
Setting device to "cpu", since torch is not built with "cuda" support...
It is recommended to use GPU if possible...
Visualizing mask overlap...

got this message for a while. Can you please confirm if the code is running as expected? Thank you.

Loss of resolution

Hello,

Thank you for implementing this great pipeline! I successfully ran it but I saw the resolution of output is much lower than the original input. The loss of resolution also happended to those unmasked areas without the watermark. I set the output dimention the same as the input dimention. Is that a way to improve the resolution and avoid changing the quality of unmasked regions?

Thanks,
Shujun

requirements.txt

Consider adding a requirements.txt file.
Dependencies I noticed:

torch
torch vision
numpy
tqdm
pillow
matplotlib

CPU usage.

When set to CPU mode this sent a 13900k into 100% and made the cpu pulse/squeal with electronic noise.
I have never had this before even stress testing.
I thought a hard drive was getting read its that loud!

Inference time

Really nice project! Thank you for making it public!

The question is: what's the time of inference on image?
Actually my goal is to remove watermark in the bottom rigth corner of the video. I've already prepared the docker container and simple code to divide video into frames. How long will it take to run on 1k frames?

I can't actually check for now because video is not currently available. (I'll post my results when I finish)

Specs:

  • i5 12600kf
  • 16gb RAM 3200MHz
  • RTX 3080 (10 GB version)

Model learns to synthesize watermark too

Hi @braindotai ,

First of all, thank you!! This is definitely a nice wrapper for the original implementation, which is quite poorly documented.

I am in the second case: watermark not available; I am using a new file of my own. However, I am getting nowhere close to your results. The model learns to reproduce the watermark.

Am I doing sth wrong? Maybe the image format of the watermark?

Cheers.

Default parameters

Hi,
Could you please give us some correct, default parameters for :

max_dim=MAX_DIM,
show_step=SHOW_STEPS,
reg_noise=REG_NOISE,
input_depth=INPUT_DEPTH,
lr=LR,
training_steps=TRAINING_STEPS

Processing was stopped out , Upsampling is deprecated

Starting training...

Completed: 0%| | 0/9 [00:00<?, ?it/s]C:\Program Files\Python36\lib\site-packages\torch\nn\modules\upsampling.py:122: UserWarning: nn.Upsampling is deprecated. Use nn.functional.interpolate instead.
warnings.warn("nn.Upsampling is deprecated. Use nn.functional.interpolate instead.")
Completed: 100%|β–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆβ–ˆ| 9/9 [00:12<00:00, 1.43s/it, Loss=0.0263]

Saving final output image to: "E:\Watermark-Removal-Pytorch\demo\frames\00000-output.jpg"

on video ?

would this model be available on videos ?

How to use the script?

The script looks great! Unfortunately I am not a developer.
May I ask how to use it?

I first downloaded it to my Mac, with GitHub Desktop.
Then I run the following in the terminal

cd /Users/MyUserName/Documents/GitHub/Watermark-Removal-Pytorch/
python inference.py

It has the following error

Traceback (most recent call last):
File "inference.py", line 2, in
from api import remove_watermark
File "/Users/MyUserName/Documents/GitHub/Watermark-Removal-Pytorch/api.py", line 63
print(f'\nSaving final output image to: "{output_path}"\n')
^
SyntaxError: invalid syntax

Could you explain a little where I did it wrong?
Thanks in advance.
I would love to try this great script.
Thanks for sharing

Need some data on performance of this application

First of all, Thank you very much for this program. It is working great for an image.

Currently, I am running this application on a windows laptop and was able to execute the program python inference.py.
I do have training-steps as 3000 which is the default. i have prepared a masked image to remove the watermark from the original image. it is taking nearly 6hrs for an image to be processed and remove the watermark.

I wanted to know if I am doing anything wrong or will it be better if I run on some computer with good CPU power.

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    πŸ–– Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. πŸ“ŠπŸ“ˆπŸŽ‰

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❀️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.