Трансформація цільового класу для задачі сегментації з використанням U-GAN

Abstract

Наведено огляд сучасних генеративних змагальних моделей (ГЗМ) для аугментації даних з фокусом на дослідження створення зображень та відповідних до них сегментаційних масок. Це особливо корисно у випадках, коли даних недостатньо, вони важкодоступні, мають конфіденційний характер, або розмітка вимагає значних ресурсів. Робота спрямована на підвищення ефективності процесу аугментації міноритарного класу за рахунок перетворення зображення з іншого класу та створення сегментаційної маски. Запропоновано новий підхід до одночасної генерації зображення та сегментаційної маски з використанням генеративної змагальної мережі, де генератором виступає U-Net. Цей генератор приймає на вхід зображення одного класу та шум, який подається як додатковий канал зображення. Генератор намагається створити зображення іншого класу, мінімізуючи зміни в початковому зображенні та додаючи ознаки іншого, разом з сегментаційною маскою, нового класу. Дискримінатор же визначає, чи є пара картинка–маска реальною чи згенерованою. У моделі для збереження оригінального вигляду вхідного зображення з мінімальними змінами використовується алгоритм, який застосовує тільки ті зміни згенерованого зображення, що вказані створеною сегментаційною маскою. Це дозволяє отримати зображення з ознаками нового класу з мінімальними змінами. Апробацію запропонованого підходу виконано на наборі даних панорамних знімків зубів, на основі якого створено набір зображень окремих зубів, частина з яких з пломбами, а частина — без них. Експериментальний набір даних включав 128 зубів без пломб і 128 з пломбами. Для створення ГЗМ взято всі зображення з пломбами та без пломб і навчено генератор перетворювати зображення без пломб в аналогічні з пломбами. Для перевірки ефективності аугментації проведено два експерименти по 50 симуляцій з різним випадковим станом для тренування сегментаційної моделі U-Net на основі ResNet-34. Перший експеримент використовував тільки реальні дані для тренування, а другий включав 64 додаткові зображення та відповідні маски, створені генератором на основі наявних зображень нульового класу. Середні значення коефіцієнта Жаккара серед всіх симуляцій для першого та другого експериментів відповідно склали 94,2 та 96,1. Це свідчить про те, що дані, згенеровані за допомогою запропонованої аугментації, дійсно сприяють підвищенню якості сегментаційної моделі та цей підхід можна комбінувати разом з іншими типами аугментацій.

The paper presents a review of modern generative adversarial models for data augmentation, focusing on research, aimed at creating images and their corresponding segmentation masks. This task is particularly useful in cases where data is insufficient, hard to access, has confidential nature, or where labeling requires significant resources. The paper is aimed at the task of augmenting the minority class by transforming an image from another class and creating a segmentation mask. New approach is proposed for the simultaneous generation of the image and segmentation mask, using a generative adversarial network with U-Net generator. This generator takes an image of one class and noise, which is fed as an additional image channel. The generator tries to create an image of another class, minimizing changes in the original image and adding features of another along with the segmentation mask of the new class. The discriminator then determines whether the picture-mask pair is real or generated. The algorithm that applies only those changes of the generated image that are indicated by the created segmentation mask used to preserve the original appearance of the input image with minimal changes. This technique allows to obtain an image with features of the new class with minimal changes. The practical implementation of the proposed approach was conducted on a dataset of panoramic dental X-rays, based on which a set of individual teeth was created, some with fillings and some without. The experimental data set included 128 teeth without fillings and 128 with fillings. The GAN is trained to transform images without fillings into similar ones with fillings using all input images. Two experiments of 50 simulations each with different random states were conducted for training the segmentation model U-Net with ResNet-34 backbone to check the effectiveness of this augmentation. The first experiment used only real data for training, while the second included 64 additional images and masks created by the generator based on existing zero-class images. The average Jaccard score among all simulations for the first and second experiments were respectively 94.2 and 96.1. This result indicates that data generated using the proposed augmentation helps improve the quality of segmentation models and this approach can be combined with other augmentation techniques.