Home

Awesome

Hierarchical-Contrastive-Learning-Corruption-Detection

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

Overall Framework

Installation

  1. Clone this repo:

    git clone https://github.com/xyfJASON/HCL.git
    cd HCL
    
  2. Create and activate a conda environment:

    conda create -n HCL python=3.10
    conda activate HCL
    
  3. 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
    
  4. Install required packages:

    pip install -r requirements.txt
    

Datasets Preparation

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

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:

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}
}