/AODnet-by-pytorch

Image Dehaze, Pytorch, An All-in-One Network for Dehazing, AOD-Net

Primary LanguagePython

AOD-Net by Pytorch

This is an implementation of AOD-Net : All-in-One Network for Dehazing on Python3, Pytorch. The model can removal hazy, smoke or even water impurities.

The repository includes:

  • Source code of AOD-Net
  • Building code for synthesized hazy images based on NYU Depth V2
  • Training code for our hazy dataset
  • Pre-trained model for AOD-Net

Requirements

Python 3.6, Pytorch 0.4.0 and other common packages

NYU Depth V2

To build synthetic hazy dataset, you'll also need:

Training Part

Dateset Setup

  1. Clone this repository
  2. Create dataset from the repository root directory 
    $ cd make_dataset
$ python create_train.py --nyu {Your NYU Depth V2 path} --dataset {Your trainset path}
  3. Random pick 3,169 pictures as validation set 
    $ python random_select.py --traindir {Your trainset path} --valdir {Your valset path}

Start to training

  1. training AOD-Net 
    $ python train.py --dataroot {Your trainset path} --valDataroot {Your valset path} --cuda

Testing Part

  1. test hazy image on AOD-Net 
    $ python test.py --input_image /test/canyon1.jpg  --model /model_pretrained/AOD_net_epoch_relu_10.pth --output_filename /result/canyon1_dehaze.jpg --cuda