[Instance Segmentation] 캐글 - Sartorius (3) UPerNet 전처리~학습~시각화

대회 링크 : https://www.kaggle.com/competitions/sartorius-cell-instance-segmentation/overview

코드 링크 : https://github.com/euiraekim/kaggle-sartorius-cell-instance-segmentation

 

 

이전 포스팅에서는 YOLOX-x 모델을 학습해 bounding box를 찾아냈고 visualization까지 해서 확인했다.

이번에는 현재 Sementic Segmentation의 sota 모델인 UPerNet모델을 학습할 것이다.

 

목표

mmsegmentation을 사용하여 UPerNet을 학습할 것이다. 전체 이미지를 학습시키지 않고 ground truth bounding box로 해당 부분을 crop하여 각각을 input image로 사용할 것이다.

 

이를 위해 먼저 mmseg 프레임워크에서 학습할 수 있는 형태로 데이터셋을 바꿔주는 전처리를 할 것이다. 그리고 학습을 하고 시각화를 할 것이다.

 

전처리

같은 mmlab에서 만든 프레임워크지만 mmdetection과 mmsegmentation은 사용 방법이 좀 다르다. (지금 생각해보면 당연) 그런데 처음에는 같은 방식으로 시도하다가 잘 안돼서 애를 많이 먹었다.

 

먼저 전처리를 해준다. 포스팅 1편에서 원본 데이터셋을 5fold의 coco format으로 만드는 전처리를 했었다. 이 포스팅에서는 이 중 젤 첫 번째 fold를 이용하여 mmsegmentation용 train set과 valid set을 만들어 볼 것이다.

 

코드는 열심히 만들어서 최상단 전체 코드의 utils/convert_to_mmseg.py에 저장해뒀고 아래와 같다.

import numpy as np
from PIL import Image
import os
from tqdm import tqdm

import mmcv
from pycocotools.coco import COCO
import pycocotools.mask as mask_utils

# 원본 데이터셋의 전체 image가 들어있는 폴더 경로
image_path = '../data/train'
# class는 cell과 background로 2가지
classes = ('cell', 'bg')
palette = [[0, 0, 0], [255, 255, 255]]

def make_dir(path):
    if not os.path.isdir(path):
        os.makedirs(path)

def convert_to_mmseg(coco_path, save_path, mode='train'):
    image_save_path = os.path.join(save_path, 'images', mode)
    ann_save_path = os.path.join(save_path, 'annotations', mode)
    make_dir(image_save_path)
    make_dir(ann_save_path)

    coco = COCO(coco_path)
    # 모든 annotation에 대하여 루프를 돌림
    for ann in tqdm(coco.loadAnns(coco.getAnnIds())):
        x_min, y_min, width, height = ann['bbox']
        ann_id = ann['id']
        img = Image.open(os.path.join(image_path, coco.loadImgs(ann['image_id'])[0]['file_name']))
        
        # 이미지에서 bbox 부분 crop 후 저장
        crop_img = img.crop((x_min, y_min, x_min + width, y_min + height))
        crop_img.save(os.path.join(image_save_path, str(ann_id)+'.png'))
        
        # 마스크를 가져오고 b-box에 맞게 crop 후 palette를 입혀 저장
        # 이 파일이 학습 시 annotation이 된다.
        mask_np = mask_utils.decode(ann['segmentation'])
        crop_mask = mask_np[int(y_min):int(y_min+height), int(x_min):int(x_min+width)]
        seg_img = Image.fromarray(crop_mask).convert('P')
        seg_img.putpalette(np.array(palette, dtype=np.uint8))
        seg_img.save(os.path.join(ann_save_path, str(ann_id)+'.png'))

        # if ann['id'] == 5:
        #     break


def main():
    convert_to_mmseg('../data/dtrain_g0.json', '../data/segmentation', 'train_g0')
    convert_to_mmseg('../data/dval_g0.json', '../data/segmentation', 'valid_g0')


if __name__ == '__main__':
    main()

 

다음 명령어를 치면 mmsegmentation을 학습하기 위한 데이터셋이 만들어진다.

cd utils
python convert_to_mmseg.py
cd ..

data 폴더 안에 segmentation 폴더가 만들어지고 그 안에 데이터셋이 저장된다. 구조는 다음과 같다.

├─segmentation
│  ├─annotations
│  │  ├─train_g0
│  │  └─valid_g0
│  └─images
│      ├─train_g0
│      └─valid_g0

 

 

UPerNet 학습

mm시리즈같은 프레임워크의 경우 누군가가 config 폴더를 띡 만들어주고 train.py만 실행해~ 라고 하면 정말 정말 쉽다. 하지만 데이터셋은 항상 custom이고 그에 맞는 하이퍼파라미터나 모델은 따로 있기에 항상 스스로 만들줄 알아야한다. 그래서 어렵다.(아직은)

 

설정 파일은 다음과 같다. 전체 코드의 configs/segmentation/upernet_kaggle.py에 저장돼있다.

