Kolejna odsłona style transfer

Technologia style transfer, pozwalająca na przeniesienie stylu jednego obrazu na inny wywołała kilka lat temu spore zamieszanie. Dla wielu osób była to pierwsza okazja kiedy usłyszały o Deep Learning. Dla innych było to odkrycie niesamowitych możliwości nowoczesnych maszyn i algorytmów. Niestety ma on wiele wad. Do najpoważniejszych można zaliczyć ogromną złożoność obliczeniową – obróbka pojedynczego zdjęcia wymaga kilku-kilkunastu minut obliczeń z wykorzystaniem najszybszych kart graficznych. Innym problemem jest częste powstawanie artefaktów i tworzenie niespójnych obrazów w których elementy stylu są przenoszone błędnie na niewłaściwe elementy obrazu (np. liście drzew na budynki,  błękit nieba na drogi itp.)

Kilka dni temu badacze z Uniwersytetu Kalifornijskiego z Merced oraz Nvidii opublikowali najnowsze wyniki swoich badań wraz z kodem źródłowym, w których  pokazali jak można przyśpieszyć technologię style transfer 60-cio krotnie i jednocześnie generować znacznie bardziej spójne obrazy. FastPhotoStyle, bo tak nazwali swoje oprogramowanie, w przeciwieństwie do istniejących algorytmów nie wymaga wielu iteracji a cały proces sprowadza się do 3 krótkich faz: stylizacji, propagacji i post-processingu. Badacze przeprowadzili też ankietę w której porównali opinie na temat wyników uzyskiwanych przez własny system oraz trzy inne, czołowe rozwiązania, w której ponad połowa respondentów uznała, że zarówno realizm jak i stopień stylizacji osiągnięty przez badaczy z Merced najbardziej im odpowiada.

W dalszej części artykułu pokażę jak możemy samodzielnie zreplikować zaprezentowane wyniki. Niezbędny będzie do tego komputer z GPU Nvidii (jeśli nie mamy takiego w domu możemy wykorzystać cloud, np. Google Compute, gdzie ceny K80 wynoszą 0.45USD a nawet 0.2USD za godzinę przy minutowym rozliczaniu!) posiadający minimum 16GB RAM.

Od dłuższego czasu do wszystkich prac wykorzystuję obrazy Dockera, dzięki czemu nie ma problemu z reprodukowaniem wyników oraz zaśmiecaniem komputera kolejnymi programami i bibliotekami, często w konfliktujących ze sobą wersjach. Tak samo będzie tym razem. Obraz który przygotowałem zawiera wszystkie niezbędne komponenty zgodnie z instrukcją przygotowaną przez autorów FastPhotoStyle. Nawiasem mówiąc instrukcja jest niezbyt dokładna i zawiera drobne błędy, np. zaleca wykorzystanie pythona2 oraz instalację cuda-9.0 podczas gdy w kodzie źródłowym znajdują się ścieżki do cuda-9.1, użycie obrazu pozwoli nam uniknąć takich pułapek. Dockerfile z dokładnym zapisem wszystkich niezbędnych kroków można znaleźć tutaj, zawiera on też kilka patchy niezbędnych do prawidłowego działania, więc jeśli chcesz odtworzyć środowisko samodzielnie koniecznie do niego zajrzyj. Do pobrania gotowego obrazu wystarczy tylko wydać polecenie:

docker pull emsi/deep-learning:image-stylization

Obraz jest wielkości 3GB więc jego pobranie może chwilę zająć. W tym czasie, o ile jeszcze nie mamy zainstalowanego nvidia-docker, możemy zapoznać się z jego instrukcją instalacji oraz upewnić się że posiadamy niezbędne drivery.

Jeśli mamy już niezbędne komponenty możemy pobrać źródła FastPhotoStyle. Zalecam pobranie źródeł z mojego repozytorium gdzie znajdują się już drobne poprawki oraz modyfikacje zapewniające bezproblemowe działanie:

git clone https://github.com/emsi/FastPhotoStyle
cd FastPhotoStyle

Ostatnim elementem układanki jest pobranie gotowych modeli przygotowanych przez autorów pracy, które należy rozpakować do katalog models w podkatalogu FastPhotoStyle. Zawartość katalogu powinna wyglądać następująco:

# ls -la models/
total 317816
drwxr-xr-x 2 root root      4096 Feb 21 21:44 .
drwxr-xr-x 7 root root      4096 Feb 24 17:00 ..
-rw-rw-r-- 1 root root     15525 Feb  2 08:32 feature_invertor_conv1_1_mask.t7
-rw-rw-r-- 1 root root   9294251 Feb  2 08:32 feature_invertor_conv2_1_mask.t7
-rw-rw-r-- 1 root root  17033905 Feb  2 08:32 feature_invertor_conv3_1_mask.t7
-rw-rw-r-- 1 root root  42739881 Feb  2 08:32 feature_invertor_conv4_1_mask.t7
-rw-rw-r-- 1 root root 119309985 Feb  2 08:32 feature_invertor_conv5_1_mask.t7
-rw-rw-r-- 1 root root     16949 Feb  2 08:32 vgg_normalised_conv1_1_mask.t7
-rw-rw-r-- 1 root root    906559 Feb  2 08:32 vgg_normalised_conv2_1_mask.t7
-rw-rw-r-- 1 root root   4451913 Feb  2 08:32 vgg_normalised_conv3_1_mask.t7
-rw-rw-r-- 1 root root  28061323 Feb  2 08:32 vgg_normalised_conv4_1_mask.t7
-rw-rw-r-- 1 root root 103581389 Feb  2 08:32 vgg_normalised_conv5_1_mask.t7

