We perform lossy image compression and compressive sensing using a BEGAN model trained on CelebA.
We show how to smoothly vary the latent dimension of the model without retraining. Some of our compressive sensing results are bananas!
For more details, see our writeup.
See preprocess_images.py
.
python preprocess_images.py --dataset celeba --input_dir $INPUT_DIR \
--output_dir $OUTPUT_DIR
See train.py
.
python train.py --run_name $RUN_NAME \
--dataset celeba \
--latent_dim 64 \
--dataset_dir $OUTPUT_DIR \
--epochs 50 \
--n_train -1 \
--output_activ elu
The file compress.py
offers both a compress()
and an uncompress()
function.
See the docstring for an explanation of the compressed file format.
input_filename = 'bananas.jpg'
compressed_filename = 'bananas.ganz'
reconstr_filename = 'bananas_reconstr.png'
# Load the image
torch_img = load_target_image(input_filename)
# Compress
x_hat, z, psnr_gan, file_size = compress(torch_img,
output_filename=compressed_filename,
skip_linear_layer=True,
no_linear_layer=False,
compression_ratio=6,
compressive_sensing=False,
n_steps=n_steps)
# Uncompress
x_hat = uncompress(compressed_filename, reconstr_filename)
After training on faces, our method achieves a PSNR of 24.76dB on a photo of bananas! (here we use M=4900 Gaussian measurements on an image of size 128x128x3, or approximately 10% measurement ratio)