Based on the improvements of Katherine Crowson (Neural style transfer in PyTorch) on the paper A Neural Algorithm of Artistic Style, This repository brings some changes to enhance the style transfer results on images that contains human faces.
Figure 1: Stylization Diagram.
First we pass the original image to the MODNet Human Segmentation Model to generate an alpha layer and remove the background. Then the image is forwarded through the FaceNet, FaceMesh and VGG19 Models for feature extraction. The VGG19 Model is also used to extract the features from desired style images. And finally, all the extracted features are joined to create the Content, Facial and Style losses.
Using FaceNet Inception Resnet Model (Tim Esler), with VGGFace2 Pretrained weights, We can implement a FaceID Loss, comparing the internal representations of the model from content image and result image with MSE (Mean Squared Error), like in the VGG Model Content Loss.
The FaceID Loss weight can be controlled through face_weight
argument.
The FaceMesh Loss works like the FaceID Loss, but it uses only the last layer output from FaceMesh Model (George Grigorev) that represents the Facial 3D Meshes. Use only for fine adjustments on relevant expression attributes like mouth opening.
The FaceMesh Loss weight can be controlled through face_mesh
argument.
The BackgroundRemoval
class uses MODNet: Trimap-Free Portrait Matting in Real Time
as backend, and uses human matting pretrained weights provided by the author.
You can remove the background from input image for better results, but it's not necessary.
The MODNet Model can be used through the remove_bg.py
script and the BackgroundRemoval
class.
The content_weight
parameter focus on general image content,
while the face_weight
parameter focus on specific facial attributes.
mesh_weight
is helpful when adjusting finer details on result faces.
Note: crop_faces
parameter is set to False in all the experiments below.
Here you can find instructions to install the project through conda. We'll create a new environment, and install the required dependencies.
First, clone the repository locally:
git clone https://github.com/thiagoambiel/PortraitStylization.git
cd PortraitStylization
Create a new conda environment called portrait_stylization
:
conda create -n portrait_stylization python=3.7
conda activate portrait_stylization
Now install the project dependencies:
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install -r requirements.txt
You can use the CLI tools remove_bg.py
and stylize.py
or import the classes
BackgroundRemoval
and StyleTransfer
. Both methods will download the
VGG19 Weights (548 MB) and FaceNet VGGFace2 Weights (107 MB)
at the first run.
Input images will be converted to sRGB when loaded, and output images have the sRGB color space. Alpha channels in the inputs will be ignored.
The StyleTransfer
and BackgroundRemoval
classes can be used on an interactive
python session or in a common script.
from PIL import Image
from style_transfer import StyleTransfer
from remove_bg import BackgroundRemoval
# Load content image.
original_image = Image.open("content.jpg")
# Load MODNet and remove content image background.
background_removal = BackgroundRemoval("./weights/modnet.pth", device="cpu")
content_image = background_removal.remove_background(
img=original_image,
bg_color="black",
)
# Load style images.
style_images = [
Image.open("style_1.jpg"),
Image.open("style_2.jpg"),
]
# Load and run style transfer module.
st = StyleTransfer(device="cpu", pooling="max")
result_image = st.stylize(
content_image=content_image,
style_images=style_images,
content_weight=0.05,
face_weight=0.25,
)
# Save result to disk.
result_image.save("out.png")
python remove_bg.py content.jpg output.jpg
-
content.jpg
: The input image to remove the background. -
output.jpg
: The path to save the result image.
You can run python remove_bg.py --help
for more info.
python stylize.py input.jpg style_1.jpg style_2.jpg -o out.png -cw 0.05 -fw 0.25
-
input.jpg
: The input image to be stylized. -
style_N.jpg
: The style images that will be used to stylize the content image, need to be at least one. -
-o out.png
(--save-path
): The path to save the result image. -
-cw 0.05
(--content-weight
): The Content Loss weight. It defines how similar the content of result image will be when compared with the content image. -
-fw 0.25
(--face-weight
): The FaceID Loss weight. It defines how similar the detected faces in result image will be when compared with the content image.
You can run python stylize.py --help
for more info.
Note: It also works with multiple faces in the same image.
Content | Style | Result |
---|---|---|
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
![]() |
Thanks to the authors of these amazing projects.
- https://github.com/crowsonkb/style-transfer-pytorch
- https://github.com/timesler/facenet-pytorch
- https://github.com/thepowerfuldeez/facemesh.pytorch
- https://github.com/ZHKKKe/MODNet
PortraitStylization is released under the MIT license. Please see the LICENSE file for more information.