-
Notifications
You must be signed in to change notification settings - Fork 10
/
bilinear_torch.py
101 lines (72 loc) · 2.77 KB
/
bilinear_torch.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
#
# Author : Alwyn Mathew
#
# Monodepth in pytorch(https://github.com/alwynmathew/monodepth-pytorch)
# Bilinear sampler in pytorch(https://github.com/alwynmathew/bilinear-sampler-pytorch)
#
from __future__ import absolute_import, division, print_function
import torch
import torch.nn.functional as F
from torch.autograd import Variable
def bilinear_sampler_1d_h(input_images, x_offset, wrap_mode='edge', name='bilinear_sampler', **kwargs):
def _repeat(x, n_repeats):
rep = x.unsqueeze(1).repeat(1, n_repeats)
return rep.view(-1)
def _interpolate(im, x, y):
# handle both texture border types
_edge_size = 0
if _wrap_mode == 'border':
_edge_size = 1
im = F.pad(im,(0,1,1,0), 'constant',0)
x = x + _edge_size
y = y + _edge_size
elif _wrap_mode == 'edge':
_edge_size = 0
else:
return None
x = torch.clamp(x, 0.0, _width_f - 1 + 2 * _edge_size)
x0_f = torch.floor(x)
y0_f = torch.floor(y)
x1_f = x0_f + 1
x0 = x0_f.type(torch.FloatTensor).cuda()
y0 = y0_f.type(torch.FloatTensor).cuda()
min_val = _width_f - 1 + 2 * _edge_size
scalar = Variable(torch.FloatTensor([min_val]).cuda())
x1 = torch.min(x1_f, scalar)
x1 = x1.type(torch.FloatTensor).cuda()
dim2 = (_width + 2 * _edge_size)
dim1 = (_width + 2 * _edge_size) * (_height + 2 * _edge_size)
base = Variable(_repeat(torch.arange(_num_batch) * dim1, _height * _width).cuda())
base_y0 = base + y0 * dim2
idx_l = base_y0 + x0
idx_r = base_y0 + x1
idx_l = idx_l.type(torch.cuda.LongTensor)
idx_r = idx_r.type(torch.cuda.LongTensor)
im_flat = im.contiguous().view(-1, _num_channels)
pix_l = torch.gather(im_flat, 0, idx_l.repeat(_num_channels).view(-1, _num_channels))
pix_r = torch.gather(im_flat, 0, idx_r.repeat(_num_channels).view(-1, _num_channels))
weight_l = (x1_f - x).unsqueeze(1)
weight_r = (x - x0_f).unsqueeze(1)
return weight_l * pix_l + weight_r * pix_r
def _transform(input_images, x_offset):
a = Variable(torch.linspace(0.0, _width_f -1.0, _width).cuda())
b = Variable(torch.linspace(0.0, _height_f -1.0, _height).cuda())
x_t = a.repeat(_height)
y_t = b.repeat(_width,1).t().contiguous().view(-1)
x_t_flat = x_t.repeat(_num_batch, 1)
y_t_flat = y_t.repeat(_num_batch, 1)
x_t_flat = x_t_flat.view(-1)
y_t_flat = y_t_flat.view(-1)
x_t_flat = x_t_flat + x_offset.contiguous().view(-1) * _width_f
input_transformed = _interpolate(input_images, x_t_flat, y_t_flat)
output = input_transformed.view(_num_batch, _num_channels, _height, _width)
return output
_num_batch = input_images.size(0)
_num_channels = input_images.size(1)
_height = input_images.size(2)
_width = input_images.size(3)
_height_f = float(_height)
_width_f = float(_width)
_wrap_mode = wrap_mode
output = _transform(input_images, x_offset)
return output