GithubHelp home page GithubHelp logo

ljzycmd / simdeblur Goto Github PK

View Code? Open in Web Editor NEW
300.0 9.0 39.0 8.72 MB

Simple framework for image and video deblurring, implemented by PyTorch

License: MIT License

Shell 0.16% Python 82.42% C++ 7.28% Cuda 10.00% C 0.14%
image-deblurring video-deblurring dbn dblrnet ifirnn strcnn simdeblur mscnn srn

simdeblur's Introduction

SimDeblur

SimDeblur (Simple Deblurring) is an open-sourced unifying training and testing framework for image and video deblurring based on PyTorch. It supports most deep-learning based state-of-the-art deblurring algorithms, and provides easy way to implement your own image or video deblurring and restoration algorithms.

Major Features

  • Modular Design

The toolbox decomposes the deblurring framework into different components and one can easily construct a customized restoration framework by combining different modules.

  • State-of-the-art

The toolbox contains most deep-learning based state-of-the-art deblurring algorithms, including MSCNN, SRN, DeblurGAN, EDVR, etc.

  • Efficient Training

SimDeblur supports distributed data-parallel training.

New Features

[2022/12/11] SimDeblur supports NAFNet (ckpt) model for image deblurring.

[2022/11/12] SimDeblur supports MIMOUnet model.

[2022/3/8] We further provide a image deblurring-based inference code, please refer to Usage section for the using.

[2022/2/18] We add PVDNet model for video deblurring. Note that it requires the pretrained BIMNet for motion estimation. Thus please modify the CKPT path of BIMNet in the source codes.

[2022/1/21] We add Restormer model. Note that it can only works on PyTorch1.8+.

[2022/1/20] We transfer some checkpoints from the open-sourced repos into SimDeblur framework! You can find them here.

[2022/1/1] Support real-world video deblurring dataset: BSD.

[2021/3/31] Support DVD, GoPro and REDS video deblurring datasets.

[2021/3/21] First release.

Surpported Methods and Benchmarks

We will gradually release the checkpoints of each model in checkpoints.md.

Dependencies and Installation

  • Python 3 (Conda is recommended)

  • Pytorch 1.5+ (with GPU, note some methods require higher version)

  • CUDA 10.1+ with NVCC (for code compilation in some models)

  1. Clone the repositry or download the zip file
     git clone https://github.com/ljzycmd/SimDeblur.git
    
  2. Install SimDeblur
    # create a pytorch env
    conda create -n simdeblur python=3.7
    conda activate simdeblur   
    # install the packages
    cd SimDeblur
    bash Install.sh  # some problems may occur due to wrong NVCC configurations for CUDA codes compiling

Usage

You can open the Colab Notebook to learn about basic usage and see the deblurring performance.

The design of SimDeblur consists of FOUR main parts as follows:

Dataset Model Scheduler Engine
Dataset-specific classes The backbone, losses, and meta_archs. Backbone is the main network, and the meta_arch is a class for model training Opeimizer, LR scheduler Trainer, and some hook functions during model training

Note that the dataset, model and scheduler can be constructed with config (EasyDict) with corresponding build_{dataset, backbone, meta_arch, scheduler, optimizer, etc.} functions. The Trainer class automatically construct all reqiured elements for model training in a general way. This means that if you want to do some specific modeling training, you may modify the training logics in corresponding meta_arch class.

0 Quick Inference

We provide a image deblurring inference code, and you can run it to deblur a blurry image as follows:

python inference_image.py CONFIG_PATH  CKPT_PATH  --img=BLUR_IMAGE_PATH  --save_path=DEBLURRED_OUT_PATH

the deblurred latent image will be stored at ./inference_resutls in default.

1 Start with Trainer

You can construct a simple training process using the default Trainer as follows (refer to the train.py for more details):

from easydict import EasyDict as edict
from simdeblur.config import build_config, merge_args
from simdeblur.engine.parse_arguments import parse_arguments
from simdeblur.engine.trainer import Trainer


args = parse_arguments()

cfg = build_config(args.config_file)
cfg = merge_args(cfg, args)
cfg.args = edict(vars(args))

trainer = Trainer(cfg)
trainer.train()

Start training with single GPU:

CUDA_VISIBLE_DEVICES=0 bash ./tools/train.sh ./configs/dbn/dbn_dvd.yaml 1

or multiple GPUs training:

CUDA_VISIBLE_DEVICES=0,1,2,3 bash ./tools/train.sh ./configs/dbn/dbn_dvd.yaml 4

For the testing, SimDeblur only supports single GPU testing and validation util now:

CUDA_VISIBLE_DEVICES=0 python test.py ./configs/dbn/dbn_dvd.yaml PATH_TO_CKPT

2 Build specific module

SimDeblur also provides you to build some specific modules, including dataset, model, loss, etc.

