Як ми тренуємо ML-моделі для створення сегментаційних масок

💡 Усі статті, обговорення, новини про AI — в одному місці. Приєднуйтесь до AI спільноти!

Спільното, привіт! Мене звати Георгій Калайчев, уже понад рік я працюю на позиції ML Engineer та є частиною R&D-команди у продуктовій IT-компанії Skylum. З огляду на профіль нашої компанії, технології Computer Vision та безпосередня робота з зображеннями є нашою першочерговою сферою інтересів та щоденним фокусом роботи.

Останні пів року я працюю над створенням якісних сегментаційних масок, які є надважливими для нашого нового продукту. Аби досягти гарних результатів, довелося спробувати чимало підходів у тренуванні різноманітних ML-моделей.

Тож у цій статті я поділюся своїми дослідженнями та висновками стосовно роботи над сегментацією обличчя.

Розкриття потенціалу сегментації зображень

Завдання сегментації з використанням штучного інтелекту полягає у розподілі зображення на різні сегменти або області, кожна з яких відповідає певному об’єкту або його частині на зображенні. А основна мета сегментації — зробити аналіз зображень більш точним і зрозумілим, дозволяючи виокремити важливі елементи зображення для подальшої обробки або аналізу.

Наведу найбільш популярні домени використання сегментації:

Медицина. Сегментація допомагає у виявленні та аналізі медичних зображень, як-то рентгенівські знімки або МРТ. Це може сприяти встановленню точного діагнозу, наприклад, при ідентифікації пухлин або інших патологій.

Нижче наведений приклад сегментації осьового зрізу мозку. Червоний, зелений і синій кольори відповідають білій речовині, сірій речовині та спинномозковій рідині.

Automotive. Сегментація використовується для розпізнавання різних об’єктів на дорозі, як-то інші автомобілі, пішоходи та дорожні знаки. Це допомагає у прийнятті рішень щодо безпечного руху транспортного засобу. Приклад такої сегментації додаю.

Агропромисловість. Тут сегментація використовується для аналізу стану полів і врожаю, виявлення хвороб рослин, аналізу їх росту й оцінки урожайності. Нижче наводжу приклад зазначеної сегментації: зеленим кольором відмічена нормальна рослина, а червоним — бур’ян.

Аерофотозйомка і дистанційне зондування. Сегментація використовується для картографування, аналізу земельних покривів, моніторингу змін навколишнього середовища й інших досліджень на базі зображень з літаків і супутників. Також додаю приклад сегментації навколишнього середовища за допомогою супутника.

Промислова автоматизація. Використання сегментації зображень для контролю якості, робототехніки й автоматизації складських операцій.

Усі приклади показують, що сегментація зображень — досить універсальна технологія, яка має потужний потенціал для застосування у багатьох сферах життя. Утім, на відміну від медицини чи промисловості, сегментація частин обличчя має більш вузькоспрямовану специфіку застосування.

Цей клас не такий популярний, як, наприклад, детекція обличчя або розпізнавання обличчя для ідентифікації людини. Тож ми взялися за більш предметне дослідження face segmentation.

Визначення завдання

Отож, для нашого нового продукту мати змогу точно сегментувати частини обличчя є критичним та одним з ключових функціоналів. Маски обличчя використовуються для:

1. Корекції кольору шкіри обличчя.
2. Розгладжування текстури шкіри.
3. Прибирання дефектів шкіри.
4. Диференціації акне та ластовиння.
5. Прибирання мішків під очима.
6. Накладання макіяжу.

Нашою метою є якомога більш детальне розуміння семантики обличчя. З огляду на поставлені задачі продукту та функціонал конкурентів (конкуруємо в розрізі якості та швидкості технологій) як оптимальний варіант ми виділили дев’ять класів. Серед них є шкіра, волосся, нижня губа, верхня губа, зуби, ліва брова, права брова, ліве око, праве око. Нижче наведу приклади зображення та відповідної маски.

Поставлену задачу можна розв’язати різними способами:

• Подавати на вхід в модель повне зображення. Головний недолік цього підходу — можливі випадки, коли розмір обличчя відповідно розміру самої фотографії є малим. Це погано, оскільки більша частини інформації, яку отримує модель, не є релевантною та не несе користі.
• Подавати на вхід в модель тільки лице (cropped face). Для реалізації такого підходу фотографія буде подаватися на вхід до моделі детекції обличчя (face detection), щоб потім подавати в модель сегментації тільки обличчя без зайвої інформації.

Звісно, використання підходу з попередньою детекцією обличчя є більш привабливим, оскільки цей підхід гарантуватиме більшу легкість при навчанні моделі сегментації обличчя та її вищу точність.

