def __init__(self, s, flagNearest=False):
super(Interpolate2D_FixedScale, self).__init__()
self.s = s
self.flagAlignCorners = GLOBAL.torch_align_corners()
self.mode = 'nearest' if flagNearest else 'bilinear'
def __init__(self,
inCh,
levels,
maxKFactor=0.5,
lastActivation=None,
flagNearest=False):
super(SpatialPyramidPooling, self).__init__()
# Global settings.
self.flagAlighCorners = GLOBAL.torch_align_corners()
self.flagReLUInplace = GLOBAL.torch_relu_inplace()
self.inCh = inCh
# Pooling levels.
self.levels = int(levels)
assert (self.levels > 0)
# Kernel size factor.
assert (0 < maxKFactor <= 1)
self.maxKFactor = maxKFactor
# Convolusion layers for the pooling levels.
self.poolingConvs = self.make_pooling_convs()
# Last activation.
self.lastActivation = lastActivation
# The interpolation mode.
self.interMode = 'nearest' if flagNearest else self.get_interpolate_mode(
)
def __init__(self,
flagMaskedLoss=True,
flagIntNearest=False):
super(LossComputer, self).__init__()
self.flagAlignCorners = modelGLOBAL.torch_align_corners()
self.flagIntNearest = flagIntNearest
self.flagMaskedLoss = flagMaskedLoss # Set True to mask the true disparity larger than self.trueDispMask.
def __init__(self, workingDir, conf, frame=None, modelName='Stereo'):
self.conf = conf # The configuration dictionary.
self.wd = workingDir
self.frame = frame
self.modelName = modelName
# NN.
self.countTrain = 0
self.countTest = 0
self.flagAlignCorners = modelGLOBAL.torch_align_corners()
self.flagIntNearest = False
self.trainIntervalAccWrite = 10 # The interval to write the accumulated values.
self.trainIntervalAccPlot = 1 # The interval to plot the accumulate values.
self.flagUseIntPlotter = False # The flag of intermittent plotter.
self.flagCPU = False
self.multiGPUs = False
self.readModelString = ""
self.readOptimizerString = ""
self.autoSaveModelLoops = 0 # The number of loops to perform an auto-saving of the model. 0 for disable.
self.autoSnapLoops = 100 # The number of loops to perform an auto-snap.
self.optimizer = None
self.flagTest = False # Should be set to True when testing.
self.flagRandomSeedSet = False
# Specified by conf.
self.trueDispMask = conf['tt']['trueDispMask']
self.model = None # make_object later.
self.dataloader = None # make_object later.
self.optType = conf['tt']['optType'] # The optimizer type. adam, sgd.
# Learning rate scheduler.
self.flagUseLRS = conf['tt']['flagUseLRS']
self.learningRate = conf['tt']['lr']
self.lrs = None # make_object later.
# True value and loss.
self.trueValueGenerator = None # make_object later.
self.lossComputer = None # make_object later.
self.testResultSubfolder = conf['tt']['testResultSubfolder']
# Test figure generator.
self.testFigGenerator = None
# Temporary values during traing and testing.
self.ctxInputs = None
self.ctxOutputs = None
self.ctxTrueValues = None
self.ctxLossValues = None
# InfoUpdaters.
self.trainingInfoUpdaters = []
self.testingInfoUpdaters = []
def __init__(self,
nConvs, inCh, interCh, outCh,
baseStride=(1,1,1), nStrides=1,
outputUpSampledFeat=False, pooling=False ):
super(DecoderBlock, self).__init__()
# Get the global settings.
self.flagAlignCorners = GLOBAL.torch_align_corners()
self.flagReLUInplace = GLOBAL.torch_relu_inplace()
# Prepare the list of strides.
assert( nConvs >= nStrides )
strideList = [baseStride] * nStrides + [(1,1,1)] * (nConvs - nStrides)
# Create the the convolusion layers.
convs = [ SepConv3DBlock( inCh, interCh, stride=strideList[0] ) ]
for i in range(1, nConvs):
convs.append( SepConv3DBlock( interCh, interCh, stride=strideList[i] ) )
self.entryConvs = WrappedModule( nn.Sequential(*convs) )
self.append_init_here( self.entryConvs )
# Classification layer.
self.classify = WrappedModule(
nn.Sequential(
cm3d.Conv3D_W( interCh, interCh,
normLayer=cm3d.FeatureNorm3D(interCh),
activation=nn.ReLU(inplace=self.flagReLUInplace) ),
cm3d.Conv3D_W(interCh, outCh, bias=True) ) )
self.append_init_here(self.classify)
# Feature up-sample setting.
self.featUpSampler = None
if outputUpSampledFeat:
self.featUpSampler = WrappedModule(
nn.Sequential(
cm3d.Interpolate3D_FixedScale(2),
cm3d.Conv3D_W( interCh, interCh//2,
normLayer=cm3d.FeatureNorm3D(interCh//2),
activation=nn.ReLU(inplace=self.flagReLUInplace) ) ) )
self.append_init_here(self.featUpSampler)
# Pooling.
if pooling:
self.spp = SPP3D( interCh, levels=4 )
self.append_init_here(self.spp)
else:
self.spp = None
def __init__(self,
nconvs,
inchannelF,
channelF,
stride=(1, 1, 1),
up=False,
nstride=1,
pool=False):
super(decoderBlock, self).__init__()
self.flagAlignCorners = GLOBAL.torch_align_corners()
self.flagReLUInplace = GLOBAL.torch_relu_inplace()
self.pool = pool
stride = [stride] * nstride + [(1, 1, 1)] * (nconvs - nstride)
self.convs = [sepConv3dBlock(inchannelF, channelF, stride=stride[0])]
for i in range(1, nconvs):
self.convs.append(
sepConv3dBlock(channelF, channelF, stride=stride[i]))
self.convs = nn.Sequential(*self.convs)
self.classify = nn.Sequential(
sepConv3d(channelF, channelF, 3, (1, 1, 1), 1),
nn.ReLU(inplace=self.flagReLUInplace),
sepConv3d(channelF, 1, 3, (1, 1, 1), 1, bias=True))
self.up = False
if up:
self.up = True
self.up = nn.Sequential(
nn.Upsample(scale_factor=(2, 2, 2),
mode='trilinear',
align_corners=self.flagAlignCorners),
sepConv3d(channelF, channelF // 2, 3, (1, 1, 1), 1,
bias=False), nn.ReLU(inplace=self.flagReLUInplace))
if pool:
self.pool_convs = torch.nn.ModuleList([
sepConv3d(channelF, channelF, 1, (1, 1, 1), 0),
sepConv3d(channelF, channelF, 1, (1, 1, 1), 0),
sepConv3d(channelF, channelF, 1, (1, 1, 1), 0),
sepConv3d(channelF, channelF, 1, (1, 1, 1), 0)
])
def __init__(self, in_channels, pool_sizes, model_name='pspnet', fusion_mode='cat', with_bn=True):
super(pyramidPooling, self).__init__()
self.flagAlignCorners = GLOBAL.torch_align_corners()
self.flagReLUInplace = GLOBAL.torch_relu_inplace()
bias = not with_bn
self.paths = []
if pool_sizes is None:
for i in range(4):
self.paths.append(conv2DBatchNormRelu(in_channels, in_channels, 1, 1, 0, bias=bias, with_bn=with_bn))
else:
for i in range(len(pool_sizes)):
self.paths.append(conv2DBatchNormRelu(in_channels, int(in_channels / len(pool_sizes)), 1, 1, 0, bias=bias, with_bn=with_bn))
self.path_module_list = nn.ModuleList(self.paths)
self.pool_sizes = pool_sizes
self.model_name = model_name
self.fusion_mode = fusion_mode
def disable_align_corners(self):
self.check_frame()
self.frame.logger.info("Use align_corners=False.")
modelGLOBAL.torch_align_corners(False)
self.flagAlignCorners = False
def enable_align_corners(self):
self.check_frame()
self.frame.logger.info("Use align_corners=True.")
modelGLOBAL.torch_align_corners(True)
self.flagAlignCorners = True
os.path.dirname(
os.path.dirname(
os.path.dirname( _CF ) ) ) ) )
print(f'Adding {_PKG_PATH} to the package search path. ')
import sys
sys.path.insert(0, _PKG_PATH)
# Import the package.
import stereo
from stereo.models.globals import GLOBAL
# Configure the global settings for testing.
# Not the right setting. Only for testing.
GLOBAL.torch_align_corners(True)
if __name__ == '__main__':
print(f'Hello, {os.path.basename(__file__)}! ')
# Show the global settings.
print(f'GLOBAL.torch_align_corners() = {GLOBAL.torch_align_corners()}')
# The dummy dictionary.
maxDisp = 192
d = dict(type='HSMNet',
maxdisp=maxDisp,
clean=-1,
featExtConfig=dict(type='UNet', initialChannels=32),
costVolConfig=dict(type='CVDiff', refIsRight=False),
dispRegConfigs=[
def __init__(self):
super(ClassifiedCostVolumeEpistemic, self).__init__()
self.flagAlignCorners = modelGLOBAL.torch_align_corners()
def __init__(self, trueDispMax, flagIntNearest=False):
super(TrueValueGenerator, self).__init__()
self.trueDispMax = trueDispMax
self.flagAlignCorners = modelGLOBAL.torch_align_corners()
self.flagIntNearest = flagIntNearest
def __init__(self,
maxDisp=192,
featExtConfig=None,
costVolConfig=None,
costPrcConfig=None,
dispRegConfigs=None,
uncertainty=False,
freeze=False):
super(CostVolPrimitive, self).__init__(freeze=freeze)
# Global setttings.
self.flagAlignCorners = GLOBAL.torch_align_corners()
# ========== Module definition. ==========
self.maxDisp = maxDisp
# Uncertainty setting.
self.uncertainty = uncertainty
# Feature extractor.
if ( featExtConfig is None ):
featExtConfig = UNet.get_default_init_args()
self.featureExtractor = make_object(FEAT_EXT, featExtConfig)
self.append_init_impl(self.featureExtractor)
# Cost volume.
if ( costVolConfig is None ):
costVolConfig = CVDiff.get_default_init_args()
self.costVol = make_object(COST_VOL, costVolConfig)
self.append_init_impl(self.costVol)
# Cost volume processing/regularization layers.
if ( costPrcConfig is None ):
costPrcConfig = C3D_SPP_CLS.get_default_init_args()
self.costProcessor = make_object(COST_PRC, costPrcConfig)
self.append_init_impl( self.costProcessor )
# Disparity regressions.
nLevels = self.featureExtractor.n_levels()
if ( dispRegConfigs is None ):
dispRegConfigs = [ ClsLinearCombination.get_default_init_args() for _ in range(nLevels)]
else:
assert( isinstance( dispRegConfigs, (tuple, list) ) ), \
f'dispRegConfigs must be a tuple or list. It is {type(dispRegConfigs)}'
assert( len(dispRegConfigs) == nLevels ), \
f'len(dispRegConfigs) = {len(dispRegConfigs)}, nLevels = {nLevels}'
self.dispRegList = nn.ModuleList()
for config in dispRegConfigs:
dispReg = make_object(DISP_REG, config)
self.dispRegList.append( dispReg )
self.append_init_impl( dispReg )
# Uncertainty computer.
self.uncertaintyComputer = ClassifiedCostVolumeEpistemic() \
if self.uncertainty else None
# Must be called at the end of __init__().
self.update_freeze()
def __init__(self,
maxdisp=192,
clean=-1,
level=1,
featExtConfig=None,
costVolConfig=None,
dispRegConfigs=None,
freeze=False):
super(HSMNet, self).__init__(freeze=freeze)
# Global setttings.
self.flagAlignCorners = GLOBAL.torch_align_corners()
# ========== Module definition. ==========
self.maxdisp = maxdisp
self.clean = clean
self.level = level
# Feature extractor.
if (featExtConfig is None):
featExtConfig = UNet.get_default_init_args()
self.featureExtractor = make_object(FEAT_EXT, featExtConfig)
self.append_init_impl(self.featureExtractor)
# Cost volume.
if (costVolConfig is None):
costVolConfig = CVDiff.get_default_init_args()
self.costVol = make_object(COST_VOL, costVolConfig)
self.append_init_impl(self.costVol)
# block 4
self.decoder6 = decoderBlock(6, 32, 32, up=True, pool=True)
self.append_init_impl(self.decoder6)
if self.level > 2:
self.decoder5 = decoderBlock(6, 32, 32, up=False, pool=True)
self.append_init_impl(self.decoder5)
else:
self.decoder5 = decoderBlock(6, 32, 32, up=True, pool=True)
self.append_init_impl(self.decoder5)
if self.level > 1:
self.decoder4 = decoderBlock(6, 32, 32, up=False)
self.append_init_impl(self.decoder4)
else:
self.decoder4 = decoderBlock(6, 32, 32, up=True)
self.decoder3 = decoderBlock(5,
32,
32,
stride=(2, 1, 1),
up=False,
nstride=1)
self.append_init_impl(self.decoder4)
self.append_init_impl(self.decoder3)
# Disparity regressions.
# self.disp_reg8 = disparityregression(self.maxdisp,16)
# self.disp_reg16 = disparityregression(self.maxdisp,16)
# self.disp_reg32 = disparityregression(self.maxdisp,32)
# self.disp_reg64 = disparityregression(self.maxdisp,64)
if (dispRegConfigs is None):
dispRegConfigs = [
disparityregression.get_default_init_args() for _ in range(4)
]
else:
assert( isinstance( dispRegConfigs, (tuple, list) ) ), \
f'dispRegConfigs must be a tuple or list. It is {type(dispRegConfigs)}'
self.disp_reg8 = make_object(DISP_REG, dispRegConfigs[0])
self.disp_reg16 = make_object(DISP_REG, dispRegConfigs[1])
self.disp_reg32 = make_object(DISP_REG, dispRegConfigs[2])
self.disp_reg64 = make_object(DISP_REG, dispRegConfigs[3])
self.append_init_impl(self.disp_reg8)
self.append_init_impl(self.disp_reg16)
self.append_init_impl(self.disp_reg32)
self.append_init_impl(self.disp_reg64)
# Must be called at the end of __init__().
self.update_freeze()
def __init__(self,
maxDisp=192,
featExtConfig=None,
costVolConfig=None,
dispRegConfigs=None,
uncertainty=False,
freeze=False):
super(HSMNet, self).__init__(freeze=freeze)
# Global setttings.
self.flagAlignCorners = GLOBAL.torch_align_corners()
# ========== Module definition. ==========
self.maxDisp = maxDisp
# Feature extractor.
if ( featExtConfig is None ):
featExtConfig = UNet.get_default_init_args()
self.featureExtractor = make_object(FEAT_EXT, featExtConfig)
self.append_init_impl(self.featureExtractor)
# Cost volume.
if ( costVolConfig is None ):
costVolConfig = CVDiff.get_default_init_args()
self.costVol = make_object(COST_VOL, costVolConfig)
self.append_init_impl(self.costVol)
# block 4
self.decoder6 = decoderBlock(6,32,32,up=True, pool=True, uncertainty=uncertainty)
self.append_init_impl(self.decoder6)
self.decoder5 = decoderBlock(6,32,32,up=True, pool=True, uncertainty=uncertainty)
self.append_init_impl(self.decoder5)
self.decoder4 = decoderBlock(6,32,32, up=True, uncertainty=uncertainty)
self.append_init_impl(self.decoder4)
self.decoder3 = decoderBlock(5,32,32, stride=(2,1,1),up=False, nstride=1, uncertainty=uncertainty)
self.append_init_impl(self.decoder3)
self.uncertainty = uncertainty
# Disparity regressions.
# self.disp_reg3 = disparityregression(self.maxDisp,16)
# self.disp_reg4 = disparityregression(self.maxDisp,16)
# self.disp_reg3 = disparityregression(self.maxDisp,32)
# self.disp_reg6 = disparityregression(self.maxDisp,64)
if ( dispRegConfigs is None ):
dispRegConfigs = [ ClsLinearCombination.get_default_init_args() for _ in range(4)]
else:
assert( isinstance( dispRegConfigs, (tuple, list) ) ), \
f'dispRegConfigs must be a tuple or list. It is {type(dispRegConfigs)}'
self.disp_reg3 = make_object(DISP_REG, dispRegConfigs[0])
self.disp_reg4 = make_object(DISP_REG, dispRegConfigs[1])
self.disp_reg5 = make_object(DISP_REG, dispRegConfigs[2])
self.disp_reg6 = make_object(DISP_REG, dispRegConfigs[3])
self.append_init_impl(self.disp_reg3)
self.append_init_impl(self.disp_reg4)
self.append_init_impl(self.disp_reg5)
self.append_init_impl(self.disp_reg6)
self.uncertaintyComputer = ClassifiedCostVolumeEpistemic() \
if self.uncertainty else None
# Must be called at the end of __init__().
self.update_freeze()