Build a dataset:

from easydict import EasyDict as edict
from simdeblur.dataset import build_dataset

# construct configs of target dataset.
# SimDeblur adopts EasyDict to store configs.
dataset_cfg = edict({
    "name": "DVD",
    "mode": "train",
    "sampling": "n_c",
    "overlapping": True,
    "interval": 1,
    "root_gt": "./dataset/DVD/quantitative_datasets",
    "num_frames": 5,
    "augmentation": {
        "RandomCrop": {
            "size": [256, 256] },
        "RandomHorizontalFlip": {
            "p": 0.5 },
        "RandomVerticalFlip": {
            "p": 0.5 },
        "RandomRotation90": {
            "p": 0.5 },
    }
})

dataset = build_dataset(dataset_cfg)

print(dataset[0])

Build a model:

from easydict import EasyDict as edict
from simdeblur.model import build_backbone

model_cfg = edict({
    "name": "DBN",
    "num_frames": 5,
    "in_channels": 3,
    "inner_channels": 64
})

model = build_backbone(model_cfg)

x = torch.randn(1, 5, 3, 256, 256)
out = model(x)

Build a loss:

from easydict import EasyDict as edict
from simdeblur.model import build_loss

criterion_cfg = {
    "name": "MSELoss",
}

criterion = build_loss()

x = torch.randn(2, 3, 256, 256)
y = torch.randn(2, 3, 256, 256)

print(criterion(x, y))

And the optimizer and lr_scheduler also can be created by the functions build_optimizer and build_lr_scheduler in the simdeblur.scheduler, etc.

Dataset Description

SimDeblur supports the most popular image and video deblurring datasets, including GOPRO, DVD, REDS, BSD. We design different data reading strategies that can meet the input requirements of different image and video deblurring models.

You can click here for more information about the design of the dataset.

To start, note that you should change the path of the dataset in related config files.

Acknowledgment

The design spirit of SimDeblur comes most from Detectron2 [1], we highly thank for this amazing open-sourced toolbox. We also thank for the paper and code collections in Awesome-Deblurring repositry [2].

[1] facebookresearch. detectron2. https://github.com/facebookresearch/detectron2

[2] subeeshvasu. Awesome-Deblurring. https://github.com/subeeshvasu/Awesome-Deblurring

Citations

If SimDeblur helps your research or work, please consider citing SimDeblur.

@misc{cao2021simdeblur,
  author       = {Mingdeng Cao},
  title        = {SimDeblur: A Simple Framwork for Image and Video Deblurring},
  howpublished = {\url{https://github.com/ljzycmd/SimDeblur}},
  year         = {2021}
}

Last, if you have any questions about SimDeblur, please feel free to open an new issue or contact me at mingdengcao [AT] gmail.com, and I will try to solve your problem. Meanwhile, any contribution to this Repo is highly welcome. Let's make SimDeblur more powerful!

simdeblur's People

Contributors

ljzycmd avatar weihaox avatar

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

simdeblur's Issues

An error is reported in the video input

Hello, I would like to ask, can't a new video be directly input and then generate a new clear video?Does it have to be framed?1.I use restormer and get an error "assert inp_img.shape[1] == 1, "Input images should contain only 1 frame."
AssertionError: Input images should contain only 1 frame.” 2. "AttributeError: 'EasyDict' object has no attribute 'model'" is also reported with dbn

Using the other supported neural network models on ColabNotebook and inference_image.py

Hello there,

I was able to follow your example which you have posted on the ColabNotebook and have successfully able to perform deblurring using the DBN model on the test images locally on my PC via JupyterNotebook with CUDA enabled on PyTorch. So the example in ColabNotebook using DBN is working well.

Next, I tried to load a different mode (i.e. the DBLRNet) to compare the results, with the code snippet below.

...
model = build_backbone(model_cfg)
ckpt = torch.load("./demo/dblrnet_dvd.pth")
model_ckpt = ckpt["model"]
model_ckpt = {k[7:]: v for k, v in model_ckpt.items()}
model.load_state_dict(model_ckpt)
model = model.to(device)
...

I then get an error from the Python below.

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_3448/3417428220.py in <module>
     19 model_ckpt = ckpt["model"]
     20 model_ckpt = {k[7:]: v for k, v in model_ckpt.items()}
---> 21 model.load_state_dict(model_ckpt)
     22 model = model.to(device)

~/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/module.py in load_state_dict(self, state_dict, strict)
   1666         if len(error_msgs) > 0:
   1667             raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
-> 1668                                self.__class__.__name__, "\n\t".join(error_msgs)))
   1669         return _IncompatibleKeys(missing_keys, unexpected_keys)
   1670 

