def train():
model = UNet(cfg.input_shape)
#编译和打印模型
model.compile(optimizer=cfg.optimizer, loss=cfg.loss, metrics=cfg.metrics)
print_summary(model=model)
#训练数据生成器G1
G1 = imageSegmentationGenerator(cfg.train_images, cfg.train_annotations,
cfg.train_batch_size, cfg.n_classes,
cfg.input_shape[0], cfg.input_shape[1],
cfg.output_shape[0], cfg.output_shape[1])
#测试数据生成器G2
if cfg.validate:
G2 = imageSegmentationGenerator(cfg.val_images, cfg.val_annotations,
cfg.val_batch_size, cfg.n_classes,
cfg.input_shape[0], cfg.input_shape[1],
cfg.output_shape[0],
cfg.output_shape[1])
#循环训练
save_index = 1
for ep in range(cfg.epochs):
#1、训练两种方式
if not cfg.validate: #只有G1
hisroy = model.fit_generator(
G1,
steps_per_epoch=cfg.train_steps_per_epoch,
workers=cfg.workers,
epochs=1,
verbose=1,
use_multiprocessing=cfg.use_multiprocessing)
else: #有G1和G2
hisroy = model.fit_generator(
G1,
steps_per_epoch=cfg.train_steps_per_epoch,
workers=cfg.workers,
epochs=1,
verbose=1,
use_multiprocessing=cfg.use_multiprocessing,
validation_data=G2,
validation_steps=cfg.validate_steps_per_epoch)
# 2、保存模型
if save_index == cfg.epochs_save:
save_index = 1
save_weights_name = 'model.{}'.format(ep)
save_weights_path = os.path.join(cfg.save_weights_path,
save_weights_name)
model.save_weights(save_weights_path)
save_index += 1
model = multi_gpu_model(model, gpus=args.gpus)
# train_mode = False
if args.mode == 'validation':
print("inference")
model.load_weights(filepath="outputs_4_channels/best_model_final.h5",
by_name=True)
model.evaluate()
else:
if args.weights is not None:
model.load_weights(filepath=args.weigths,
by_name=True,
skip_mismatch=True)
optimizer = optim.adam(lr=args.lr, decay=5e-4)
model.compile(optimizer=optimizer,
loss=bce_dice_loss,
metrics=[dice_coeff])
log_filename = os.path.join(args.outputs, 'logger_model_train.csv')
csv_log = CSVLogger(log_filename, separator=',', append=True)
lr = ReduceLROnPlateau(monitor='val_loss',
factor=0.8,
patience=5,
verbose=2,
mode='min')
ckpt_path = os.path.join(
args.outputs, "best_weight_model_{epoch:03d}_{val_loss:.4f}.h5")
ckpt = ModelCheckpoint(monitor='val_loss',
filepath=ckpt_path,
def train_model(X_train, X_test, y_train, y_test, save_dir):
"""
Run train model process.
Args:
X_train (list): List of strings X_train files patches.
X_test (list): List of strings X_test files patches.
y_train (list): List of strings y_train files patches.
y_test (list): List of strings y_test files patches.
save_dir (pathlib.PosixPath): Path for save results.
"""
len_train, len_test = len(X_train), len(X_test)
# Create datasets from data.
train_dataset = tf.data.Dataset.from_tensor_slices(
(X_train, y_train)).map(load_data, num_parallel_calls=AUTOTUNE)
test_dataset = tf.data.Dataset.from_tensor_slices(
(X_test, y_test)).map(load_data, num_parallel_calls=AUTOTUNE)
# Prepare data for training.
train_dataset = train_dataset.map(normalize, num_parallel_calls=AUTOTUNE)
train_dataset = train_dataset.shuffle(len_train * 2)
train_dataset = train_dataset.repeat()
train_dataset = train_dataset.batch(BATCH_SIZE)
train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
test_dataset = test_dataset.map(normalize, num_parallel_calls=AUTOTUNE)
test_dataset = test_dataset.batch(BATCH_SIZE)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# TRAIN MODEL.
# Init paths.
save_dir.mkdir(parents=True, exist_ok=True)
path_weights = (save_dir / "weights.hdf5").as_posix()
path_history = (save_dir / "training.csv").as_posix()
path_model = (save_dir / "model").as_posix()
path_predict_fulls = (save_dir / "predict_fulls.csv").as_posix()
path_predict_means = (save_dir / "predict_means.csv").as_posix()
# Init callbacks.
callback_save_best_weights = tf.keras.callbacks.ModelCheckpoint(
filepath=path_weights,
monitor="val_loss",
mode="min",
save_best_only=True,
verbose=0,
)
callback_csv_logger = tf.keras.callbacks.CSVLogger(
filename=path_history,
separator=",",
append=False,
)
# Init optimizer.
optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
# Init model.
model = UNet(
input_shape=(IMAGE_NET_W, IMAGE_NET_H, CLASSES_COUNT),
out_channels=CLASSES_COUNT,
filters=[16, 32, 64, 128, 256],
)
model.compile(
optimizer=optimizer,
loss=tf.keras.losses.MSE,
metrics=[tf.keras.metrics.Accuracy()],
)
# Fake predict for build model and summary.
model.predict(next(iter(test_dataset))[0])
model.summary()
# Load weights.
if MODEL_LOAD_WEIGHTS:
model.load_weights(path_weights)
# Train model.
if MODEL_RUN_TRAIN:
model.fit(
train_dataset,
validation_data=test_dataset,
epochs=EPOCHS,
steps_per_epoch=len_train // BATCH_SIZE,
validation_steps=len_test // BATCH_SIZE,
callbacks=[callback_save_best_weights, callback_csv_logger],
)
# Evaluate model.
evaluate = model.evaluate(test_dataset)
print(tabulate(zip(model.metrics_names, evaluate), tablefmt="fancy_grid"))
# Make prediction.
predict_dataset = model.predict(test_dataset)
# Save model.
model.save(path_model)
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# COMPARE WITH REAL COORDINATES.
X_test_names = [pathlib.Path(x).stem for x in X_test]
result_array = np.zeros((len_test, CLASSES_COUNT))
for k, predict_array in enumerate(predict_dataset):
for channel_type, channel_number in zip(CLASSES, range(CLASSES_COUNT)):
# Get channel.
predict_channel = predict_array[:, :, channel_number]
# Get predicted coordinates with max value.
maxes = np.argwhere(predict_channel.max() == predict_channel)
x_p, y_p = np.mean(maxes, axis=0).astype(np.int)
# Rescale predicted coordinates.
x_p, y_p = rescale_coords(
(x_p, y_p),
src_size=(IMAGE_NET_W, IMAGE_NET_H),
out_size=(IMAGE_SRC_W, IMAGE_SRC_H),
)
# Load text coordinates.
filename = PWD_COORDS / channel_type / f"{X_test_names[k]}.txt"
x_t, y_t = np.fromfile(filename, sep=",")
# Compute distance.
# ! NOTE: in this dataset, 1 cm == 11 pix.
distance = np.sqrt((x_t - x_p)**2 + (y_t - y_p)**2) / 11.0
# Save result.
result_array[k][channel_number] = distance
# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# DUMP RESULTS.
# Dump full results to CSV.
results = pd.DataFrame(columns=["filename", *CLASSES])
for n, filename in enumerate(X_test_names):
# Write rows of matrix to CSV.
named_results = dict(zip(CLASSES, list(result_array[n])))
named_results["filename"] = filename
results = results.append(named_results, ignore_index=True)
results.to_csv(path_predict_fulls, index=False, sep=";")
# Dump mean results to CSV.
results = pd.DataFrame(columns=["class", "mean", "max"])
means = np.mean(result_array, axis=0)
maxes = np.max(result_array, axis=0)
for class_label, val_mean, val_max in zip(CLASSES, means, maxes):
results = results.append(
{
"class": class_label,
"mean": val_mean,
"max": val_max
},
ignore_index=True,
)
results.to_csv(path_predict_means, index=False, sep=";")
print(results)
