用於腦部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]))
參考文獻
