Segmentation

Добрый - всем, зовут Андрей(27). Постараюсь коротко. Почему программирование? По образованию - бакалавр электромеханик, профессию знаю. Отработал 2 года на должности инженера-энергетика в буровой компании вполне успешно, вместо повышения написал заявление - сгорел, да не по мне оказалось это всё. Нравится создавать, находить решения сложных задач, с ПК в обнимку с сознательных лет. Выбор очевиден. Вначале (полгода назад), всерьёз думал записаться на курсы от Я или подобные. Начитался отзывов, поговорил с участниками и понял что с получением информацией проблем нет. Так нашел сайт, там получил базу по Python и с ним уже начал свой путь (сейчас там постепенно изучаю всё, что связано с ML). Сразу заинтересовало машинное обучение, CV в частности. Придумал себе задачу и вот здесь (по мне, так отличный способ учиться).

import nrrdimport torchimport torchvision.transforms as tfclass 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 tudclass 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

На выходе имеем следующее:

*эти размеры менялись, в зависимости от эксперимента, подробности - дальше.

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,721Trainable params: 7,702,721Non-trainable params: 0----------------------------------------------------------------Input size (MB): 0.08Forward/backward pass size (MB): 7434.08Params size (MB): 29.38Estimated Total Size (MB): 7463.54"""

Это размер глубины. Был подобран так чтобы:

• количество данных(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,321Trainable params: 22,400,321Non-trainable params: 0----------------------------------------------------------------Input size (MB): 0.61Forward/backward pass size (MB): 15974.12Params size (MB): 85.45Estimated Total Size (MB): 16060.18----------------------------------------------------------------"""

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,075Trainable params: 430,075Non-trainable params: 0----------------------------------------------------------------Input size (MB): 1.32Forward/backward pass size (MB): 5744.38Params size (MB): 1.64Estimated 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*(выше).

Первоначальная идея при переходе на 3D архитектуру была в том чтобы делить данные не слайсами (как в данном посте) (1512, 120, 120) --> 21*(1, 72, 120, 120), а кубиками ~х*(30, 30, 30) или около того (результат этой попытки не был сохранен оп понятным причинам). Опытным путем понял 2 вещи: чем большими порциями ты подаешь 3-х мерные объекты, тем лучше результат(для моего конкретного случая); и нужно больше изучать теорию того, с чем работаешь.

Параметры сети подобраны так, что обучение 1 epochs на моей "старушке" занимает ~13сек, а размер конечной модели не превышает 2мб (прошлая>80мб). Время рабочего цикла примерно равно 1 epochs. Однако стоит понимать, это обучение и работа на данных достаточно маленького размера.

import torchfrom tqdm import tqdmfrom _loss_f import LossFunctionclass 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))

import torchclass 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. Ну а дальше всё без изменений.

import nrrd# читает в numpyread = 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_cubesimport nrrdimport numpy as npfrom stl import meshpath = '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 cubestl = three_d_creator(datas)

Пользовался этим способом, но иногда файлы "ломались" в процессе сохранения и не открывались. А на те, которые открывались, ругался встроенный в Win 10 3D Builder и постоянно пытался там что-то исправить. Так же еще придется "прикрутить" к коду модуль для просмотра 3D объектов без их сохранения. Решение "из коробки" дальше.

from vedo import Volume, show, writeprediction = '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 = 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,036Trainable params: 108,036Non-trainable params: 0----------------------------------------------------------------Input size (MB): 3.96Forward/backward pass size (MB): 12170.30Params size (MB): 0.41Estimated Total Size (MB): 12174.66----------------------------------------------------------------    """

*Ввиду ещё большего уменьшения параметров сети(фильтр[9, 18, 36, 72]), удалось уместить объект в память видеокарты целиком - 9*(168, 120, 120)