Готови ли сте да пренесете вашия LipSync PyTorch модел на следващото ниво и да го конвертирате в CoreML модел?

Страхотен! В тази публикация в блога ще ви преведем през стъпките за преобразуване на модела и подготовката му за интегриране на приложение за iOS. Така че нека се потопим!

Първо, нека започнем с малко основна информация. Моделът, с който ще работим, се нарича Wav2Lip, който може да се използва за генериране на видеоклипове, синхронизирани с устни. Можете да намерите повече подробности и демонстрация на официалния уебсайт на адрес https://bhaasha.iiit.ac.in/lipsync/.

За да започнете, ето ръководство стъпка по стъпка за преобразуване на модела:

Стъпка 1: Клонирайте хранилището Започнете с клониране на хранилището Wav2Lip от GitHub. Можете да го намерите на адрес https://github.com/Rudrabha/Wav2Lip. Веднъж клониран, отидете до папката Wav2Lip.

Стъпка 2: Подгответе зависимостите За да сте сигурни, че всичко работи гладко, уверете се, че имате инсталирани всички необходими зависимости. Ако изпълнявате това на вашата локална машина с процесор, готови сте. Ако обаче предпочитате да използвате Google Colab, хранилището предоставя бележник Colab за тази цел.

Стъпка 3: Изтеглете предварително обучените модели Има два предварително обучени модела, които трябва да изтеглите. Първо изтеглете модела за разпознаване на лица от https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth и го запазете в папка face_detection/detection/sfd/. След това изтеглете основните предварително обучени тегла на модела от https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/Eb3LEzbfuKlJiR600lQWRxgBIY27JZg80f7V9jtMfbNDaQ?e=TBFBVW. Можете също да намерите алтернативни връзки за изтегляне на страницата GitHub.

Стъпка 4: Стартирайте модела Сега е време да стартирате модела и да генерирате демонстрационен видеоклип. Отворете терминал или команден ред и изпълнете следната команда:

python inference.py --checkpoint_path <ckpt> --face <video.mp4> --audio <an-audio-source>

Не забравяйте да замените <ckpt> с пътя към изтеглените предварително обучени тегла на модела, <video.mp4> с видеото, което искате да синхронизирате, и <an-audio-source> с аудио файл (за предпочитане във формат WAV), който е с продължителност 20 секунди или по-малко.

След като сте доволни от генерираното демонстрационно видео и сте потвърдили, че моделът работи според очакванията, е време да продължите с преобразуването на CoreML модела.

Преобразуване на CoreML модел:

Стъпка 1: Импортирайте необходимите библиотеки В горната част на файла inference.py добавете следния оператор за импортиране:

from coremltools import ct

Ако не сте инсталирали coremltools, можете да го направите, като стартирате pip install coremltools във вашия терминал или команден ред.

Стъпка 2: Намерете входните тензори в кода Във файла inference.py намерете ред #260. Това е мястото, където ще видите входните тензори, които моделът очаква, по-специално тензорите за изображение и звук. Под тези редове ще добавите кода за преобразуване на CoreML модела.

Ето пример за това как трябва да изглежда кодът за преобразуване:

img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)        
traced_model = torch.jit.trace(model, (mel_batch, img_batch))    
print("This is the traced model:")    
print(traced_model)    
coreml_model = ct.converters.convert(
    traced_model,
    inputs=[
        ct.TensorType(name='mel_batch', shape=mel_batch.shape),
        ct.TensorType(name='img_batch', shape=img_batch.shape)
    ],
    minimum_deployment_target=ct.target.iOS15
)
coreml_model.save("model.mlpackage")

Кодът по-горе създава TorchScript представяне на оригиналния модел на PyTorch, проследява го и след това го преобразува в CoreML модел. Използваме функцията convert от coremltools и указваме входните форми, които моделът очаква. В този случай предоставяме формите на mel_batch и img_batch тензорите.

Също така задаваме minimum_deployment_target=ct.target.iOS15 да генерира ML програма вместо модел на невронна мрежа. Програмата ML е с по-малък размер и е отделена от своите тежести, което я прави предпочитан тип, ако сте насочени към iOS 15 и по-нова версия. Това намалява размера на модела поне наполовина в сравнение с модел на невронна мрежа.