RuntimeError: Error(s) in loading state_dict for DBN:
	Missing key(s) in state_dict: "F0.0.weight", "F0.0.bias", "F0.1.weight", "F0.1.bias", "F0.1.running_mean", "F0.1.running_var", "D1.0.weight", "D1.0.bias", "D1.1.weight", "D1.1.bias", "D1.1.running_mean", "D1.1.running_var", "F1_1.0.weight", "F1_1.0.bias", "F1_1.1.weight", "F1_1.1.bias", "F1_1.1.running_mean", "F1_1.1.running_var",
...

I have also attempted to run the inference_image.py script to do the same thing with following command in Linux.
python inference_image.py ./configs/dblrnet/dblrnet_dvd.yaml ./demo/dblrnet_dvd.pth --img=./datasets/input/00000.jpg

This resulted in the following error below.

Using checkpoint loaded from ./demo/dblrnet_dvd.pth for testing.
Traceback (most recent call last):
  File "inference_image.py", line 81, in <module>
    inference()
  File "inference_image.py", line 70, in inference
    outputs = arch.postprocess(arch.model(arch.preprocess(input_image)))
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/emui/sandbox-git/SimDeblur/simdeblur/model/backbone/dblrnet/dblrnet.py", line 52, in forward
    l2 = self.L_in(x)
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/container.py", line 204, in forward
    input = module(input)
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1190, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 613, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/emui/anaconda3/envs/simdeblur/lib/python3.7/site-packages/torch/nn/modules/conv.py", line 609, in _conv_forward
    input, weight, bias, self.stride, self.padding, self.dilation, self.groups
RuntimeError: Calculated padded input size per channel: (1 x 722 x 1282). Kernel size: (3 x 3 x 3). Kernel size can't be greater than actual input size

Do you know what could be causing the issue here? It would be really nice if you could provide step-by-step examples on how to use the scripts to deblur images/videos along with test inputs and expected output, so that we know we have properly setup the SimDeblur on our machines locally.

P.S. It would be useful to also add in the README.md on the instructions of how to install all the required dependencies for SimDeblur, e.g. the Python packages and the CUDA AI libraries and toolkit on Linux. Thanks again for the great work. 👍

frames folder example results in black images

Hi, and thank you for making this code available!

I am running:

python video_infer.py D:\\DeBlur\\SimDeblur-main\\SimDeblur-main\\configs\\dbn\\dbn_gopro.yaml D:\\DeBlur\\SimDeblur-main\\SimDeblur-main\\saves\\checkpoints\\DBN\\dbn_ckpt.pth --frames_folder_path=datasets/input --save_path=deblur

where video_infer.py is:


import os
import sys
import argparse
import numpy as np
import cv2
import torch
import torch.nn as nn
import torch.nn.functional as F
from easydict import EasyDict as edict

from simdeblur.config import build_config
from simdeblur.model import build_backbone
from simdeblur.dataset.frames_folder import FramesFolder

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("config", type=str, help="The config .yaml file of deblurring model. ")
    parser.add_argument("ckpt", type=str, help="The trained checkpoint of the selected deblurring model. ")
    parser.add_argument("--frames_folder_path", type=str, help="The video frames folder path. ")
    parser.add_argument("--save_path", type=str, help="The output deblurred path")

    args = parser.parse_args()

    return args


def frames_foler_demo():
    args = parse_args()
    config = build_config(args.config)
    config.args = args
    model = build_backbone(config.model).cuda()

    ckpt = torch.load(args.ckpt, map_location="cuda:0")

    model_ckpt = ckpt["model"]
    model_ckpt = {k[7:]: v for k, v in model_ckpt.items()}
    model.load_state_dict(model_ckpt)

    data_config = edict({
        "root_input": "D:\\DeBlur\\SimDeblur-main\\SimDeblur-main\\datasets\\input",
        "num_frames": 5,
        "overlapping": True,
        "sampling": "n_c"
    })
    frames_data = FramesFolder(data_config)
    frames_dataloader = torch.utils.data.DataLoader(frames_data, 1)

    model.eval()
    with torch.no_grad():
        for i, batch_data in enumerate(frames_dataloader):
            out = model(batch_data["input_frames"].cuda())
            print(batch_data["gt_names"], out.shape)
            save_image_path = f"./deblurred_frame{i}.png"
            cv2.imwrite(save_image_path, out[0].cpu().permute(1, 2, 0).numpy())


if __name__ == "__main__":
    frames_foler_demo()

The code runs, and prints:

Cannot inport EDVR modules!!!
[('00002.jpg',)] torch.Size([1, 3, 1080, 1920])
[('00003.jpg',)] torch.Size([1, 3, 1080, 1920])
....

But all the resulting frames are just solid black. What might be happening here?

Thanks!

Unable to execute inference_image script

Hi,
Great repository, just loved it. I am trying to execute inference_image.py with the dbn architecture as the backbone but I'm getting stuck on this:
Executing the script produces the following:

