이미지 분류, 객체 탐지와 더불어 고급 기술인 이미지 분할과 관련한 유용한 사이트에 대해 알아 볼 것이다.

"Segmentation Models"

공식 사이트

설치

pip install -U segmentation-models-pytorch

모델 생성

import segmentation_models_pytorch as smp

model = smp.Unet(
    encoder_name="resnet34",        # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
    encoder_weights="imagenet",     # use `imagenet` pre-trained weights for encoder initialization
    in_channels=1,                  # model input channels (1 for gray-scale images, 3 for RGB, etc.)
    classes=3,                      # model output channels (number of classes in your dataset)
)

segmentation 모델 중 하나인 Unet을 사용한 것이다.

이걸로 모델 생성은 끝... 정말 쉽다.
필자는 이 사이트 존재 유무를 몰라 모델을 코드로 작성하고 사용했다...

생성된 모델은 torch.nn.Module이라 기존에 사용하던 방식으로 모델을 대입하면 끝난다.

모델 종류와 각 모델 파라미터 설명

공식 사이트

나와있는 모델은 총 9가지로 Unet, Unet++, MAnet, Linknet, FPN, PSPNet, PAN, DeepLabV3, DeepLabV3+ 가 된다. 이들 중 마음에 드는 모델을 사용하면 되고 각 모델마다 파라미터가 다를 수 있으니 그것만 주의하면 된다.

그 밖의 인코더 선택이나 손실함수 관련된 것들이 모두 저 사이트에 있으니 저 사이트만 사용한다면 진짜 간편하게 segmentation model을 학습시키고 사용할 수 있다.

ImageNet

위 '모델 생성' 코드를 보면, Unet() 함수의 파라미터 중에

encoder_weights="imagenet"

가 있는 것을 볼 수 있다. 여기서 imagenet은 오픈 소스로 풀린 대규모 데이터셋이다. 불러오고 사용하는 방법은 이 블로그 주소로 들어가면 알 수 있다.