1
2
3
# 구글 드라이브 사용 시 연결
from google.colab import drive
drive.mount('/content/drive')
1
Mounted at /content/drive
0. Library Call
1
2
3
4
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
1
2
3
4
5
6
7
8
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import matplotlib.pyplot as plt
from torchvision import transforms, datasets
1. Data Load % Data Formatting
1
2
3
4
5
6
7
8
9
10
11
12
13
dir_data = './UNet'
name_label = 'train-labels.tif'
name_input = 'train-volume.tif'
img_label = Image.open(os.path.join(dir_data, name_label))
img_input = Image.open(os.path.join(dir_data, name_input))
ny, nx = img_label.size
nframe = img_label.n_frames
print(ny, nx)
print(nframe)
1
2
512 512
30
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
# Train / Valid / Test
nframe_train = 24
nframe_val = 3
nframe_test = 3
dir_save_train = os.path.join(dir_data, 'train')
dir_save_val = os.path.join(dir_data, 'val')
dir_save_test = os.path.join(dir_data, 'test')
if not os.path.exists(dir_save_train):
os.makedirs(dir_save_train)
if not os.path.exists(dir_save_val):
os.makedirs(dir_save_val)
if not os.path.exists(dir_save_test):
os.makedirs(dir_save_test)
1
2
3
# 전체 이미지 30개를 랜덤하게 섞어줌
id_frame = np.arange(nframe)
np.random.shuffle(id_frame)
1
2
3
4
5
6
7
8
9
10
11
12
# Train image를 npy 파일로 저장
offset_nframe = 0
for i in range(nframe_train):
img_label.seek(id_frame[i + offset_nframe])
img_input.seek(id_frame[i + offset_nframe])
label_ = np.asarray(img_label)
input_ = np.asarray(img_input)
np.save(os.path.join(dir_save_train, 'label_%03d.npy' % i), label_)
np.save(os.path.join(dir_save_train, 'input_%03d.npy' % i), input_)
1
2
3
4
5
6
7
8
9
10
11
12
13
# Valid image를 npy 파일로 저장
offset_nframe = nframe_train
for i in range(nframe_val):
img_label.seek(id_frame[i + offset_nframe])
img_input.seek(id_frame[i + offset_nframe])
label_ = np.asarray(img_label)
input_ = np.asarray(img_input)
np.save(os.path.join(dir_save_val, 'label_%03d.npy' % i), label_)
np.save(os.path.join(dir_save_val, 'input_%03d.npy' % i), input_)
1
2
3
4
5
6
7
8
9
10
11
12
# Test image를 npy 파일로 저장
offset_nframe = nframe_train + nframe_val
for i in range(nframe_test):
img_label.seek(id_frame[i + offset_nframe])
img_input.seek(id_frame[i + offset_nframe])
label_ = np.asarray(img_label)
input_ = np.asarray(img_input)
np.save(os.path.join(dir_save_test, 'label_%03d.npy' % i), label_)
np.save(os.path.join(dir_save_test, 'input_%03d.npy' % i), input_)
1
2
3
4
5
6
7
8
9
10
# 위에서 생성한 이미지 시각화
plt.subplot(122)
plt.imshow(label_, cmap='gray')
plt.title('label')
plt.subplot(121)
plt.imshow(input_, cmap='gray')
plt.title('input')
plt.show()
1
2
3
4
5
6
7
8
9
10
11
# 픽셀 값의 분포
plt.subplot(122)
plt.hist(label_.flatten(), bins=20)
plt.title('label')
plt.subplot(121)
plt.hist(input_.flatten(), bins=20)
plt.title('input')
plt.tight_layout()
plt.show()
2. U-Net Architecture
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
# 네트워크 구축
class UNet(nn.Module):
def __init__(self):
super(UNet, self).__init__()
# Convolution + BatchNormalization + Relu 정의 (반복적으로 사용)
def CBR2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True):
layers = []
layers += [nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
kernel_size=kernel_size, stride=stride, padding=padding,
bias=bias)]
layers += [nn.BatchNorm2d(num_features=out_channels)]
layers += [nn.ReLU()]
cbr = nn.Sequential(*layers)
return cbr
# 수축 경로(Contracting path) -> Encoder
self.enc1_1 = CBR2d(in_channels=1, out_channels=64)
self.enc1_2 = CBR2d(in_channels=64, out_channels=64)
self.pool1 = nn.MaxPool2d(kernel_size=2)
self.enc2_1 = CBR2d(in_channels=64, out_channels=128)
self.enc2_2 = CBR2d(in_channels=128, out_channels=128)
self.pool2 = nn.MaxPool2d(kernel_size=2)
self.enc3_1 = CBR2d(in_channels=128, out_channels=256)
self.enc3_2 = CBR2d(in_channels=256, out_channels=256)
self.pool3 = nn.MaxPool2d(kernel_size=2)
self.enc4_1 = CBR2d(in_channels=256, out_channels=512)
self.enc4_2 = CBR2d(in_channels=512, out_channels=512)
self.pool4 = nn.MaxPool2d(kernel_size=2)
self.enc5_1 = CBR2d(in_channels=512, out_channels=1024)
# 확장 경로(Expansive path) -> Decoder
self.dec5_1 = CBR2d(in_channels=1024, out_channels=512)
self.unpool4 = nn.ConvTranspose2d(in_channels=512, out_channels=512,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec4_2 = CBR2d(in_channels=2 * 512, out_channels=512)
self.dec4_1 = CBR2d(in_channels=512, out_channels=256)
self.unpool3 = nn.ConvTranspose2d(in_channels=256, out_channels=256,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec3_2 = CBR2d(in_channels=2 * 256, out_channels=256)
self.dec3_1 = CBR2d(in_channels=256, out_channels=128)
self.unpool2 = nn.ConvTranspose2d(in_channels=128, out_channels=128,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec2_2 = CBR2d(in_channels=2 * 128, out_channels=128)
self.dec2_1 = CBR2d(in_channels=128, out_channels=64)
self.unpool1 = nn.ConvTranspose2d(in_channels=64, out_channels=64,
kernel_size=2, stride=2, padding=0, bias=True)
self.dec1_2 = CBR2d(in_channels=2 * 64, out_channels=64)
self.dec1_1 = CBR2d(in_channels=64, out_channels=64)
self.fc = nn.Conv2d(in_channels=64, out_channels=1, kernel_size=1, stride=1, padding=0, bias=True)
# forward 함수 정의하기
def forward(self, x):
enc1_1 = self.enc1_1(x)
enc1_2 = self.enc1_2(enc1_1)
pool1 = self.pool1(enc1_2)
enc2_1 = self.enc2_1(pool1)
enc2_2 = self.enc2_2(enc2_1)
pool2 = self.pool2(enc2_2)
enc3_1 = self.enc3_1(pool2)
enc3_2 = self.enc3_2(enc3_1)
pool3 = self.pool3(enc3_2)
enc4_1 = self.enc4_1(pool3)
enc4_2 = self.enc4_2(enc4_1)
pool4 = self.pool4(enc4_2)
enc5_1 = self.enc5_1(pool4)
dec5_1 = self.dec5_1(enc5_1)
unpool4 = self.unpool4(dec5_1)
cat4 = torch.cat((unpool4, enc4_2), dim=1)
dec4_2 = self.dec4_2(cat4)
dec4_1 = self.dec4_1(dec4_2)
unpool3 = self.unpool3(dec4_1)
cat3 = torch.cat((unpool3, enc3_2), dim=1)
dec3_2 = self.dec3_2(cat3)
dec3_1 = self.dec3_1(dec3_2)
unpool2 = self.unpool2(dec3_1)
cat2 = torch.cat((unpool2, enc2_2), dim=1)
dec2_2 = self.dec2_2(cat2)
dec2_1 = self.dec2_1(dec2_2)
unpool1 = self.unpool1(dec2_1)
cat1 = torch.cat((unpool1, enc1_2), dim=1)
dec1_2 = self.dec1_2(cat1)
dec1_1 = self.dec1_1(dec1_2)
x = self.fc(dec1_1)
return x
3. DataLoader and Transform
3-1. DataLoader
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
# DataLoader
class Dataset(torch.utils.data.Dataset):
def __init__(self, data_dir, transform=None):
self.data_dir = data_dir
self.transform = transform
lst_data = os.listdir(self.data_dir)
lst_label = [f for f in lst_data if f.startswith('label')]
lst_input = [f for f in lst_data if f.startswith('input')]
lst_label.sort()
lst_input.sort()
self.lst_label = lst_label
self.lst_input = lst_input
def __len__(self):
return len(self.lst_label)
def __getitem__(self, index):
label = np.load(os.path.join(self.data_dir, self.lst_label[index]))
input = np.load(os.path.join(self.data_dir, self.lst_input[index]))
# 정규화
label = label/255.0
input = input/255.0
# 이미지와 레이블의 차원 = 2일 경우(채널이 없을 경우, 흑백 이미지), 새로운 채널(축) 생성
if label.ndim == 2:
label = label[:, :, np.newaxis]
if input.ndim == 2:
input = input[:, :, np.newaxis]
data = {'input': input, 'label': label}
# transform이 정의되어 있다면 transform을 거친 데이터를 불러옴
if self.transform:
data = self.transform(data)
return data
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# DataLoader 작동 확인
dataset_train = Dataset(data_dir=dir_save_train)
data = dataset_train.__getitem__(0) # Index0에 해당하는 한 개의 이미지 불러오기
input = data['input']
label = data['label']
# 불러온 이미지 시각화
plt.subplot(122)
plt.imshow(label.reshape(512,512), cmap='gray')
plt.title('label')
plt.subplot(121)
plt.imshow(input.reshape(512,512), cmap='gray')
plt.title('input')
plt.show()
3-2. Transform
- ToTensor : numpy에서 tensor로 데이터 type을 변경
- Normalization : 데이터 정규화
- RandomFlip : 랜덤 좌우 및 상하 반전
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
# Transform
class ToTensor(object):
def __call__(self, data):
label, input = data['label'], data['input']
label = label.transpose((2, 0, 1)).astype(np.float32)
input = input.transpose((2, 0, 1)).astype(np.float32)
data = {'label': torch.from_numpy(label), 'input': torch.from_numpy(input)}
return data
# Normalization
class Normalization(object):
def __init__(self, mean=0.5, std=0.5):
self.mean = mean
self.std = std
def __call__(self, data):
label, input = data['label'], data['input']
input = (input - self.mean) / self.std
data = {'label': label, 'input': input}
return data
# 상하/좌우 반전
class RandomFlip(object):
def __call__(self, data):
label, input = data['label'], data['input']
if np.random.rand() > 0.5:
label = np.fliplr(label)
input = np.fliplr(input)
if np.random.rand() > 0.5:
label = np.flipud(label)
input = np.flipud(input)
data = {'label': label, 'input': input}
return data
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# Transform 작동 확인
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), RandomFlip(), ToTensor()])
dataset_train = Dataset(data_dir=dir_save_train, transform=transform)
data = dataset_train.__getitem__(0) # Index0에 해당하는 한 개의 이미지
input = data['input']
label = data['label']
# 불러온 이미지 시각화
plt.subplot(122)
plt.hist(label.flatten(), bins=20)
plt.title('label')
plt.subplot(121)
plt.hist(input.flatten(), bins=20)
plt.title('input')
plt.tight_layout()
plt.show()
4. Network Save and Load
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
# Network save
def save(ckpt_dir, net, optim, epoch):
if not os.path.exists(ckpt_dir):
os.makedirs(ckpt_dir)
torch.save({'net': net.state_dict(), 'optim': optim.state_dict()},
"%s/model_epoch%d.pth" % (ckpt_dir, epoch))
# Network load
def load(ckpt_dir, net, optim):
if not os.path.exists(ckpt_dir):
epoch = 0
return net, optim, epoch
ckpt_lst = os.listdir(ckpt_dir)
ckpt_lst.sort(key=lambda f: int(''.join(filter(str.isdigit, f))))
dict_model = torch.load('%s/%s' % (ckpt_dir, ckpt_lst[-1]))
net.load_state_dict(dict_model['net'])
optim.load_state_dict(dict_model['optim'])
epoch = int(ckpt_lst[-1].split('epoch')[1].split('.pth')[0])
return net, optim, epoch
5. Model Training
1
2
3
4
5
6
7
8
9
# 훈련 파라미터 설정하기
lr = 1e-3
batch_size = 4
num_epoch = 50
base_dir = './drive/MyDrive/net'
data_dir = dir_data
ckpt_dir = os.path.join(base_dir, "checkpoint")
log_dir = os.path.join(base_dir, "log")
1
2
3
4
5
6
7
8
# 훈련을 위한 Transform과 DataLoader
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), RandomFlip(), ToTensor()])
dataset_train = Dataset(data_dir=os.path.join(data_dir, 'train'), transform=transform)
loader_train = DataLoader(dataset_train, batch_size=batch_size, shuffle=True, num_workers=0)
dataset_val = Dataset(data_dir=os.path.join(data_dir, 'val'), transform=transform)
loader_val = DataLoader(dataset_val, batch_size=batch_size, shuffle=False, num_workers=0)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
# 네트워크 생성하기
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net = UNet().to(device)
# 손실함수 정의하기
fn_loss = nn.BCEWithLogitsLoss().to(device)
# Optimizer 설정하기
optim = torch.optim.Adam(net.parameters(), lr=lr)
# 그밖에 부수적인 variables 설정하기
num_data_train = len(dataset_train)
num_data_val = len(dataset_val)
num_batch_train = np.ceil(num_data_train / batch_size)
num_batch_val = np.ceil(num_data_val / batch_size)
# 그 밖에 부수적인 functions 설정하기
fn_tonumpy = lambda x: x.to('cpu').detach().numpy().transpose(0, 2, 3, 1)
fn_denorm = lambda x, mean, std: (x * std) + mean
fn_class = lambda x: 1.0 * (x > 0.5)
# Tensorboard 를 사용하기 위한 SummaryWriter 설정
writer_train = SummaryWriter(log_dir=os.path.join(log_dir, 'train'))
writer_val = SummaryWriter(log_dir=os.path.join(log_dir, 'val'))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
# 네트워크 학습시키기
st_epoch = 0
# 학습한 모델이 있을 경우 모델 로드하기
net, optim, st_epoch = load(ckpt_dir=ckpt_dir, net=net, optim=optim)
for epoch in range(st_epoch + 1, num_epoch + 1):
net.train()
loss_arr = []
for batch, data in enumerate(loader_train, 1):
# forward pass
label = data['label'].to(device)
input = data['input'].to(device)
output = net(input)
# backward pass
optim.zero_grad()
loss = fn_loss(output, label)
loss.backward()
optim.step()
# 손실 함수 계산
loss_arr += [loss.item()]
print("TRAIN: EPOCH %04d / %04d | BATCH %04d / %04d | LOSS %.4f" %
(epoch, num_epoch, batch, num_batch_train, np.mean(loss_arr)))
# Tensorboard 저장하기
label = fn_tonumpy(label)
input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
output = fn_tonumpy(fn_class(output))
writer_train.add_image('label', label, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')
writer_train.add_image('input', input, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')
writer_train.add_image('output', output, num_batch_train * (epoch - 1) + batch, dataformats='NHWC')
writer_train.add_scalar('loss', np.mean(loss_arr), epoch)
with torch.no_grad():
net.eval()
loss_arr = []
for batch, data in enumerate(loader_val, 1):
# forward pass
label = data['label'].to(device)
input = data['input'].to(device)
output = net(input)
# 손실함수 계산하기
loss = fn_loss(output, label)
loss_arr += [loss.item()]
print("VALID: EPOCH %04d / %04d | BATCH %04d / %04d | LOSS %.4f" %
(epoch, num_epoch, batch, num_batch_val, np.mean(loss_arr)))
# Tensorboard 저장하기
label = fn_tonumpy(label)
input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
output = fn_tonumpy(fn_class(output))
writer_val.add_image('label', label, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')
writer_val.add_image('input', input, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')
writer_val.add_image('output', output, num_batch_val * (epoch - 1) + batch, dataformats='NHWC')
writer_val.add_scalar('loss', np.mean(loss_arr), epoch)
# epoch 50마다 모델 저장하기
if epoch % 50 == 0:
save(ckpt_dir=ckpt_dir, net=net, optim=optim, epoch=epoch)
writer_train.close()
writer_val.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
TRAIN: EPOCH 0001 / 0050 | BATCH 0001 / 0006 | LOSS 0.6968
TRAIN: EPOCH 0001 / 0050 | BATCH 0002 / 0006 | LOSS 0.6383
TRAIN: EPOCH 0001 / 0050 | BATCH 0003 / 0006 | LOSS 0.6072
TRAIN: EPOCH 0001 / 0050 | BATCH 0004 / 0006 | LOSS 0.5800
TRAIN: EPOCH 0001 / 0050 | BATCH 0005 / 0006 | LOSS 0.5547
TRAIN: EPOCH 0001 / 0050 | BATCH 0006 / 0006 | LOSS 0.5346
VALID: EPOCH 0001 / 0050 | BATCH 0001 / 0001 | LOSS 0.5948
...
TRAIN: EPOCH 0050 / 0050 | BATCH 0002 / 0006 | LOSS 0.1701
TRAIN: EPOCH 0050 / 0050 | BATCH 0003 / 0006 | LOSS 0.1784
TRAIN: EPOCH 0050 / 0050 | BATCH 0004 / 0006 | LOSS 0.1781
TRAIN: EPOCH 0050 / 0050 | BATCH 0005 / 0006 | LOSS 0.1779
TRAIN: EPOCH 0050 / 0050 | BATCH 0006 / 0006 | LOSS 0.1774
VALID: EPOCH 0050 / 0050 | BATCH 0001 / 0001 | LOSS 0.1970
6. Evaluation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# Model Test
transform = transforms.Compose([Normalization(mean=0.5, std=0.5), ToTensor()])
dataset_test = Dataset(data_dir=os.path.join(data_dir, 'test'), transform=transform)
loader_test = DataLoader(dataset_test, batch_size=batch_size, shuffle=False, num_workers=8)
# 그밖에 부수적인 variables 설정하기
num_data_test = len(dataset_test)
num_batch_test = np.ceil(num_data_test / batch_size)
# 결과 디렉토리 생성하기
result_dir = os.path.join(base_dir, 'result')
if not os.path.exists(result_dir):
os.makedirs(os.path.join(result_dir, 'png'))
os.makedirs(os.path.join(result_dir, 'numpy'))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
net, optim, st_epoch = load(ckpt_dir=ckpt_dir, net=net, optim=optim)
with torch.no_grad():
net.eval()
loss_arr = []
for batch, data in enumerate(loader_test, 1):
# forward pass
label = data['label'].to(device)
input = data['input'].to(device)
output = net(input)
# 손실함수 계산하기
loss = fn_loss(output, label)
loss_arr += [loss.item()]
print("TEST: BATCH %04d / %04d | LOSS %.4f" %
(batch, num_batch_test, np.mean(loss_arr)))
# Tensorboard 저장하기
label = fn_tonumpy(label)
input = fn_tonumpy(fn_denorm(input, mean=0.5, std=0.5))
output = fn_tonumpy(fn_class(output))
# 테스트 결과 저장하기
for j in range(label.shape[0]):
id = num_batch_test * (batch - 1) + j
plt.imsave(os.path.join(result_dir, 'png', 'label_%04d.png' % id), label[j].squeeze(), cmap='gray')
plt.imsave(os.path.join(result_dir, 'png', 'input_%04d.png' % id), input[j].squeeze(), cmap='gray')
plt.imsave(os.path.join(result_dir, 'png', 'output_%04d.png' % id), output[j].squeeze(), cmap='gray')
np.save(os.path.join(result_dir, 'numpy', 'label_%04d.npy' % id), label[j].squeeze())
np.save(os.path.join(result_dir, 'numpy', 'input_%04d.npy' % id), input[j].squeeze())
np.save(os.path.join(result_dir, 'numpy', 'output_%04d.npy' % id), output[j].squeeze())
print("AVERAGE TEST: BATCH %04d / %04d | LOSS %.4f" %
(batch, num_batch_test, np.mean(loss_arr)))
1
2
TEST: BATCH 0001 / 0001 | LOSS 0.2014
AVERAGE TEST: BATCH 0001 / 0001 | LOSS 0.2014
1
2
3
4
5
6
7
8
9
lst_data = os.listdir(os.path.join(result_dir, 'numpy'))
lst_label = [f for f in lst_data if f.startswith('label')]
lst_input = [f for f in lst_data if f.startswith('input')]
lst_output = [f for f in lst_data if f.startswith('output')]
lst_label.sort()
lst_input.sort()
lst_output.sort()
1
2
3
4
5
id = 0
label = np.load(os.path.join(result_dir,"numpy", lst_label[id]))
input = np.load(os.path.join(result_dir,"numpy", lst_input[id]))
output = np.load(os.path.join(result_dir,"numpy", lst_output[id]))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# 결과 시각화 -> input / label / output(predict)
plt.figure(figsize=(8,6))
plt.subplot(131)
plt.imshow(input, cmap='gray')
plt.title('input')
plt.subplot(132)
plt.imshow(label, cmap='gray')
plt.title('label')
plt.subplot(133)
plt.imshow(output, cmap='gray')
plt.title('output')
plt.show()