Набір даних для навчання моделі

Зазвичай для розв’язання подібних завдань використовуються відкриті набори даних, які можна знайти у вільному доступі. В нашому випадку використання відкритих даних не є розв’язанням проблеми, оскільки, як я вказував вище, точність та якість масок є надважливими для майбутнього продукту.

Ми використовували два джерела даних:

— дані, які розмітили наші анотатори;
— синтетичні дані, які були штучно згенеровані.

Велика кількість експериментів показала, що саме використання комбінації реальних і штучних даних дає кращий результат, ніж використання кожного з них окремо.

Переваги даних, які розмітили анотатори:

• З даними працюють реальні люди, яким можна гнучко й досить специфічно пояснити, як при розмітці опрацьовувати тонкі деталі обличчя. Наприклад, при роботі з маскою губ анотатори концентрувалися на точному контурі куточків губ, аби модель їх не губила.
• Реалістичність даних, яка потрібна для нашого майбутнього продукту. Ми плануємо працювати з професійними фотографами, а вони оперують даними великого розширення, отримати які можна виключно з реальних даних.

Переваги синтетичних даних:

• Вони можуть бути згенеровані у великій кількості, що дозволяє розширити набір даних і покращити загальну продуктивність моделі.
• Можуть бути налаштовані таким чином, щоб охопити рідкісні або важкодоступні випадки, що часто відсутні у реальних наборах даних.
• Можуть допомогти зменшити залежність від ручної анотації, що економить час і ресурси.

Отож, синтетичні дані — це зручний та швидкий спосіб отримати необмежену кількість даних з обмеженою якістю. Водночас професіональні анотатори здатні видавати вищу якість розмітки з урахуванням усіх особливих потреб і поставлених перед ними задач.

За нашим досвідом, якраз використання синтетичних даних у поєднанні з реальними даними дозволяє досягти більш збалансованого та надійного навчання моделі. Своєю чергою це сприяє покращенню її точності та загальної ефективності.

Навчання моделі

Навчання моделі сегментації обличчя — це процес, що містить кілька вкрай важливих компонентів, і відбувається у контексті навчання з вчителем. Основна ціль полягає у тому, щоб навчити модель ефективно розрізняти різні області на зображенні. До прикладу, відрізняти губи від інших фрагментів обличчя.

У процесі розробки потрібної моделі ми експериментували з різноманітними архітектурами, включно з популярними енкодерами (ResNet, Inception, MobileNet, FastViT різних версій та інші) та декодерами (класичний U-Net, FPN, DeepLab різних версій та інші). Деякі з них у комбінаціях працювали гірше, деякі — давали гарні результати у плані безпосередньої точності. Однак крім точності, для нашого продукту важливо також отримати високу швидкість та мати змогу конвертувати модель у різні фреймворки.

Врешті-решт ми зупинилися на двох моделях — FastViT й FPN. І ось чому.

Encoder. Як енкодер використали модель FastViT, оскільки вона швидка, потужна, оптимізована програмістами Apple для всієї екосистеми продуктів компанії та досить гнучка у розмірності. Застосування FastViT дозволило нам ефективно стиснути вхідне зображення до простору з меншою розмірністю, використовуючи переваги трансформерів для глибокого аналізу контексту в зображеннях.

Також варто зазначити, що при конвертації найбільшої моделі FastViT у CoreML, ONNX або OpenVINO ваги Torch, отримані під час навчання, у форматі float32 стискаються із 170 Мб до 67.8 Мб без зміни формату моделі. Це дозволяє використовувати архітектуру з більшою кількістю параметрів та «займати» місця як слабкіша модель.

Decoder. Щодо декодера, зупинилися на моделі FPN, адже вона якраз розроблена для завдань комп’ютерного зору і часто використовується у задачах сегментації зображень. Модель є відносно легкою архітектурою, що ефективно працює з об’єктами різного розміру та дозволяє точно визначати межі об’єктів.

Застосування FPN як декодера дозволило точно відновлювати вихідне зображення зі стиснутого представлення, забезпечуючи високу точність сегментації.

Тепер детальніше розкажу про архітектури, від яких команді через особливість поставленої задачі та високих вимог до продуктивності довелося відмовитися.

Енкодери

ResNet

Ця архітектура вважається класичною для багатьох задач комп’ютерного зору. Вона забезпечує високу точність завдяки residual-блокам, які допомагають уникати проблеми зникнення градієнта у глибоких мережах. Однак для наших задач ResNet виявився занадто важким і повільним.

Inception

Відрізняється можливістю одночасного використання різних фільтрів, що дає змогу моделі виявляти різноманітні патерни. І хоча вона показала хороші результати в плані точності, але все ж гірші за ResNet та FastViT.

