Segmentasi semantik dengan U-Net

Deep Learning untuk Gambar dengan PyTorch

Michal Oleszak

Machine Learning Engineer

Segmentasi semantik

Tidak membedakan instance berbeda dari kelas yang sama
Berguna untuk citra medis atau analisis citra satelit
Arsitektur populer: U-Net

Arsitektur U-Net

Diagram arsitektur U-Net

Encoder:

Lapisan konvolusi dan pooling
Downsampling: mengurangi dimensi spasial sambil menambah kedalaman

Arsitektur U-Net

Diagram arsitektur U-Net

Decoder:

Simetris terhadap encoder
Meng-upsample feature map dengan konvolusi tertransposisi

Arsitektur U-Net

Diagram arsitektur U-Net

Skip connection:

Tautan dari encoder ke decoder
Mempertahankan detail yang hilang saat downsampling

Konvolusi tertransposisi

Diagram konvolusi tertransposisi

Meng-upsample feature map di decoder: menambah tinggi dan lebar sambil mengurangi kedalaman
Proses konvolusi tertransposisi:
1. Sisipkan nol di antara atau di sekitar feature map input
2. Lakukan konvolusi biasa pada input yang dipadding nol

Konvolusi tertransposisi di PyTorch

import torch.nn as nn

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

U-Net: definisi lapisan

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)

Encoder:

Blok konvolusi

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)
)

Lapisan pooling

Decoder:
- Konvolusi tertransposisi
- Blok konvolusi

U-Net: metode 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)

Umpankan input melalui blok konvolusi encoder dan lapisan pooling
Decoder dan skip connection:
- Lewatkan representasi melalui konvolusi tertransposisi
- Konkatenasi dengan keluaran encoder terkait
- Lewatkan melalui blok konvolusi
- Ulangi untuk semua langkah decoder
Kembalikan keluaran langkah decoder terakhir

Menjalankan inferensi

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()

Gambar mobil asli

Mask semantik ditumpuk di atas gambar mobil

Ayo berlatih!

Deep Learning untuk Gambar dengan PyTorch