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]