Official repository of "Joint RGB and Near-infrared Image Compression via Common-Unique Feature Disentanglement"

Environment

compressai==1.2.4
Pytorch==2.0.1

Training

First train the CFD module:

python train_AGNet.py

Then use code in temp.ipynb to mix model weights, and choose to train TCC module

CUDA_VISIBLE_DEVICES=0,1 python -W ignore train_union.py --dataset ~/database/worldstart/ --cuda --patch-size 512 512 --test-patch-size 512 512 -b 4 --lambda 0.0932 --exp-name exp_worldstart --save -lr 5e-5 -n 0 --save-checkpoint compress_best_loss.pth.tar --no-optim --N 192 --patience 5

Then create a new model union_best_loss.pth.tar from compress_best_loss.pth.tar using code in temp.ipynb, and train union， first froze compress part then train together

# first train with compress part forzen
CUDA_VISIBLE_DEVICES=0,1 python -W ignore train_union_restore.py --dataset ~/database/worldstart/ --cuda --patch-size 512 512 --test-patch-size 512 512 -b 4 --lambda 0.0932 --exp-name exp_worldstart --save -lr 5e-5 -n 0 --save-checkpoint union_best_loss.pth.tar --no-optim --N 192 --patience 5
# Then train together
CUDA_VISIBLE_DEVICES=0,1 python -W ignore train_union.py --dataset ~/database/worldstart/ --cuda --patch-size 512 512 --test-patch-size 512 512 -b 1 --lambda 0.0932 --exp-name exp_worldstart --save -lr 5e-5 -n 0 --save-checkpoint union_best_loss.pth.tar --no-optim --N 192 --patience 5

Testing

CUDA_VISIBLE_DEVICES=0 python eval.py --cuda --dataset ~/database/worldstart --entropy-estimation --lambda 0.0025 --N 192 --load-checkpoint exp_worldstart/lambda_0.0025/union_best_loss.pth.tar