用于脑部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)
加载在脑部MRI体积数据集上预训练的U-Net模型，用于异常分割 kaggle.com/mateuszbuda/lgg-mri-segmentation 预训练模型需要3个输入通道、1个输出通道，并且第一层有32个特征。
模型描述
该U-Net模型包含四个级别的块，每个块在编码部分包含两个卷积层（带批量归一化和ReLU激活函数）和一个最大池化层，在解码部分则包含上采样卷积层。每个块中的卷积滤波器数量分别为32、64、128和256。瓶颈层有512个卷积滤波器。从编码层到解码部分的相应层使用了跳跃连接。输入图像是来自造影前、FLAIR和造影后序列的3通道脑部MRI切片。输出是一个与输入图像大小相同的单通道异常区域概率图。可以通过阈值处理将其转换为二进制分割掩码，如下例所示。
示例
预训练模型的输入图像应具有3个通道，并调整为256×256像素，并按体积进行z-score归一化。
# 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]))
参考文献