Cannot inport EDVR modules!!!
Cannot import STFAN modules!!!
Using checkpoint loaded from ./checkpoints/dbn_ckpt.pth for testing.
Traceback (most recent call last):
File "inference_image.py", line 81, in
inference()
File "inference_image.py", line 70, in inference
outputs = arch.postprocess(arch.model(arch.preprocess(input_image)))
File "/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/content/drive/MyDrive/SimDeblur/simdeblur/model/backbone/dbn/dbn.py", line 137, in forward
central_frame = x[:, self.num_frames // 2]
IndexError: index 2 is out of bounds for dimension 1 with size 1

I tried printing the dimensions of the image and it came out to be: torch.Size([1, 1, 3, 385, 1504])
and self.num_frames is 5 for my case.
Don't know how to resolve the issue. Please help out.

Have you trained the model with your own script?

Hello, thanks for your contributions.
Have you trained the model with your own script, or you just copied or wrote the code for that method without training by yourself?
If not, how to verify the correctness of your tool?
Thanks.

AttributeError: 'EasyDict' object has no attribute 'model'

Here is my argument: python inference_image.py configs/base_dvd.yaml demo/dbn_ckpt.pth --img=2016-3A-00000.png --save_path=inference_results

Traceback (most recent call last):
File "C:\Users\Test\Documents\Image_Scripts\SimDeblur\inference_image.py", line 81, in
inference()
File "C:\Users\Test\Documents\Image_Scripts\SimDeblur\inference_image.py", line 46, in inference
arch = build_meta_arch(cfg)
File "C:\Users\Test\Documents\Image_Scripts\SimDeblur\simdeblur\model\build.py", line 38, in build_meta_arch
ret = META_ARCH_REGISTRY.get(name)(cfg)
File "C:\Users\Test\Documents\Image_Scripts\SimDeblur\simdeblur\model\meta_arch\plain_cnn.py", line 40, in init
self.model = self.build_model(cfg).to(self.device)
File "C:\Users\Test\Documents\Image_Scripts\SimDeblur\simdeblur\model\meta_arch\base_arch.py", line 76, in build_model
model = build_backbone(cfg.model)
AttributeError: 'EasyDict' object has no attribute 'model'

Colab Questions

Thank you so much for sharing the this!
Forgive me, because I'm not super technical but I was hoping you could walk me through how to use the Colab notebook to test my own image sequence. When I "run all" in Colab it shows a single output comparison jpeg. Does it save the entire deblurred image sequence anywhere? All I can see are the input images when I search the folder structure. How would I deblur and entire image sequence and save the resulting images? Thanks in advance!

Integrating blur decomposition works and real-world benchmarks

Hi, great work!

Any interests to integrate blur decomposition works into this general framework, such as:
Learning to Extract a Video Sequence from a Single Motion-Blurred Image (CVPR 2018) (Image)
Bringing Alive Blurred Moments (CVPR 2019) (Image)
Learning to Extract Flawless Slow Motion From Blurry Videos (CVPR 2019) (Video)
Blurry Video Frame Interpolation (CVPR 2020) (Video)

Also, real-world benchmarks such as:
RealBlur (Image)
BSD (Video)

If you would like to do it, maybe I can offer some help :)

Email: [email protected]
WeChat: zzh-tech

Checkpoints for ESTRNN

Are there plans to add pre-trained model for ESTRNN? The ones provided on the official ESTRNN github produce artifacts in my tests so far. Would be interested to know if there are plans to add models trained of REDS or GoPro.

Comparison of results

Hello, I have three questions I would like to consult, 1. After trying some models, I feel that the effect of video deblurring is not very good, are there any improvement measures? 2. Can't reason about your own video directly? I want to enter a video path directly? 3. Compare the deblurring effect of nafnet and the source code in the framework, I feel that the effect of the source code is much better than the effect in the framework, have you done a test comparison?

No object names xxx found in backbone registry!

I cannot run mscnn and srn model with the error:

Traceback (most recent call last):
  File "train.py", line 28, in <module>
    main()
  File "train.py", line 23, in main
    trainer = Trainer(cfg)
  File "/deblur/SimDeblur/simdeblur/engine/trainer.py", line 69, in __init__
    self.model = self.build_model(cfg).to(self.device)
  File "/deblur/SimDeblur/simdeblur/engine/trainer.py", line 305, in build_model
    model = build_backbone(cfg.model)
  File "/SimDeblur/simdeblur/model/build.py", line 26, in build_backbone
    return build(cfg, BACKBONE_REGISTRY)
  File "/SimDeblur/simdeblur/model/build.py", line 21, in build
    ret = registry.get(name)(**args)
  File "/SimDeblur/utils/registry.py", line 34, in get
    raise KeyError("No object names {} found in {} registry!".format(name, self._name))
KeyError: 'No object names MSResNet found in backbone registry!'

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.