A neural network writen in PyTorch to unfog an image or stbilize the colour balance by unhazing.
The Project tries to replicate the architecture outlined in the project AOD-Net.
AOD-Net stands for All in One Dehazing Network which is a lightweight image dehazing network that is fast to train and using for inference.
Preparation:
- Create folder "data".
- Create 2 subfolders, "images" and "data" with original images and foggy images respectively.
- Install the requirements of the project by typing
sudo pip3 install -r requirements.txt
Training:
- To traning the model run:
python3 train.py
-
Open train.py file for details of the parameters and parser arguments.
-
Random validation samples will be saved to "samples" directory after every epoch and the model will be saved as "net.pth" in the "snapshots" directory. A pretrained model has been provided in the snapshots directory.
- Copy the images to be unfogged in the "test_images" directory and run
python3 main.py
The net will save the unfogged images in the "results" directory.
Results:
The models performs well in removing the fog and haziness in the foreground of the input images. for removing complete fog it would require some sort of Generative nets to replace the foggy patches where nothing is seen.
Some of the results are:
