We have provided the Pytorch Implementation for our course project Image Segmentation Enhanced Style Transfer. We have used Image Segmentation to help us improve the performance of conventional Style Transfer, especially for problem that the elements within the picture will interfere with each other.
The Implementation is based on CycleGAN Pytorch. We have added the part in Image Segmentation and several new files under the root repository, which is the key element to our image segmentation enhancing.
The poster of our project is in ECCG.
The process of our proposed ECCG contains two essential parts:
(1) Image Segmentation. We directly used the pretrained model from FastFCN You can download the models from checkpoint and then put them into the folder checkpoint
.
(2) Style Transfer for each segmentation. For simplicity, we used the pretrained model horse2zebra_pretrained
and winter2summer_yosemite_pretrained
from CycleGAN Pytorch. These models should also be put in the folder checkpoint
. You can download the checkpoint using the following command(from CycleGAN).
bash ./scripts/download_cyclegan_model.sh horse2zebra
- The pretrained model is saved at
./checkpoints/{name}_pretrained/latest_net_G.pth
. Check here for all the available CycleGAN models.
After downloading these two models, you can directly run the following command:
python ECCG.py --image $path_to_image --target_path $path_to_results
After the program is finished, there will be six images in the target_path
folder:
File_name | meaning |
---|---|
xxx_fakehorse2zebra_pretrained.png | Baseline1: fake picture from pretrained model 1 |
xxx_fakewinter2summer_yosemite_pretrained.png | Baseline2: fake picture from pretrained model 2 |
xxx_fake0.png | ECCG result 1 with animal from pretrained model 1 |
xxx_fake1.png | ECCG result 2 with animal from pretrained model 1 and background from pretrained model 2 |
xxx_fakepoisson0.png | ECCG result 3: poisson editing on xxx_fake0.png |
xxx_fakepoisson1.png | ECCG result 4: poisson editing on xxx_fake1.png |
Besides, you can also use L0-smooth (from the paper: Image Smoothing via
python main.py --image $path_to_image --target_path $path_to_results --with_L0
As we can see in the results, our proposed methods have significantly improve the performance of CycleGAN, especially in the elements apart from the main animal.