3D teeth instance segmentation. В темноте, но не один
3D сегментация зубов от поиска данных до конечного результата. Почти.
Дисклеймер
Данная статья не является обучающей в любом понимании этого термина и носит сугубо информативный характер. Автор статьи не несет ответственности за время, потраченное на ее чтение.
Об автореДобрый — всем, зовут Андрей (27). Постараюсь коротко. Почему программирование? По образованию — бакалавр электромеханик, профессию знаю. Отработал 2 года на должности инженера-энергетика в буровой компании вполне успешно, вместо повышения написал заявление — сгорел, да не по мне оказалось это всё. Нравится создавать, находить решения сложных задач, с ПК в обнимку с сознательных лет. Выбор очевиден. Вначале (полгода назад), всерьёз думал записаться на курсы от Я или подобные. Начитался отзывов, поговорил с участниками и понял что с получением информацией проблем нет. Так нашел сайт, там получил базу по Python и с ним уже начал свой путь (сейчас там постепенно изучаю всё, что связано с ML). Сразу заинтересовало машинное обучение, CV в частности. Придумал себе задачу и вот здесь (по мне, так отличный способ учиться).
1. Введение
В результате нескольких неудачных попыток, пришел к решению использовать 2 легковесные модели для получения желаемого результата. 1-ая сегментирует все зубы как [1, 0] категорию, а вторая делит их на категории[0, 8]. Но начнем по порядку.
2. Поиск и подготовка данных
Потратив не один вечер на поиск данных для работы, пришел в выводу что в свободном доступе челюсть в хорошем качестве и формате (*.stl, *.nrrd и т.д.) не получится. Лучшее, что мне попалось — это тестовый образец головы пациента после хирургической операции на челюсти в программе 3D Slicer.
Очевидно, мне не нужна голова целиком, поэтому обрезал исходник в той же программе до размера 163×112*120рх (в данном посте {x*y*z = ш-г-в} и 1рх — 0,5 мм), оставив только зубы и сопутствующие челюстно-лицевые части.
Уже больше похоже на то что нужно, дальше — интереснее. Теперь нужно создать маски всех необходимых нам объектов. Для тех, кто уже работал с этим — «autothreshold» не то чтобы совсем не работает, просто лишнего много, думаю, исправление заняло бы столько же времени, сколько и разметка вручную (через маски).
— Пиксели (срезы слева)? — Вспоминаем размер изображенияРазмечал часов 12~14. И да, тот факт что я не сразу разметил каждый зуб как категорию стоил мне еще порядка 4 часов. В итоге у нас есть данные, с которыми у же можно работать.
Конечный вариант маски. Smooth 0.5. (сглаживание в обучении не использовалось)Должен добавить, даже на мой (без опыта) взгляд, этих данных очень мало для обучения и последующей полноценной работы нейронной сети. На данном этапе, единственное что пришло в голову, повернуть имеющиеся данные N-раз и соединить, random-crop использовать не стал.
Код подготовки данныхimport nrrd
import torch
import torchvision.transforms as tf
class DataBuilder:
def __init__(self,
data_path,
list_of_categories,
num_of_chunks: int = 0,
augmentation_coeff: int = 0,
num_of_classes: int = 0,
normalise: bool = False,
fit: bool = True,
data_format: int = 0,
save_data: bool = False
):
self.data_path = data_path
self.number_of_chunks = num_of_chunks
self.augmentation_coeff = augmentation_coeff
self.list_of_cats = list_of_categories
self.num_of_cls = num_of_classes
self.normalise = normalise
self.fit = fit
self.data_format = data_format
self.save_data = save_data
def forward(self):
data = self.get_data()
data = self.fit_data(data) if self.fit else data
data = self.pre_normalize(data) if self.normalise else data
data = self.data_augmentation(data, self.augmentation_coeff) if self.augmentation_coeff != 0 else data
data = self.new_chunks(data, self.number_of_chunks) if self.number_of_chunks != 0 else data
data = self.category_splitter(data, self.num_of_cls, self.list_of_cats) if self.num_of_cls != 0 else data
torch.save(data, self.data_path[-14:]+'.pt') if self.save_data else None
return torch.unsqueeze(data, 1)
def get_data(self):
if self.data_format == 0:
return torch.from_numpy(nrrd.read(self.data_path)[0])
elif self.data_format == 1:
return torch.load(self.data_path).cpu()
elif self.data_format == 2:
return torch.unsqueeze(self.data_path, 0).cpu()
else:
print('Available types are: "nrrd", "tensor" or "self.tensor(w/o load)"')
@staticmethod
def fit_data(some_data):
data = torch.movedim(some_data, (1, 0), (0, -1))
data_add_x = torch.nn.ZeroPad2d((5, 0, 0, 0))
data = data_add_x(data)
data = torch.movedim(data, -1, 0)
data_add_z = torch.nn.ZeroPad2d((0, 0, 8, 0))
return data_add_z(data)
@staticmethod
def pre_normalize(some_data):
min_d, max_d = torch.min(some_data), torch.max(some_data)
return (some_data - min_d) / (max_d - min_d)
@staticmethod
def data_augmentation(some_data, aug_n):
torch.manual_seed(17)
tr_data = []
for e in range(aug_n):
transform = tf.RandomRotation(degrees=(20*e, 20*e))
for image in some_data:
image = torch.unsqueeze(image, 0)
image = transform(image)
tr_data.append(image)
return tr_data
def new_chunks(self, some_data, n_ch):
data = torch.stack(some_data, 0) if self.augmentation_coeff != 0 else some_data
data = torch.squeeze(data, 1)
chunks = torch.chunk(data, n_ch, 0)
return torch.stack(chunks)
@staticmethod
def category_splitter(some_data, alpha, list_of_categories):
data, _ = torch.squeeze(some_data, 1).to(torch.int64), alpha
for i in list_of_categories:
data = torch.where(data < i, _, data)
_ += 1
return data - alpha
Имейте ввиду что это финальная версия кода подготовки данных для 3D U-net. Форвард:
Загружаем дату (в зависимости от типа).
Добавляем 0 по краям чтобы подогнать размер до 168×120*120 (вместо исходных 163×112*120). *пригодится дальше.
Нормализуем входящие данные в 0…1 (исходные ~-2000…16000).
Поворачиваем N-раз и соединяем.
Полученные данные режем на равные части чтобы забить память видеокарты по максимуму (в моем случае это 1, 1, 72, 120, 120).
Эта часть распределяет по категориям 28 имеющихся зубов и фон для облегчения обучения моделей (см. Введение):
одну категорию для 1-ой;
на 9 категорий (8+фон) для 2-ой.
import torch.utils.data as tud
class ToothDataset(tud.Dataset):
def __init__(self, images, masks):
self.images = images
self.masks = masks
def __len__(self): return len(self.images)
def __getitem__(self, index):
if self.masks is not None:
return self.images[index, :, :, :, :],\
self.masks[index, :, :, :, :]
else:
return self.images[index, :, :, :, :]
def get_loaders(images, masks,
batch_size: int = 1,
num_workers: int = 1,
pin_memory: bool = True):
train_ds = ToothDataset(images=images,
masks=masks)
data_loader = tud.DataLoader(train_ds,
batch_size=batch_size,
shuffle=False,
num_workers=num_workers,
pin_memory=pin_memory)
return data_loader
На выходе имеем следующее:
Semantic | Instance | Predictions | |
Data | (27*, 1, 56*, 120,120)[0…1] | (27*, 1, 56*, 120,120) [0, 1] | (1, 1, 168, 120, 120)[0…1] |
Masks | (27*, 1, 56*, 120,120)[0, 1] | (27*, 1, 56*, 120,120)[0, 8] | - |
*эти размеры менялись, в зависимости от эксперимента, подробности — дальше.
3. Выбор и настройка моделей обучения
Цель работы — обучение. Поэтому взял наиболее простую и понятную для себя модель нейросети архитектуры U-Net. Код не выкладываю, можно посмотреть тут.
2D U-NetПодробно рассказывать не буду, информации в достатке в сети. Метод оптимизации — Adam, функция расчета потерь Dice-loss (implement), спусков/подъемов 4, фильтры [64, 128, 256, 512] (знаю, много, об этом — позже). Обучал в среднем 60–80 epochs на эксперимент. Transfer learning не использовал.
model.summary ()model = UNet(dim=2, in_channels=1, out_channels=1, n_blocks=4, start_filters=64).to(device)
print(summary(model, (1, 168, 120)))
"""
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 168, 120] 640
ReLU-2 [-1, 64, 168, 120] 0
BatchNorm2d-3 [-1, 64, 168, 120] 128
Conv2d-4 [-1, 64, 168, 120] 36,928
ReLU-5 [-1, 64, 168, 120] 0
BatchNorm2d-6 [-1, 64, 168, 120] 128
MaxPool2d-7 [-1, 64, 84, 60] 0
DownBlock-8 [[-1, 64, 84, 60], [-1, 64, 168, 120]] 0
Conv2d-9 [-1, 128, 84, 60] 73,856
ReLU-10 [-1, 128, 84, 60] 0
BatchNorm2d-11 [-1, 128, 84, 60] 256
Conv2d-12 [-1, 128, 84, 60] 147,584
ReLU-13 [-1, 128, 84, 60] 0
BatchNorm2d-14 [-1, 128, 84, 60] 256
MaxPool2d-15 [-1, 128, 42, 30] 0
DownBlock-16 [[-1, 128, 42, 30], [-1, 128, 84, 60]] 0
Conv2d-17 [-1, 256, 42, 30] 295,168
ReLU-18 [-1, 256, 42, 30] 0
BatchNorm2d-19 [-1, 256, 42, 30] 512
Conv2d-20 [-1, 256, 42, 30] 590,080
ReLU-21 [-1, 256, 42, 30] 0
BatchNorm2d-22 [-1, 256, 42, 30] 512
MaxPool2d-23 [-1, 256, 21, 15] 0
DownBlock-24 [[-1, 256, 21, 15], [-1, 256, 42, 30]] 0
Conv2d-25 [-1, 512, 21, 15] 1,180,160
ReLU-26 [-1, 512, 21, 15] 0
BatchNorm2d-27 [-1, 512, 21, 15] 1,024
Conv2d-28 [-1, 512, 21, 15] 2,359,808
ReLU-29 [-1, 512, 21, 15] 0
BatchNorm2d-30 [-1, 512, 21, 15] 1,024
DownBlock-31 [[-1, 512, 21, 15], [-1, 512, 21, 15]] 0
ConvTranspose2d-32 [-1, 256, 42, 30] 524,544
ReLU-33 [-1, 256, 42, 30] 0
BatchNorm2d-34 [-1, 256, 42, 30] 512
Concatenate-35 [-1, 512, 42, 30] 0
Conv2d-36 [-1, 256, 42, 30] 1,179,904
ReLU-37 [-1, 256, 42, 30] 0
BatchNorm2d-38 [-1, 256, 42, 30] 512
Conv2d-39 [-1, 256, 42, 30] 590,080
ReLU-40 [-1, 256, 42, 30] 0
BatchNorm2d-41 [-1, 256, 42, 30] 512
UpBlock-42 [-1, 256, 42, 30] 0
ConvTranspose2d-43 [-1, 128, 84, 60] 131,200
ReLU-44 [-1, 128, 84, 60] 0
BatchNorm2d-45 [-1, 128, 84, 60] 256
Concatenate-46 [-1, 256, 84, 60] 0
Conv2d-47 [-1, 128, 84, 60] 295,040
ReLU-48 [-1, 128, 84, 60] 0
BatchNorm2d-49 [-1, 128, 84, 60] 256
Conv2d-50 [-1, 128, 84, 60] 147,584
ReLU-51 [-1, 128, 84, 60] 0
BatchNorm2d-52 [-1, 128, 84, 60] 256
UpBlock-53 [-1, 128, 84, 60] 0
ConvTranspose2d-54 [-1, 64, 168, 120] 32,832
ReLU-55 [-1, 64, 168, 120] 0
BatchNorm2d-56 [-1, 64, 168, 120] 128
Concatenate-57 [-1, 128, 168, 120] 0
Conv2d-58 [-1, 64, 168, 120] 73,792
ReLU-59 [-1, 64, 168, 120] 0
BatchNorm2d-60 [-1, 64, 168, 120] 128
Conv2d-61 [-1, 64, 168, 120] 36,928
ReLU-62 [-1, 64, 168, 120] 0
BatchNorm2d-63 [-1, 64, 168, 120] 128
UpBlock-64 [-1, 64, 168, 120] 0
Conv2d-65 [-1, 1, 168, 120] 65
================================================================
Total params: 7,702,721
Trainable params: 7,702,721
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.08
Forward/backward pass size (MB): 7434.08
Params size (MB): 29.38
Estimated Total Size (MB): 7463.54
"""![Эксп.№1 2D U-Net, подача изображений покадрово, плоскость [x, z]](https://habrastorage.org/getpro/habr/upload_files/d24/21a/246/d2421a24672cfc2928bf98a22ebd3bc7.gif)
Эксп.№1
2D U-Net, подача изображений покадрово, плоскость [x, z]Определенно, это — зубы. Только кроме зубов есть много всего, нам ненужного. Подробнее о трансформации numpy — *.stl в Главе 6. Посмотрим ещё раз на фактический размер и качество изображений, которые попадают на вход нейросети:
![Слева на право: 1. Не видно[x, y]. 2. Немного лучше[x, z]. 3.Ещё лучше[y, z]](https://habrastorage.org/getpro/habr/upload_files/75b/d08/ffb/75bd08ffbdeac83181abf79355378872.png)
Слева на право:
1. Не видно[x, y]. 2. Немного лучше[x, z]. 3.Ещё лучше[y, z]Если сам не видишь на 100% где там начался зуб, а где нет, то как тогда эту работу выполнит нейросеть? Как минимум, необходимо изменить плоскость подачи изображения.
Проведя не один день разбираясь в том, как можно улучшить сложившуюся ситуацию, пришел к тому, что можно составлять каскад и сетей, поочередно обрабатывающих изображение, аналогично работе фильтров грубой и тонкой очистки.
![Эксп.№2 Каскад 2-ух 2D U-Net, подача изображений покадрово, плоскость [y, z]](https://habrastorage.org/getpro/habr/upload_files/eef/335/e4e/eef335e4ebb6f3a3d23709cf81948e40.png)
Эксп.№2
Каскад 2-ух 2D U-Net, подача изображений покадрово, плоскость [y, z]Прогресс виден, однако вместе с помехами пропадают и части зубов, дальнейшее обучение тому подтверждение:
![Эксп.№3 Каскад 2-ух 2D U-Net, подача изображений покадрово плоскость [y, z] с увеличением времени обучения на 50%](https://habrastorage.org/getpro/habr/upload_files/3ce/fc5/3d5/3cefc53d5fee0fdf2693b95e6db91152.png)
Эксп.№3
Каскад 2-ух 2D U-Net, подача изображений покадрово плоскость [y, z]
с увеличением времени обучения на 50%Ввиду последних событий было принято решение о переходе на 3D архитектуру нейронной сети. Переподготовил входные данные, а именно разделил на части размером (24*, 120, 120). Почему так? — изначально большая модель обучения (~22 млн. параметров). Моя видеокарта (1063gtx) не могла физически вместить больше.
24*Это размер глубины. Был подобран так чтобы:
количество данных (1512, 120, 120) делится нацело на это число — получается 63;
в свою очередь получившийся batch size (24, 120, 120) — максимум, вмещающийся в память видеокарты с текущими параметрами сети;
само это число (24) делилось на количество спусков/подъемов так же нацело (имеется в виду соответствие выражению 24/2/2/2=3 и 3×2*2×2=24, где количество делений/умножений на 2 соответствует количеству спусков/подъемов минус 1);
то же самое не только для глубины данных, но и длинны и ширины. Подробнее в .summary ()
model = UNet(dim=3, in_channels=1, out_channels=1, n_blocks=4, start_filters=64).to(device)
print(summary(model, (1, 24, 120, 120)))
"""
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv3d-1 [-1, 64, 24, 120, 120] 1,792
ReLU-2 [-1, 64, 24, 120, 120] 0
BatchNorm3d-3 [-1, 64, 24, 120, 120] 128
Conv3d-4 [-1, 64, 24, 120, 120] 110,656
ReLU-5 [-1, 64, 24, 120, 120] 0
BatchNorm3d-6 [-1, 64, 24, 120, 120] 128
MaxPool3d-7 [-1, 64, 12, 60, 60] 0
DownBlock-8 [[-1, 64, 12, 60, 60], [-1, 64, 24, 120, 120]] 0
Conv3d-9 [-1, 128, 12, 60, 60] 221,312
ReLU-10 [-1, 128, 12, 60, 60] 0
BatchNorm3d-11 [-1, 128, 12, 60, 60] 256
Conv3d-12 [-1, 128, 12, 60, 60] 442,496
ReLU-13 [-1, 128, 12, 60, 60] 0
BatchNorm3d-14 [-1, 128, 12, 60, 60] 256
MaxPool3d-15 [-1, 128, 6, 30, 30] 0
DownBlock-16 [[-1, 128, 6, 30, 30], [-1, 128, 12, 60, 60]] 0
Conv3d-17 [-1, 256, 6, 30, 30] 884,992
ReLU-18 [-1, 256, 6, 30, 30] 0
BatchNorm3d-19 [-1, 256, 6, 30, 30] 512
Conv3d-20 [-1, 256, 6, 30, 30] 1,769,728
ReLU-21 [-1, 256, 6, 30, 30] 0
BatchNorm3d-22 [-1, 256, 6, 30, 30] 512
MaxPool3d-23 [-1, 256, 3, 15, 15] 0
DownBlock-24 [[-1, 256, 3, 15, 15], [-1, 256, 6, 30, 30]] 0
Conv3d-25 [-1, 512, 3, 15, 15] 3,539,456
ReLU-26 [-1, 512, 3, 15, 15] 0
BatchNorm3d-27 [-1, 512, 3, 15, 15] 1,024
Conv3d-28 [-1, 512, 3, 15, 15] 7,078,400
ReLU-29 [-1, 512, 3, 15, 15] 0
BatchNorm3d-30 [-1, 512, 3, 15, 15] 1,024
DownBlock-31 [[-1, 512, 3, 15, 15], [-1, 512, 3, 15, 15]] 0
ConvTranspose3d-32 [-1, 256, 6, 30, 30] 1,048,832
ReLU-33 [-1, 256, 6, 30, 30] 0
BatchNorm3d-34 [-1, 256, 6, 30, 30] 512
Concatenate-35 [-1, 512, 6, 30, 30] 0
Conv3d-36 [-1, 256, 6, 30, 30] 3,539,200
ReLU-37 [-1, 256, 6, 30, 30] 0
BatchNorm3d-38 [-1, 256, 6, 30, 30] 512
Conv3d-39 [-1, 256, 6, 30, 30] 1,769,728
ReLU-40 [-1, 256, 6, 30, 30] 0
BatchNorm3d-41 [-1, 256, 6, 30, 30] 512
UpBlock-42 [-1, 256, 6, 30, 30] 0
ConvTranspose3d-43 [-1, 128, 12, 60, 60] 262,272
ReLU-44 [-1, 128, 12, 60, 60] 0
BatchNorm3d-45 [-1, 128, 12, 60, 60] 256
Concatenate-46 [-1, 256, 12, 60, 60] 0
Conv3d-47 [-1, 128, 12, 60, 60] 884,864
ReLU-48 [-1, 128, 12, 60, 60] 0
BatchNorm3d-49 [-1, 128, 12, 60, 60] 256
Conv3d-50 [-1, 128, 12, 60, 60] 442,496
ReLU-51 [-1, 128, 12, 60, 60] 0
BatchNorm3d-52 [-1, 128, 12, 60, 60] 256
UpBlock-53 [-1, 128, 12, 60, 60] 0
ConvTranspose3d-54 [-1, 64, 24, 120, 120] 65,600
ReLU-55 [-1, 64, 24, 120, 120] 0
BatchNorm3d-56 [-1, 64, 24, 120, 120] 128
Concatenate-57 [-1, 128, 24, 120, 120] 0
Conv3d-58 [-1, 64, 24, 120, 120] 221,248
ReLU-59 [-1, 64, 24, 120, 120] 0
BatchNorm3d-60 [-1, 64, 24, 120, 120] 128
Conv3d-61 [-1, 64, 24, 120, 120] 110,656
ReLU-62 [-1, 64, 24, 120, 120] 0
BatchNorm3d-63 [-1, 64, 24, 120, 120] 128
UpBlock-64 [-1, 64, 24, 120, 120] 0
Conv3d-65 [-1, 1, 24, 120, 120] 65
================================================================
Total params: 22,400,321
Trainable params: 22,400,321
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.61
Forward/backward pass size (MB): 15974.12
Params size (MB): 85.45
Estimated Total Size (MB): 16060.18
----------------------------------------------------------------
"""![Эксп.№4 3D U-Net, подача объемом, плоскость [y, z], время*0,38](https://habrastorage.org/getpro/habr/upload_files/618/cdd/83d/618cdd83db3c6855c2f916a11303c960.png)
Эксп.№4
3D U-Net, подача объемом, плоскость [y, z],
время*0,38С учетом сокращенного на ~60% времени обучения (25 epochs) результат меня устроил, продолжаем.
![Эксп.№5 3D U-Net, подача объемом, плоскость [y, z], 65 epochs ~ 1,5 часа](https://habrastorage.org/getpro/habr/upload_files/7c7/e67/19a/7c7e6719a026fb11df22cd32eab03620.png)
Эксп.№5
3D U-Net, подача объемом, плоскость [y, z],
65 epochs ~ 1,5 часа Особых потерь в искомых зонах не заметил. Решил продолжать, однако результат дальнейшего обучения мы уже где то видели (эксп.№3) — значительное уменьшение искомых зон и появление артефактов:
![Эксп.№6 3D U-Net, подача объемом, плоскость [x, z], 105 epochs ~ 2,1 часа](https://habrastorage.org/getpro/habr/upload_files/bc8/372/a07/bc8372a07b3aa13e874b0f9abfc4d21a.png)
Эксп.№6
3D U-Net, подача объемом, плоскость [x, z],
105 epochs ~ 2,1 часа «Научный» перебор параметров в течении недели принес результат. Уменьшил количество параметров сети до ~400к (от первоначальных ~22 м) путем уменьшения фильтра [18, 32, 64, 128] и спуска/подъема до 3. Изменил метод оптимизации на RSMProp. Уменьшение количества параметров нейросети позволило увеличить объем входных данных в три раза (1, 1, 72*, 120, 120). Посмотрим результат?
model.summary ()model = UNet(dim=3, in_channels=1, out_channels=1, n_blocks=3, start_filters=18).to(device)
print(summary(model, (1, 1, 72, 120, 120)))
"""
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv3d-1 [-1, 18, 72, 120, 120] 504
ReLU-2 [-1, 18, 72, 120, 120] 0
BatchNorm3d-3 [-1, 18, 72, 120, 120] 36
Conv3d-4 [-1, 18, 72, 120, 120] 8,766
ReLU-5 [-1, 18, 72, 120, 120] 0
BatchNorm3d-6 [-1, 18, 72, 120, 120] 36
MaxPool3d-7 [-1, 18, 36, 60, 60] 0
DownBlock-8 [[-1, 18, 36, 60, 60], [-1, 18, 24, 120, 120]] 0
Conv3d-9 [-1, 36, 36, 60, 60] 17,532
ReLU-10 [-1, 36, 36, 60, 60] 0
BatchNorm3d-11 [-1, 36, 36, 60, 60] 72
Conv3d-12 [-1, 36, 36, 60, 60] 35,028
ReLU-13 [-1, 36, 36, 60, 60] 0
BatchNorm3d-14 [-1, 36, 36, 60, 60] 72
MaxPool3d-15 [-1, 36, 18, 30, 30] 0
DownBlock-16 [[-1, 36, 18, 30, 30], [-1, 36, 36, 60, 60]] 0
Conv3d-17 [-1, 72, 18, 30, 30] 70,056
ReLU-18 [-1, 72, 18, 30, 30] 0
BatchNorm3d-19 [-1, 72, 18, 30, 30] 144
Conv3d-20 [-1, 72, 18, 30, 30] 140,040
ReLU-21 [-1, 72, 18, 30, 30] 0
BatchNorm3d-22 [-1, 72, 18, 30, 30] 144
DownBlock-23 [[-1, 72, 18, 30, 30], [-1, 72, 18, 30, 30]] 0
ConvTranspose3d-24 [-1, 36, 36, 60, 60] 20,772
ReLU-25 [-1, 36, 36, 60, 60] 0
BatchNorm3d-26 [-1, 36, 36, 60, 60] 72
Concatenate-27 [-1, 72, 36, 60, 60] 0
Conv3d-28 [-1, 36, 36, 60, 60] 70,020
ReLU-29 [-1, 36, 36, 60, 60] 0
BatchNorm3d-30 [-1, 36, 36, 60, 60] 72
Conv3d-31 [-1, 36, 36, 60, 60] 35,028
ReLU-32 [-1, 36, 36, 60, 60] 0
BatchNorm3d-33 [-1, 36, 36, 60, 60] 72
UpBlock-34 [-1, 36, 36, 60, 60] 0
ConvTranspose3d-35 [-1, 18, 72, 120, 120] 5,202
ReLU-36 [-1, 18, 72, 120, 120] 0
BatchNorm3d-37 [-1, 18, 72, 120, 120] 36
Concatenate-38 [-1, 36, 72, 120, 120] 0
Conv3d-39 [-1, 18, 72, 120, 120] 17,514
ReLU-40 [-1, 18, 72, 120, 120] 0
BatchNorm3d-41 [-1, 18, 72, 120, 120] 36
Conv3d-42 [-1, 18, 72, 120, 120] 8,766
ReLU-43 [-1, 18, 72, 120, 120] 0
BatchNorm3d-44 [-1, 18, 72, 120, 120] 36
UpBlock-45 [-1, 18, 72, 120, 120] 0
Conv3d-46 [-1, 1, 72, 120, 120] 19
================================================================
Total params: 430,075
Trainable params: 430,075
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 1.32
Forward/backward pass size (MB): 5744.38
Params size (MB): 1.64
Estimated Total Size (MB): 5747.34
----------------------------------------------------------------
"""Некоторые из вас подумают, исходные данные (168, 120, 120), а часть (72, 120, 120). Назревает вопрос, как делить. Всё просто, во 2 главе мы увеличивали размер наших данных и затем делили их на части, соответствующие объему памяти видеокарты. Я увеличил данные в 9 раз (1512, 120, 120) т.е. повернул на 9 различных углов относительно одной оси, а затем разделил на 21(batch size) часть по (72, 120, 120). Так же 72 соответствует всем условиям, описанным в 24*(выше).
![Эксп.№7 3D U-Net, подача объемом, плоскость [x, z], Маска (слева) и готовая сегментация (справа), оптимизированные параметры сети, время обучения(65 epochs) ~ 14мин.](https://habrastorage.org/getpro/habr/upload_files/9d5/497/b04/9d5497b04c3c3a80a2c6275e5148cb2c.png)
Эксп.№7
3D U-Net, подача объемом, плоскость [x, z],
Маска (слева) и готовая сегментация (справа),
оптимизированные параметры сети,
время обучения (65 epochs) ~ 14 мин.Результат вполне удовлетворительный, есть недочеты (вроде «похудевших» зубов). Возможно, исправим их в другом посте. Для этапа semantic segmentation я думаю мы сделали достаточно, теперь необходимо задать категории.
О размере подаваемых данныхПервоначальная идея при переходе на 3D архитектуру была в том чтобы делить данные не слайсами (как в данном посте) (1512, 120, 120) --> 21*(1, 72, 120, 120), а кубиками ~х*(30, 30, 30) или около того (результат этой попытки не был сохранен оп понятным причинам). Опытным путем понял 2 вещи: чем большими порциями ты подаешь 3-х мерные объекты, тем лучше результат (для моего конкретного случая); и нужно больше изучать теорию того, с чем работаешь.
Параметры сети подобраны так, что обучение 1 epochs на моей «старушке» занимает ~13сек, а размер конечной модели не превышает 2 мб (прошлая>80 мб). Время рабочего цикла примерно равно 1 epochs. Однако стоит понимать, это обучение и работа на данных достаточно маленького размера.
Для разделения на категории пришлось немного повозиться с функцией расчета ошибки и визуализацией данных. Первоначально поставил себе задачу разделить на 8 категорий + фон. О loss function и визуализации поговорим подробнее.
Код training loopimport torch
from tqdm import tqdm
from _loss_f import LossFunction
class TrainFunction:
def __init__(self,
data_loader,
device_for_training,
model_name,
model_name_pretrained,
model,
optimizer,
scale,
learning_rate: int = 1e-2,
num_epochs: int = 1,
transfer_learning: bool = False,
binary_loss_f: bool = True
):
self.data_loader = data_loader
self.device = device_for_training
self.model_name_pretrained = model_name_pretrained
self.semantic_binary = binary_loss_f
self.num_epochs = num_epochs
self.model_name = model_name
self.transfer = transfer_learning
self.optimizer = optimizer
self.learning_rate = learning_rate
self.model = model
self.scale = scale
def forward(self):
print('Running on the:', torch.cuda.get_device_name(self.device))
self.model.load_state_dict(torch.load(self.model_name_pretrained)) if self.transfer else None
optimizer = self.optimizer(self.model.parameters(), lr=self.learning_rate)
for epoch in range(self.num_epochs):
self.train_loop(self.data_loader, self.model, optimizer, self.scale, epoch)
torch.save(self.model.state_dict(), 'models/' + self.model_name+str(epoch+1)
+ '_epoch.pth') if (epoch + 1) % 10 == 0 else None
def train_loop(self, loader, model, optimizer, scales, i):
loop, epoch_loss = tqdm(loader), 0
loop.set_description('Epoch %i' % (self.num_epochs - i))
for batch_idx, (data, targets) in enumerate(loop):
data, targets = data.to(device=self.device, dtype=torch.float), \
targets.to(device=self.device, dtype=torch.long)
optimizer.zero_grad()
*тут секрет*
with torch.cuda.amp.autocast():
predictions = model(data)
loss = LossFunction(predictions, targets,
device_for_training=self.device,
semantic_binary=self.semantic_binary
).forward()
scales.scale(loss).backward()
scales.step(optimizer)
scales.update()
epoch_loss += (1 - loss.item())*100
loop.set_postfix(loss=loss.item())
print('Epoch-acc', round(epoch_loss / (batch_idx+1), 2))
4. Функция расчета ошибки
Мне в целом понравилось как проявляет себя Dice-loss в сегментации, только 'проблема' в том что он работает с форматом данных [0, 1]. Однако, если предварительно разделить данные на категории (а так же привести к формату [0, 1]), и пропускать пары (имеется ввиду «предсказание» и «маска» только одной категории) в стандартную Dice-loss функцию, то это может сработать.
Код categorical_dice_lossimport torch
class LossFunction:
def __init__(self,
prediction,
target,
device_for_training,
semantic_binary: bool = True,
):
self.prediction = prediction
self.device = device_for_training
self.target = target
self.semantic_binary = semantic_binary
def forward(self):
if self.semantic_binary:
return self.dice_loss(self.prediction, self.target)
return self.categorical_dice_loss(self.prediction, self.target)
@staticmethod
def dice_loss(predictions, targets, alpha=1e-5):
intersection = 2. * (predictions * targets).sum()
denomination = (torch.square(predictions) + torch.square(targets)).sum()
dice_loss = 1 - torch.mean((intersection + alpha) / (denomination + alpha))
return dice_loss
def categorical_dice_loss(self, prediction, target):
pr, tr = self.prepare_for_multiclass_loss_f(prediction, target)
target_categories, losses = torch.unique(tr).tolist(), 0
for num_category in target_categories:
categorical_target = torch.where(tr == num_category, 1, 0)
categorical_prediction = pr[num_category][:][:][:]
losses += self.dice_loss(categorical_prediction, categorical_target).to(self.device)
return losses / len(target_categories)
@staticmethod
def prepare_for_multiclass_loss_f(prediction, target):
prediction_prepared = torch.squeeze(prediction, 0)
target_prepared = torch.squeeze(target, 0)
target_prepared = torch.squeeze(target_prepared, 0)
return prediction_prepared, target_prepared
Тут просто, но всё равно объясню «categorical_dice_loss»:
подготовка данных (убираем ненужные в данном расчете измерения);
получения списка категорий, которые содержит каждый batch масок;
для каждой категории берем «прогноз» и «маску» соответствующих категорий, приводим значения к формату [0, 1] и пропускаем через стандартную Dice-loss;
складывая результаты и деля на количество категорий, получаем усредненное значение для каждого batct. Ну, а дальше всё без изменений.
Так же, думаю, помог бы перевод данных к one-hot формату, но только не в момент формирования основного дата сета (раздует в размере), а непосредственно перед расчетом ошибки, но я не проверял. Кто в курсе, напишите, пожалуйста, буду рад. Результат работы данной функции будет в Главе (5).
5. Визуализация данных
Так и хочется добавить »…как отдельный вид искусства». Начну с того что прочитать *.nrrd оказалось самым простым.
Кодimport nrrd
# читает в numpy
read = nrrd.read(data_path)
data, meta_data = read[0], read[1]
print(data.shape, np.max(data), np.min(data), meta_data, sep="\n")
(163, 112, 120)
14982
-2254
OrderedDict([('type', 'short'), ('dimension', 3), ('space', 'left-posterior-superior'), ('sizes', array([163, 112, 120])), ('space directions', array([[-0.5, 0. , 0. ],
[ 0. , -0.5, 0. ],
[ 0. , 0. , 0.5]])), ('kinds', ['domain', 'domain', 'domain']), ('endian', 'little'), ('encoding', 'gzip'), ('space origin', array([131.57200623, 80.7661972 , 32.29940033]))])Дальше — сложнее, как обратно перевести? Получается, если я ничего не путаю, необходимо из числа получить вершины и грани, между которыми образуется поверхность.
Неправильный путь
Иными словами, чтобы сделать куб нам необходимо 8 вершин и 12 треугольных поверхностей. В этом и состояла первая идея (до применения специальных библиотек) — заменить все пиксели (числа в 3-х мерной матрице) на такие кубики. Код я не сохранил, но смысл прост, рисуем куб на месте «пикселя» со сдвигом -1 по трем направлениям, потом следующий и т.д.
Выглядит это так же бредово, как и звучитОтрицательный результат — тоже результат, продолжаем. На этом этапе я уже понял, что без сторонних библиотек мне не обойтись. Первой попыткой в была пара Skimage и Stl.
from skimage.measure import marching_cubes
import nrrd
import numpy as np
from stl import mesh
path = 'some_path.nrrd'
data = nrrd.read(path)[0]
def three_d_creator(some_data):
vertices, faces, volume, _ = marching_cubes(some_data)
cube = mesh.Mesh(np.full(faces.shape[0], volume.shape[0], dtype=mesh.Mesh.dtype))
for i, f in enumerate(faces):
for j in range(3):
cube.vectors[i][j] = vertices[f[j]]
cube.save('name.stl')
return cube
stl = three_d_creator(datas)Пользовался этим способом, но иногда файлы «ломались» в процессе сохранения и не открывались. А на те, которые открывались, ругался встроенный в Win 10 3D Builder и постоянно пытался там что-то исправить. Так же еще придется «прикрутить» к коду модуль для просмотра 3D объектов без их сохранения. Решение «из коробки» дальше.
На момент написания статью пользуюсь v3do. Коротко, быстро, удобно и можно сразу осмотреть модель.
Код перевода npy в stl и вывода объекта на дисплейfrom vedo import Volume, show, write
prediction = 'some_data_path.npy'
def show_save(data, save=False):
data_multiclass = Volume(data, c='Set2', alpha=(0.1, 1), alphaUnit=0.87, mode=1)
data_multiclass.addScalarBar3D(nlabels=9)
show([(data_multiclass, "Multiclass teeth segmentation prediction")], bg='black', N=1, axes=1).close()
write(data_multiclass.isosurface(), 'some_name_.stl') if save else None
show_save(prediction, save=True)Названия функций говорят сами за себя.
Пришло время увидеть конечный результат всего вышесказанного. Томить не буду:
model.summary ()model = UNet(dim=3, in_channels=1, out_channels=9, n_blocks=3, start_filters=9).to(device)
print(summary(model, (1, 168*, 120, 120)))
"""
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv3d-1 [-1, 9, 168, 120, 120] 252
ReLU-2 [-1, 9, 168, 120, 120] 0
BatchNorm3d-3 [-1, 9, 168, 120, 120] 18
Conv3d-4 [-1, 9, 168, 120, 120] 2,196
ReLU-5 [-1, 9, 168, 120, 120] 0
BatchNorm3d-6 [-1, 9, 168, 120, 120] 18
MaxPool3d-7 [-1, 9, 84, 60, 60] 0
DownBlock-8 [[-1, 9, 84, 60, 60], [-1, 9, 168, 120, 120]] 0
Conv3d-9 [-1, 18, 84, 60, 60] 4,392
ReLU-10 [-1, 18, 84, 60, 60] 0
BatchNorm3d-11 [-1, 18, 84, 60, 60] 36
Conv3d-12 [-1, 18, 84, 60, 60] 8,766
ReLU-13 [-1, 18, 84, 60, 60] 0
BatchNorm3d-14 [-1, 18, 84, 60, 60] 36
MaxPool3d-15 [-1, 18, 42, 30, 30] 0
DownBlock-16 [[-1, 18, 18, 42, 30], [-1, 18, 84, 60, 60]] 0
Conv3d-17 [-1, 36, 42, 30, 30] 17,532
ReLU-18 [-1, 36, 42, 30, 30] 0
BatchNorm3d-19 [-1, 36, 42, 30, 30] 72
Conv3d-20 [-1, 36, 42, 30, 30] 35,028
ReLU-21 [-1, 36, 42, 30, 30] 0
BatchNorm3d-22 [-1, 36, 42, 30, 30] 72
DownBlock-23 [[-1, 36, 42, 30, 30], [-1, 36, 42, 30, 30]] 0
ConvTranspose3d-24 [-1, 18, 84, 60, 60] 5,202
ReLU-25 [-1, 18, 84, 60, 60] 0
BatchNorm3d-26 [-1, 18, 84, 60, 60] 36
Concatenate-27 [-1, 36, 84, 60, 60] 0
Conv3d-28 [-1, 18, 84, 60, 60] 17,514
ReLU-29 [-1, 18, 84, 60, 60] 0
BatchNorm3d-30 [-1, 18, 84, 60, 60] 36
Conv3d-31 [-1, 18, 84, 60, 60] 8,766
ReLU-32 [-1, 18, 84, 60, 60] 0
BatchNorm3d-33 [-1, 18, 84, 60, 60] 36
UpBlock-34 [-1, 18, 84, 60, 60] 0
ConvTranspose3d-35 [-1, 9, 168, 120, 120] 1,305
ReLU-36 [-1, 9, 168, 120, 120] 0
BatchNorm3d-37 [-1, 9, 168, 120, 120] 18
Concatenate-38 [-1, 18, 168, 120, 120] 0
Conv3d-39 [-1, 9, 168, 120, 120] 4,383
ReLU-40 [-1, 9, 168, 120, 120] 0
BatchNorm3d-41 [-1, 9, 168, 120, 120] 18
Conv3d-42 [-1, 9, 168, 120, 120] 2,196
ReLU-43 [-1, 9, 168, 120, 120] 0
BatchNorm3d-44 [-1, 9, 168, 120, 120] 18
UpBlock-45 [-1, 9, 168, 120, 120] 0
Conv3d-46 [-1, 9, 168, 120, 120] 90
================================================================
Total params: 108,036
Trainable params: 108,036
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 3.96
Forward/backward pass size (MB): 12170.30
Params size (MB): 0.41
Estimated Total Size (MB): 12174.66
----------------------------------------------------------------
"""*Ввиду ещё большего уменьшения параметров сети (фильтр[9, 18, 36, 72]), удалось уместить объект в память видеокарты целиком — 9*(168, 120, 120)
6. After words
Думал, что закончил, а оказалось — только начал. Тут еще есть над чем поработать. Мне, в целом, 2 этап не нравится, хоть он и работает. Зачем заново переопределять каждый пиксель, когда мне нужен целый регион? А если, образно, есть 28 разделенных регионов, зачем мне пытаться определить их все, не проще ли определить один зуб и завязать это всё на «условный» ориентированный/неориентированный граф? Или вместо U-net использовать GCNN и вместо Pytorch — Pytorch3D? Пятна, думаю, можно убрать с помощью выравнивания данных внутри bounding box (ведь один зуб может принадлежать только 1 категории). Но, возможно, это вопросы для следующей публикации.
Прототип (набросок)
Пример неориентированного графа на 28 категорий с «разделителями»Отдельное спасибо моей жене — Алёне, за особую поддержку во время этого «погружения в темноту».
Благодарю всех за внимание. Конструктивная критика и предложения, как исправлений, так и новых проектов — приветствуются.
