/
utils.py
112 lines (94 loc) · 4.69 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
from __future__ import division
import shutil
import numpy as np
import torch
from path import Path
import datetime
from collections import OrderedDict
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
def save_path_formatter(args, parser):
def is_default(key, value):
return value == parser.get_default(key)
args_dict = vars(args)
data_folder_name = str(Path(args_dict['data']).normpath().name)
folder_string = [data_folder_name]
if not is_default('epochs', args_dict['epochs']):
folder_string.append('{}epochs'.format(args_dict['epochs']))
keys_with_prefix = OrderedDict()
keys_with_prefix['epoch_size'] = 'epoch_size'
keys_with_prefix['sequence_length'] = 'seq'
keys_with_prefix['rotation_mode'] = 'rot_'
keys_with_prefix['padding_mode'] = 'padding_'
keys_with_prefix['batch_size'] = 'b'
keys_with_prefix['lr'] = 'lr'
keys_with_prefix['photo_loss_weight'] = 'p'
keys_with_prefix['mask_loss_weight'] = 'm'
keys_with_prefix['smooth_loss_weight'] = 's'
for key, prefix in keys_with_prefix.items():
value = args_dict[key]
if not is_default(key, value):
folder_string.append('{}{}'.format(prefix, value))
save_path = Path(','.join(folder_string))
timestamp = datetime.datetime.now().strftime("%m-%d-%H:%M")
return save_path/timestamp
def high_res_colormap(low_res_cmap, resolution=1000, max_value=1):
# Construct the list colormap, with interpolated values for higer resolution
# For a linear segmented colormap, you can just specify the number of point in
# cm.get_cmap(name, lutsize) with the parameter lutsize
x = np.linspace(0, 1, low_res_cmap.N)
low_res = low_res_cmap(x)
new_x = np.linspace(0, max_value, resolution)
high_res = np.stack([np.interp(new_x, x, low_res[:, i]) for i in range(low_res.shape[1])], axis=1)
return ListedColormap(high_res)
def opencv_rainbow(resolution=1000):
# Construct the opencv equivalent of Rainbow
opencv_rainbow_data = (
(0.000, (1.00, 0.00, 0.00)),
(0.400, (1.00, 1.00, 0.00)),
(0.600, (0.00, 1.00, 0.00)),
(0.800, (0.00, 0.00, 1.00)),
(1.000, (0.60, 0.00, 1.00))
)
return LinearSegmentedColormap.from_list('opencv_rainbow', opencv_rainbow_data, resolution)
COLORMAPS = {'rainbow': opencv_rainbow(),
'magma': high_res_colormap(cm.get_cmap('magma')),
'bone': cm.get_cmap('bone', 10000)}
def log_output_tensorboard(writer, prefix, index, suffix, n_iter, depth, disp, warped, diff, mask):
disp_to_show = tensor2array(disp[0], max_value=None, colormap='magma')
depth_to_show = tensor2array(depth[0], max_value=None)
writer.add_image('{} Dispnet Output Normalized{}/{}'.format(prefix, suffix, index), disp_to_show, n_iter)
writer.add_image('{} Depth Output Normalized{}/{}'.format(prefix, suffix, index), depth_to_show, n_iter)
# log warped images along with explainability mask
if (warped is None) or (diff is None):
return
for j, (warped_j, diff_j) in enumerate(zip(warped, diff)):
whole_suffix = '{} {}/{}'.format(suffix, j, index)
warped_to_show = tensor2array(warped_j)
diff_to_show = tensor2array(0.5*diff_j)
writer.add_image('{} Warped Outputs {}'.format(prefix, whole_suffix), warped_to_show, n_iter)
writer.add_image('{} Diff Outputs {}'.format(prefix, whole_suffix), diff_to_show, n_iter)
if mask is not None:
mask_to_show = tensor2array(mask[0, j], max_value=1, colormap='bone')
writer.add_image('{} Exp mask Outputs {}'.format(prefix, whole_suffix), mask_to_show, n_iter)
def tensor2array(tensor, max_value=None, colormap='rainbow'):
tensor = tensor.detach().cpu()
if max_value is None:
max_value = tensor[tensor < np.inf].max().item()
if tensor.ndimension() == 2 or tensor.size(0) == 1:
norm_array = tensor.squeeze().numpy()/max_value
norm_array[norm_array == np.inf] = np.nan
array = COLORMAPS[colormap](norm_array).astype(np.float32)
array = array.transpose(2, 0, 1)[:3]
elif tensor.ndimension() == 3:
assert(tensor.size(0) == 3)
array = 0.5 + tensor.numpy()*0.5
return array
def save_checkpoint(save_path, dispnet_state, exp_pose_state, is_best, filename='checkpoint.pth.tar'):
file_prefixes = ['dispnet', 'exp_pose']
states = [dispnet_state, exp_pose_state]
for (prefix, state) in zip(file_prefixes, states):
torch.save(state, save_path/'{}_{}'.format(prefix, filename))
if is_best:
for prefix in file_prefixes:
shutil.copyfile(save_path/'{}_{}'.format(prefix, filename), save_path/'{}_model_best.pth.tar'.format(prefix))