HAOWEI960421/PEN-A-Parallel-Enhancement-Network-of-Compressed-Video

Parallel Enhancement Net

• Parallel Enhancement Net
  • 0. Background
  • 1. Pre-request
    • 1.1. Environment
    • 1.2. DCNv2
  • 2. Train
  • 3. Test

0. Background

PyTorch implementation of Parallel Enhancement Net

1. Pre-request

1.1. Environment

• Ubuntu 20.04/18.04
• CUDA 10.1
• PyTorch 1.6
• Packages: tqdm, lmdb, pyyaml, opencv-python, scikit-image

Suppose that you have installed CUDA 10.1, then:

git clone --depth=1 https://github.com/RyanXingQL/STDF-PyTorch 
cd STDF-PyTorch/
conda create -n stdf python=3.7 -y
conda activate stdf
python -m pip install torch==1.6.0+cu101 torchvision==0.7.0+cu101 -f https://download.pytorch.org/whl/torch_stable.html
python -m pip install tqdm lmdb pyyaml opencv-python scikit-image

1.2. DCNv2

Build DCNv2.

cd ops/dcn/
bash build.sh

(Optional) Check if DCNv2 works.

python simple_check.py

The DCNv2 source files here is different from the open-sourced version due to incompatibility. [issue]

1.3 change dataset

Modify the make_peak_label.py to creat label file. Change the path of hdf5 file in /dataset/ChallengeDataset.py

2. Train

Change the path of training data in option_R3_ChallengeDataset.yml. Run: CUDA_VISIBLE_DEVICES=0,2 python -m torch.distributed.launch --nproc_per_node=2 --master_port=12354 train.py --opt_path option_R3_ChallengeDataset.yml

3. Test

Run: CUDA_VISIBLE_DEVICES=2 python test.py --opt_path option_R3_ChallengeDataset.yml