[VOC PASCAL 2012] Semantic segmentation 하기 - 2

이전 포스팅 내용이 궁금하시면 아래의 링크를 참고하시면 감사하겠습니다!
   - [VOC PASCAL 2012] Semantic segmentation 하기 -1

 

지난 번 포스팅에서는 VOC 데이터셋을 준비하고 데이터를 분석하였다.

이번 포스팅에서는 데이터셋에 전처리와 증강을 적용하는 코드를 알아보고 학습에 사용할 코드를 작성해 볼 것이다.

 

우선, 지난 번 포스팅에서 잘못된 부분이 있어서 그 부분에 대해 먼저 짚고 넘어가겠다.

이전 포스팅에서는 학습과 검증 데이터셋의 이미지에 바로 평균과 표준편차를 구했지만, 이 부분은 resizing을 적용한
데이터에 대해 계산해야한다. 때문에 코드를 아래와 같이 수정해야한다.

 

import torch
from tqdm import tqdm
from Dataset import _VOCdataset
import torchvision.transforms.functional as F

# 데이터셋을 초기화합니다.
ds_train = _VOCdataset(mode='train')
ds_valid = _VOCdataset(mode='valid')

def calculate_mean_std(ds):
    sum_rgb = torch.tensor([0.0, 0.0, 0.0])
    sum_rgb_squared = torch.tensor([0.0, 0.0, 0.0])
    n_pixels = 0
    
    for i in tqdm(range(len(ds)), desc="Calculating mean and std"):
        sample = ds.__getitem__(i)
        img = sample['image']
        img = F.resize(img, (224,224))
        sum_rgb += img.float().sum(dim=[1,2])
        sum_rgb_squared += (img.float() ** 2).sum(dim=[1,2])
        n_pixels += img.shape[1] * img.shape[2]
    
    mean = sum_rgb / n_pixels
    std = (sum_rgb_squared / n_pixels - mean ** 2) ** 0.5
    return mean, std

mean_train, std_train = calculate_mean_std(ds_train)
mean_valid, std_valid = calculate_mean_std(ds_valid)

print(f"Training set mean: {mean_train}, std: {std_train}")
print(f"Validation set mean: {mean_valid}, std: {std_valid}")
=======================================================================================
Calculating mean and std: 100%|██████████| 1464/1464 [00:02<00:00, 534.56it/s]
Calculating mean and std: 100%|██████████| 1449/1449 [00:03<00:00, 461.32it/s]
Training set mean: tensor([0.4567, 0.4431, 0.4086]), std: tensor([0.2680, 0.2649, 0.2806])
Validation set mean: tensor([0.4568, 0.4387, 0.4014]), std: tensor([0.2650, 0.2634, 0.2794])

 

이번 학습에서는 이미지의 크기를 $224 \times 224$를 적용할 것이기 때문에, 저 크기로 변환한 크기에 대해 계산해준다.

약간의 차이는 있지만 지난 값과 거의 비슷하다.

---

1. 전처리 및 데이터 증강

 

이번에 사용할 전처리는 Normalization과 Resizing, 증강 기법은 Horizontal Flip이다. 해당 기능을 수행할 코드는 아래와 같다.

 

import torchvision.transforms.functional as F
import random
import torch


class Random_processing(object):
    def __init__(self, mean=False, std=False):
        '''
        mean : defalut=False, format : [0.0, 0.0, 0.0]
        std : default=False, format : [0.0, 0.0, 0.0]
        '''
        self.mean = mean
        self.std = std

    def __call__(self, samples):
        img, lbl = samples['image'], samples['label']
        lbl = torch.where(lbl == 255, torch.tensor(21, dtype=lbl.dtype), lbl)

        img = F.resize(img, [224, 224], interpolation=F.InterpolationMode.BILINEAR)
        lbl = F.resize(lbl, [224, 224], interpolation=F.InterpolationMode.NEAREST)

        if random.random() > 0.5:
            img, lbl = F.hflip(img), F.hflip(lbl)

        if self.mean and self.std:
            img = F.normalize(img, self.mean, self.std)

        sample = {'image': img, 'label': lbl}

        return sample

 

Resizing과 Normalization을 적용하기 앞서서, 255로 할당된 class를 21로 바꿔주었다.

 

---

2. 모델 구성

 

이번에 비교할 모델은 Vgg16 을 backbone으로 사용하는 FCN-8s 모델과 UNet 모델이다. 각 모델의 대한 자세한 정보는

아래의 링크를 참고하면 좋을 것 같다.

• U-Net 논문 리뷰
  https://cvad.tistory.com/10

[논문리뷰] U-Net : Convolutional Networks for Biomedical Image Segmentation

• FCN 논문 리뷰
  https://cvad.tistory.com/13

[논문리뷰] Fully Convolutional Networks for Semantic Segmentation

 

Vgg 16 모델은 Imagenet1k로 사전학습된 모델을 사용하였다. 이제 각 모델의 코드를 알아보자.

먼저, FCN 모델의 코드는 아래와 같다.

 

import torch.nn as nn
import torchvision.models as models