# model settings
norm_cfg = dict(type='SyncBN', requires_grad=True)
backbone_norm_cfg = dict(type='LN', requires_grad=True)

checkpoint_file = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_tiny_patch4_window7_224_20220317-1cdeb081.pth'  # noqa

model = dict(
    type='EncoderDecoder',
    pretrained=None,
    backbone=dict(
        type='SwinTransformer',
        init_cfg=dict(type='Pretrained', checkpoint=checkpoint_file),
        pretrain_img_size=224,
        embed_dims=96,
        patch_size=4,
        window_size=7,
        mlp_ratio=4,
        depths=[2, 2, 6, 2],
        num_heads=[3, 6, 12, 24],
        strides=(4, 2, 2, 2),
        out_indices=(0, 1, 2, 3),
        qkv_bias=True,
        qk_scale=None,
        patch_norm=True,
        drop_rate=0.,
        attn_drop_rate=0.,
        drop_path_rate=0.3,
        use_abs_pos_embed=False,
        act_cfg=dict(type='GELU'),
        norm_cfg=backbone_norm_cfg),
    decode_head=dict(
        type='UPerHead',
        in_channels=[96, 192, 384, 768],
        in_index=[0, 1, 2, 3],
        pool_scales=(1, 2, 3, 6),
        channels=512,
        dropout_ratio=0.1,
        num_classes=2,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0)),
    auxiliary_head=dict(
        type='FCNHead',
        in_channels=384,
        in_index=2,
        channels=256,
        num_convs=1,
        concat_input=False,
        dropout_ratio=0.1,
        num_classes=2,
        norm_cfg=norm_cfg,
        align_corners=False,
        loss_decode=dict(
            type='CrossEntropyLoss', use_sigmoid=False, loss_weight=0.4)),
    # model training and testing settings
    train_cfg=dict(),
    test_cfg=dict(mode='whole'))

# dataset settings
dataset_type = 'CustomDataset'
data_root = 'data/segmentation/'
img_norm_cfg = dict(
    mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375], to_rgb=True)
img_scale = (128, 128)
classes = ('cell', 'bg')
palette = [[255, 255, 255], [0, 0, 0]]
train_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='LoadAnnotations', reduce_zero_label=True),
    dict(type='Resize', img_scale=img_scale),
    dict(type='RandomFlip', prob=0.5),
    dict(type='PhotoMetricDistortion'),
    dict(type='Normalize', **img_norm_cfg),
    dict(type='DefaultFormatBundle'),
    dict(type='Collect', keys=['img', 'gt_semantic_seg']),
]
test_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(
        type='MultiScaleFlipAug',
        img_scale=img_scale,
        # img_ratios=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
        flip=False,
        transforms=[
            dict(type='Resize', keep_ratio=True),
            dict(type='RandomFlip'),
            dict(type='Normalize', **img_norm_cfg),
            dict(type='ImageToTensor', keys=['img']),
            dict(type='Collect', keys=['img']),
        ])
]
data = dict(
    samples_per_gpu=64,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        classes=classes,
        palette=palette,
        data_root=data_root,
        img_dir='images/train_g0',
        ann_dir='annotations/train_g0',
        img_suffix='.png',
        pipeline=train_pipeline),
    val=dict(
        type=dataset_type,
        classes=classes,
        palette=palette,
        data_root=data_root,
        img_dir='images/valid_g0',
        ann_dir='annotations/valid_g0',
        img_suffix='.png',
        pipeline=test_pipeline),
    test=dict(
        type=dataset_type,
        classes=classes,
        palette=palette,
        data_root=data_root,
        img_dir='images/valid_g0',
        ann_dir='annotations/valid_g0',
        img_suffix='.png',
        pipeline=test_pipeline))

# AdamW optimizer, no weight decay for position embedding & layer norm
# in backbone
optimizer = dict(
    # _delete_=True,
    type='AdamW',
    lr=0.00006 / 64,
    betas=(0.9, 0.999),
    weight_decay=0.01,
    paramwise_cfg=dict(
        custom_keys={
            'absolute_pos_embed': dict(decay_mult=0.),
            'relative_position_bias_table': dict(decay_mult=0.),
            'norm': dict(decay_mult=0.)
        }))
optimizer_config = dict()
# learning policy
lr_config = dict(policy='poly', power=0.9, min_lr=1e-4, by_epoch=False)
# runtime settings
runner = dict(type='IterBasedRunner', max_iters=30000)
checkpoint_config = dict(by_epoch=False, interval=200)
evaluation = dict(interval=1000, metric='mIoU', pre_eval=True)

# yapf:disable
log_config = dict(
    interval=100,
    # _delete_=True,
    policy='poly',
    warmup='linear',
    warmup_iters=1500,
    warmup_ratio=1e-6,
    power=1.0,
    min_lr=0.0,
    by_epoch=False,
    hooks=[
        dict(type='TextLoggerHook', by_epoch=False),
        # dict(type='TensorboardLoggerHook')
    ])
# yapf:enable
dist_params = dict(backend='nccl')
log_level = 'INFO'
load_from = None
resume_from = None
workflow = [('train', 1)]
cudnn_benchmark = True

 