След като стартирате кода за преобразуване, трябва да видите файла model.mlpackage, записан в главната директория на хранилището на Wav2Lip.

Сега нека да преминем към изследването на конвертирания модел чрез Xcode. Ето как изглежда:

Нека заредим CoreML модела от диска и го отпечатаме, за да видим как изглеждат входните и изходните тензори:

import coremltools as ct

model = ct.models.MLModel('/path/to/model/Wav2Lip.mlpackage')
print(model)

Страхотен! Нека разгледаме отблизо:

input {
  name: "mel_batch"
  type {
    multiArrayType {
      shape: 102
      shape: 1
      shape: 80
      shape: 16
      dataType: FLOAT32
    }
  }
}
input {
  name: "img_batch"
  type {
    multiArrayType {
      shape: 102
      shape: 6
      shape: 96
      shape: 96
      dataType: FLOAT32
    }
  }
}
output {
  name: "var_1049"
  type {
    multiArrayType {
      dataType: FLOAT32
    }
  }
}
metadata {
  userDefined {
    key: "com.github.apple.coremltools.source"
    value: "torch==2.0.1"
  }
  userDefined {
    key: "com.github.apple.coremltools.version"
    value: "7.0b1"
  }
}

Оригиналният модел извършва значителна предварителна и последваща обработка на видео кадрите и данните от аудио спектрограмата извън самите слоеве на модела. Ето разбивка на процеса:

1- Разпознаване на лица: Моделът използва библиотека за разпознаване на лица, за да разпознае лица в дадено видео. Той запазва тези изображения на лица като масив от 96x96 пиксела. Първоначално видеото се преоразмерява до 256x256 пиксела преди разпознаване на лица. Тази стъпка за откриване на лице се изпълнява от друга библиотека дори преди моделът да получи входа.

2- Предварителна обработка на аудио: Аудио файлът преминава през специфични стъпки за предварителна обработка, които можете да разгледате, като разгледате кода във файла на Python.

3- Входни данни на модела: Моделът приема масив от лица с изображения от 96x96 пиксела и 6 канала (не обичайните 3) като един от своите входове. Броят на лицата съответства на броя на кадрите.

4- Аудио вход: Моделът също така очаква тензор, съдържащ предварително обработените аудио данни.

5- Резултат от модела: Моделът генерира масив от пренастроени лица с изображения 96x96, където областите на устата са пренастроени, за да изглеждат така, сякаш са синхронизирани с устни със съответстващите данни на аудио спектрограмата. Резултатът е масив от пренастроени лица със същия брой кадри като входа.

6- Последваща обработка: Python файлът извършва последваща обработка, като реконструира изходните лицеви рамки с оригиналните видео кадри и добавя оригиналния аудио файл към видеото. Тези стъпки се извършват извън модела, което води до крайно видео от 256x256 пиксела.

Последни бележки

Това прави ли модела добър, особено за обработка в реално време? Простият отговор е не, не мисля така. Отделните стъпки за предварителна и последваща обработка в идеалния случай трябва да бъдат интегрирани като вътрешни слоеве в оригиналния модел, с оптимизации за подобряване на производителността.

Важно е да се отбележи, че текущото внедряване добавя сложност и може да направи предизвикателство да съберете всичко заедно и да работите ефективно.

Имайте предвид, че вашите изисквания и цели може да се различават от другите. Ето защо исках да ви покажа лесен начин за превключване на преобразуването на модели за машинно обучение.

За оптимална производителност обаче се препоръчва да имате модел с висока производителност, който работи бързо за обработка в реално време и обработва повечето от предварителната и последващата обработка вътрешно.

Като разберете включените сложности и проучите различни подходи, можете да намерите най-доброто решение, което отговаря на вашите специфични нужди.

Успех с преобразуването на LipSync PyTorch в CoreML модел и се надявам, че тази публикация в блога ви е била полезна в пътуването!

Къде да изтеглите? Изтеглете напълно конвертирания CoreML модел тук.

За всякакви въпроси или запитвания, моля, свържете се с нас по имейл на [email protected] Любезно молим само за сериозни запитвания.