3D U-Net для медицинской сегментации

Эта модель представляет собой 3D U-Net, обученную на медицинских данных для задач сегментации.

Архитектура модели

Тип: 3D U-Net (MONAI)
Входные каналы: 1
Выходные каналы: 2 (фон + целевой класс)
Каналы: (16, 32, 64, 128, 256, 512)
Размер патча: 256x256x256
Функция потерь: Dice Loss

Использование

import os
import torch
from monai.networks.nets import UNet
from huggingface_hub import hf_hub_download

# Отключение предупреждения о symlinks на Windows
os.environ['HF_HUB_DISABLE_SYMLINKS_WARNING'] = '1'

# Загрузка модели
model = UNet(
    spatial_dims=3, 
    in_channels=1, 
    out_channels=2,
    channels=(16, 32, 64, 128, 256, 512), 
    strides=(2, 2, 2, 2, 2, 2),
    num_res_units=2, 
    dropout=0.2
)

# Загрузка весов безопасно (исправляет FutureWarning)
model_path = hf_hub_download(repo_id="Nikitapyscript/MRISegment_model", filename="best_metric_model.pth")
model.load_state_dict(torch.load(model_path, map_location='cpu', weights_only=True))
model.eval()

Предобработка данных

Модель ожидает входные данные в следующем формате:

Нормализация интенсивности: [-1000, 4000] -> [0, 1]
Размер входа: (1, 256, 256, 256)
Тип данных: torch.float32

Обучение

Модель была обучена с использованием:

Оптимизатор: Adam (lr=1e-4, weight_decay=1e-5)
Scheduler: ReduceLROnPlateau
Аугментации: случайные отражения, повороты, изменения интенсивности
Метрика: Dice Score

Требования

torch
monai
numpy