[Kaggle] CT Medical Image - (2) DataLoader 만들기

CT Medical Image Classification

(0) 서론, 액션플랜   
(1) DICOM 파일 array로 전환 및 시각화

 

 

0. 필요한 라이브러리 임포트

 

from dicom_file import dicom_2_array

from random import random
import pandas as pd
import numpy as np

import torch, torchvision
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.utils.data import Sampler

 

1. df 불러오기

 

df = pd.read_csv("final_df.csv")

df.loc[df['Contrast']==True,['Contrast']] = 1
df.loc[df['Contrast']==False,['Contrast']] = 0

df['Contrast'] = df['Contrast'].astype('int')

train_images = df['dicom_name'].values
train_labels = df['Contrast'].values

 

저번 시간에 만든 final_df.csv 를 이용한다. 

 

Contrast가 Bool type으로 들어있기 때문에 위와 같이 0과 1의 값을 가지는 이진변수로 변환해준다. 

 

추후에 만들 customdataset과 customdatasampler, 그리고 customdataloader를 만들기 위해서 

train_images와 train_labels를 만든다. 

 

실은 저번에 train_images로 한 번에 넣기 위해서 pixel_array column을 추가해놓기는 했었는데, 다시 불러오니까 이게 string으로 저장되어 있고 \n 과 같은 개행문자를 포함하고 있었기 때문에 그냥 dicom 파일의 위치를 train_images로 받아들이고 customdataset에서 이를 array로 바꿔주었다. 

 

 

2. CustomDataset 생성

 

class CustomDataset(Dataset):
    def __init__(self,train_images,train_labels):
    
        self.x_data = [dicom_2_array('dicom_dir/' + path) for path in train_images]
        self.y_data = [[i] for i in train_labels]

    def __len__(self):

        return len(self.x_data)

    def __getitem__(self, idx):

        x = torch.FloatTensor(self.x_data[idx])
        y = torch.LongTensor(self.y_data[idx])

        return x,y

 

3. CustomSampler 생성

 

class customsampler(Sampler):

    def __init__(self, data, positive_num, negative_num):

        self.data = data
        self.total = len(self.data)
        self.positive_num = positive_num
        self.negative_num = negative_num

        self.batch_size = self.positive_num + self.negative_num
        self.iter = self.total // self.batch_size

        self.positive_list = len([i for i, predict in enumerate(self.data) if predict == 1 ])

        self.idx = list(range(self.total)) 

    def __iter__(self):
        batch = list()

        for i in range(self.iter):

            temp = np.concatenate((random.sample(self.idx[:self.positive_list], self.positive_num),
                                random.sample(self.idx[self.positive_list:], self.negative_num)))

            random.shuffle(temp)
            batch.extend(temp)
            
        return iter(batch)

    def __len__(self) -> int:
        return self.iter * self.batch_size

    def get_num_batch(self) -> int:
        return self.iter

 

데이터 셋을 열어보니 0~49 번은 Contrast == True 였고 50~99 번은 Contrast = False였다. 

따라서 적절한 학습을 위해서는 적절히 섞는 샘플링이 필요하므로 customsampler를 생성한다. 

 

기타 object detection의 경우 background 사진이 많다는 점을 반영해서 negative sample을 더 늘리기는 하지만, 본 프로젝트에서는 그런 것과는 깊은 관련을 찾기는 어려울 것 같아서 단순히 50:50의 비율을 이용하였다. 

 

4. 그리고 CustomDataLoader 까지 

 

dataset = CustomDataset(train_images, train_labels)
sampler = customsampler(train_labels, 8,8)
dataloader = DataLoader(dataset, batch_size=16,sampler=sampler)

total_batch = len(dataloader)

 

후 다했다. 

 

다 됐는지 확인하기 위해서 아래와 같이 출력해 보았다. 

 

for num, data in enumerate(dataloader):
    image, label = data
    print(image,label)
    if num ==0 : break

 

배치 사이즈가 16개 이기 때문에 중간에 잘랐다. 

아래는 우선 image 

 

아래는 label이다. 

 

 

custom dataset, loader를 성공적으로 잘 만들었다. 

다음부터는 본격적인 모델이다! 실은 계획에 따르면 시작하기 좀 이른감이 있지만 어서 시작하고 어서 끝내도록 하자. 시험이 얼마 안 남은고로.. 고로..