-
Notifications
You must be signed in to change notification settings - Fork 1
/
utils.py
executable file
·86 lines (70 loc) · 3.51 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
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
from img_show import img_show_singleimage_numpy,visualize_flow,show_flow, tensor2array_flow
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 tensor2array(tensor, max_value=None, colormap='rainbow'):
tensor = tensor.detach().cpu()
if max_value is None:
max_value = tensor.max().item()
if tensor.ndimension() == 2 or tensor.size(0) == 1:
norm_array = tensor.squeeze().numpy()/max_value
np.set_printoptions(threshold=np.inf)
#print(norm_array)
array = COLORMAPS[colormap](norm_array).astype(np.float32)
array = array.transpose(2, 0, 1)
elif tensor.ndimension() == 3:
assert(tensor.size(0) == 3)
array = 0.45 + tensor.numpy()*0.225
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))
# TODO: depth
def log_output_tensorboard(writer, prefix, index, suffix, n_iter, depth, disp, flow_forward, flow_backward): # TODO: add flow
# depth for tgt
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
for i, (flow, b_f) in enumerate(zip(flow_forward,flow_backward)):
whole_suffix = 'flow:{} {}/{}'.format(suffix, i, index)
flow_to_show = tensor2array_flow(flow)
flowb_to_show = tensor2array_flow(b_f)
writer.add_image('{} flow Outputs {}'.format(prefix, whole_suffix), flow_to_show, n_iter)
writer.add_image('{} flowback Outputs {}'.format(prefix, whole_suffix), flowb_to_show, n_iter)