/AODnet-by-pytorch

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

Primary LanguagePython

AOD-Net by Pytorch

This project implements the AOD-Net : All-in-One Network for Dehazing for image dehazing using Python and PyTorch. The model is capable of removing haze, smoke, and water impurities from images."

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