[Instance Segmentation] 캐글 - Sartorius (1) 데이터 전처리

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

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

 

목표

Sartorius의 1등 솔루션을 참고하여 내 스타일대로 구현한다. 이번 포스팅에서는 데이터셋을 살펴보고 이를 학습에 적합한 형태로 변환할 것이다.

 

데이터 살펴보기

대회 링크에서 데이터를 다운받으면 다음과 같은 폴더가 생긴다. 프로젝트 루트 폴더 아래 data 폴더를 만들고 그 아래에 넣어두자.

• LIVE_Cell_dataset_2021 : 이 대회의 선행연구의 데이터셋이다. 이 대회에 포함된 cell class가 일부 포함되어 있고 다른 종류의 cell class들도 포함되어 있다. 사용한다면 대회의 본 데이터셋을 학습시키기 전 pretrain 데이터로 사용될 예정이다.
• test : 테스트 PNG 이미지 파일들이다. (3개)
• train : 학습 PNG 이미지 파일들이다. (606개)
• train_semi_supervised : 라벨링이 되어있지 않은 이미지들이다. 사용하지 않을 예정이다.
• sample_submission.csv : 샘플 제출 파일이다.
• train.csv : train 이미지들의 라벨링 파일이다.

모든 이미지의 크기는 704 x 520으로 동일하다.

live cell의 이미지는 3727개(train+val)이고, 대회의 이미지는 총 606개다. live cell로 pretrain한 후 대회의 이미지를 학습시키는 방법이 좋을 수 있겠다.

 

live cell의 cell 종류는 8개이지만 주어진 coco format의 데이터셋에는 category가 나뉘어있지 않고 cell의 종류가 file name에 써있다. 이에 대한 처리를 해야할 것 같다. 또한 대회의 데이터셋은 class가 3개고, coco format이 아니므로 변환해줘야한다.

 

LIVECell 데이터 확인 / 전처리

먼저 이 중 랜덤하게 3장의 이미지를 뽑아 원본 이미지와 마스킹된 이미지를 확인해보자.

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

%matplotlib inline

coco = COCO('./data/LIVECell_dataset_2021/livecell_coco_train.json')
img_dir = './data/LIVECell_dataset_2021/images/livecell_train_val_images'

#이미지 아이디 랜덤하게 3개 샘플링
imgIds = random.sample(coco.getImgIds(), 3)
for imgid in imgIds:
    img = coco.imgs[imgid]
    image = np.array(Image.open(os.path.join(img_dir, img['file_name'])))
    plt.imshow(Image.open(os.path.join(img_dir, img['file_name'])))
    plt.show()

    anns_ids = coco.getAnnIds(imgIds=imgid)
    anns = coco.loadAnns(anns_ids)
    plt.imshow(Image.open(os.path.join(img_dir, img['file_name'])))
    coco.showAnns(anns)
    plt.show()

 

이 데이터셋은 이미 coco format이긴 하지만 cell class가 category에 각각 나뉘어있지 않고 file name에 들어가 있는 형태다. 8개의 class로 나눠 저장해줘야한다. 그에 대한 코드를 utils/conver_to_coco_livecell.py에 입력해준다. 코드는 아래와 같고 그 아래있는 명령어로 코드를 실행한다.

import copy
import os

import mmcv
from tqdm import tqdm
from pycocotools.coco import COCO

CATEGORIES = (
    'shsy5y', 'a172', 'bt474', 'bv2', 'huh7', 'mcf7', 'skov3', 'skbr3'
)
CAT2IDX = {cat: idx for idx, cat in enumerate(CATEGORIES)}


def init_coco():
    return {
        'info': {},
        'categories':
            [{
                'id': idx,
                'name': cat,
            } for cat, idx in CAT2IDX.items()]
    }


def to_multiclass(ann_file):
    img_infos = []
    ann_infos = []
    img_id = 0
    ann_id = 0
    coco = COCO(ann_file)
    for _id in tqdm(coco.getImgIds()):
        img_info = copy.deepcopy(coco.loadImgs(_id)[0])
        del img_info['original_filename']
        del img_info['url']
        img_info['id'] = img_id
        filename = img_info['file_name']
        cat = filename.split('_')[0].lower()

        img_infos.append(img_info)
        for ann_info in coco.loadAnns(coco.getAnnIds(_id)):
            ann_info = copy.deepcopy(ann_info)
            ann_info['image_id'] = img_id
            ann_info['id'] = ann_id
            ann_info['category_id'] = CAT2IDX[cat]
            ann_infos.append(ann_info)
            ann_id += 1
        img_id += 1

    coco = init_coco()
    coco['images'] = img_infos
    coco['annotations'] = ann_infos
    return coco

def main():
    inputs = ['../data/LIVECell_dataset_2021/livecell_coco_train.json', '../data/LIVECell_dataset_2021/livecell_coco_val.json', '../data/LIVECell_dataset_2021/livecell_coco_test.json']
    outputs = ['../data/LIVECell_dataset_2021/train_8class.json', '../data/LIVECell_dataset_2021/val_8class.json', '../data/LIVECell_dataset_2021/test_8class.json']
    for i in range(len(inputs)):
        mmcv.dump(to_multiclass(inputs[i]), outputs[i])