각자의 환경에 맞는 버전으로 mmsegmentation을 설치하고

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu111/torch1.10/index.html
pip install mmsegmentation

다음 명령어를 입력해 학습을 한다.

python utils/segmentation/train.py configs/segmentation/upernet_kaggle.py

 

 

시각화

개별 crop image에 대한 시각화를 해보자. 본 시각화 과정은 detection 포스팅 때의 과정과 거의 같다.

 

먼저 config 파일과 가중치 파일을 이용해 모델을 불러오자.

from mmseg.apis import init_segmentor, inference_segmentor

config_path = './work_dirs/upernet_kaggle/upernet_kaggle.py'
checkpoint_path = './work_dirs/upernet_kaggle/iter_1000.pth'

model = init_segmentor(config_path, checkpoint_path, device='cuda:0')

 

모든 valid 이미지에 대하여 inference하자. 이미지가 너무 많으므로 500개를 랜덤 샘플링했다.

import os
import time
import random
from tqdm import tqdm

image_path = './data/segmentation/images/valid_g0'
image_list = os.listdir(image_path)
image_list = random.sample(image_list, 500)

start = time.time()

# inference_detector의 인자로 string(file경로), ndarray가 단일 또는 list형태로 입력 될 수 있음. 
result = []
for image in tqdm(image_list):
    result.append(inference_segmentor(model, os.path.join(image_path, image)))
    
end = time.time()
print(str((end-start)/len(image_list)) + ' sec per image')

###
0.03536668872833252 sec per image
###

 

각각 원본과 모델의 결과가 시각화된 이미지파일을 각각 저장하자.

import shutil

result_path = './result/upernet_kaggle'

for i in tqdm(range(len(image_list))):
    image = os.path.join(image_path, image_list[i])

    model.show_result(image, result[i],
                        out_file=os.path.join(result_path, f'img_{i}_result.jpg'), opacity=0.5)
    # 원본 이미지 저장
    shutil.copy(image, os.path.join(result_path, f'img_{i}_src.jpg'))

 

그리고 예시 결과 사진을 몇 개 가져와보았다. 왼쪽이 원본, 오른쪽은 모델이 배경이라고 판단한 부분을 검정색으로 칠한 이미지다.

보다시피 성능은 조금 눈물나는 수준이지만 목적에 맞게 따로 더 하이퍼파라미터 튜닝 등을 하지는 않을 예정이다. 그 이외에 목적한 바에 따라 얻은 것들이 많다.

 

 

전체 이미지 시각화

이번에는 crop 이미지들에 대한 결과들을 모아서 crop하기 전 원본 사진에 입혀 시각화해 볼 차례다.

 

모델을 가져온다. crop 이미지 시각화할 때랑 같다. 이어서 진행한다면 이미 모델이 메모리에 있으므로 아래 코드는 필요없다.

from mmseg.apis import init_segmentor, inference_segmentor

config_path = './work_dirs/upernet_kaggle/upernet_kaggle.py'
checkpoint_path = './work_dirs/upernet_kaggle/iter_1000.pth'

model = init_segmentor(config_path, checkpoint_path, device='cuda:0')

 

필요 패키지들을 임포트하고 원본 image들의 경로, valid 부분 coco 경로, valid crop image 경로를 설정하고 COCO 오브젝트를 불러온다.

import os
from tqdm import tqdm
from PIL import Image
import numpy as np
from pycocotools.coco import COCO
import matplotlib.pyplot as plt
import random

image_path = './data/train'
coco_path = './data/dval_g0.json'
crop_image_path = './data/segmentation/images/valid_g0'

val_coco = COCO(coco_path)

 

전체 원본 이미지 중 랜덤하게 10개를 뽑아서 원본 이미지, 모델이 추론한 mask를 입힌 이미지를 입힌 이미지를 plot한다.

imgIds = val_coco.getImgIds()
imgIds = random.sample(imgIds, 10)
for imgId in tqdm(imgIds):
    img_file = val_coco.loadImgs(imgId)[0]['file_name']
    image_np = np.array(Image.open(os.path.join(image_path, img_file)).convert('RGB'))

    # 원본 이미지
    plt.imshow(image_np)
    plt.show()

    anns = val_coco.loadAnns(val_coco.getAnnIds(imgIds=imgId))

    mask_np = np.zeros(image_np.shape[:2])
    for ann in anns:
        x_min, y_min, width, height = ann['bbox']
        crop_file_path = os.path.join(crop_image_path, str(ann['id']) + '.png')
        result = inference_segmentor(model, crop_file_path)

        # result[0] = 0은 class가 0(cell)인 부분
        # mask_np에 cell 부분만 1씩 더해주는 코드
        mask_np[int(y_min):int(y_min+height), int(x_min):int(x_min+width)] += result[0] == 0
    
    # mask가 0보다 큰 부분은 255, 0인 부분은 그대로 0으로 채움
    mask_np = np.where(mask_np>0, 255, 0)
    plt.imshow(mask_np)
    plt.show()