Hierarchical-Contrastive-Learning-Corruption-Detection

Codebase of ICCV 2023 paper "Hierarchical Contrastive Learning for Pattern-Generalizable Image Corruption Detection".

Overall Framework

Installation

Clone this repo:

git clone https://github.com/xyfJASON/HCL.git
cd HCL

Create and activate a conda environment:

conda create -n HCL python=3.10
conda activate HCL

Install pytorch from https://pytorch.org, for example:

pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url https://download.pytorch.org/whl/cu116

Install required packages:
```
pip install -r requirements.txt
```

Datasets Preparation

We use CelebA-HQ, FFHQ, ImageNet and Places365 datasets in our experiments.

CelebA-HQ: Download CelebAMask-HQ dataset from the official repo. Unzip the downloaded file. Put ./scripts/celebahq_map_index.py under the unzipped directory and run the script. It will change the name of the images to match the index in the original CelebA dataset.
FFHQ: Download FFHQ dataset from the official repo. We only need the images1024x1024 part. No post-processing is needed.
ImageNet: Download ImageNet dataset (ILSVRC 2012) from the official website. No post-processing is needed.
Places365: Download Places365-Standard from the official website. No post-processing is needed.

We use NVIDIA Irregular Mask Dataset in our experiments, which can be downloaded from the official website. To avoid heavy computation in transforming the masks during training, we use this dataset in a way similar to EdgeConnect. To do so, put ./scripts/make_irregular_dataset.py under the unzipped dataset and run the script. It will augment and split the original "testing set" into a new training split and a test split.

After downloading and processing the datasets, set the dataroot in the configuration files to your downloaded path.

Pretrained Weights

The pretrained model weights and configuration files are provided in the following links:

FFHQ (256x256): pretrained weight | config file
ImageNet (256x256): pretrained weight | config file
Places (256x256): pretrained weight | config file

Training

The training process contains two phases:

First phase training:

torchrun --nproc_per_node 4 train.py -c ./configs/xxxxx.yml -p 1

Second phase training:

torchrun --nproc_per_node 4 train.py -c ./configs/xxxxx.yml -p 2 \
    --train.pretrained /path/to/checkpoint/of/first/phase.pt \
    --train.n_steps 100000

Downstream tasks finetuning (e.g., watermark removal on LOGO-30K dataset):

torchrun --nproc_per_node 2 train.py -c ./configs/xxxxx.yml -p 2 \
    --downstream \
    --data.name LOGO-30K \
    --data.dataroot /path/to/dataroot/ \
    --train.pretrained /path/to/pretrained/checkpoint.pt

Testing

Evaluate:

torchrun --nproc_per_node 4 test.py evaluate -c /path/to/config/file.yml \
    --test.pretrained /path/to/checkpoint/of/trained/model.pt \
    --mask.dir_path /path/to/irregular/mask/dataset/test/split

Sample:

python test.py sample -c /path/to/config/file.yml \
    --test.pretrained /path/to/checkpoint/of/trained/model.pt \
    --test.n_samples {number of samples} \
    --test.save_dir /directory/to/save/the/results/ \
    --mask.dir_path /path/to/irregular/mask/dataset/test/split

Citation

If you find this work useful for your research, please cite:

@InProceedings{Feng_2023_ICCV,
    author    = {Feng, Xin and Xu, Yifeng and Lu, Guangming and Pei, Wenjie},
    title     = {Hierarchical Contrastive Learning for Pattern-Generalizable Image Corruption Detection},
    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
    month     = {October},
    year      = {2023},
    pages     = {12076-12085}
}

xyfJASON / HCL