class Vgg16_backbone(nn.Module):
    def __init__(self):
        super(Vgg16_backbone, self).__init__()
        self.layers = [0,5,10,17,24,31]
        self.weights = models.VGG16_Weights.DEFAULT
        self.vgg16 = models.vgg16(weights=self.weights)
        self.features = self.vgg16.features

    def forward(self, x):
        output = {}
        for i in range(len(self.layers)-1):
            for layer in range(self.layers[i], self.layers[i+1]):
                x = self.features[layer](x)
            output[f'x{i+1}'] = x
        return output
    
class O_FCN8s(nn.Module):
    def __init__(self, n_class, backbone=Vgg16_backbone()):
        super().__init__()
        self.n_class = n_class
        self.backbone = backbone
        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512, 512, 3, 2, 1, dilation=1, output_padding=1)
        self.deconv2 = nn.ConvTranspose2d(512, 256, 3, 2, 1, dilation=1, output_padding=1)
        self.deconv3 = nn.ConvTranspose2d(256, 128, 3, 2, 1, dilation=1, output_padding=1)
        self.deconv4 = nn.ConvTranspose2d(128, 64, 3, 2, 1, dilation=1, output_padding=1)
        self.deconv5 = nn.ConvTranspose2d(64, 32, 3, 2, 1, dilation=1, output_padding=1)
        self.classifier = nn.Conv2d(32, self.n_class, kernel_size=1)

    def forward(self, x):
        out_backbone = self.backbone(x)
        
        pool5 = out_backbone['x5']
        pool4 = out_backbone['x4']
        pool3 = out_backbone['x3']

        segmap = self.relu(self.deconv1(pool5))
        segmap = segmap + pool4
        segmap = self.relu(self.deconv2(segmap))
        segmap = segmap + pool3
        segmap = self.relu(self.deconv3(segmap))
        segmap = self.relu(self.deconv4(segmap))
        segmap = self.relu(self.deconv5(segmap))
        segmap = self.classifier(segmap)

        return segmap

 

FCN-8s 모델이기 때문에, vgg16의 4번째 pooling layer와 5번째 pooling layer의 feature map을 concat 해줬다.

 

다음은 U-Net모델이다. 코드는 아래의 wandb의 예제코드를 사용하였다.

https://wandb.ai/ishandutta/semantic_segmentation_unet/reports/Semantic-Segmentation-with-UNets-in-PyTorch--VmlldzoyMzA3MTk1

Semantic Segmentation with UNets in PyTorch

 

import torch
import torch.nn as nn
from torchvision.models import vgg16_bn

def conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True)
    )


def up_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
        nn.ReLU(inplace=True)
    )


class UNet(nn.Module):
    def __init__(self, pretrained=True, out_channels=22):
        super().__init__()


        self.encoder = vgg16_bn(pretrained=pretrained).features
        self.block1 = nn.Sequential(*self.encoder[:6])
        self.block2 = nn.Sequential(*self.encoder[6:13])
        self.block3 = nn.Sequential(*self.encoder[13:20])
        self.block4 = nn.Sequential(*self.encoder[20:27])
        self.block5 = nn.Sequential(*self.encoder[27:34])


        self.bottleneck = nn.Sequential(*self.encoder[34:])
        self.conv_bottleneck = conv(512, 1024)


        self.up_conv6 = up_conv(1024, 512)
        self.conv6 = conv(512 + 512, 512)
        self.up_conv7 = up_conv(512, 256)
        self.conv7 = conv(256 + 512, 256)
        self.up_conv8 = up_conv(256, 128)
        self.conv8 = conv(128 + 256, 128)
        self.up_conv9 = up_conv(128, 64)
        self.conv9 = conv(64 + 128, 64)
        self.up_conv10 = up_conv(64, 32)
        self.conv10 = conv(32 + 64, 32)
        self.conv11 = nn.Conv2d(32, out_channels, kernel_size=1)
        
    def forward(self, x):
        block1 = self.block1(x)
        block2 = self.block2(block1)
        block3 = self.block3(block2)
        block4 = self.block4(block3)
        block5 = self.block5(block4)


        bottleneck = self.bottleneck(block5)
        x = self.conv_bottleneck(bottleneck)


        x = self.up_conv6(x)
        x = torch.cat([x, block5], dim=1)
        x = self.conv6(x)


        x = self.up_conv7(x)
        x = torch.cat([x, block4], dim=1)
        x = self.conv7(x)


        x = self.up_conv8(x)
        x = torch.cat([x, block3], dim=1)
        x = self.conv8(x)


        x = self.up_conv9(x)
        x = torch.cat([x, block2], dim=1)
        x = self.conv9(x)


        x = self.up_conv10(x)
        x = torch.cat([x, block1], dim=1)
        x = self.conv10(x)


        x = self.conv11(x)


        return x

 

다음 포스팅부터는 학습 코드를 작성하는 과정에 대해 다루겠다.

 

포스팅에 대한 질문이나 잘못된 사항있다면 댓글 남겨주시면 감사하겠습니다.