Phân đoạn ngữ nghĩa với U-Net

Deep Learning cho Ảnh với PyTorch

Michal Oleszak

Machine Learning Engineer

Phân đoạn ngữ nghĩa

Không phân biệt các thể hiện khác nhau của cùng một lớp
Hữu ích cho ảnh y khoa hoặc ảnh vệ tinh
Kiến trúc phổ biến: U-Net

Kiến trúc U-Net

Sơ đồ kiến trúc U-Net

Bộ mã hoá:

Lớp tích chập và pooling
Downsampling: giảm kích thước không gian, tăng độ sâu

Kiến trúc U-Net

Sơ đồ kiến trúc U-Net

Bộ giải mã:

Đối xứng với bộ mã hoá
Upsample đặc trưng bằng tích chập chuyển vị

Kiến trúc U-Net

Sơ đồ kiến trúc U-Net

Skip connection:

Liên kết từ bộ mã hoá sang bộ giải mã
Giữ chi tiết mất đi khi downsampling

Tích chập chuyển vị

Sơ đồ tích chập chuyển vị

Upsample đặc trưng trong bộ giải mã: tăng chiều cao, chiều rộng và giảm độ sâu
Quy trình tích chập chuyển vị:
1. Chèn các số 0 vào giữa/bao quanh bản đồ đặc trưng đầu vào
2. Thực hiện tích chập thường trên đầu vào đã đệm 0

Tích chập chuyển vị trong PyTorch

import torch.nn as nn

upsample = nn.ConvTranspose2d(
    in_channels=in_channels,
    out_channels=out_channels,
    kernel_size=2,
    stride=2,
)

U-Net: định nghĩa các lớp

class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()


        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)
        self.enc4 = self.conv_block(256, 512)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)


        self.upconv3 = nn.ConvTranspose2d(512, 256,
                                          kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(256, 128,
                                          kernel_size=2, stride=2)
        self.upconv1 = nn.ConvTranspose2d(128, 64,
                                          kernel_size=2, stride=2)


        self.dec1 = self.conv_block(512, 256)
        self.dec2 = self.conv_block(256, 128)
        self.dec3 = self.conv_block(128, 64)
        self.out = nn.Conv2d(64, out_channels, kernel_size=1)

Bộ mã hoá:

Khối tích chập

def conv_block(self, in_channels, out_channels):
return nn.Sequential(
  nn.Conv2d(in_channels, out_channels),
  nn.ReLU(inplace=True),
  nn.Conv2d(out_channels, out_channels),
  nn.ReLU(inplace=True)
)

Lớp pooling

Bộ giải mã:
- Tích chập chuyển vị
- Khối tích chập

U-Net: phương thức forward

def forward(self, x):

    x1 = self.enc1(x)
    x2 = self.enc2(self.pool(x1))
    x3 = self.enc3(self.pool(x2))
    x4 = self.enc4(self.pool(x3))


    x = self.upconv3(x4)

    x = torch.cat([x, x3], dim=1)

    x = self.dec1(x)


    x = self.upconv2(x)
    x = torch.cat([x, x2], dim=1)
    x = self.dec2(x)

    x = self.upconv1(x)
    x = torch.cat([x, x1], dim=1)
    x = self.dec3(x)


    return self.out(x)

Truyền đầu vào qua các khối tích chập và lớp pooling của bộ mã hoá
Bộ giải mã và skip connection:
- Truyền đầu vào đã mã hoá qua tích chập chuyển vị
- Nối với đầu ra bộ mã hoá tương ứng
- Truyền qua khối tích chập
- Lặp lại cho mọi bước giải mã
Trả về đầu ra của bước giải mã cuối

Chạy suy luận

model = UNet()
model.eval()


image = Image.open("car.jpg")
transform = transforms.Compose([transforms.ToTensor()])
image_tensor = transform(image).unsqueeze(0)


with torch.no_grad():
    prediction = model(image_tensor).squeeze(0)


plt.imshow(prediction[1, :, :])
plt.show()

Ảnh ô tô gốc

Mặt nạ ngữ nghĩa chồng lên ảnh ô tô

Hãy luyện tập!

Deep Learning cho Ảnh với PyTorch