MobileNet

Декодери

Щодо декодерів, які відіграють ключову роль в архітектурі сегментаційних моделей, адже відповідають за відновлення просторової інформації після процесу енкодінгу, довелося відмовитися від таких:

U-Net

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

DeepLab

Ця серія моделей була розроблена компанією Google. У нашому дослідженні ми зосереджували увагу саме на версії DeepLabV3. Важливим складником DeepLabV3 є Atrous Spatial Pyramid Pooling (ASPP). Цей модуль дозволяє глибше інтегрувати контекстну інформацію на різних масштабах, використовуючи атріальні згортки з різними факторами диляції.

Завдяки ASPP модель може краще розрізняти об’єкти на різних масштабах, що сприяє більш точній сегментації. У деяких випадках вона показала кращі результати, ніж U-Net, особливо на великих наборах даних. Проте її складність призводила до зниження швидкості роботи, що стало недоліком для нашого продукту.

А сам навчальний процес, власне, включав такі етапи:

Підготовка даних. Перед навчанням моделі зображення були відповідним чином підготовлені та нормалізовані, щоб забезпечити універсальність даних і виключити зміщення у вхідних даних. Також були застосовані різні методи аугментації зображень при навчанні.

Функція втрат. Використання FocalLoss дозволило зменшити проблему дисбалансу класів, зосереджуючись більше на складних класах або тих, які рідко зустрічаються, що є типовою проблемою у задачах сегментації.

Оптимізація. Оптимізатор AdamW був обраний для кращого управління вагами моделі, що дозволяє вдало мінімізувати функцію втрат. Це дає більш стабільні та кращі результати навчання.

Регулювання швидкості навчання. Скедулер OneCycleLr був використаний для адаптивної модифікації швидкості навчання у процесі тренувань, що допомогло швидше досягати оптимальних параметрів й уникати проблеми локальних мінімумів.

Аналіз результатів

Після тренувань модель була ретельно оцінена за допомогою метрик для класифікації (Precision, Recall, F1-score) і сегментаційних метрик (IOU та Dice Coefficient).

Precision

Визначає частку правильних позитивних прогнозів серед усіх позитивних прогнозів, зроблених моделлю. Метрика важлива для завдань, коли помилкові позитивні результати можуть мати серйозні наслідки. Наприклад, помилкове віднесення не шкірних зон до шкіри може призвести до появи артефактів та небажаних результатів.

Recall

Визначає частку правильних позитивних прогнозів серед усіх дійсних позитивних зразків. Нею послуговуються, коли необхідно мінімізувати помилкові негативні результати. Ця метрика показує, чи ми дійсно добираємо усі позитивні класи, приміром, до маски шкіри, і чи ми достатньо повно відносимо пікселі до класу шкіри.

F1-score

Використовується для оцінки загальної продуктивності моделі. Це особливо важливо, коли потрібно знайти певний баланс між Precision і Recall, тим самим гармонізувавши результати.

IOU (Intersection Over Union)

Ключова метрика для задач сегментації зображень, яка визначає частку перетину площі між передбачуваними та справжніми масками об’єкта. Вона показує наскільки точно модель передбачає межі об’єктів. Саме показник IOU на рівні 87% був для нашої команди підтвердженням високої якості роботи моделі.

Dice Coefficient

Також дуже популярна метрика для задач сегментації. Вона особливо корисна у випадку великого дисбалансу класів. Це пов’язано з тим, що Dice більш чутливий до невеликих сегментів і може краще підкреслювати важливість дрібних об’єктів. Додам, що Dice Coefficient часто використовується у сфері сегментації медичних зображень, де важливо точно визначити межі органів чи патологій. Враховуючи це, зазначена метрика може бути основною, наприклад, у задачі сегментації акне на шкірі обличчя.

Підсумки

Попри те, що кожна метрика має свої особливості, найнадійнішим способом оцінити роботу моделі залишається мануальне тестування. Важливо, аби тестувальники перевірили на реальних кейсах продукту ефективність роботи розробленої моделі. Їхній позитивний відгук стосовно якості роботи не менш важливий за високий показник IOU. Особливо, якщо валідація моделі під час навчання проводиться не тільки на реальних, а й на штучних даних.

Підсумовуючи, всі вищеописані метрики підтвердили високу ефективність обраних архітектур і методик навчання. До слова, наші ключові партнери — професійні фотографи — уже встигли дати позитивні відгуки про роботу майбутнього продукту. Зокрема, щодо якості та швидкості.

Тож наостанок наведу приклади роботи нашої моделі на реальних даних.