def load_data(dataset_path, contour_type, crop_size):
contours_path = os.path.join(dataset_path, 'contours')
images_path = os.path.join(dataset_path, 'images')
contours = load_all_contours(contours_path, contour_type, shuffle=True)
return export_all_contours(
contours,
images_path,
crop_size=crop_size,
sax_series=get_SAX_SERIES(),
)
#!/usr/bin/env python2.7
import re, sys, os
import shutil, cv2
import numpy as np
from train_sunnybrook_unetres import read_contour, map_all_contours, export_all_contours
from helpers import reshape, get_SAX_SERIES, draw_result
from unet_model_inv import unet_model_inv, dice_coef_each
SAX_SERIES = get_SAX_SERIES()
SUNNYBROOK_ROOT_PATH = 'D:\cardiac_data\Sunnybrook'
VAL_CONTOUR_PATH = os.path.join(
SUNNYBROOK_ROOT_PATH, 'Sunnybrook Cardiac MR Database ContoursPart2',
'ValidationDataContours')
VAL_IMG_PATH = os.path.join(SUNNYBROOK_ROOT_PATH,
'Sunnybrook Cardiac MR Database DICOMPart2',
'ValidationDataDICOM')
VAL_OVERLAY_PATH = os.path.join(SUNNYBROOK_ROOT_PATH,
'Sunnybrook Cardiac MR Database OverlayPart2',
'ValidationDataOverlay')
ONLINE_CONTOUR_PATH = os.path.join(
SUNNYBROOK_ROOT_PATH, 'Sunnybrook Cardiac MR Database ContoursPart1',
'OnlineDataContours')
ONLINE_IMG_PATH = os.path.join(SUNNYBROOK_ROOT_PATH,
'Sunnybrook Cardiac MR Database DICOMPart1',
'OnlineDataDICOM')
ONLINE_OVERLAY_PATH = os.path.join(
SUNNYBROOK_ROOT_PATH, 'Sunnybrook Cardiac MR Database OverlayPart1',
'OnlineDataOverlay')
SAVE_VAL_PATH = os.path.join(SUNNYBROOK_ROOT_PATH, 'Sunnybrook_val_submission')
def train(image_path, contour_path, *, contour_type, crop_size, batch_size,
seed, epoch_count):
contours = load_all_contours(contour_path, contour_type, shuffle=True)
loaded_train_x, loaded_train_y = export_all_contours(
contours,
image_path,
crop_size=crop_size,
sax_series=get_SAX_SERIES(),
)
input_shape = (crop_size, crop_size, 1)
num_classes = 2
m = fcn_model
model = m(input_shape, num_classes)
kwargs = dict(
rotation_range=180,
zoom_range=0.0,
width_shift_range=0.0,
height_shift_range=0.0,
horizontal_flip=True,
vertical_flip=True,
)
image_datagen = ImageDataGenerator(**kwargs)
mask_datagen = ImageDataGenerator(**kwargs)
image_generator = image_datagen.flow(
loaded_train_x,
shuffle=False,
batch_size=batch_size,
seed=seed,
)
mask_generator = mask_datagen.flow(
loaded_train_y,
shuffle=False,
batch_size=batch_size,
seed=seed,
)
train_generator = zip(image_generator, mask_generator)
max_iter = (len(contours) / batch_size) * epoch_count
curr_iter = 0
base_lr = K.eval(model.optimizer.lr)
learning_rate = lr_poly_decay(model,
base_lr,
curr_iter,
max_iter,
power=0.5)
for epoch in range(1, epoch_count + 1):
print()
print('Main Epoch {:d}'.format(epoch))
print('Learning rate: {:6f}'.format(learning_rate))
train_result = []
iter_count = len(img_train) // batch_size
for img, mask in islice(train_generator, iter_count):
res = model.train_on_batch(img, mask)
curr_iter += 1
learning_rate = lr_poly_decay(model,
base_lr,
curr_iter,
max_iter,
power=0.5)
train_result.append(res)
train_result = np.asarray(train_result)
train_result = np.mean(train_result, axis=0).round(decimals=10)
print('Train result {}:'.format(str(model.metrics_names)))
print('{}'.format(str(train_result)))
print()
print('Evaluating dev set ...')
result = model.evaluate(img_dev, mask_dev, batch_size=32)
result = np.round(result, decimals=10)
print()
print('Dev set result {}:'.format(str(model.metrics_names)))
print('{}'.format(str(result)))
predict_and_save(
model,
img_dev[:10],
output_dir=os.path.join(DIR_DATA, 'predictions'),
)
save_file = '_'.join([
'sunnybrook',
contour_type,
'epoch',
str(epoch),
]) + '.h5'
logs_dir = os.path.join(DIR_DATA, 'model_logs')
os.makedirs(logs_dir, exist_ok=True)
save_path = os.path.join(logs_dir, save_file)
print()
print('Saving model weights to {}'.format(save_path))
model.save_weights(save_path)
