def main(): parser = argparse.ArgumentParser( description= 'py, train_data_txt, train_data_ture_txt, validation_data_txt, outdir') parser.add_argument('--train_data_txt', '-i1', default='', help='train data list') parser.add_argument('--train_ground_truth_txt', '-i2', default='', help='train ground truth list') parser.add_argument('--validation_data_txt', '-i3', default='', help='validation data list') parser.add_argument('--validation_ground_truth_txt', '-i4', default='', help='validation ground truth list') parser.add_argument('--outdir', '-i5', default='', help='outdir') args = parser.parse_args() # check folder if not (os.path.exists(args.outdir)): os.mkdir(args.outdir) # define batch_size = 3 epoch = 2500 # load train data train_data = io.load_matrix_data(args.train_data_txt, 'float32') train_data = np.expand_dims(train_data, axis=4) # load train ground truth train_truth = io.load_matrix_data(args.train_ground_truth_txt, 'float32') train_truth = np.expand_dims(train_truth, axis=4) # load validation data val_data = io.load_matrix_data(args.validation_data_txt, 'float32') val_data = np.expand_dims(val_data, axis=4) # load validation ground truth val_truth = io.load_matrix_data(args.validation_ground_truth_txt, 'float32') val_truth = np.expand_dims(val_truth, axis=4) print(' number of training: {}'.format(len(train_data))) print('size of traning: {}'.format(train_data.shape)) print(' number of validation: {}'.format(len(val_data))) print('size of validation: {}'.format(val_data.shape)) image_size = [] image_size.extend([ list(train_data.shape)[1], list(train_data.shape)[2], list(train_data.shape)[3] ]) # set network network = Autoencoder(*image_size) model = network.model() model.summary() model.compile(optimizer='Nadam', loss=losses.mean_squared_error, metrics=['mse']) # set data_set train_steps, train_data = batch_iter(train_data, train_truth, batch_size) valid_steps, val_data = batch_iter(val_data, val_truth, batch_size) # fit network model_checkpoint = ModelCheckpoint(os.path.join( args.outdir, 'weights.{epoch:02d}-{val_loss:.2f}.hdf5'), verbose=1) history = model.fit_generator(train_data, steps_per_epoch=train_steps, epochs=epoch, validation_data=val_data, validation_steps=valid_steps, verbose=1, callbacks=[model_checkpoint]) plot_history(history, args.outdir)
def predict(): parser = argparse.ArgumentParser( description='py, test_data_list, name_list, outdir') parser.add_argument('--test_data_list', '-i1', default='', help='test data') parser.add_argument('--name_list', '-i2', default='', help='name list') parser.add_argument('--model', '-i3', default='', help='model') parser.add_argument('--outdir', '-i4', default='', help='outdir') args = parser.parse_args() if not (os.path.exists(args.outdir)): os.mkdir(args.outdir) # load name_list name_list = [] with open(args.name_list) as paths_file: for line in paths_file: line = line.split() if not line: continue name_list.append(line[:]) print('number of test data : {}'.format(len(name_list))) test_data = io.load_matrix_data(args.test_data_list, 'float32') test_data = np.expand_dims(test_data, axis=4) print(test_data.shape) image_size = [] image_size.extend([ list(test_data.shape)[1], list(test_data.shape)[2], list(test_data.shape)[3] ]) print(image_size) # set network network = Autoencoder(*image_size) model = network.model() model.load_weights(args.model) preds = model.predict(test_data, 1) preds = preds[:, :, :, :, 0] print(preds.shape) for i in range(preds.shape[0]): # EUDT eudt_image = sitk.GetImageFromArray(preds[i]) eudt_image.SetSpacing([1, 1, 1]) eudt_image.SetOrigin([0, 0, 0]) # label label = np.where(preds[i] > 0, 0, 1) label_image = sitk.GetImageFromArray(label) label_image.SetSpacing([1, 1, 1]) label_image.SetOrigin([0, 0, 0]) io.write_mhd_and_raw( eudt_image, '{}.mhd'.format(os.path.join(args.outdir, 'EUDT', *name_list[i]))) io.write_mhd_and_raw( label_image, '{}.mhd'.format(os.path.join(args.outdir, 'label', *name_list[i])))