Awesome
Big GANs Are Watching You
This repository is the official implementation of Big GANs Are Watching You: Towards Unsupervised Object Segmentation with Off-the-Shelf Generative Models by Andrey Voynov, Stanislav Morozov, and Artem Babenko.
Schematic representation of our approach.
Core steps:
- Find the BigBiGAN latent direction responsible for background darkening and foreground lightening;
- Sample real images embeddings and generate segmentation masks with their shifts;
- Train U-net with that synthetic data.
example of variation along the latent direction
Requirements
python 3.6 or later
torch>=1.4
torchvision>=0.5
tensorboardX
scikit-image
scipy
h5py
at least 12Gb of GPU-memory for batch size 95 (though smaller batch gives competitive result)
For results visualization:
jupyter
Optional for BigBiGAN tf to torch conversion and BigBiGAN embeddings evaluation:
tensorflow_gpu==1.15.2
tensorflow_hub
parse
if troubles – check the authors packages versions in requirements.txt
Training
First run download_all.sh
this scrips downloads:
- BigBiGAN weights of the original tf-hub model (https://tfhub.dev/deepmind/bigbigan-resnet50/1) ported to pytorch;
- precalculated BigBiGAN-embeddings for some popular datasets;
- background darkening and foreground lightening latent directions coefficients.
note that the ImageNet-BigBiGAN-embeddings takes 551Mb.
To train the U-net segmentation model, run this command:
python train_segmentation.py \
--out "results" \
--gan_weights "BigGAN/weights/BigBiGAN_x1.pth" \
--z "embeddings/BigBiGAN_ImageNet_z.npy"
--bg_direction "BigGAN/weights/bg_direction.pth" \
--val_images_dir __path_to_images_dir__ \
--val_masks_dir __path_to_masks_dir__ \
Custom Data
Once you want to use your own images BigBiGAN-embeddings, please run bigbigan_embeddings.ipynb
notebook and provide your own data path.
Evaluation
To check the synthetic data / segmentation samples / model metrics, please run
inspection.ipynb
notebook.
Pre-trained Models
You can download pretrained models here:
- U-net trained on ImageNet-BigBiGAN embeddings using default parameters.
- U-net trained on CUB-200-2011-BigBiGAN embeddings using default parameters.
- U-net trained on Flowers-BigBiGAN embeddings using default parameters.
Results
Our model achieves state-of-the-art results among fully unsupervised methods on many datasets. Check Tables 1, 2, 3 of the original paper for more details. In the tables below we refer our approach as BigBiGAN and E-BigBiGAN.
Object Detection
Saliency Detection
Here are some samples of the model performing on DUTS dataset:
BigBiGAN weights conversion
Once the original BigBiGAN weigths are distributed within tfhub, we also provide the conversion code. See bigbigan2biggan.ipynb
. This code is based on this script: https://github.com/ajbrock/BigGAN-PyTorch/blob/master/TFHub/converter.py . Note that BigBiGAN and BigGAN almost share the architecture and we use the author's officially unofficial PyTorch BigGAN implementation.
Licenses
- Unsupervised Discovery of Latent Directions: https://github.com/anvoynov/GANLatentDiscovery
- BigBiGAN weights: https://tfhub.dev/deepmind/bigbigan-resnet50/1
- BigGAN pytorch: https://github.com/ajbrock/BigGAN-PyTorch
- U-Net model code is based on: https://github.com/milesial/Pytorch-UNet
- prefetch_generator: https://github.com/justheuristic/prefetch_generator