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 个卷积核。从编码层到解码层的对应层使用了跳跃连接。输入图像是来自造影前、FLAIR 和造影后序列的 3 通道脑部 MRI 切片。输出是一个单通道的异常区域概率图，其大小与输入图像相同。可以通过阈值处理将其转换为二值分割掩码，如下例所示。

示例

预训练模型的输入图像应具有 3 个通道，并被调整大小到 256x256 像素，并按体积进行 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]))

参考文献