def train(self, training_data, target_station, verbose=1):
print(self.session_name)
# print(training_data, target_station, self.vol_size, self.strides)
dataset = Dataset(self._builder_class, training_data, target_station, self.vol_size, self.strides)
(X_train, Y_train) = dataset.construct()
model = self.model(summary=False)
model.fit(X_train, Y_train, epochs=self.epochs, batch_size=self.batch_size, callbacks=self.callbacks(),
validation_split=0.05, verbose=verbose, shuffle=True)
def test(self, test_data, target_station):
"""
:param target_station:
:param test_data:
"""
model_path = '/home/jamesneill/vessel_seg/session/{}/models/{}.h5'.format(self.__class__.__name__,
self.session_name)
save_path = '/home/jamesneill/vessel_seg/session/{}/results/{}.pkl'.format(self.__class__.__name__,
self.session_name)
custom_objects = {
'dsc': metrics.dsc,
'iou': metrics.iou,
'dsc_loss': losses.dsc,
'prelu': activations.prelu
}
print('Testing: {}'.format(model_path))
# Load model with custom objects suchas dice loss, mIoU etc.
model = load_model(model_path, custom_objects=custom_objects)
self.strides = (None, None, 1)
dataset = Dataset(self._builder_class, test_data, target_station, self.vol_size, self.strides, testing=True)
(X_test, _) = dataset.construct()
output_shape = X_test.shape
padding = []
padding.append((0, 0))
for axis in X_test.shape[1:-1]:
padding.append((0, 16 - (axis % 16)))
padding.append((0, 0))
padding = tuple(padding)
print(padding)
X_test = np.pad(X_test, padding, 'constant', constant_values=(0, 0))
rebuilt = model.predict(X_test)
rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
rebuilt = rebuilt[0:output_shape[0], 0:output_shape[1], 0:output_shape[2]]
rebuilt = np.swapaxes(rebuilt, 2, 0)
print('rebuilt.shape', rebuilt.shape)
# index = 0
# for x in range(0, (output_shape[0] / strides[0]) / 2):
# for y in range(0, (output_shape[1] / strides[1]) / 2):
# for z in range(0, (output_shape[2] / strides[2]) / 2):
# rebuilt[index] = predict_test[index]
#
# # rebuilt[strides[0] * x:vol_size[0] + (strides[0] * x),
# # strides[1] * y:vol_size[1] + (strides[1] * y),
# # strides[2] * z:vol_size[2] + (strides[2] * z)] = predict_test[index]
#
#
# # strides[2] * z:vol_size[2] + (strides[2] * z)] += predict_test[index]
# index += 1
# rebuilt -= np.min(rebuilt)
# rebuilt /= np.max(rebuilt)
# rebuilt -= np.mean(rebuilt) - 0.5
# print('np.max(rebuilt)', np.max(rebuilt))
# rebuilt /= np.max(rebuilt)
classification_thresh = 0.5
rebuilt[rebuilt < classification_thresh] = 0
rebuilt[rebuilt >= classification_thresh] = 1
print('Saving output..')
pickle.dump(rebuilt, open(save_path, 'wb'))
imsave('{}.tiff'.format(save_path), rebuilt.astype(np.int16))
print('Done!')
