2018年11月8日 更新

# U-Net：セグメンテーションに特化したネットワーク

47,699 view 1
class double_conv(nn.Module):
'''(conv => BN => ReLU) * 2'''
def __init__(self, in_ch, out_ch):
super(double_conv, self).__init__()
self.conv = nn.Sequential(
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True)
)

def forward(self, x):
x = self.conv(x)
return x

class inconv(nn.Module):
def __init__(self, in_ch, out_ch):
super(inconv, self).__init__()
self.conv = double_conv(in_ch, out_ch)

def forward(self, x):
x = self.conv(x)
return x

class down(nn.Module):
def __init__(self, in_ch, out_ch):
super(down, self).__init__()
self.mpconv = nn.Sequential(
nn.MaxPool2d(2),
double_conv(in_ch, out_ch)
)

def forward(self, x):
x = self.mpconv(x)
return x

class up(nn.Module):
def __init__(self, in_ch, out_ch, bilinear=True):
super(up, self).__init__()

#  would be a nice idea if the upsampling could be learned too,
#  but my machine do not have enough memory to handle all those weights
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_ch//2, in_ch//2, 2, stride=2)

self.conv = double_conv(in_ch, out_ch)

def forward(self, x1, x2):
x1 = self.up(x1)
diffX = x1.size()[2] - x2.size()[2]
diffY = x1.size()[3] - x2.size()[3]
x2 = F.pad(x2, (diffX // 2, int(diffX / 2),
diffY // 2, int(diffY / 2)))
x = torch.cat([x2, x1], dim=1)
x = self.conv(x)
return x

class outconv(nn.Module):
def __init__(self, in_ch, out_ch):
super(outconv, self).__init__()
self.conv = nn.Conv2d(in_ch, out_ch, 1)

def forward(self, x):
x = self.conv(x)
return x

class UNet(nn.Module):
def __init__(self, n_channels, n_classes):
super(UNet, self).__init__()
self.inc = inconv(n_channels, 64)
self.down1 = down(64, 128)
self.down2 = down(128, 256)
self.down3 = down(256, 512)
self.down4 = down(512, 512)
self.up1 = up(1024, 256)
self.up2 = up(512, 128)
self.up3 = up(256, 64)
self.up4 = up(128, 64)
self.outc = outconv(64, n_classes)

def forward(self, x):
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
x = self.outc(x)
return x
U-Net ネットワークの定義.py

Dice 係数を算出するクラスの定義を行います．
class DiceCoeff(Function):
"""Dice coeff for individual examples"""

def forward(self, input, target):
self.save_for_backward(input, target)
eps = 0.0001
self.inter = torch.dot(input.view(-1), target.view(-1))
self.union = torch.sum(input) + torch.sum(target) + eps

t = (2 * self.inter.float() + eps) / self.union.float()
return t

# This function has only a single output, so it gets only one gradient

input, target = self.saved_variables

grad_input = grad_output * 2 * (target * self.union + self.inter) \
/ self.union * self.union

def dice_coeff(input, target):
"""Dice coeff for batches"""
if input.is_cuda:
s = torch.FloatTensor(1).cuda().zero_()
else:
s = torch.FloatTensor(1).zero_()

for i, c in enumerate(zip(input, target)):
s = s + DiceCoeff().forward(c[0], c[1])

return s / (i + 1)
Dice 係数を算出するクラスを定義

ネットワークの学習を行います．
net = UNet(n_channels=3, n_classes=1).cuda()

optimizer = optim.SGD(
net.parameters(),
lr=0.1,
momentum=0.9,
weight_decay=0.0005
)

criterion = nn.BCELoss()

for epoch in range(args['epoch']):

#---- Train section
epoch_loss = 0
for i, b in enumerate(batch(train, args['batch_size'])):
img = np.array([i[0] for i in b]).astype(np.float32)
mask = np.array([i[1] for i in b])

img = torch.from_numpy(img).cuda()

loss.backward()
optimizer.step()
epoch_loss += loss.item()

if i%10 == 0:
print('{}/{} ---- loss: {}'.format(i, int(len_train/args['batch_size']), loss.item()))

print('Epoch finished ! Loss: {}'.format(epoch_loss / len_train))

#---- Val section
val_dice = 0
for j, b in enumerate(val):
img =  torch.from_numpy(b[0]).unsqueeze(0).cuda()

if j%10 == 0:
print('val: {}/{}'.format(j, len_val))

torch.save(net.state_dict(), '{}CP{}.pth'.format(args['dir_checkpoint'], epoch + 1))
print('Checkpoint {} saved !'.format(epoch + 1))
print('Validation Dice Coeff: {}'.format(val_dice / len_val))

テストデータを用いて，学習したネットワークを評価します．
file_img_test = os.listdir(args['dir_img_test'])
random.shuffle(file_img_test)

for i, file in enumerate(file_img_test):
img_original = Image.open(args['dir_img_test']+file)
img = img_original

w = img.size[0]
h = img.size[1]

newW = int(w * args['scale'])
newH = int(h * args['scale'])

img = img.resize((newW, newH))
img = img.crop((0, 0, newW, newH))
img = np.array(img, dtype=np.float32)
img = img / 255

img_left = img[:, :newH]
img_right = img[:, -newH:]

img_left = np.transpose(img_left, axes=[2, 0, 1])
img_right = np.transpose(img_right, axes=[2, 0, 1])

img_left = torch.from_numpy(img_left).unsqueeze(0).cuda()
img_right = torch.from_numpy(img_right).unsqueeze(0).cuda()

tf = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize(h),
transforms.ToTensor()
])

d = dcrf.DenseCRF2D(w, h, 2)
U = U.reshape((2, -1))
U = np.ascontiguousarray(U)
img = np.ascontiguousarray(np.array(img_original).astype(np.uint8))

d.setUnaryEnergy(U)

fig = plt.figure(figsize=(27, 7))

ax1.imshow(img_original)
ax1.set_title('input', fontsize=28)

ax2.set_title('output', fontsize=28)

ax3.imshow(img_original)
ax3.set_title('input and output', fontsize=28)

plt.show()
テスト.py

24 件

## 関連する記事 こんな記事も人気です♪

#### Deep learningで画像認識③〜ネオコグニトロンとは？〜

Deep learningは、画像認識において大きな成功を収めています。そこで用いられる多層ネットワークは、畳み込みニューラルネットワーク（convolutional neural network：CNN）と呼ばれており、画像認識に適した独特の構造を持っています。