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