Generate Portrait on browser with U²-Net

This time we will use U²-Net to generate a portrait. This is the one that was buzzing a while ago.

It seems that several methods have been proposed for generating portraits using machine learning, including APDrawingGAN, which is referenced by U²-Net. Since portrait generation is one of the image generation methods, I kind of think of it as an extension of GAN. In other words, as APDrawingGAN does, it creates a portrait by fusing an image reconstructed from the features of the whole face, such as where the eyes and nose are located, and an image reconstructed from the individual features of the eyes and nose. However, U²-Net takes a completely different approach: U²-Net actually comes out of research in the field of Salient Object Detection, which is the separation of the object of interest in an image from the background. Therefore, there is no mention of portraits in the main body of the paper at all.

If you refer to the U²-Net repository, you’ll see that the authors found the new way to use their output just after submitting paper. They say, “If we use portraits drawn by imitating photos as teacher data, we can create a model for segmentation. I have both thinkings “Really?” and “maybe that’s true”. But it’s actually working, so I can only say, “Wow!”. However, it was possible because of the U²-Net network, which is able to capture the multi-scale and multi-layer features, and I think it would not work well with segmentation models that use other backends. (If you want to know a little more about the architecture and other details of U²-Net, please read the paper.)

In this article, I would like to try to run a function to generate portraits in a browser using such an interesting model with an unexpected background.

U²-Net is currently only available as a PyTorch implementation. So in order to run it in Tensorflowjs, we need to convert the model. The method is detailed in PINTO’s article on converting PyTorch models to Tensorflowjs. Basically, I’m not going to explain much here because I’m converting according to this article by PINTO. Please refer to the article by PINTO. In the previous article, I explained how to run the conversion process in docker, so please refer to that article as well. After conversion, incorporate it into Tensorflowjs and you are done.

In this section, we will assume that you have already built a docker environment referring to the previous article.

(1) Clone the U²-Net repository from github and download the trained model.

$ git clone https://github.com/NathanUA/U-2-Net.git
$ cd U-2-Net/
$ mkdir ./saved_models/u2net_portrait/
$ curl -sc /tmp/cookie "https://drive.google.com/uc?export=download&id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW" > /dev/null
$ CODE="$(awk '/_warning_/ {print $NF}' /tmp/cookie)"  
$ curl -Lb /tmp/cookie "https://drive.google.com/uc?export=download&confirm=${CODE}&id=1IG3HdpcRiDoWNookbncQjeaPN28t90yW" -o saved_models/u2net_portrait/u2net_portrait.pth

(2) Next, we will perform the transformation process. The model name --model-name U2Net and the file name of the trained model —weights saved_models/u2net_portrait.pth are different, but the basic process is the same as the one for U²-Net described in the previous article. The output file name of pytorch_to_onnx.py is the same as in the previous article, so please refer to the previous article for further processing. However, I think that an error occurs when openvino2tensorflow is executed and SavedModel cannot be created. This is probably because SavedModel should be less than 2GB (issue). A frozen model will be generated, so please convert this model to tensorflowjs. The total file size is over 150MB.

$ export PYTHONPATH=/U-2-Net
$ SIZE=512
$ python3 /opt/intel/openvino_2021/deployment_tools/tools/model_downloader/pytorch_to_onnx.py  
--import-module model.u2net   
--model-name U2NET   
--input-shape 1,3,${SIZE},${SIZE}   
--weights saved_models/u2net_portrait/u2net_portrait.pth   
--output-file u2netp_${SIZE}x${SIZE}.onnx   --input-names "x"   
--output-names "a/F.sigmoid(d0)"

(3) Once the conversion is complete, you can place it on your website and load it with Tensorflowjs for use. The basic usage is the same as for other models. Please note that normalization and thresholding of input and output are required to generate images with high accuracy. Please refer to the source code available in the repository below for the thresholds I use.

Now let’s use it to generate a portrait. As usual, I will vary the size of the input image to see the processing time and quality. Note that we are experimenting here on a Linux PC with GeForce GTX 1660 & Chrome.

This is what it looks like at 320×320. It’s a pretty solid portrait. Processing time is about 1.5 seconds per frame on the GPU. It’s heavy.