if __name__ == '__main__':
    main()

cd utils
python convert_to_coco_livecell.py
cd ..

 

캐글 데이터 확인 / 전처리

livecell과 마찬가지로 마스킹된 데이터를 확인해보기 위해서는 coco format으로 변환해줘야한다. 그냥 확인해도 되지만 굳이 코드를 한 번 더 짤 필요는 없겠다.

우선 데이터를 coco format으로 변환한다. segmentation mask는 RLE(Run Length Encoding)된 형태로 되어있다. 이를 mask로 바꾸고 다시 바이너리 RLE 형태로 바꿔줘야한다.

 

또한 valid dataset을 만들어야하기 때문에 stratifiedKFold를 이용하여 각 class의 비율에 맞게 5fold로 나눠주어야 한다.

 

이 과정을 담은게 아래 코드다.

 

import pandas as pd
import numpy as np
from sklearn.model_selection import KFold, StratifiedKFold

import mmcv
import pycocotools.mask as mask_utils
from tqdm import tqdm

CATEGORIES = ('shsy5y', 'astro', 'cort')
CAT2IDX = {cat: idx for idx, cat in enumerate(CATEGORIES)}
IMG_HEIGHT = 520
IMG_WIDTH = 704


def init_coco():
    return {
        'info': {},
        'categories':
            [{
                'id': idx,
                'name': cat,
            } for cat, idx in CAT2IDX.items()]
    }


def krle2mask(rle, height, width):
    s = rle.split()
    starts, lengths = [
        np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])
    ]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(height * width, dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(height, width)

def df2coco(df):
    img_infos = []
    ann_infos = []
    img_id = 0
    ann_id = 0
    for img_name, img_df in tqdm(df.groupby('id'), total=df['id'].nunique()):
        img_info = dict(
            id=img_id,
            width=IMG_WIDTH,
            height=IMG_HEIGHT,
            file_name=f'{img_name}.png',
        )
        for kaggle_rle, cell_type in zip(
            img_df['annotation'], img_df['cell_type']
        ):
            mask = krle2mask(kaggle_rle, IMG_HEIGHT, IMG_WIDTH)
            mask = np.asfortranarray(mask)
            rle = mask_utils.encode(mask)
            rle['counts'] = rle['counts'].decode()
            bbox = mask_utils.toBbox(rle).tolist()
            ann_info = dict(
                id=ann_id,
                image_id=img_id,
                category_id=CAT2IDX[cell_type],
                iscrowd=0,
                segmentation=rle,
                area=bbox[2] * bbox[3],
                bbox=bbox,
            )
            ann_infos.append(ann_info)
            ann_id += 1
        img_infos.append(img_info)
        img_id += 1

    coco = init_coco()
    coco['images'] = img_infos
    coco['annotations'] = ann_infos
    return coco


def main():
    df = pd.read_csv('../data/train.csv')
    dtrainval = df2coco(df)
    mmcv.dump(dtrainval, '../data/dtrainval.json')

    all_samples = np.array(sorted(set(df['sample_id'])))
    sample2celltype = dict(zip(df['sample_id'], df['cell_type']))
    cell_types = [sample2celltype[_] for _ in all_samples]
    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=0)
    splits = list(skf.split(all_samples, cell_types))
    for fold, (train_inds, val_inds) in enumerate(splits):
        train_samples = all_samples[train_inds]
        train_df = df[df['sample_id'].isin(train_samples)]
        train_coco = df2coco(train_df)
        mmcv.dump(train_coco, f'../data/dtrain_g{fold}.json')

        val_samples = all_samples[val_inds]
        val_df = df[df['sample_id'].isin(val_samples)]
        val_coco = df2coco(val_df)
        mmcv.dump(val_coco, f'../data/dval_g{fold}.json')


if __name__ == '__main__':
    main()

cd utils
python convert_to_coco_kaggle.py
cd ..

 

 

지금까지의 과정을 거치면 다음과 같은 폴더구조가 만들어진다. checkpoints 폴더는 학습할 때 pretrained model을 넣어두기 위해 미리 만들었다.

data

data/LIVECell_dataset_2021

 

kaggle 데이터를 coco format으로 만들었으므로 이에 대한 시각화도 몇 개 해보자. 위 livecell 시각화와 동일한 코드에 경로만 바꾸었다.

coco = COCO('./data/dtrainval.json')
img_dir = './data/train'

#이미지 아이디 랜덤하게 3개 샘플링
imgIds = random.sample(coco.getImgIds(), 3)
for imgid in imgIds:
    img = coco.imgs[imgid]
    image = np.array(Image.open(os.path.join(img_dir, img['file_name'])))
    plt.imshow(Image.open(os.path.join(img_dir, img['file_name'])), cmap='gray')
    plt.show()

    anns_ids = coco.getAnnIds(imgIds=imgid)
    anns = coco.loadAnns(anns_ids)
    plt.imshow(Image.open(os.path.join(img_dir, img['file_name'])), cmap='gray')
    coco.showAnns(anns)
    plt.show()