Примечание:
Эта статья также доступна здесь. (на японском языке)
https://zenn.dev/wok/articles/0007_u2net-portrait-drawing

Во-первых

Я хочу легко запускать модели машинного обучения в браузере. Поэтому я работаю над библиотекой для запуска различных моделей в tensorflowjs (см. этот репозиторий). В этой статье я попробую запустить функцию генерации портрета из фото в браузере.

Например, вы можете сгенерировать изображение справа из изображения слева, вот так. (Как обычно, изображение слева было позаимствовано из Pakutaso. Это Кавамура-сан. Она всегда милая. )

Немного о U²-Net

На этот раз мы будем использовать U²-Net для создания портрета. Это тот, который гудел некоторое время назад.

Кажется, было предложено несколько методов создания портретов с использованием машинного обучения, в том числе APDrawingGAN, на который ссылается U²-Net. Поскольку генерация портретов — это один из методов создания изображений, я думаю о нем как о расширении GAN. Другими словами, как и в случае с APDrawingGAN, он создает портрет, объединяя изображение, реконструированное по чертам лица в целом, например по расположению глаз и носа, и изображение, реконструированное по отдельным чертам глаз и носа. Однако U²-Net использует совершенно другой подход: U²-Net фактически является результатом исследований в области обнаружения существенных объектов, то есть отделения интересующего объекта на изображении от фона. Поэтому о портретах в основной части статьи вообще не упоминается.

Если вы обратитесь к репозиторию U²-Net, вы увидите, что авторы нашли новый способ использования своих результатов сразу после отправки статьи. Они говорят: «Если мы будем использовать портреты, нарисованные путем имитации фотографий, в качестве данных учителя, мы сможем создать модель для сегментации. У меня есть оба мнения. «Правда?» и «может быть, это правда». Но это на самом деле работает, поэтому я могу сказать только «Вау!». Однако это стало возможным благодаря сети U²-Net, которая способна захватывать многомасштабные и многоуровневые функции, и я думаю, что это не будет хорошо работать с моделями сегментации, использующими другие серверные части. (Если вы хотите узнать немного больше об архитектуре и других деталях U²-Net, прочтите статью.)

В этой статье я хотел бы попробовать запустить функцию для генерации портретов в браузере, используя вот такую ​​интересную модель с неожиданным фоном.

Использование модели U²-Net в браузере

U²-Net в настоящее время доступен только в виде реализации PyTorch. Итак, чтобы запустить его в Tensorflowjs, нам нужно преобразовать модель. Метод подробно описан в статье ПИНТО о преобразовании моделей PyTorch в Tensorflowjs. В принципе, я не буду здесь много объяснять, потому что я конвертирую в соответствии с этой статьей ПИНТО. Пожалуйста, обратитесь к статье ПИНТО. В предыдущей статье я объяснил, как запустить процесс конвертации в докере, поэтому, пожалуйста, обратитесь к этой статье. После преобразования включите его в Tensorflowjs, и все готово.

Преобразовать и внедрить

В этом разделе мы будем предполагать, что вы уже построили Docker-окружение со ссылкой на предыдущую статью.

(1) Клонируйте репозиторий U²-Net с github и загрузите обученную модель.

$ 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) Далее мы выполним процесс преобразования. Имя модели --model-name U2Net и имя файла обученной модели —weights saved_models/u2net_portrait.pth различаются, но основной процесс такой же, как и для U²-Net, описанный в предыдущей статье. Имя выходного файла pytorch_to_onnx.py такое же, как и в предыдущей статье, поэтому для дальнейшей обработки обратитесь к предыдущей статье. Однако я думаю, что ошибка возникает, когда выполняется openvino2tensorflow и не может быть создана SavedModel. Вероятно, это связано с тем, что SavedModel должен быть меньше 2 ГБ (проблема). Будет создана замороженная модель, поэтому преобразуйте эту модель в tensorflowjs. Общий размер файла составляет более 150 МБ.

$ 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) После завершения преобразования вы можете разместить его на своем веб-сайте и загрузить с помощью Tensorflowjs для использования. Основное использование такое же, как и для других моделей. Обратите внимание, что нормализация и пороговое значение ввода и вывода необходимы для создания изображений с высокой точностью. Пожалуйста, обратитесь к исходному коду, доступному в репозитории ниже, для пороговых значений, которые я использую.

Бегать

Теперь воспользуемся им для создания портрета. Как обычно, я буду изменять размер входного изображения, чтобы увидеть время обработки и качество. Обратите внимание, что здесь мы экспериментируем на ПК с Linux, GeForce GTX 1660 и Chrome.

Вот как это выглядит в разрешении 320х320. Это довольно солидный портрет. Время обработки составляет около 1,5 секунды на кадр на графическом процессоре. Это тяжелый.

Пробовал делать 256х256. Это довольно сломано. Время обработки составляет около 1,1 секунды.

192x192 немного непрактично. Время обработки составляло около 0,4 секунды на кадр.

Сделайте 512x512. Первоначальное обучение, похоже, было выполнено в разрешении 512x512, так что оно действительно очень хорошо сделано. Глаза выглядят очень привлекательно. К сожалению, в моей среде графическому процессору не хватило памяти, чтобы справиться с этим, поэтому мне пришлось позволить центральному процессору справиться с этим, что занимало около 4 минут на кадр. Обработка каждого кадра процессором занимает около 4 минут.

Я также попробовал 1024x1024 для удовольствия. Это также работает очень хорошо. Впрочем, 512x512 тоже было хорошо, поэтому у людей могут быть разные мнения о том, какой из них лучше. Кстати, это тоже было сделано с CPU. Это было сделано после того, как я проснулся после того, как вздремнул. Я думаю, что это, вероятно, заняло около часа.

Это все. Качество портретов достаточно высокое, но время обработки мне кажется слишком долгим.

Репозиторий и демо

Эта демонстрация хранится в репозитории ниже. Вы также можете запустить демо на следующем сайте.

  • Репозиторий
    https://github.com/w-okada/image-analyze-workers
  • Демонстрация
    https://flect-lab-web.s3-us-west-2.amazonaws.com/P01_wokers/t09_u2net-portrait/index.html
    * Примечание о демонстрации.< br /> (1) Это очень медленно. Подождите около минуты, пока появится первое изображение.
    (2) Я настроил его для работы на процессоре с разрешением 192x192. Контроллер в правом верхнем углу позволяет переключаться на другое разрешение и графический процессор.
    (3) Установите для параметра useTFWasmBackend значение Выкл. и нажмите кнопку перезагрузки модели, чтобы войти в режим графического процессора.
    (4) Выберите разрешение. из раскрывающегося меню modelPath и нажмите кнопку перезагрузки модели, чтобы обработать изображение с этим разрешением.
    (5) Вы можете изменить изображение, которое будет вводиться из ввода. Также можно использовать камеру.
    (6) При увеличении разрешения до 512x512 и выше лучше не использовать режим GPU. Существует риск переполнения памяти.

Ну наконец то

Я попытался запустить функцию создания портрета в браузере, и мне кажется, что это немного сложно запустить в режиме реального времени даже на машине с графическим процессором. Если нагрузка такая высокая, возможно, лучше использовать сервер, а не браузер.