用于脑部MRI的U-Net
import torch
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet',
in_channels=3, out_channels=1, init_features=32, pretrained=True)
加载一个 U-Net 模型，该模型已在一个脑部 MRI 数据集上进行预训练，用于异常分割 kaggle.com/mateuszbuda/lgg-mri-segmentation 预训练模型需要 3 个输入通道、1 个输出通道，并且第一层有 32 个特征。
模型描述
这个 U-Net 模型包含四个级别的块，每个块在编码部分包含两个带有批量归一化和 ReLU 激活函数的卷积层，以及一个最大池化层；在解码部分则包含上采样卷积层。每个块中的卷积滤波器数量分别为 32、64、128 和 256。瓶颈层有 512 个卷积滤波器。从编码层到解码部分相应的层使用跳跃连接。输入图像是一个 3 通道脑部 MRI 切片，分别来自对比前、FLAIR 和对比后序列。输出是一个单通道异常区域概率图，其大小与输入图像相同。通过阈值化可以将其转换为二元分割掩码，如下面的示例所示。
示例
预训练模型的输入图像应具有 3 个通道，并调整为 256×256 像素，并按体积进行 Z 值标准化。
# Download an example image
import urllib
url, filename = ("https://github.com/mateuszbuda/brain-segmentation-pytorch/raw/master/assets/TCGA_CS_4944.png", "TCGA_CS_4944.png")
try: urllib.URLopener().retrieve(url, filename)
except: urllib.request.urlretrieve(url, filename)
import numpy as np
from PIL import Image
from torchvision import transforms
input_image = Image.open(filename)
m, s = np.mean(input_image, axis=(0, 1)), np.std(input_image, axis=(0, 1))
preprocess = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=m, std=s),
])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)
if torch.cuda.is_available():
input_batch = input_batch.to('cuda')
model = model.to('cuda')
with torch.no_grad():
output = model(input_batch)
print(torch.round(output[0]))
参考文献
