Joint network for specular highlight detection and adversarial generation of specular-free images trained with polarimetric data

Official repository for the article @ Neurocomputing Journal (Elsevier, ScienceDirect)

Citation

@article{Anwer2023,
  doi = {10.1016/j.neucom.2023.126769},
  url = {https://doi.org/10.1016/j.neucom.2023.126769},
  year = {2023},
  month = nov,
  publisher = {Elsevier {BV}},
  volume = {559},
  pages = {126769},
  author = {Atif Anwer and Samia Ainouz and Naufal M. Saad and Syed Saad Azhar Ali and Fabrice Meriaudeau},
  title = {Joint network for specular highlight detection and adversarial generation of specular-free images trained with polarimetric data},
  journal = {Neurocomputing}
}

The article can be viewed and downloaded from Neurocomputing (Elsevier ScienceDirect).

(back to top)

Introduction

This repository is the code of our paper 'Joint network for specular highlight detection and adversarial generation of specular-free images trained with polarimetric data'.

SHMGAN consists of a single generator–discriminator pair trained simultaneously using polarimetric images. The proposed GAN uses a dynamically generated attention mask based on a specularity segmentation network (SpecSeg) without requiring additional manual input. The network is able to learn the illumination variation between the four polarimetric images and a pseudo-diffuse image. Once trained, SHMGAN is able to generate specular-free image from a single RGB image as input; without requiring any additional external labels.

The developed network can be used for network training and testing. Please cite the paper if you use them and find them useful.

SHMGAN Generator with SpecSeg

Full article @ Neurocomputing Journal

Requirements

• Tested with Python 3.8, Tensorflow 2.8 on local system (Nvidia RTX 3070, 8GB VRAM, 32GB Ram and CUDA 11.8)
• Tested with Python 3.7, Tensorflow 2.8+ on Google Colab
• Uses Tensorflow Function API for a custom training loop for a multi-input GAN (see the paper for details)
• Uses the Segmentation Models library
• Training also uses Comet.ml for logging the GAN training/generation progress and generating the confusion matrices (Optional, and commented out by default).

Datasets

The datasets used in the paper can be found at the following Github Repos (property of the respective authors):

1. PSD Dataset, Wu et al.
2. SIHQ Dataset, Fu et al.

Usage

For training on your dataset, please run main.py with the required paths to the datsaet.

(back to top)

License

Distributed under the GNU General Public License v3.0 License, for research training and/or testing purposes. Pleas cite our paper if you use this code or network. For more details, please see Choose a license.

(back to top)