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!')