def run_pwc(self, input_dict, x1_raw, x2_raw, k1, k2):
output_dict = {}
# on the bottom level are original images
x1_pyramid = self.feature_pyramid_extractor(x1_raw) + [x1_raw]
x2_pyramid = self.feature_pyramid_extractor(x2_raw) + [x2_raw]
# outputs
disps_1 = []
disps_2 = []
for l, (x1, x2) in enumerate(zip(x1_pyramid, x2_pyramid)):
# warping
if l == 0:
x2_warp = x2
x1_warp = x1
else:
disp_1 = interpolate2d_as(disp_1, x1, mode="bilinear")
disp_2 = interpolate2d_as(disp_2, x1, mode="bilinear")
x1_out = self.upconv_layers[l - 1](x1_out)
x2_out = self.upconv_layers[l - 1](x2_out)
# disparity estimator
if l == 0:
x1_out, disp_1 = self.disp_estimators[l](x1)
x2_out, disp_2 = self.disp_estimators[l](x2)
else:
x1_out, disp_1 = self.disp_estimators[l](torch.cat(
[x1, x1_out, disp_1], dim=1))
x2_out, disp_2 = self.disp_estimators[l](torch.cat(
[x2, x2_out, disp_2], dim=1))
# upsampling or post-processing
disp_1 = self.sigmoid(disp_1) * 0.3
disp_2 = self.sigmoid(disp_2) * 0.3
disps_1.append(disp_1)
disps_2.append(disp_2)
if l == self.output_level:
break
x1_rev = x1_pyramid[::-1]
output_dict['disp_l1'] = upsample_outputs_as(disps_1[::-1], x1_rev)
output_dict['disp_l2'] = upsample_outputs_as(disps_2[::-1], x1_rev)
return output_dict
def forward(self, x, sceneflow, disp, k1, input_size):
_, _, h_x, w_x = x.size()
disp = interpolate2d_as(disp, x) * w_x
local_scale = torch.zeros_like(input_size)
local_scale[:, 0] = h_x
local_scale[:, 1] = w_x
pts1, k1_scale = pixel2pts_ms(k1, disp, local_scale / input_size)
_, _, coord1 = pts2pixel_ms(k1_scale, pts1, sceneflow, [h_x, w_x])
grid = coord1.transpose(1, 2).transpose(2, 3)
x_warp = tf.grid_sample(x, grid)
mask = torch.ones_like(x, requires_grad=False)
mask = tf.grid_sample(mask, grid)
mask = (mask >= 1.0).float()
return x_warp * mask
def run_pwc(self, input_dict, x1_raw, x2_raw, k1, k2):
output_dict = {}
# on the bottom level are original images
x1_pyramid = self.feature_pyramid_extractor(x1_raw) + [x1_raw]
x2_pyramid = self.feature_pyramid_extractor(x2_raw) + [x2_raw]
# outputs
sceneflows_f = []
sceneflows_b = []
disps_1 = []
disps_2 = []
for l, (x1, x2) in enumerate(zip(x1_pyramid, x2_pyramid)):
# warping
if l == 0:
x2_warp = x2
x1_warp = x1
else:
flow_f = interpolate2d_as(flow_f, x1, mode="bilinear")
flow_b = interpolate2d_as(flow_b, x1, mode="bilinear")
disp_l1 = interpolate2d_as(disp_l1, x1, mode="bilinear")
disp_l2 = interpolate2d_as(disp_l2, x1, mode="bilinear")
x1_out = self.upconv_layers[l - 1](x1_out)
x2_out = self.upconv_layers[l - 1](x2_out)
x2_warp = self.warping_layer_sf(
x2, flow_f, disp_l1, k1, input_dict['aug_size']
) # becuase K can be changing when doing augmentation
x1_warp = self.warping_layer_sf(x1, flow_b, disp_l2, k2,
input_dict['aug_size'])
# correlation
out_corr_f = Correlation.apply(x1, x2_warp, self.corr_params)
out_corr_b = Correlation.apply(x2, x1_warp, self.corr_params)
out_corr_relu_f = self.leakyRELU(out_corr_f)
out_corr_relu_b = self.leakyRELU(out_corr_b)
# monosf estimator
if l == 0:
x1 = self.camconv(x1, x1_raw, k1)
x2 = self.camconv(x2, x2_raw, k2)
x1_out, flow_f, disp_l1 = self.flow_estimators[l](torch.cat(
[out_corr_relu_f, x1], dim=1))
x2_out, flow_b, disp_l2 = self.flow_estimators[l](torch.cat(
[out_corr_relu_b, x2], dim=1))
else:
x1 = self.camconv(x1)
x2 = self.camconv(x2)
x1_out, flow_f_res, disp_l1 = self.flow_estimators[l](
torch.cat([out_corr_relu_f, x1, x1_out, flow_f, disp_l1],
dim=1))
x2_out, flow_b_res, disp_l2 = self.flow_estimators[l](
torch.cat([out_corr_relu_b, x2, x2_out, flow_b, disp_l2],
dim=1))
flow_f = flow_f + flow_f_res
flow_b = flow_b + flow_b_res
# upsampling or post-processing
if l != self.output_level:
disp_l1 = self.sigmoid(disp_l1) * 0.3
disp_l2 = self.sigmoid(disp_l2) * 0.3
sceneflows_f.append(flow_f)
sceneflows_b.append(flow_b)
disps_1.append(disp_l1)
disps_2.append(disp_l2)
else:
flow_res_f, disp_l1 = self.context_networks(
torch.cat([x1_out, flow_f, disp_l1], dim=1))
flow_res_b, disp_l2 = self.context_networks(
torch.cat([x2_out, flow_b, disp_l2], dim=1))
flow_f = flow_f + flow_res_f
flow_b = flow_b + flow_res_b
sceneflows_f.append(flow_f)
sceneflows_b.append(flow_b)
disps_1.append(disp_l1)
disps_2.append(disp_l2)
break
x1_rev = x1_pyramid[::-1]
output_dict['flow_f'] = upsample_outputs_as(sceneflows_f[::-1], x1_rev)
output_dict['flow_b'] = upsample_outputs_as(sceneflows_b[::-1], x1_rev)
output_dict['disp_l1'] = upsample_outputs_as(disps_1[::-1], x1_rev)
output_dict['disp_l2'] = upsample_outputs_as(disps_2[::-1], x1_rev)
return output_dict
def run_pwc(self, input_dict, x1_raw, x2_raw, k1, k2):
output_dict = {}
# on the bottom level are original images
x1_pyramid = self.feature_pyramid_extractor(x1_raw) + [x1_raw]
x2_pyramid = self.feature_pyramid_extractor(x2_raw) + [x2_raw]
# outputs
flows_f = []
flows_b = []
for l, (x1, x2) in enumerate(zip(x1_pyramid, x2_pyramid)):
# warping
if l == 0:
x2_warp = x2
x1_warp = x1
else:
flow_f = interpolate2d_as(flow_f, x1, mode="bilinear")
flow_b = interpolate2d_as(flow_b, x1, mode="bilinear")
x1_out = self.upconv_layers[l - 1](x1_out)
x2_out = self.upconv_layers[l - 1](x2_out)
x2_warp = self.warping_layer(x2, flow_f)
x1_warp = self.warping_layer(x1, flow_b)
# correlation
out_corr_f = Correlation.apply(x1, x2_warp, self.corr_params)
out_corr_b = Correlation.apply(x2, x1_warp, self.corr_params)
out_corr_relu_f = self.leakyRELU(out_corr_f)
out_corr_relu_b = self.leakyRELU(out_corr_b)
# flow estimator
if l == 0:
x1_out, flow_f = self.flow_estimators[l](torch.cat(
[out_corr_relu_f, x1], dim=1))
x2_out, flow_b = self.flow_estimators[l](torch.cat(
[out_corr_relu_b, x2], dim=1))
else:
x1_out, flow_f_res = self.flow_estimators[l](torch.cat(
[out_corr_relu_f, x1, x1_out, flow_f], dim=1))
x2_out, flow_b_res = self.flow_estimators[l](torch.cat(
[out_corr_relu_b, x2, x2_out, flow_b], dim=1))
flow_f = flow_f + flow_f_res
flow_b = flow_b + flow_b_res
# upsampling or post-processing
if l != self.output_level:
flows_f.append(flow_f)
flows_b.append(flow_b)
else:
flow_res_f = self.context_networks(
torch.cat([x1_out, flow_f], dim=1))
flow_res_b = self.context_networks(
torch.cat([x2_out, flow_b], dim=1))
flow_f = flow_f + flow_res_f
flow_b = flow_b + flow_res_b
flows_f.append(flow_f)
flows_b.append(flow_b)
break
x1_rev = x1_pyramid[::-1]
output_dict['flow_f'] = upsample_outputs_as(flows_f[::-1], x1_rev)
output_dict['flow_b'] = upsample_outputs_as(flows_b[::-1], x1_rev)
return output_dict
def run_pwc(self, input_dict, x1_raw, x2_raw, k1, k2):
output_dict = {}
x1_pyramid = self.feature_pyramid_extractor(x1_raw) + [x1_raw]
x2_pyramid = self.feature_pyramid_extractor(x2_raw) + [x2_raw]
# outputs
flows_f = []
flows_b = []
disps_l1 = []
disps_l2 = []
for l, (x1, x2) in enumerate(zip(x1_pyramid, x2_pyramid)):
# warping
if l == 0:
x2_warp = x2
x1_warp = x1
else:
sf_f = interpolate2d_as(sf_f, x1, mode="bilinear")
sf_b = interpolate2d_as(sf_b, x1, mode="bilinear")
dp_f = interpolate2d_as(dp_f, x1, mode="bilinear")
dp_b = interpolate2d_as(dp_b, x1, mode="bilinear")
sf_f_out = self.upconv_layers_sf[l - 1](sf_f_out)
sf_b_out = self.upconv_layers_sf[l - 1](sf_b_out)
dp_f_out = self.upconv_layers_dp[l - 1](dp_f_out)
dp_b_out = self.upconv_layers_dp[l - 1](dp_b_out)
x2_warp = self.warping_layer_sf(
x2, sf_f, dp_f, k1, input_dict['aug_size']
) # becuase K can be changing when doing augmentation
x1_warp = self.warping_layer_sf(x1, sf_b, dp_b, k2,
input_dict['aug_size'])
# correlation
out_corr_f = Correlation.apply(x1, x2_warp, self.corr_params)
out_corr_b = Correlation.apply(x2, x1_warp, self.corr_params)
out_corr_relu_f = self.leakyRELU(out_corr_f)
out_corr_relu_b = self.leakyRELU(out_corr_b)
# flow estimator
if l == 0:
sf_f_out, sf_f, dp_f_out, dp_f = self.sf_estimators[l](
torch.cat([out_corr_relu_f, x1], dim=1))
sf_b_out, sf_b, dp_b_out, dp_b = self.sf_estimators[l](
torch.cat([out_corr_relu_b, x2], dim=1))
else:
sf_f_out, sf_f_res, dp_f_out, dp_f = self.sf_estimators[l](
torch.cat(
[out_corr_relu_f, x1, sf_f_out, sf_f, dp_f_out, dp_f],
dim=1))
sf_b_out, sf_b_res, dp_b_out, dp_b = self.sf_estimators[l](
torch.cat(
[out_corr_relu_b, x2, sf_b_out, sf_b, dp_b_out, dp_b],
dim=1))
sf_f = sf_f + sf_f_res
sf_b = sf_b + sf_b_res
# upsampling or post-processing
if l != self.output_level:
dp_f = self.sigmoid(dp_f) * 0.3
dp_b = self.sigmoid(dp_b) * 0.3
flows_f.append(sf_f)
flows_b.append(sf_b)
disps_l1.append(dp_f)
disps_l2.append(dp_b)
else:
sf_res_f, dp_f = self.context_networks(
torch.cat([sf_f_out, sf_f], dim=1),
torch.cat([dp_f_out, dp_f], dim=1))
sf_res_b, dp_b = self.context_networks(
torch.cat([sf_b_out, sf_b], dim=1),
torch.cat([dp_b_out, dp_b], dim=1))
sf_f = sf_f + sf_res_f
sf_b = sf_b + sf_res_b
flows_f.append(sf_f)
flows_b.append(sf_b)
disps_l1.append(dp_f)
disps_l2.append(dp_b)
break
x1_rev = x1_pyramid[::-1]
output_dict['flow_f'] = upsample_outputs_as(flows_f[::-1], x1_rev)
output_dict['flow_b'] = upsample_outputs_as(flows_b[::-1], x1_rev)
output_dict['disp_l1'] = upsample_outputs_as(disps_l1[::-1], x1_rev)
output_dict['disp_l2'] = upsample_outputs_as(disps_l2[::-1], x1_rev)
return output_dict
def upsample_outputs_as(input_list, ref_list):
output_list = []
for ii in range(0, len(input_list)):
output_list.append(interpolate2d_as(input_list[ii], ref_list[ii]))
return output_list