基於深度學習的CT影象肺結節自動檢測技術四—資料增強—定義神經網路並訓練
阿新 • • 發佈：2019-01-05
# -- coding: utf-8 --
#訓練影象分割網路(u-net)模型
import csv
import glob
import random
import cv2
import numpy
import os
from typing import List, Tuple
from keras.optimizers import SGD
from keras.layers import Input, Convolution2D, MaxPooling2D,
UpSampling2D, merge, 資料增強
# 隨機縮放影象的函式，用於資料增廣(augmentation)
def random_scale_img(img, xy_range, lock_xy=False):
if random.random() > xy_range.chance:
return img
if not isinstance 測試前影象準備—影象資料型別轉換—通道歸一化（灰度圖）
#影象型別轉換，通道歸一化（灰度圖）
def prepare_image_for_net(img):
img = img.astype(numpy.float)
img /= 255.
if len(img.shape) == 3:
img = img.reshape(img.shape[-3], img.shape[-2], img.shape[-1])
else:
img = img.reshape(1, img.shape[-2], img.shape[-1], 1)
return img
# 訓練集圖片預處理函式
def get_train_holdout_files():
train_path = TRAIN_LIST
val_path = VAL_LIST
train_res = []
holdout_res = []
with open(train_path, 'r') as f:
reader = csv.reader(f)
samples_train = list(reader)
random.shuffle(samples_train)
with open(val_path, 'r') as f:
reader = csv.reader(f)
samples_holdout = list(reader)
random.shuffle(samples_holdout)
for img_path in samples_train:
if len(img_path) == 0:
print('space line, skip')
continue
overlay_path = img_path[0].replace("_img.png", "_mask.png")
train_res.append((img_path[0], overlay_path))
for img_path in samples_holdout:
if len(img_path) == 0:
print('space line, skip')
continue
overlay_path = img_path[0].replace("_img.png", "_mask.png")
holdout_res.append((img_path[0], overlay_path))
print("Train count: ", len(train_res), ", holdout count: ", len(holdout_res))
return train_res, holdout_res
定義神經網路並訓練
# unet模型損失函式
def dice_coef(y_true, y_pred):
y_true_f = K.flatten(y_true)
y_pred_f = K.flatten(y_pred)
intersection = K.sum(y_true_f * y_pred_f)
return (2. * intersection + 100) / (K.sum(y_true_f) + K.sum(y_pred_f) + 100)
# unet模型損失函式
def dice_coef_np(y_true, y_pred):
y_true_f = y_true.flatten()
y_pred_f = y_pred.flatten()
intersection = numpy.sum(y_true_f * y_pred_f)
return (2. * intersection + 100) / (numpy.sum(y_true_f) + numpy.sum(y_pred_f) + 100)
# unet模型損失函式
def dice_coef_loss(y_true, y_pred):
return -dice_coef(y_true, y_pred)
# 將每個epoch結束後的驗證結果儲存為圖片
class DumpPredictions(Callback):
def __init__(self, dump_filelist: List[Tuple[str, str]], model_type):
super(DumpPredictions, self).__init__()
self.dump_filelist = dump_filelist
self.batch_count = 0
if not os.path.exists(TRAIN_TEMP_DIR):
os.mkdir(TRAIN_TEMP_DIR)
for file_path in glob.glob(TRAIN_TEMP_DIR + "*.*"):
os.remove(file_path)
self.model_type = model_type
def on_epoch_end(self, epoch, logs=None):
model = self.model # type: Model
generator = image_generator(self.dump_filelist, 1, train_set=False)
for i in range(0, 10):
x, y = next(generator)
y_pred = model.predict(x, batch_size=1)
x = x.swapaxes(0, 3)
x = x[0]
# print(x.shape, y.shape, y_pred.shape)
x *= 255.
x = x.reshape((x.shape[0], x.shape[0])).astype(numpy.uint8)
y *= 255.
y = y.reshape((y.shape[1], y.shape[2])).astype(numpy.uint8)
y_pred *= 255.
y_pred = y_pred.reshape((y_pred.shape[1], y_pred.shape[2])).astype(numpy.uint8)
cv2.imwrite(TRAIN_TEMP_DIR + "img_{0:03d}_{1:02d}_i.png".format(epoch, i), x)
cv2.imwrite(TRAIN_TEMP_DIR + "img_{0:03d}_{1:02d}_o.png".format(epoch, i), y)
cv2.imwrite(TRAIN_TEMP_DIR + "img_{0:03d}_{1:02d}_p.png".format(epoch, i), y_pred)
# 訓練集資料載入器
def image_generator(batch_files, batch_size, train_set):
global ELASTIC_INDICES
while True:
if train_set:
random.shuffle(batch_files)
img_list = []
overlay_list = []
ELASTIC_INDICES = None
for batch_file_idx, batch_file in enumerate(batch_files):
images = []
img = cv2.imread(batch_file[0], cv2.IMREAD_GRAYSCALE)
images.append(img)
overlay = cv2.imread(batch_file[1], cv2.IMREAD_GRAYSCALE)
if train_set:
if random.randint(0, 100) > 50:
for img_index, img in enumerate(images):
images[img_index] = elastic_transform(img, 128, 15)
overlay = elastic_transform(overlay, 128, 15)
if True:
augmented = images + [overlay]
augmented = random_rotate_img(augmented, 0.8, -20, 20)
augmented = random_flip_img(augmented, 0.5, 0.5)
augmented = random_translate_img(augmented, XYRange(-30, 30, -30, 30, 0.8))
images = augmented[:-1]
overlay = augmented[-1]
for index, img in enumerate(images):
img = prepare_image_for_net(img)
images[index] = img
overlay = prepare_image_for_net(overlay)
images3d = numpy.vstack(images)
images3d = images3d.swapaxes(0, 3)
img_list.append(images3d)
overlay_list.append(overlay)
if len(img_list) >= batch_size:
x = numpy.vstack(img_list)
y = numpy.vstack(overlay_list)
# if len(img_list) >= batch_size:
yield x, y
img_list = []
overlay_list = []
# 實現unet的網路結構，並載入預訓練好的權重
def get_unet(learn_rate=0.0001