def __init__(self, name_model, maxdisp): super(WSMCnetELB, self).__init__(name_model, maxdisp) # modules of [feature_extraction and cost_compute] from Submodules.FeatureExtraction import FeatureExtraction from Submodules.CostComputeL import CostCompute self.feature_extraction = FeatureExtraction(planes=32) self.cost_compute = CostCompute(inplanes=64, planes=self.F3, C=self.C) # init weight Submodules.weight_init(self.modules())
def get_matching_costs(self, left, right): # feature extraction bn = left.shape[0] features = torch.cat([left, right], dim=0) features = self.feature_extraction(features) fL, fR = features[:bn], features[bn:] # compute matching cost stride = self.S shift = (self.maxdisp // (stride << 2)) + 1 cost = Submodules.cost_volume_gen(fL, fR, shift, stride) costs = self.cost_compute(cost) return costs
def forward(self, left, right): # compute matching costs costs = self.get_matching_costs(left, right) # disparty regression h, w = left.shape[-2:] disps = [] for tcost in costs: tcost = self.upsample(tcost.squeeze(1))[..., :h, :w] disps.append(Submodules.disp_regression(tcost, self.disp_step)) # return logger.debug('training: %s \n' % str(self.training)) return disps if self.training else disps[-1]
def forward(self, left, right): # compute matching costs costs = self.get_matching_costs(left, right) # disparty regression h, w = left.shape[-2:] costs = [ self.upsample(tcost.squeeze(1))[..., :h, :w] for tcost in costs ] disps = [ Submodules.disp_regression_nearby(tcost, self.disp_step) for tcost in costs ] if (self.training): return costs, disps else: return disps[-1]