clausmichele / ViDeNN

ViDeNN - Deep Blind Video Denoising
MIT License
232 stars 35 forks source link

How to use your own checkpoint #12

Closed ojasviyadav closed 3 years ago

ojasviyadav commented 4 years ago

I am trying to locate the checkpoint created after you train the model on your own data.

I have tried to test the model with the checkpoint that gets stored during training for Spatial-CNN as well as Temp3-CNN folder.

I'm just a little unsure where exactly does the checkpoint get stored in the format that can be used with main_VideNN.py to test the denoising.

clausmichele commented 4 years ago

Can you tell me step by step the procedure you have followed?

ojasviyadav commented 4 years ago
  1. I trained the model on default settings on the dataset that is recommended. First I trained the spatial model and then I trained the temporal model. For this I used the instructions given in the README. I am assuming that the checkpoints of this training did get stored in the folders Spatial-CNN and Temp3-CNN.

  2. I then noticed that the denoise.sh script, which is used for testing, actually uses the checkpoints that were included in the repository already, namely ckpt_videnn. I am assuming this checkpoint stores the pretrained model.

  3. I would like to use the checkpoints that were created in step 1, instead of using the checkpoint of the pretrained model, to test the model on the CBD dataset. For that I would first like to confirm which is the appropriate checkpoint, is it the one saved when training the Spatial model or the one saved when training the temporal model? Or is there some other checkpoint which I should be using for testing the model on a video?

clausmichele commented 4 years ago

You are right! the README misses the last part of training that creates a single checkpoint from the spatial and temporal ones. I'll update it as soon as I can.

clausmichele commented 4 years ago

So, the solution was already inside the code! If you run main_ViDeNN.pywithout giving the --checkpoint_dir parameter, it will load the last spatial checkpoint in ./Spatial-CNN/ckpt_awgn and the last temporal checkpoint in ./Temp3-CNN/ckpt. (You can modify those two defaults folder in model_ViDeNN.py)

Then, if you want to use the denoise.sh script and you have the checkpoints in the correct folders, just modify denoise.sh: python main_ViDeNN.py --use_gpu=1 --checkpoint_dir=ckpt_videnn --save_dir='./data/denoised' --test_dir='./data' to python main_ViDeNN.py --use_gpu=1 --save_dir='./data/denoised' --test_dir='./data'