Segmentare semantică cu U-Net

Deep Learning pentru imagini cu PyTorch

Michal Oleszak

Machine Learning Engineer

Segmentare semantică

Nu distinge instanțele diferite ale aceleiași clase
Util pentru imagistică medicală sau analiza imaginilor satelitare
Arhitectură populară: U-Net

Arhitectura U-Net

Diagramă arhitectură U-Net

Encoder:

Straturi convoluționale și de pooling
Downsampling: reduce dimensiunile spațiale, crescând adâncimea

Arhitectura U-Net

Diagramă arhitectură U-Net

Decoder:

Simetric față de encoder
Upsamplează hărțile de caracteristici cu convoluții transpuse

Arhitectura U-Net

Diagramă arhitectură U-Net

Conexiuni skip:

Legături de la encoder la decoder
Păstrează detaliile pierdute la downsampling

Convoluție transpusă

Diagramă convoluție transpusă

Upsamplează hărțile de caracteristici în decoder: crește înălțimea și lățimea, reducând adâncimea
Procesul de convoluție transpusă:
1. Inserare de zerouri între sau în jurul hărții de caracteristici de intrare
2. Aplicare convoluție standard pe intrarea cu zero-padding

Convoluție transpusă în PyTorch

import torch.nn as nn

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

U-Net: definiții straturi

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:

Blocuri convoluționale

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

Strat de pooling

Decoder:
- Convoluții transpuse
- Blocuri convoluționale

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

Intrarea parcurge blocurile convoluționale și straturile de pooling ale encoderului
Decoder și conexiuni skip:
- Intrarea codificată parcurge convoluția transpusă
- Concatenare cu ieșirea corespunzătoare din encoder
- Parcurgere bloc convoluțional
- Repetare pentru toți pașii decoderului
Returnează ieșirea ultimului pas al decoderului

Rularea inferenței

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

Imaginea originală a mașinii

Mască semantică suprapusă pe imaginea mașinii

Să exersăm!

Deep Learning pentru imagini cu PyTorch