for horizon_path in config['horizon-path']:
horizon = Horizon(horizon_path, geometry=geometry)
hm = HorizonMetrics(horizon)
prefix = '' if config['add-prefix'] is False else horizon.name + '_'
with open(os.path.join(config['savedir'], f'{prefix}metrics_info.txt'),
'w') as result_txt:
_ = horizon.evaluate(
compute_metric=False,
printer=lambda msg: print(msg, file=result_txt))
for metric_name in config['metrics']:
kwargs = copy(LOCAL_KWARGS) if metric_name.startswith(
'local') else copy(SUPPORT_KWARGS)
kwargs = {} if metric_name.startswith('insta') else kwargs
savepath = os.path.join(config['savedir'],
f'{prefix}{metric_name}')
metric = hm.evaluate(metric_name,
**kwargs,
plot=True,
show_plot=False,
plot_kwargs={'figsize': (20, 20)},
savepath=savepath + '.png')
if config['save-txt']:
save_point_cloud(metric, savepath + '.txt', geometry)
print(
f'{metric_name} avg value: {""*20} {np.nanmean(metric):5.5}',
file=result_txt)
loss = loss_d1(coarse_gt, y_coarse) + args.alpha * loss_d1(dense_gt, y_detail)
total_loss += loss.item()
iter_count += 1
mean_loss = total_loss / iter_count
print("Validation epoch {}/{}, loss is {}".format(epoch, args.epochs, mean_loss))
# Save model
if mean_loss < minimum_loss:
partial_input, coarse_gt, dense_gt = val_dataset[random.randint(0, len(val_dataset))]
partial_input = partial_input.to(DEVICE)
coarse_gt = coarse_gt.to(DEVICE)
dense_gt = dense_gt.to(DEVICE)
partial_input = partial_input.permute(0, 2, 1)
v, y_coarse, y_detail = network(partial_input)
y_coarse = y_coarse.permute(0, 2, 1)
y_detail = y_detail.permute(0, 2, 1)
save_point_cloud(partial_input.numpy(), 'visual/partial_input.pcd')
save_point_cloud(dense_gt.numpy(), 'visual/gt.pcd')
save_point_cloud(y_detail.numpy(), 'visual/output.pcd')
best_epoch = epoch
minimum_loss = mean_loss
torch.save(network.state_dict(), + 'model/trained_model_fold.pth')
# Argname, description, dtype, default
ARGS = [
('cube-path', 'path to the seismic cube in HDF5 format', str, None),
('savedir', 'path to save files to', str, '_placeholder_'),
('metrics', 'metrics to compute for quality map creation', str, ['support_hellinger']),
('add-prefix', 'whether to prepend cube name to the saved file names', str2bool, True),
('save-txt', 'whether to save point cloud of quality map to disk', str2bool, False),
]
if __name__ == '__main__':
config = make_config(MSG, ARGS, os.path.basename(__file__).split('.')[0])
if config['savedir'] == '_placeholder_':
config['savedir'] = os.path.dirname(config['cube-path'])
geometry = SeismicGeometry(config['cube-path'])
geometry.make_quality_map([0.1, 0.15, 0.2, 0.4, 0.5], config['metrics'])
prefix = '' if config['add-prefix'] is False else geometry.short_name + '_'
safe_mkdir(config['savedir'])
plot_image(geometry.quality_map, cmap='Reds',
xlabel='INLINE_3D', ylabel='CROSSLINE_3D', title='',
savepath=os.path.join(config['savedir'], f'{prefix}quality_map.png'))
if config['save-txt']:
save_point_cloud(geometry.quality_map,
os.path.join(config['savedir'], f'{prefix}quality_map.txt'),
geometry)