bigwangyudong / uwenhancement Goto Github PK

Python 97.87% MATLAB 2.13%

uwenhancement's Introduction

Underwater Image Enhancement Toolbox

This is an open source underwater image enhancement toolbox based on PyTorch. Meanwhile, this is the code of the implementation of the UIEC^2-Net. If you use our code, please consider citing our paper. Thanks.

Documentation: https://uwenhancement.readthedocs.io/

Introduction

This is a code that summarizes deep learning-based underwater image enhancement. We converted relevant Tensorflow code to PyTorch, and transcoded the Tensorflow checkpoint provided by the author so that it can be read by PyTorch. We confirmed the reliability of transcoding, and uploaded the PyTorch checkpoint models(.pth file) to Google Drive and Baidu Cloud. (You can download at Model Zoo)

This is convenient for everyone to use the same environment to run experiments. We currently support UIEC^2-Net, UWCNN, WaterNet, and we will continue to update. (More details in Model Zoo)

Highlights

TTA x8 (Test Time Argument x8 images)
fp16
Visualize Network Structure
Training process monitoring (Use Visdom package)
Training log records, and loss visualize(Use TensorboardX package)
Tensorflow to PyTorch.
...

You can get more details at useful_tools.

TODO:

We will support GAN soon.
In the future we will support PIxel-Level related tasks, such as Dehazing, Salient Object Detection(SOD) etc.

Installation

Please refer getting_start.md for installation.

Prepare Data

Using write_txt.py prepare train.txt and test.txt. All datasets should put in DATA folders, the directory structure is:

.
├── ...
├── DATA
│   ├── Train
│   │   ├── train
│   │   ├── gt
│   │   └── train.txt
│   ├── Test
│   │   ├── test
│   │   ├── gt(if need)
│   │   └── test.txt
└── ...

If you want to run WaterNet, you should prepare GC, WB, HE images and put in DATA/Train/WaterNet and DATA/Train/WaterNet, the directory structure is:

.
├── WaterNet
│    ├── ce
│    ├── gc
│    └── wb

Prepare Training and Testing

Clone the repo.

git clone https://github.com/BIGWangYuDong/UWEnhancement.git UW

put the checkpoint(*.pth) in checkpoints/
change args.config and args.load_from root path.
change config/XXX.py data root path, and some other important options.

If you want to use our toolbox to add some New Models, Data Pipelines, loss functions or some other operation, We provide Tutorials.

Testing

python test.py
# or 
python test_UWModels/test_XXX.py

Training

python -m visdom.server			# open this in a single terminal
python train.py --gpus XX 		# XX means use numbers of gpus

Contact Us

If you have any questions, please contact us ([email protected] or [email protected]).

Bibliography entry for citation:

@article{wang2021uiec,
title = {{UIEC\^{}2-Net}: CNN-based underwater image enhancement using two color space},
journal = {Signal Processing: Image Communication},
volume = {96},
pages = {116250},
year = {2021},
issn = {0923-5965},
doi = {https://doi.org/10.1016/j.image.2021.116250},
url = {https://www.sciencedirect.com/science/article/pii/S0923596521001004},
author = {Yudong Wang and Jichang Guo and Huan Gao and Huihui Yue},
}

uwenhancement's People

Contributors

Stargazers

Watchers

uwenhancement's Issues

Using test_WaterNet.py for evaluating UIEB dataset

Hello, when I used the test_WaterNet.py for evaluating UIEB dataset on google colab, I met a trouble.

The info seems like told me I had not read the dataset into the DataLoader, however, I had changed the important configs as you mentioned in README.md:

test_UWModels/test_WaterNet.py:

parser.add_argument('--config',type=str,
default='/content/drive/MyDrive/UWEnhancement-master/config/WaterNet.py',
help='train config file path')

parser.add_argument('--load_from',
default='/content/drive/MyDrive/UWEnhancement-master/checkpoints/WaterNet.pth',
help='the dir to save logs and models,')

config/WaterNet.py

data_root_test = '/content/drive/MyDrive/UWEnhancement-master/DATA/Test/'

test_pipeling = [dict(type='LoadImageFromFile', gt_type='color', get_gt=False),
dict(type='LoadWaterNetImage', path='/content/drive/MyDrive/UWEnhancement-master/DATA/Test/WaterNet/')

test=dict( # load data in test process
type=dataset_type,
ann_file=data_root_test + test_ann_file_path,
img_prefix=data_root_test + 'WaterNet/',
gt_prefix=data_root_test + 'gt/'

and my file trees are shown above:

Moreover, the code of reading dataset is well encapsulated so that I have no idea how to solve this trouble. If you could do me a favor, it will be appreciated!

delete developed codes
add trainer

where is importlib._setup()

First of all, thank you very much for sharing! When I run train.py, it has the error shown in the image below.And I tried to run _bootstrap_external.py.But I can't find importlib._setup()