Teraz już możemy uruchomić nasz kontener. Musimy tylko pamiętać o zamontowaniu katalogu ze źródłami FastPhotoStyle jako volumen. Zakładając że znajdujmy się w katalogu ze źródłami FastPhotoStyle wystarczy wydać komendę:

nvidia-docker run --rm -ti -h FastPhotoStyle -v $(pwd):/FastPhotoStyle emsi/deep-learning:image-stylization /bin/bash

W jej wyniku powinniśmy otrzymać prompt w kontenerze. Teraz wystarczy już tylko przejść do odpowiedniego katalogu i sprawdzić czy demo działa prawidłowo:

root@FastPhotoStyle:/FastPhotoStyle# cd
root@FastPhotoStyle:~# cd /FastPhotoStyle/
root@FastPhotoStyle:/FastPhotoStyle# ./demo.py 
Content image ./images/content/in81.png
Style image ./images/style/in81.png
Content map ./images/contentSeg/in81.png
Style map ./images/styleSeg/in81.png
Output image ./results/in81.png
Elapsed time in stylization: 4.529068
Elapsed time in propagation: 17.120014
Elapsed time in post processing: 0.569869
root@FastPhotoStyle:/FastPhotoStyle#

Jak widać całość procesu trwała nieco ponad 20 sekund a w jego wyniku utworzony został plik ./results/in81.png o mniej więcej poniższej zawartości:

Jest on oczywiście efektem nałożenia na lewy z poniższych obrazków stylu z prawego poniższego zdjęcia.

Dysponując działającym narzędziem możemy przeprowadzić podobną operację z użyciem własnych obrazków, a także wykorzystać maski do wskazania które partie poszczególnych zdjęć odpowiadają sobie nawzajem. Służą do tego odpowiednie argumenty skryptu:

./demo.py --help
usage: demo.py [-h] [--vgg1 VGG1] [--vgg2 VGG2] [--vgg3 VGG3] [--vgg4 VGG4]
               [--vgg5 VGG5] [--decoder5 DECODER5] [--decoder4 DECODER4]
               [--decoder3 DECODER3] [--decoder2 DECODER2]
               [--decoder1 DECODER1] [--content-image CONTENT_IMAGE]
               [--style-image STYLE_IMAGE] [--content-map CONTENT_MAP]
               [--style-map STYLE_MAP] [--output-image OUTPUT_IMAGE]

Photorealistic Image Stylization

optional arguments:
  -h, --help            show this help message and exit
  --vgg1 VGG1           Path to the VGG conv1_1
  --vgg2 VGG2           Path to the VGG conv2_1
  --vgg3 VGG3           Path to the VGG conv3_1
  --vgg4 VGG4           Path to the VGG conv4_1
  --vgg5 VGG5           Path to the VGG conv5_1
  --decoder5 DECODER5   Path to the decoder5
  --decoder4 DECODER4   Path to the decoder4
  --decoder3 DECODER3   Path to the decoder3
  --decoder2 DECODER2   Path to the decoder2
  --decoder1 DECODER1   Path to the decoder1
  --content-image CONTENT_IMAGE
                        Path to content image
  --style-image STYLE_IMAGE
                        Path to style image
  --content-map CONTENT_MAP
                        Path to content image
  --style-map STYLE_MAP
                        Path to style image
  --output-image OUTPUT_IMAGE
                        Path to style image

Przykładowo w celu uzyskania odwrotnego efektu, czyli nałożenia letniego stylu na zimowe zdjęcie należy wydać komendę:

./demo.py --content-image ./images/style/in81.png --content-map ./images/styleSeg/in81.png --style-image ./images/content/in81.png --style-map ./images/contentSeg/in81.png --output-image ./results/in81-reverse.png

Efekt jego działania wygląda następująco:

Do tworzenia odpowiednich masek możemy wykorzystać labelme, jego autor przygotował odpowiedni obraz dockera który możemy uruchomić z katalogu FastPhotoStyle poprzez wydanie komendy:

docker run -ti -v $(pwd):/FastPhotoStyle -v /tmp/.X11-unix:/tmp/.X11-unix -e DISPLAY=:0 -e QT_X11_NO_MITSHM=1 wkentaro/labelme /bin/bash

Następnie sam labelme możemy uruchomić po prostu poprzez wpisanie z konsoli labelme. Analogicznie wygląda opisana w instrukcji FastPhotoStyle kwestia konwersji plików .json zawierających opis utworzonych przez nas masek do plików .png. Musimy tylko pamiętać aby samo uruchomienie labelme_json_to_dataset wykonywać z kontenera labelme a wyniki naszych działań zapisywać do zamontowanego z hosta volumenu /FastPhotStyle.

Na koniec przedstawię jeszcze jeden efekt działania FastPhotoStyle na zdjęciu gabloty zawierającej makietę bitwy z czasów Drugiej Wojny Światowej z nałożonym stylem obrazu Mazurovskiego:

p.s. W artykule jako obraz nagłówkowy wykorzystano zdjęcie z pracy Adobe: https://arxiv.org/abs/1703.07511, https://github.com/luanfujun/deep-photo-styletransfer.