From a27fad6e54dffcdeb7e1433d84b65386df2da4d7 Mon Sep 17 00:00:00 2001 From: reveLATONG Date: Wed, 30 Jul 2025 10:35:08 +0800 Subject: [PATCH 1/4] merged igev_plusplus --- cfgs/igevpp/igev_rt_sceneflow_amp.yaml | 80 ++++ cfgs/igevpp/igevpp_sceneflow_amp.yaml | 86 +++++ stereo/modeling/__init__.py | 6 +- stereo/modeling/models/igevpp/extractor.py | 355 ++++++++++++++++++ stereo/modeling/models/igevpp/geometry.py | 88 +++++ .../modeling/models/igevpp/igevpp_stereo.py | 285 ++++++++++++++ stereo/modeling/models/igevpp/submodule.py | 266 +++++++++++++ stereo/modeling/models/igevpp/trainer.py | 13 + stereo/modeling/models/igevpp/update.py | 172 +++++++++ stereo/modeling/models/igevpp/utils.py | 134 +++++++ stereo/modeling/models/igevpp_rt/extractor.py | 355 ++++++++++++++++++ stereo/modeling/models/igevpp_rt/geometry.py | 34 ++ .../modeling/models/igevpp_rt/igev_stereo.py | 235 ++++++++++++ stereo/modeling/models/igevpp_rt/submodule.py | 265 +++++++++++++ stereo/modeling/models/igevpp_rt/trainer.py | 13 + stereo/modeling/models/igevpp_rt/update.py | 133 +++++++ stereo/modeling/models/igevpp_rt/utils.py | 212 +++++++++++ 17 files changed, 2731 insertions(+), 1 deletion(-) create mode 100644 cfgs/igevpp/igev_rt_sceneflow_amp.yaml create mode 100644 cfgs/igevpp/igevpp_sceneflow_amp.yaml create mode 100644 stereo/modeling/models/igevpp/extractor.py create mode 100644 stereo/modeling/models/igevpp/geometry.py create mode 100644 stereo/modeling/models/igevpp/igevpp_stereo.py create mode 100644 stereo/modeling/models/igevpp/submodule.py create mode 100644 stereo/modeling/models/igevpp/trainer.py create mode 100644 stereo/modeling/models/igevpp/update.py create mode 100644 stereo/modeling/models/igevpp/utils.py create mode 100644 stereo/modeling/models/igevpp_rt/extractor.py create mode 100644 stereo/modeling/models/igevpp_rt/geometry.py create mode 100644 stereo/modeling/models/igevpp_rt/igev_stereo.py create mode 100644 stereo/modeling/models/igevpp_rt/submodule.py create mode 100644 stereo/modeling/models/igevpp_rt/trainer.py create mode 100644 stereo/modeling/models/igevpp_rt/update.py create mode 100644 stereo/modeling/models/igevpp_rt/utils.py diff --git a/cfgs/igevpp/igev_rt_sceneflow_amp.yaml b/cfgs/igevpp/igev_rt_sceneflow_amp.yaml new file mode 100644 index 00000000..ff5e92a8 --- /dev/null +++ b/cfgs/igevpp/igev_rt_sceneflow_amp.yaml @@ -0,0 +1,80 @@ +DATA_CONFIG: + DATA_INFOS: + - DATASET: SceneFlowDataset + DATA_SPLIT: { + TRAINING: ./data/SceneFlow/sceneflow_finalpass_train.txt, + EVALUATING: ./data/SceneFlow/sceneflow_finalpass_test.txt, + TESTING: ./data/SceneFlow/sceneflow_finalpass_test.txt + } + DATA_PATH: /mnt/d/prr/depth/double/dataset/sceneflow + RETURN_RIGHT_DISP: false + + DATA_TRANSFORM: + TRAINING: + - { NAME: StereoColorJitter, BRIGHTNESS: [ 0.6, 1.4 ], CONTRAST: [ 0.6, 1.4 ], SATURATION: [ 0.6, 1.4 ], HUE: 0.5, ASYMMETRIC_PROB: 0.2 } + - { NAME: RandomErase, PROB: 0.5, MAX_TIME: 2, BOUNDS: [ 50, 100 ] } + - { NAME: RandomCrop, SIZE: [ 320, 736 ], Y_JITTER: false } + - { NAME: TransposeImage } + - { NAME: ToTensor } + - { NAME: NormalizeImage, MEAN: [ 0.485, 0.456, 0.406 ], STD: [ 0.229, 0.224, 0.225 ] } + EVALUATING: + - { NAME: RightTopPad, SIZE: [ 544, 960 ] } + - { NAME: TransposeImage } + - { NAME: ToTensor } + +MODEL: + NAME: IGEVPP_rt + MAX_DISP: &max_disp 192 + HIDDEN_DIMS: [ 128, 128, 128 ] + HIDDEN_DIM: 96 + N_GRU_LAYERS: 3 + N_DOWNSAMPLE: 2 + SLOW_FAST_GRU: True + CORR_LEVELS: 2 + CORR_RADIUS: 4 + TRAIN_ITERS: 22 + VALID_ITERS: 8 + FIND_UNUSED_PARAMETERS: false + CKPT: -1 + PRETRAINED_MODEL: '' + +OPTIMIZATION: + FREEZE_BN: true + SYNC_BN: true + AMP: true + BATCH_SIZE_PER_GPU: 8 + NUM_EPOCHS: 30 + + OPTIMIZER: + NAME: AdamW + LR: &lr 0.0002 + WEIGHT_DECAY: 1.0e-05 + EPS: 1.0e-08 + + SCHEDULER: + NAME: OneCycleLR + MAX_LR: *lr + PCT_START: 0.01 + ON_EPOCH: False + + CLIP_GRAD: + TYPE: value + CLIP_VALUE: 0.1 + +EVALUATOR: + BATCH_SIZE_PER_GPU: 4 + MAX_DISP: *max_disp + METRIC: + - d1_all + - epe + - thres_1 + - thres_2 + - thres_3 + +TRAINER: + EVAL_INTERVAL: 1 + CKPT_SAVE_INTERVAL: 1 + MAX_CKPT_SAVE_NUM: 30 + LOGGER_ITER_INTERVAL: 10 + TRAIN_VISUALIZATION: True + EVAL_VISUALIZATION: True diff --git a/cfgs/igevpp/igevpp_sceneflow_amp.yaml b/cfgs/igevpp/igevpp_sceneflow_amp.yaml new file mode 100644 index 00000000..7839be51 --- /dev/null +++ b/cfgs/igevpp/igevpp_sceneflow_amp.yaml @@ -0,0 +1,86 @@ +DATA_CONFIG: + DATA_INFOS: + - DATASET: SceneFlowDataset + DATA_SPLIT: { + TRAINING: ./data/SceneFlow/sceneflow_finalpass_train.txt, + EVALUATING: ./data/SceneFlow/sceneflow_finalpass_test.txt, + TESTING: ./data/SceneFlow/sceneflow_finalpass_test.txt + } + RETURN_RIGHT_DISP: false + + DATA_TRANSFORM: + TRAINING: + - { NAME: RandomCrop, SIZE: [ 320, 736 ], Y_JITTER: false } + - { NAME: TransposeImage } + - { NAME: ToTensor } + EVALUATING: + - { NAME: RightTopPad, SIZE: [ 544, 960 ] } + - { NAME: TransposeImage } + - { NAME: ToTensor } + +MODEL: + NAME: IGEVPP + MAX_DISP: &max_disp 768 + HIDDEN_DIMS: [ 128, 128, 128 ] + + S_DISP_RANGE: 48 + S_DISP_INTERVAL: 1 + M_DISP_RANGE: 96 + M_DISP_INTERVAL: 2 + L_DISP_RANGE: 192 + L_DISP_INTERVAL: 4 + + N_GRU_LAYERS: 3 + N_DOWNSAMPLE: 2 + SLOW_FAST_GRU: True + CORR_LEVELS: 2 + CORR_RADIUS: 4 + TRAIN_ITERS: 22 + VALID_ITERS: 32 + FIND_UNUSED_PARAMETERS: false + CKPT: -1 + PRETRAINED_MODEL: '' + +OPTIMIZATION: + BATCH_SIZE_PER_GPU: 4 + FREEZE_BN: true + SYNC_BN: true + AMP: true + NUM_EPOCHS: 64 + + OPTIMIZER: + NAME: AdamW + LR: 0.0002 + WEIGHT_DECAY: 0.00001 + EPS: 0.00000001 + + SCHEDULER: + NAME: MultiStepLR + GAMMA: 0.5 + MILESTONES: [ 10, 12, 14, 16 ] + ON_EPOCH: True + WARMUP: + WARM_STEPS: 100 + + CLIP_GRAD: + TYPE: norm + MAX_NORM: 35 + NORM_TYPE: 2 + +EVALUATOR: + BATCH_SIZE_PER_GPU: 10 + MAX_DISP: *max_disp + METRIC: + - d1_all + - epe + - thres_1 + - thres_2 + - thres_3 + +TRAINER: + EVAL_INTERVAL: 1 + CKPT_SAVE_INTERVAL: 1 + MAX_CKPT_SAVE_NUM: 30 + LOGGER_ITER_INTERVAL: 10 + TRAIN_VISUALIZATION: True + EVAL_VISUALIZATION: True diff --git a/stereo/modeling/__init__.py b/stereo/modeling/__init__.py index a4970a82..fa7e584a 100644 --- a/stereo/modeling/__init__.py +++ b/stereo/modeling/__init__.py @@ -13,6 +13,8 @@ from .models.sttr.trainer import Trainer as STTRTrainer from .models.lightstereo.trainer import Trainer as LightStereoTrainer from .models.stereobase.trainer import Trainer as StereoBaseGRUTrainer +from .models.igevpp.trainer import Trainer as IGEVPPTrainer +from .models.igevpp_rt.trainer import Trainer as IGEVPPrtTrainer # from .models.iinet.trainer import Trainer as IINetTrainer ''' @@ -45,7 +47,9 @@ 'StereoBaseGRU': StereoBaseGRUTrainer, # 'FoundationStereo': FoundationStereoTrainer, # 'IInet': IINetTrainer, - 'NMRF': NMRFTrainer + 'NMRF': NMRFTrainer, + 'IGEVPP': IGEVPPTrainer, + 'IGEVPP_rt': IGEVPPrtTrainer, } diff --git a/stereo/modeling/models/igevpp/extractor.py b/stereo/modeling/models/igevpp/extractor.py new file mode 100644 index 00000000..61c4bdb3 --- /dev/null +++ b/stereo/modeling/models/igevpp/extractor.py @@ -0,0 +1,355 @@ +import timm + +from .submodule import * + + +class ResidualBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes) + self.norm2 = nn.InstanceNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.Sequential() + + if stride == 1 and in_planes == planes: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) + + def forward(self, x): + y = x + y = self.conv1(y) + y = self.norm1(y) + y = self.relu(y) + y = self.conv2(y) + y = self.norm2(y) + y = self.relu(y) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BottleneckBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(BottleneckBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0) + self.conv2 = nn.Conv2d(planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not stride == 1: + self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes // 4) + self.norm2 = nn.BatchNorm2d(planes // 4) + self.norm3 = nn.BatchNorm2d(planes) + if not stride == 1: + self.norm4 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes // 4) + self.norm2 = nn.InstanceNorm2d(planes // 4) + self.norm3 = nn.InstanceNorm2d(planes) + if not stride == 1: + self.norm4 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + self.norm3 = nn.Sequential() + if not stride == 1: + self.norm4 = nn.Sequential() + + if stride == 1: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + y = self.relu(self.norm3(self.conv3(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BasicEncoder(nn.Module): + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0, downsample=3): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + self.downsample = downsample + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + + # output convolution + self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) + + self.dropout = None + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, dual_inp=False): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.conv2(x) + + if self.training and self.dropout is not None: + x = self.dropout(x) + + if is_list: + x = x.split(split_size=batch_dim, dim=0) + + return x + + +class MultiBasicEncoder(nn.Module): + def __init__(self, output_dim=[128], norm_fn='batch', dropout=0.0, downsample=3): + super(MultiBasicEncoder, self).__init__() + self.norm_fn = norm_fn + self.downsample = downsample + + # self.norm_111 = nn.BatchNorm2d(128, affine=False, track_running_stats=False) + # self.norm_222 = nn.BatchNorm2d(128, affine=False, track_running_stats=False) + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + self.layer4 = self._make_layer(128, stride=2) + self.layer5 = self._make_layer(128, stride=2) + + output_list = [] + + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[2], 3, padding=1)) + output_list.append(conv_out) + + self.outputs04 = nn.ModuleList(output_list) + + output_list = [] + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[1], 3, padding=1)) + output_list.append(conv_out) + + self.outputs08 = nn.ModuleList(output_list) + + output_list = [] + for dim in output_dim: + conv_out = nn.Conv2d(128, dim[0], 3, padding=1) + output_list.append(conv_out) + + self.outputs16 = nn.ModuleList(output_list) + + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + else: + self.dropout = None + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, dual_inp=False, num_layers=3): + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + if dual_inp: + v = x + x = x[:(x.shape[0] // 2)] + + outputs04 = [f(x) for f in self.outputs04] + if num_layers == 1: + return (outputs04, v) if dual_inp else (outputs04,) + + y = self.layer4(x) + outputs08 = [f(y) for f in self.outputs08] + + if num_layers == 2: + return (outputs04, outputs08, v) if dual_inp else (outputs04, outputs08) + + z = self.layer5(y) + outputs16 = [f(z) for f in self.outputs16] + + return (outputs04, outputs08, outputs16, v) if dual_inp else (outputs04, outputs08, outputs16) + + +class SubModule(nn.Module): + def __init__(self): + super(SubModule, self).__init__() + + def weight_init(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.Conv3d): + n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm3d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + +class Feature(SubModule): + def __init__(self): + super(Feature, self).__init__() + pretrained = True + model = timm.create_model('mobilenetv2_100', pretrained=pretrained, features_only=True) + layers = [1, 2, 3, 5, 6] + chans = [16, 24, 32, 96, 160] + self.conv_stem = model.conv_stem + self.bn1 = model.bn1 + self.act1 = model.act1 + + self.block0 = torch.nn.Sequential(*model.blocks[0:layers[0]]) + self.block1 = torch.nn.Sequential(*model.blocks[layers[0]:layers[1]]) + self.block2 = torch.nn.Sequential(*model.blocks[layers[1]:layers[2]]) + self.block3 = torch.nn.Sequential(*model.blocks[layers[2]:layers[3]]) + self.block4 = torch.nn.Sequential(*model.blocks[layers[3]:layers[4]]) + + self.deconv32_16 = Conv2x_IN(chans[4], chans[3], deconv=True, concat=True) + self.deconv16_8 = Conv2x_IN(chans[3] * 2, chans[2], deconv=True, concat=True) + self.deconv8_4 = Conv2x_IN(chans[2] * 2, chans[1], deconv=True, concat=True) + self.conv4 = BasicConv_IN(chans[1] * 2, chans[1] * 2, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + x = self.act1(self.bn1(self.conv_stem(x))) + x2 = self.block0(x) + x4 = self.block1(x2) + x8 = self.block2(x4) + x16 = self.block3(x8) + x32 = self.block4(x16) + + x16 = self.deconv32_16(x32, x16) + x8 = self.deconv16_8(x16, x8) + x4 = self.deconv8_4(x8, x4) + x4 = self.conv4(x4) + return [x4, x8, x16, x32] diff --git a/stereo/modeling/models/igevpp/geometry.py b/stereo/modeling/models/igevpp/geometry.py new file mode 100644 index 00000000..a9f9fc14 --- /dev/null +++ b/stereo/modeling/models/igevpp/geometry.py @@ -0,0 +1,88 @@ +import torch +import torch.nn.functional as F + +from .utils import bilinear_sampler + + +class Combined_Geo_Encoding_Volume: + def __init__(self, geo_volume0, geo_volume1, geo_volume2, init_fmap1, init_fmap2, radius=4, num_levels=2): + self.num_levels = num_levels + self.radius = radius + # self.geo_volume_pyramid = [] + self.init_corr_pyramid = [] + self.geo_volume0_pyramid = [] + + # all pairs correlation + init_corr = Combined_Geo_Encoding_Volume.corr(init_fmap1, init_fmap2) + + b, h, w1, _, w2 = init_corr.shape + b, c, d0, h, w = geo_volume0.shape + d1 = geo_volume1.shape[2] + d2 = geo_volume2.shape[2] + geo_volume0 = geo_volume0.permute(0, 3, 4, 1, 2).reshape(b * h * w, c, 1, d0) + self.geo_volume1 = geo_volume1.permute(0, 3, 4, 1, 2).reshape(b * h * w, c, 1, d1) + self.geo_volume2 = geo_volume2.permute(0, 3, 4, 1, 2).reshape(b * h * w, c, 1, d2) + + init_corr = init_corr.reshape(b * h * w1, 1, 1, w2) + self.init_corr_pyramid.append(init_corr) + self.geo_volume0_pyramid.append(geo_volume0) + for i in range(self.num_levels - 1): + geo_volume0 = F.avg_pool2d(geo_volume0, [1, 2], stride=[1, 2]) + self.geo_volume0_pyramid.append(geo_volume0) + + init_corr = F.avg_pool2d(init_corr, [1, 2], stride=[1, 2]) + self.init_corr_pyramid.append(init_corr) + + def __call__(self, disp, coords): + r = self.radius + b, _, h, w = disp.shape + init_corr_pyramid = [] + geo_feat0_pyramid = [] + dx = torch.linspace(-r, r, 2 * r + 1) + dx = dx.view(1, 1, 2 * r + 1, 1).to(disp.device) + + x1 = dx + disp.reshape(b * h * w, 1, 1, 1) / 2 + y0 = torch.zeros_like(x1) + disp_lvl1 = torch.cat([x1, y0], dim=-1) + geo_feat1 = bilinear_sampler(self.geo_volume1, disp_lvl1) + geo_feat1 = geo_feat1.view(b, h, w, -1) + + x2 = dx + disp.reshape(b * h * w, 1, 1, 1) / 4 + y0 = torch.zeros_like(x2) + disp_lvl2 = torch.cat([x2, y0], dim=-1) + geo_feat2 = bilinear_sampler(self.geo_volume2, disp_lvl2) + geo_feat2 = geo_feat2.view(b, h, w, -1) + + for i in range(self.num_levels): + geo_volume0 = self.geo_volume0_pyramid[i] + x0 = dx + disp.reshape(b * h * w, 1, 1, 1) / 2 ** i + y0 = torch.zeros_like(x0) + disp_lvl0 = torch.cat([x0, y0], dim=-1) + geo_feat0 = bilinear_sampler(geo_volume0, disp_lvl0) + geo_feat0 = geo_feat0.view(b, h, w, -1) + geo_feat0_pyramid.append(geo_feat0) + + init_corr = self.init_corr_pyramid[i] + init_x0 = coords.reshape(b * h * w, 1, 1, 1) / 2 ** i - disp.reshape(b * h * w, 1, 1, 1) / 2 ** i + dx + init_coords_lvl = torch.cat([init_x0, y0], dim=-1) + init_corr = bilinear_sampler(init_corr, init_coords_lvl) + init_corr = init_corr.view(b, h, w, -1) + init_corr_pyramid.append(init_corr) + + init_corr = torch.cat(init_corr_pyramid, dim=-1).permute(0, 3, 1, 2).contiguous().float() + geo_feat0 = torch.cat(geo_feat0_pyramid, dim=-1) + geo_feat0 = geo_feat0.permute(0, 3, 1, 2).contiguous().float() + geo_feat1 = geo_feat1.permute(0, 3, 1, 2).contiguous().float() + geo_feat2 = geo_feat2.permute(0, 3, 1, 2).contiguous().float() + + return geo_feat0, geo_feat1, geo_feat2, init_corr + + @staticmethod + def corr(fmap1, fmap2): + B, D, H, W1 = fmap1.shape + _, _, _, W2 = fmap2.shape + fmap1 = fmap1.view(B, D, H, W1) + fmap2 = fmap2.view(B, D, H, W2) + corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2) + corr = corr.reshape(B, H, W1, 1, W2).contiguous() + return corr diff --git a/stereo/modeling/models/igevpp/igevpp_stereo.py b/stereo/modeling/models/igevpp/igevpp_stereo.py new file mode 100644 index 00000000..4f345fbe --- /dev/null +++ b/stereo/modeling/models/igevpp/igevpp_stereo.py @@ -0,0 +1,285 @@ +from .extractor import MultiBasicEncoder, Feature +from .geometry import Combined_Geo_Encoding_Volume +from .submodule import * +from .update import BasicMultiUpdateBlock +from .utils import Map + + +class hourglass(nn.Module): + def __init__(self, in_channels): + super(hourglass, self).__init__() + + self.conv1 = nn.Sequential( + BasicConv(in_channels, in_channels * 2, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv2 = nn.Sequential( + BasicConv(in_channels * 2, in_channels * 4, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv3 = nn.Sequential( + BasicConv(in_channels * 4, in_channels * 6, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 6, in_channels * 6, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv3_up = BasicConv(in_channels * 6, in_channels * 4, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv2_up = BasicConv(in_channels * 4, in_channels * 2, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv1_up = BasicConv(in_channels * 2, 8, deconv=True, is_3d=True, bn=False, + relu=False, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.agg_0 = nn.Sequential( + BasicConv(in_channels * 8, in_channels * 4, is_3d=True, kernel_size=1, padding=0, stride=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, kernel_size=3, padding=1, stride=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, kernel_size=3, padding=1, stride=1), ) + + self.agg_1 = nn.Sequential( + BasicConv(in_channels * 4, in_channels * 2, is_3d=True, kernel_size=1, padding=0, stride=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, kernel_size=3, padding=1, stride=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, kernel_size=3, padding=1, stride=1)) + + self.feature_att_8 = FeatureAtt(in_channels * 2, 64) + self.feature_att_16 = FeatureAtt(in_channels * 4, 192) + self.feature_att_32 = FeatureAtt(in_channels * 6, 160) + self.feature_att_up_16 = FeatureAtt(in_channels * 4, 192) + self.feature_att_up_8 = FeatureAtt(in_channels * 2, 64) + + def forward(self, x, features): + conv1 = self.conv1(x) + conv1 = self.feature_att_8(conv1, features[1]) + + conv2 = self.conv2(conv1) + conv2 = self.feature_att_16(conv2, features[2]) + + conv3 = self.conv3(conv2) + conv3 = self.feature_att_32(conv3, features[3]) + + conv3_up = self.conv3_up(conv3) + conv2 = torch.cat((conv3_up, conv2), dim=1) + conv2 = self.agg_0(conv2) + conv2 = self.feature_att_up_16(conv2, features[2]) + + conv2_up = self.conv2_up(conv2) + conv1 = torch.cat((conv2_up, conv1), dim=1) + conv1 = self.agg_1(conv1) + conv1 = self.feature_att_up_8(conv1, features[1]) + + conv = self.conv1_up(conv1) + + return conv + + +class IGEVPPStereo(nn.Module): + def __init__(self, args): + super().__init__() + self.args = args + + context_dims = args.HIDDEN_DIMS + + self.cnet = MultiBasicEncoder(output_dim=[args.HIDDEN_DIMS, context_dims], norm_fn="batch", + downsample=args.N_DOWNSAMPLE) + self.update_block = BasicMultiUpdateBlock(self.args, hidden_dims=args.HIDDEN_DIMS) + self.context_zqr_convs = nn.ModuleList( + [nn.Conv2d(context_dims[i], args.HIDDEN_DIMS[i] * 3, 3, padding=3 // 2) for i in + range(self.args.N_GRU_LAYERS)]) + self.feature = Feature() + + self.stem_2 = nn.Sequential( + BasicConv(3, 32, kernel_size=3, stride=2, padding=1), + nn.Conv2d(32, 32, 3, 1, 1, bias=False), + nn.InstanceNorm2d(32), nn.ReLU() + ) + self.stem_4 = nn.Sequential( + BasicConv(32, 48, kernel_size=3, stride=2, padding=1), + nn.Conv2d(48, 48, 3, 1, 1, bias=False), + nn.InstanceNorm2d(48), nn.ReLU() + ) + self.spx = nn.Sequential(nn.ConvTranspose2d(2 * 32, 9, kernel_size=4, stride=2, padding=1), ) + self.spx_2 = Conv2x(64, 32, True) + self.spx_4 = nn.Sequential( + BasicConv(96, 64, kernel_size=3, stride=1, padding=1), + nn.Conv2d(64, 64, 3, 1, 1, bias=False), + nn.InstanceNorm2d(64), nn.ReLU() + ) + + self.spx_2_gru = Conv2x(64, 32, True) + self.spx_gru = nn.Sequential(nn.ConvTranspose2d(2 * 32, 9, kernel_size=4, stride=2, padding=1), ) + + self.conv = BasicConv(96, 96, kernel_size=3, padding=1, stride=1) + self.desc = nn.Conv2d(96, 96, kernel_size=1, padding=0, stride=1) + self.patch0 = nn.Conv3d(8, 8, kernel_size=(2, 1, 1), stride=(2, 1, 1), bias=False) + self.patch1 = nn.Conv3d(8, 8, kernel_size=(4, 1, 1), stride=(4, 1, 1), bias=False) + self.cost_agg0 = hourglass(8) + self.cost_agg1 = hourglass(8) + self.cost_agg2 = hourglass(8) + self.classifier = nn.Conv3d(8, 1, 3, 1, 1, bias=False) + self.disp_conv = nn.Sequential( + BasicConv(3, 64, kernel_size=1, stride=1, padding=0), + BasicConv(64, 64, kernel_size=3, stride=1, padding=1), + ) + self.selective_conv = nn.Sequential( + BasicConv(96 + 64, 128, kernel_size=1, stride=1, padding=0), + BasicConv(128, 128, kernel_size=3, stride=1, padding=1), + nn.Conv2d(128, 3, 3, 1, 1, bias=False), + ) + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def upsample_disp(self, disp, mask_feat_4, stem_2x): + xspx = self.spx_2_gru(mask_feat_4, stem_2x) + spx_pred = self.spx_gru(xspx) + spx_pred = F.softmax(spx_pred, 1) + up_disp = context_upsample(disp * 4., spx_pred).unsqueeze(1) + return up_disp + + def forward(self, data): + image1 = data['left'] + image2 = data['right'] + """ Estimate disparity between pair of frames """ + test_mode = not self.training + image1 = (2 * (image1 / 255.0) - 1.0).contiguous() + image2 = (2 * (image2 / 255.0) - 1.0).contiguous() + + features_left = self.feature(image1) + features_right = self.feature(image2) + stem_2x = self.stem_2(image1) + stem_4x = self.stem_4(stem_2x) + stem_2y = self.stem_2(image2) + stem_4y = self.stem_4(stem_2y) + features_left[0] = torch.cat((features_left[0], stem_4x), 1) + features_right[0] = torch.cat((features_right[0], stem_4y), 1) + + match_left = self.desc(self.conv(features_left[0])) + match_right = self.desc(self.conv(features_right[0])) + all_disp_volume = build_gwc_volume(match_left, match_right, self.args.MAX_DISP // 4, 8) + + disp_volume0 = all_disp_volume[:, :, :self.args.S_DISP_RANGE] + disp_volume1 = self.patch0(all_disp_volume[:, :, :self.args.M_DISP_RANGE]) + disp_volume2 = self.patch1(all_disp_volume) + + geo_encoding_volume0 = self.cost_agg0(disp_volume0, features_left) + geo_encoding_volume1 = self.cost_agg1(disp_volume1, features_left) + geo_encoding_volume2 = self.cost_agg2(disp_volume2, features_left) + + cost_volume0 = self.classifier(geo_encoding_volume0) + prob_volume0 = F.softmax(cost_volume0.squeeze(1), dim=1) + agg_disp0 = disparity_regression(prob_volume0, self.args.S_DISP_RANGE, self.args.S_DISP_INTERVAL) + + cost_volume1 = self.classifier(geo_encoding_volume1) + prob_volume1 = F.softmax(cost_volume1.squeeze(1), dim=1) + agg_disp1 = disparity_regression(prob_volume1, self.args.M_DISP_RANGE, self.args.M_DISP_INTERVAL) + + cost_volume2 = self.classifier(geo_encoding_volume2) + prob_volume2 = F.softmax(cost_volume2.squeeze(1), dim=1) + agg_disp2 = disparity_regression(prob_volume2, self.args.L_DISP_RANGE, self.args.L_DISP_INTERVAL) + + disp_feature = self.disp_conv(torch.cat([agg_disp0, agg_disp1, agg_disp2], dim=1)) + selective_weights = torch.sigmoid(self.selective_conv(torch.cat([features_left[0], disp_feature], dim=1))) + cnet_list = self.cnet(image1, num_layers=self.args.N_GRU_LAYERS) + net_list = [torch.tanh(x[0]) for x in cnet_list] + inp_list = [torch.relu(x[1]) for x in cnet_list] + inp_list = [list(conv(i).split(split_size=conv.out_channels // 3, dim=1)) for i, conv in + zip(inp_list, self.context_zqr_convs)] + + geo_block = Combined_Geo_Encoding_Volume + geo_fn = geo_block(geo_encoding_volume0.float(), geo_encoding_volume1.float(), geo_encoding_volume2.float(), + match_left.float(), match_right.float(), radius=self.args.CORR_RADIUS) + b, c, h, w = match_left.shape + coords = torch.arange(w).float().to(match_left.device).reshape(1, 1, w, 1).repeat(b, h, 1, 1) + disp = agg_disp0 + iter_preds = [] + + # GRUs iterations to update disparity + iters = self.args.VALID_ITERS if test_mode else self.args.TRAIN_ITERS + for itr in range(iters): + disp = disp.detach() + geo_feat0, geo_feat1, geo_feat2, init_corr = geo_fn(disp, coords) + # with autocast(enabled=self.args.mixed_precision, dtype=getattr(torch, self.args.precision_dtype, torch.float16)): + net_list, mask_feat_4, delta_disp = self.update_block(net_list, inp_list, geo_feat0, geo_feat1, geo_feat2, + init_corr, selective_weights, disp, + iter16=self.args.N_GRU_LAYERS == 3, + iter08=self.args.N_GRU_LAYERS >= 2) + + disp = disp + delta_disp + if test_mode and itr < iters - 1: + continue + + # upsample predictions + disp_up = self.upsample_disp(disp, mask_feat_4, stem_2x) + iter_preds.append(disp_up) + + if test_mode: + return disp_up + + xspx = self.spx_4(features_left[0]) + xspx = self.spx_2(xspx, stem_2x) + spx_pred = self.spx(xspx) + spx_pred = F.softmax(spx_pred, 1) + agg_disp0 = context_upsample(agg_disp0 * 4., spx_pred.float()) + agg_disp1 = context_upsample(agg_disp1 * 4., spx_pred.float()) + agg_disp2 = context_upsample(agg_disp2 * 4., spx_pred.float()) + agg_preds = [agg_disp0, agg_disp1, agg_disp2] + + return {'init_disp': agg_preds, + 'disp_preds': iter_preds, + 'disp_pred': iter_preds[-1]} + + def get_loss(self, model_pred, input_data): + disp_gt = input_data["disp"] + mask = (disp_gt < self.max_disp) & (disp_gt > 0) + valid = mask.float() + max_disp0 = 192 + max_disp1 = 384 + max_disp = 700 + + disp_gt = disp_gt.unsqueeze(1) + mag = torch.sum(disp_gt ** 2, dim=1).sqrt() + valid = ((valid >= 0.5) & (mag < self.max_disp)).unsqueeze(1) + assert valid.shape == disp_gt.shape, [valid.shape, disp_gt.shape] + assert not torch.isinf(disp_gt[valid.bool()]).any() + disp_loss = 0.0 + mag = torch.sum(disp_gt**2, dim=1).sqrt() + mask0 = ((valid >= 0.5) & (mag < max_disp0)).unsqueeze(1) + mask1 = ((valid >= 0.5) & (mag < max_disp1)).unsqueeze(1) + mask = ((valid >= 0.5) & (mag < max_disp)).unsqueeze(1) + + disp_init_pred = model_pred['init_disp'] + disp_loss += 1.0 * F.smooth_l1_loss(disp_init_pred[0][mask0.bool()], disp_gt[mask0.bool()], reduction='mean') + disp_loss += 0.5 * F.smooth_l1_loss(disp_init_pred[1][mask1.bool()], disp_gt[mask1.bool()], reduction='mean') + disp_loss += 0.2 * F.smooth_l1_loss(disp_init_pred[2][mask.bool()], disp_gt[mask.bool()], reduction='mean') + + # gru loss + loss_gamma = 0.9 + disp_preds = model_pred['disp_preds'] + n_predictions = len(disp_preds) + assert n_predictions >= 1 + for i in range(n_predictions): + adjusted_loss_gamma = loss_gamma ** (15 / (n_predictions - 1)) + i_weight = adjusted_loss_gamma ** (n_predictions - i - 1) + i_loss = (disp_preds[i] - disp_gt).abs() + assert i_loss.shape == mask.shape, [i_loss.shape, mask.shape, disp_gt.shape, disp_preds[i].shape] + disp_loss += i_weight * i_loss[mask.bool()].mean() + + # epe = torch.sum((disp_preds[-1] - disp_gt) ** 2, dim=1).sqrt() + # epe = epe.view(-1)[mask.view(-1)] + # + # metrics = { + # 'epe': epe.mean().item(), + # '1px': (epe < 1).float().mean().item(), + # '3px': (epe < 3).float().mean().item(), + # '5px': (epe < 5).float().mean().item(), + # } + + loss_info = {'scalar/train/loss_disp': disp_loss.item()} + return disp_loss, loss_info diff --git a/stereo/modeling/models/igevpp/submodule.py b/stereo/modeling/models/igevpp/submodule.py new file mode 100644 index 00000000..c2e404b0 --- /dev/null +++ b/stereo/modeling/models/igevpp/submodule.py @@ -0,0 +1,266 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np + + +class BasicConv(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, **kwargs): + super(BasicConv, self).__init__() + + self.relu = relu + self.use_bn = bn + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_bn: + x = self.bn(x) + if self.relu: + x = nn.LeakyReLU()(x) # , inplace=True) + return x + + +class Conv2x(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, + relu=True, keep_dispc=False): + super(Conv2x, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, + stride=stride, padding=padding) + else: + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, + stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = BasicConv(out_channels * 2, out_channels * mul, False, is_3d, bn, relu, kernel_size=3, + stride=1, padding=1) + else: + self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, + padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate( + x, + size=(rem.shape[-2], rem.shape[-1]), + mode='nearest') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + +class BasicConv_IN(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs): + super(BasicConv_IN, self).__init__() + + self.relu = relu + self.use_in = IN + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_in: + x = self.IN(x) + if self.relu: + x = nn.LeakyReLU()(x) # , inplace=True) + return x + + +class Conv2x_IN(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, + relu=True, keep_dispc=False): + super(Conv2x_IN, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, + stride=stride, padding=padding) + else: + self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, + stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = BasicConv_IN(out_channels * 2, out_channels * mul, False, is_3d, IN, relu, kernel_size=3, + stride=1, padding=1) + else: + self.conv2 = BasicConv_IN(out_channels, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, + padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate( + x, + size=(rem.shape[-2], rem.shape[-1]), + mode='nearest') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + +def groupwise_correlation(fea1, fea2, num_groups): + B, C, H, W = fea1.shape + assert C % num_groups == 0 + channels_per_group = C // num_groups + cost = (fea1 * fea2).view([B, num_groups, channels_per_group, H, W]).mean(dim=2) + assert cost.shape == (B, num_groups, H, W) + return cost + + +def build_gwc_volume(refimg_fea, targetimg_fea, maxdisp, num_groups): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = groupwise_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], + num_groups) + else: + volume[:, :, i, :, :] = groupwise_correlation(refimg_fea, targetimg_fea, num_groups) + volume = volume.contiguous() + return volume + + +def norm_correlation(fea1, fea2): + cost = torch.mean( + ((fea1 / (torch.norm(fea1, 2, 1, True) + 1e-05)) * (fea2 / (torch.norm(fea2, 2, 1, True) + 1e-05))), dim=1, + keepdim=True) + return cost + + +def build_norm_correlation_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = norm_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i]) + else: + volume[:, :, i, :, :] = norm_correlation(refimg_fea, targetimg_fea) + volume = volume.contiguous() + return volume + + +def correlation(fea1, fea2): + cost = torch.sum((fea1 * fea2), dim=1, keepdim=True) + return cost + + +def build_correlation_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i]) + else: + volume[:, :, i, :, :] = correlation(refimg_fea, targetimg_fea) + volume = volume.contiguous() + return volume + + +def build_concat_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 2 * C, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :C, i, :, :] = refimg_fea[:, :, :, :] + volume[:, C:, i, :, i:] = targetimg_fea[:, :, :, :-i] + else: + volume[:, :C, i, :, :] = refimg_fea + volume[:, C:, i, :, :] = targetimg_fea + volume = volume.contiguous() + return volume + + +def disparity_regression(prob, maxdisp, interval): + assert len(prob.shape) == 4 + disp_values = torch.arange(0, maxdisp, interval, dtype=prob.dtype, device=prob.device) + disp_values = disp_values.view(1, maxdisp//interval, 1, 1) + return torch.sum(prob * disp_values, 1, keepdim=True) + + +class FeatureAtt(nn.Module): + def __init__(self, cv_chan, feat_chan): + super(FeatureAtt, self).__init__() + + self.feat_att = nn.Sequential( + BasicConv(feat_chan, feat_chan // 2, kernel_size=1, stride=1, padding=0), + nn.Conv2d(feat_chan // 2, cv_chan, 1)) + + def forward(self, cv, feat): + ''' + ''' + feat_att = self.feat_att(feat).unsqueeze(2) + cv = torch.sigmoid(feat_att) * cv + return cv + + +def context_upsample(disp_low, up_weights): + ### + # cv (b,1,h,w) + # sp (b,9,4*h,4*w) + ### + b, c, h, w = disp_low.shape + + disp_unfold = F.unfold(disp_low.reshape(b, c, h, w), 3, 1, 1).reshape(b, -1, h, w) + disp_unfold = F.interpolate(disp_unfold, (h * 4, w * 4), mode='nearest').reshape(b, 9, h * 4, w * 4) + + disp = (disp_unfold * up_weights).sum(1) + + return disp diff --git a/stereo/modeling/models/igevpp/trainer.py b/stereo/modeling/models/igevpp/trainer.py new file mode 100644 index 00000000..e560e128 --- /dev/null +++ b/stereo/modeling/models/igevpp/trainer.py @@ -0,0 +1,13 @@ +# @Time : 2024/2/9 11:39 +# @Author : zhangchenming +from stereo.modeling.trainer_template import TrainerTemplate +from .igevpp_stereo import IGEVPPStereo + +__all__ = { + 'IGEVPP': IGEVPPStereo, +} + +class Trainer(TrainerTemplate): + def __init__(self, args, cfgs, local_rank, global_rank, logger, tb_writer): + model = __all__[cfgs.MODEL.NAME](cfgs.MODEL) + super().__init__(args, cfgs, local_rank, global_rank, logger, tb_writer, model) diff --git a/stereo/modeling/models/igevpp/update.py b/stereo/modeling/models/igevpp/update.py new file mode 100644 index 00000000..1d94a618 --- /dev/null +++ b/stereo/modeling/models/igevpp/update.py @@ -0,0 +1,172 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class FlowHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=2): + super(FlowHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class DispHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=1): + super(DispHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class ConvGRU(nn.Module): + def __init__(self, hidden_dim, input_dim, kernel_size=3): + super(ConvGRU, self).__init__() + self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, kernel_size, padding=kernel_size // 2) + + def forward(self, h, cz, cr, cq, *x_list): + x = torch.cat(x_list, dim=1) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz(hx) + cz) + r = torch.sigmoid(self.convr(hx) + cr) + q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1)) + cq) + h = (1 - z) * h + z * q + return h + + +class SepConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192 + 128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convr1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convq1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + + self.convz2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convr2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convq2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + + def forward(self, h, *x): + # horizontal + x = torch.cat(x, dim=1) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + return h + +class GeoEncoder(nn.Module): + def __init__(self, geo_planes): + super(GeoEncoder, self).__init__() + self.convg1 = nn.Conv2d(geo_planes, 128, 1, padding=0) + self.convg2 = nn.Conv2d(128, 96, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, geo): + return self.convg2(self.relu(self.convg1(geo))) + +class BasicDispEncoder(nn.Module): + def __init__(self, args): + super(BasicDispEncoder, self).__init__() + self.args = args + geo_planes = (2*args.CORR_RADIUS + 1) * 2 + 96 + self.convc1 = nn.Conv2d(geo_planes, 128, 1, padding=0) + self.convc2 = nn.Conv2d(128, 96, 3, padding=1) + self.convd1 = nn.Conv2d(1, 32, 7, padding=3) + self.convd2 = nn.Conv2d(32, 32, 3, padding=1) + self.conv = nn.Conv2d(96+32, 128-1, 3, padding=1) + + def forward(self, disp, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + disp_ = F.relu(self.convd1(disp)) + disp_ = F.relu(self.convd2(disp_)) + cor_disp = torch.cat([cor, disp_], dim=1) + out = F.relu(self.conv(cor_disp)) + return torch.cat([out, disp], dim=1) + + +def pool2x(x): + return F.avg_pool2d(x, 3, stride=2, padding=1) + + +def pool4x(x): + return F.avg_pool2d(x, 5, stride=4, padding=1) + + +def interp(x, dest): + original_dtype = x.dtype + x_fp32 = x.float() + interp_args = {'mode': 'bilinear', 'align_corners': True} + with torch.cuda.amp.autocast(enabled=False): + output_fp32 = F.interpolate(x_fp32, dest.shape[2:], **interp_args) + if original_dtype != torch.float32: + output = output_fp32.to(original_dtype) + else: + output = output_fp32 + return output + + +class BasicMultiUpdateBlock(nn.Module): + def __init__(self, args, hidden_dims=[]): + super().__init__() + self.args = args + self.geo_encoder0 = GeoEncoder(geo_planes=2 * (2*args.CORR_RADIUS + 1) * 8) + self.geo_encoder1 = GeoEncoder(geo_planes=(2*args.CORR_RADIUS + 1) * 8) + self.geo_encoder2 = GeoEncoder(geo_planes=(2*args.CORR_RADIUS + 1) * 8) + self.encoder = BasicDispEncoder(args) + encoder_output_dim = 128 + + self.gru04 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.N_GRU_LAYERS > 1)) + self.gru08 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.N_GRU_LAYERS == 3) + hidden_dims[2]) + self.gru16 = ConvGRU(hidden_dims[0], hidden_dims[1]) + self.disp_head = DispHead(hidden_dims[2], hidden_dim=256, output_dim=1) + factor = 2**self.args.N_DOWNSAMPLE + + self.mask_feat_4 = nn.Sequential( + nn.Conv2d(hidden_dims[2], 64, 3, padding=1), + nn.ReLU(inplace=True)) + + def forward(self, net, inp, geo_feat0=None, geo_feat1=None, geo_feat2=None, init_corr=None, selective_weights=None, disp=None, iter04=True, iter08=True, iter16=True, update=True): + + if iter16: + net[2] = self.gru16(net[2], *(inp[2]), pool2x(net[1])) + if iter08: + if self.args.N_GRU_LAYERS > 2: + net[1] = self.gru08(net[1], *(inp[1]), pool2x(net[0]), interp(net[2], net[1])) + else: + net[1] = self.gru08(net[1], *(inp[1]), pool2x(net[0])) + if iter04: + geo_feat0 = self.geo_encoder0(geo_feat0) + geo_feat1 = self.geo_encoder1(geo_feat1) + geo_feat2 = self.geo_encoder2(geo_feat2) + geo_feat = selective_weights[:,0:1]*geo_feat0 + selective_weights[:,1:2]*geo_feat1 + selective_weights[:,2:3]*geo_feat2 + geo_feat = torch.cat([geo_feat, init_corr], dim=1) + disp_features = self.encoder(disp, geo_feat) + if self.args.N_GRU_LAYERS > 1: + net[0] = self.gru04(net[0], *(inp[0]), disp_features, interp(net[1], net[0])) + else: + net[0] = self.gru04(net[0], *(inp[0]), disp_features) + + if not update: + return net + + delta_disp = self.disp_head(net[0]) + mask_feat_4 = self.mask_feat_4(net[0]) + return net, mask_feat_4, delta_disp diff --git a/stereo/modeling/models/igevpp/utils.py b/stereo/modeling/models/igevpp/utils.py new file mode 100644 index 00000000..ef709bd5 --- /dev/null +++ b/stereo/modeling/models/igevpp/utils.py @@ -0,0 +1,134 @@ +import numpy as np +import torch +import torch.nn.functional as F +from scipy import interpolate + + +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + + def __init__(self, dims, mode='sintel', divis_by=8): + self.ht, self.wd = dims[-2:] + pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by + pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by + if mode == 'sintel': + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, pad_ht - pad_ht // 2] + else: + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht] + + def pad(self, *inputs): + assert all((x.ndim == 4) for x in inputs) + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + assert x.ndim == 4 + ht, wd = x.shape[-2:] + c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + +def forward_interpolate(flow): + flow = flow.detach().cpu().numpy() + dx, dy = flow[0], flow[1] + + ht, wd = dx.shape + x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + + x1 = x0 + dx + y1 = y0 + dy + + x1 = x1.reshape(-1) + y1 = y1.reshape(-1) + dx = dx.reshape(-1) + dy = dy.reshape(-1) + + valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + x1 = x1[valid] + y1 = y1[valid] + dx = dx[valid] + dy = dy[valid] + + flow_x = interpolate.griddata( + (x1, y1), dx, (x0, y0), method='nearest', fill_value=0) + + flow_y = interpolate.griddata( + (x1, y1), dy, (x0, y0), method='nearest', fill_value=0) + + flow = np.stack([flow_x, flow_y], axis=0) + return torch.from_numpy(flow).float() + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + + # print("$$$55555", img.shape, coords.shape) + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (W - 1) - 1 + + # print("######88888", xgrid) + assert torch.unique(ygrid).numel() == 1 and H == 1 # This is a stereo problem + + grid = torch.cat([xgrid, ygrid], dim=-1) + # print("###37777", grid.shape) + img = F.grid_sample(img, grid, align_corners=True) + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + + +def coords_grid(batch, ht, wd): + coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def upflow8(flow, mode='bilinear'): + new_size = (8 * flow.shape[2], 8 * flow.shape[3]) + return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True) + + +def gauss_blur(input, N=5, std=1): + B, D, H, W = input.shape + x, y = torch.meshgrid(torch.arange(N).float() - N // 2, torch.arange(N).float() - N // 2) + unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2)) + weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4) + weights = weights.view(1, 1, N, N).to(input) + output = F.conv2d(input.reshape(B * D, 1, H, W), weights, padding=N // 2) + return output.view(B, D, H, W) + + +class Map(dict): + """ + Example: + m = Map({'first_name': 'Eduardo'}, last_name='Pool', age=24, sports=['Soccer']) + """ + def __init__(self, *args, **kwargs): + super(Map, self).__init__(*args, **kwargs) + for arg in args: + if isinstance(arg, dict): + for k, v in arg.items(): + self[k] = v + + if kwargs: + for k, v in kwargs.items(): + self[k] = v + + def __getattr__(self, attr): + return self.get(attr) + + def __setattr__(self, key, value): + self.__setitem__(key, value) + + def __setitem__(self, key, value): + super(Map, self).__setitem__(key, value) + self.__dict__.update({key: value}) + + def __delattr__(self, item): + self.__delitem__(item) + + def __delitem__(self, key): + super(Map, self).__delitem__(key) + del self.__dict__[key] diff --git a/stereo/modeling/models/igevpp_rt/extractor.py b/stereo/modeling/models/igevpp_rt/extractor.py new file mode 100644 index 00000000..61c4bdb3 --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/extractor.py @@ -0,0 +1,355 @@ +import timm + +from .submodule import * + + +class ResidualBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(ResidualBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, stride=stride) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes) + self.norm2 = nn.BatchNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes) + self.norm2 = nn.InstanceNorm2d(planes) + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + if not (stride == 1 and in_planes == planes): + self.norm3 = nn.Sequential() + + if stride == 1 and in_planes == planes: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm3) + + def forward(self, x): + y = x + y = self.conv1(y) + y = self.norm1(y) + y = self.relu(y) + y = self.conv2(y) + y = self.norm2(y) + y = self.relu(y) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BottleneckBlock(nn.Module): + def __init__(self, in_planes, planes, norm_fn='group', stride=1): + super(BottleneckBlock, self).__init__() + + self.conv1 = nn.Conv2d(in_planes, planes // 4, kernel_size=1, padding=0) + self.conv2 = nn.Conv2d(planes // 4, planes // 4, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(planes // 4, planes, kernel_size=1, padding=0) + self.relu = nn.ReLU(inplace=True) + + num_groups = planes // 8 + + if norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) + self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes // 4) + self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + if not stride == 1: + self.norm4 = nn.GroupNorm(num_groups=num_groups, num_channels=planes) + + elif norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(planes // 4) + self.norm2 = nn.BatchNorm2d(planes // 4) + self.norm3 = nn.BatchNorm2d(planes) + if not stride == 1: + self.norm4 = nn.BatchNorm2d(planes) + + elif norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(planes // 4) + self.norm2 = nn.InstanceNorm2d(planes // 4) + self.norm3 = nn.InstanceNorm2d(planes) + if not stride == 1: + self.norm4 = nn.InstanceNorm2d(planes) + + elif norm_fn == 'none': + self.norm1 = nn.Sequential() + self.norm2 = nn.Sequential() + self.norm3 = nn.Sequential() + if not stride == 1: + self.norm4 = nn.Sequential() + + if stride == 1: + self.downsample = None + + else: + self.downsample = nn.Sequential( + nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride), self.norm4) + + def forward(self, x): + y = x + y = self.relu(self.norm1(self.conv1(y))) + y = self.relu(self.norm2(self.conv2(y))) + y = self.relu(self.norm3(self.conv3(y))) + + if self.downsample is not None: + x = self.downsample(x) + + return self.relu(x + y) + + +class BasicEncoder(nn.Module): + def __init__(self, output_dim=128, norm_fn='batch', dropout=0.0, downsample=3): + super(BasicEncoder, self).__init__() + self.norm_fn = norm_fn + self.downsample = downsample + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + + # output convolution + self.conv2 = nn.Conv2d(128, output_dim, kernel_size=1) + + self.dropout = None + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, dual_inp=False): + + # if input is list, combine batch dimension + is_list = isinstance(x, tuple) or isinstance(x, list) + if is_list: + batch_dim = x[0].shape[0] + x = torch.cat(x, dim=0) + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.conv2(x) + + if self.training and self.dropout is not None: + x = self.dropout(x) + + if is_list: + x = x.split(split_size=batch_dim, dim=0) + + return x + + +class MultiBasicEncoder(nn.Module): + def __init__(self, output_dim=[128], norm_fn='batch', dropout=0.0, downsample=3): + super(MultiBasicEncoder, self).__init__() + self.norm_fn = norm_fn + self.downsample = downsample + + # self.norm_111 = nn.BatchNorm2d(128, affine=False, track_running_stats=False) + # self.norm_222 = nn.BatchNorm2d(128, affine=False, track_running_stats=False) + + if self.norm_fn == 'group': + self.norm1 = nn.GroupNorm(num_groups=8, num_channels=64) + + elif self.norm_fn == 'batch': + self.norm1 = nn.BatchNorm2d(64) + + elif self.norm_fn == 'instance': + self.norm1 = nn.InstanceNorm2d(64) + + elif self.norm_fn == 'none': + self.norm1 = nn.Sequential() + + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=1 + (downsample > 2), padding=3) + self.relu1 = nn.ReLU(inplace=True) + + self.in_planes = 64 + self.layer1 = self._make_layer(64, stride=1) + self.layer2 = self._make_layer(96, stride=1 + (downsample > 1)) + self.layer3 = self._make_layer(128, stride=1 + (downsample > 0)) + self.layer4 = self._make_layer(128, stride=2) + self.layer5 = self._make_layer(128, stride=2) + + output_list = [] + + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[2], 3, padding=1)) + output_list.append(conv_out) + + self.outputs04 = nn.ModuleList(output_list) + + output_list = [] + for dim in output_dim: + conv_out = nn.Sequential( + ResidualBlock(128, 128, self.norm_fn, stride=1), + nn.Conv2d(128, dim[1], 3, padding=1)) + output_list.append(conv_out) + + self.outputs08 = nn.ModuleList(output_list) + + output_list = [] + for dim in output_dim: + conv_out = nn.Conv2d(128, dim[0], 3, padding=1) + output_list.append(conv_out) + + self.outputs16 = nn.ModuleList(output_list) + + if dropout > 0: + self.dropout = nn.Dropout2d(p=dropout) + else: + self.dropout = None + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') + elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)): + if m.weight is not None: + nn.init.constant_(m.weight, 1) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + def _make_layer(self, dim, stride=1): + layer1 = ResidualBlock(self.in_planes, dim, self.norm_fn, stride=stride) + layer2 = ResidualBlock(dim, dim, self.norm_fn, stride=1) + layers = (layer1, layer2) + + self.in_planes = dim + return nn.Sequential(*layers) + + def forward(self, x, dual_inp=False, num_layers=3): + + x = self.conv1(x) + x = self.norm1(x) + x = self.relu1(x) + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + if dual_inp: + v = x + x = x[:(x.shape[0] // 2)] + + outputs04 = [f(x) for f in self.outputs04] + if num_layers == 1: + return (outputs04, v) if dual_inp else (outputs04,) + + y = self.layer4(x) + outputs08 = [f(y) for f in self.outputs08] + + if num_layers == 2: + return (outputs04, outputs08, v) if dual_inp else (outputs04, outputs08) + + z = self.layer5(y) + outputs16 = [f(z) for f in self.outputs16] + + return (outputs04, outputs08, outputs16, v) if dual_inp else (outputs04, outputs08, outputs16) + + +class SubModule(nn.Module): + def __init__(self): + super(SubModule, self).__init__() + + def weight_init(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.Conv3d): + n = m.kernel_size[0] * m.kernel_size[1] * m.kernel_size[2] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm3d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + +class Feature(SubModule): + def __init__(self): + super(Feature, self).__init__() + pretrained = True + model = timm.create_model('mobilenetv2_100', pretrained=pretrained, features_only=True) + layers = [1, 2, 3, 5, 6] + chans = [16, 24, 32, 96, 160] + self.conv_stem = model.conv_stem + self.bn1 = model.bn1 + self.act1 = model.act1 + + self.block0 = torch.nn.Sequential(*model.blocks[0:layers[0]]) + self.block1 = torch.nn.Sequential(*model.blocks[layers[0]:layers[1]]) + self.block2 = torch.nn.Sequential(*model.blocks[layers[1]:layers[2]]) + self.block3 = torch.nn.Sequential(*model.blocks[layers[2]:layers[3]]) + self.block4 = torch.nn.Sequential(*model.blocks[layers[3]:layers[4]]) + + self.deconv32_16 = Conv2x_IN(chans[4], chans[3], deconv=True, concat=True) + self.deconv16_8 = Conv2x_IN(chans[3] * 2, chans[2], deconv=True, concat=True) + self.deconv8_4 = Conv2x_IN(chans[2] * 2, chans[1], deconv=True, concat=True) + self.conv4 = BasicConv_IN(chans[1] * 2, chans[1] * 2, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + x = self.act1(self.bn1(self.conv_stem(x))) + x2 = self.block0(x) + x4 = self.block1(x2) + x8 = self.block2(x4) + x16 = self.block3(x8) + x32 = self.block4(x16) + + x16 = self.deconv32_16(x32, x16) + x8 = self.deconv16_8(x16, x8) + x4 = self.deconv8_4(x8, x4) + x4 = self.conv4(x4) + return [x4, x8, x16, x32] diff --git a/stereo/modeling/models/igevpp_rt/geometry.py b/stereo/modeling/models/igevpp_rt/geometry.py new file mode 100644 index 00000000..170699d4 --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/geometry.py @@ -0,0 +1,34 @@ +import torch +import torch.nn.functional as F + +from .utils import bilinear_sampler + + +class Geo_Encoding_Volume: + def __init__(self, geo_volume, num_levels=2, radius=4): + self.num_levels = num_levels + self.radius = radius + self.geo_volume_pyramid = [] + b, c, d, h, w = geo_volume.shape + geo_volume = geo_volume.permute(0, 3, 4, 1, 2).reshape(b*h*w, c, 1, d) + self.geo_volume_pyramid.append(geo_volume) + for i in range(self.num_levels-1): + geo_volume = F.avg_pool2d(geo_volume, [1,2], stride=[1,2]) + self.geo_volume_pyramid.append(geo_volume) + + def __call__(self, disp): + r = self.radius + b, _, h, w = disp.shape + out_pyramid = [] + dx = torch.linspace(-r, r, 2*r+1) + dx = dx.view(1, 1, 2*r+1, 1).to(disp.device) + for i in range(self.num_levels): + geo_volume = self.geo_volume_pyramid[i] + x0 = dx + disp.reshape(b*h*w, 1, 1, 1) / 2**i + y0 = torch.zeros_like(x0) + disp_lvl = torch.cat([x0,y0], dim=-1) + geo_volume = bilinear_sampler(geo_volume, disp_lvl) + geo_volume = geo_volume.view(b, h, w, -1) + out_pyramid.append(geo_volume) + out = torch.cat(out_pyramid, dim=-1) + return out.permute(0, 3, 1, 2).contiguous().float() diff --git a/stereo/modeling/models/igevpp_rt/igev_stereo.py b/stereo/modeling/models/igevpp_rt/igev_stereo.py new file mode 100644 index 00000000..bc79f98c --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/igev_stereo.py @@ -0,0 +1,235 @@ +from .extractor import MultiBasicEncoder, Feature +from .geometry import Geo_Encoding_Volume +from .submodule import * +from .update import BasicUpdateBlock +from .utils import Map + +class hourglass(nn.Module): + def __init__(self, in_channels): + super(hourglass, self).__init__() + + self.conv1 = nn.Sequential( + BasicConv(in_channels, in_channels * 2, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv2 = nn.Sequential( + BasicConv(in_channels * 2, in_channels * 4, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv3 = nn.Sequential( + BasicConv(in_channels * 4, in_channels * 6, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=2, dilation=1), + BasicConv(in_channels * 6, in_channels * 6, is_3d=True, bn=True, relu=True, kernel_size=3, + padding=1, stride=1, dilation=1)) + + self.conv3_up = BasicConv(in_channels * 6, in_channels * 4, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv2_up = BasicConv(in_channels * 4, in_channels * 2, deconv=True, is_3d=True, bn=True, + relu=True, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.conv1_up = BasicConv(in_channels * 2, 8, deconv=True, is_3d=True, bn=False, + relu=False, kernel_size=(4, 4, 4), padding=(1, 1, 1), stride=(2, 2, 2)) + + self.agg_0 = nn.Sequential( + BasicConv(in_channels * 8, in_channels * 4, is_3d=True, kernel_size=1, padding=0, stride=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, kernel_size=3, padding=1, stride=1), + BasicConv(in_channels * 4, in_channels * 4, is_3d=True, kernel_size=3, padding=1, stride=1), ) + + self.agg_1 = nn.Sequential( + BasicConv(in_channels * 4, in_channels * 2, is_3d=True, kernel_size=1, padding=0, stride=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, kernel_size=3, padding=1, stride=1), + BasicConv(in_channels * 2, in_channels * 2, is_3d=True, kernel_size=3, padding=1, stride=1)) + + self.feature_att_8 = FeatureAtt(in_channels * 2, 64) + self.feature_att_16 = FeatureAtt(in_channels * 4, 192) + self.feature_att_32 = FeatureAtt(in_channels * 6, 160) + self.feature_att_up_16 = FeatureAtt(in_channels * 4, 192) + self.feature_att_up_8 = FeatureAtt(in_channels * 2, 64) + + def forward(self, x, features): + conv1 = self.conv1(x) + conv1 = self.feature_att_8(conv1, features[1]) + + conv2 = self.conv2(conv1) + conv2 = self.feature_att_16(conv2, features[2]) + + conv3 = self.conv3(conv2) + conv3 = self.feature_att_32(conv3, features[3]) + + conv3_up = self.conv3_up(conv3) + conv2 = torch.cat((conv3_up, conv2), dim=1) + conv2 = self.agg_0(conv2) + conv2 = self.feature_att_up_16(conv2, features[2]) + + conv2_up = self.conv2_up(conv2) + conv1 = torch.cat((conv2_up, conv1), dim=1) + conv1 = self.agg_1(conv1) + conv1 = self.feature_att_up_8(conv1, features[1]) + + conv = self.conv1_up(conv1) + + return conv + + +class IGEVStereo(nn.Module): + def __init__(self, args): + super().__init__() + self.args = args + self.max_disp = args.MAX_DISP + + context_dim = args.HIDDEN_DIM + + self.cnet = BasicConv(96, context_dim, kernel_size=3, stride=1, padding=1) + self.update_block = BasicUpdateBlock(self.args, hidden_dim=args.HIDDEN_DIM) + self.hnet = nn.Sequential(BasicConv(96, args.HIDDEN_DIM, kernel_size=3, stride=1, padding=1), + nn.Conv2d(args.HIDDEN_DIM, args.HIDDEN_DIM, 3, 1, 1, bias=False)) + + + self.context_zqr_conv = nn.Conv2d(context_dim, context_dim*3, 3, padding=3//2) + + self.feature = Feature() + + self.stem_2 = nn.Sequential( + BasicConv_IN(3, 32, kernel_size=3, stride=2, padding=1), + nn.Conv2d(32, 32, 3, 1, 1, bias=False), + nn.InstanceNorm2d(32), nn.ReLU() + ) + self.stem_4 = nn.Sequential( + BasicConv_IN(32, 48, kernel_size=3, stride=2, padding=1), + nn.Conv2d(48, 48, 3, 1, 1, bias=False), + nn.InstanceNorm2d(48), nn.ReLU() + ) + + self.spx = nn.Sequential(nn.ConvTranspose2d(2 * 32, 9, kernel_size=4, stride=2, padding=1), ) + self.spx_2 = Conv2x_IN(24, 32, True) + self.spx_4 = nn.Sequential( + BasicConv_IN(96, 24, kernel_size=3, stride=1, padding=1), + nn.Conv2d(24, 24, 3, 1, 1, bias=False), + nn.InstanceNorm2d(24), nn.ReLU() + ) + + self.spx_2_gru = Conv2x(32, 32, True) + self.spx_gru = nn.Sequential(nn.ConvTranspose2d(2 * 32, 9, kernel_size=4, stride=2, padding=1), ) + + self.conv = BasicConv_IN(96, 96, kernel_size=3, padding=1, stride=1) + self.desc = nn.Conv2d(96, 96, kernel_size=1, padding=0, stride=1) + + # self.corr_stem = BasicConv(8, 8, is_3d=True, kernel_size=3, stride=1, padding=1) + # self.corr_feature_att = FeatureAtt(8, 96) + self.cost_agg = hourglass(8) + self.classifier = nn.Conv3d(8, 1, 3, 1, 1, bias=False) + + def freeze_bn(self): + for m in self.modules(): + if isinstance(m, nn.BatchNorm2d): + m.eval() + + def upsample_disp(self, disp, mask_feat_4, stem_2x): + xspx = self.spx_2_gru(mask_feat_4, stem_2x) + spx_pred = self.spx_gru(xspx) + spx_pred = F.softmax(spx_pred, 1) + up_disp = context_upsample(disp * 4., spx_pred).unsqueeze(1) + return up_disp + + def forward(self, data): + image1 = data['left'] + image2 = data['right'] + """ Estimate disparity between pair of frames """ + test_mode = not self.training + image1 = (2 * (image1 / 255.0) - 1.0).contiguous() + image2 = (2 * (image2 / 255.0) - 1.0).contiguous() + features_left = self.feature(image1) + features_right = self.feature(image2) + stem_2x = self.stem_2(image1) + stem_4x = self.stem_4(stem_2x) + stem_2y = self.stem_2(image2) + stem_4y = self.stem_4(stem_2y) + + features_left[0] = torch.cat((features_left[0], stem_4x), 1) + features_right[0] = torch.cat((features_right[0], stem_4y), 1) + + match_left = self.desc(self.conv(features_left[0])) + match_right = self.desc(self.conv(features_right[0])) + + gwc_volume = build_gwc_volume(match_left, match_right, self.args.MAX_DISP // 4, 8) + geo_encoding_volume = self.cost_agg(gwc_volume, features_left) + + # Init disp from geometry encoding volume + prob = F.softmax(self.classifier(geo_encoding_volume).squeeze(1), dim=1) + init_disp = disparity_regression(prob, self.args.MAX_DISP // 4) # 1, 1, 80, 184 + + del prob, gwc_volume + + if not test_mode: + xspx = self.spx_4(features_left[0]) + xspx = self.spx_2(xspx, stem_2x) # 1, 64, 160, 368 + spx_pred = self.spx(xspx) + spx_pred = F.softmax(spx_pred, 1) # 1, 9, 320, 736 + + hidden = self.hnet(features_left[0]) + net = torch.tanh(hidden) + context = self.cnet(features_left[0]) + context = list(self.context_zqr_conv(context).split(split_size=self.args.HIDDEN_DIM, dim=1)) + + geo_block = Geo_Encoding_Volume + geo_fn = geo_block(geo_encoding_volume.float(), radius=self.args.CORR_RADIUS, num_levels=self.args.CORR_LEVELS) + # b, c, h, w = match_left.shape + disp = init_disp + disp_preds = [] + + # GRUs iterations to update disparity + iters = self.args.VALID_ITERS if test_mode else self.args.TRAIN_ITERS + for itr in range(iters): + disp = disp.detach() + geo_feat = geo_fn(disp) + net, mask_feat_4, delta_disp = self.update_block(net, context, geo_feat, disp) + disp = disp + delta_disp + if test_mode and itr < iters - 1: + continue + + # upsample predictions + disp_up = self.upsample_disp(disp, mask_feat_4, stem_2x) + disp_preds.append(disp_up) + + if test_mode: + return {'disp_pred': disp_up} + + init_disp = context_upsample(init_disp * 4., spx_pred.float()).unsqueeze(1) + + return {'init_disp': init_disp, + 'disp_preds': disp_preds, + 'disp_pred': disp_preds[-1]} + + def get_loss(self, model_pred, input_data): + disp_gt = input_data["disp"] + mask = (disp_gt < self.max_disp) & (disp_gt > 0) + valid = mask.float() + + disp_gt = disp_gt.unsqueeze(1) + mag = torch.sum(disp_gt ** 2, dim=1).sqrt() + valid = ((valid >= 0.5) & (mag < self.max_disp)).unsqueeze(1) + assert valid.shape == disp_gt.shape, [valid.shape, disp_gt.shape] + assert not torch.isinf(disp_gt[valid.bool()]).any() + + disp_init_pred = model_pred['init_disp'] + disp_loss = 1.0 * F.smooth_l1_loss(disp_init_pred[valid.bool()], disp_gt[valid.bool()], reduction='mean') + + # gru loss + loss_gamma = 0.9 + disp_preds = model_pred['disp_preds'] + n_predictions = len(disp_preds) + assert n_predictions >= 1 + for i in range(n_predictions): + adjusted_loss_gamma = loss_gamma ** (15 / (n_predictions - 1)) + i_weight = adjusted_loss_gamma ** (n_predictions - i - 1) + i_loss = (disp_preds[i] - disp_gt).abs() + assert i_loss.shape == valid.shape, [i_loss.shape, valid.shape, disp_gt.shape, disp_preds[i].shape] + disp_loss += i_weight * i_loss[valid.bool()].mean() + + loss_info = {'scalar/train/loss_disp': disp_loss.item()} + return disp_loss, loss_info diff --git a/stereo/modeling/models/igevpp_rt/submodule.py b/stereo/modeling/models/igevpp_rt/submodule.py new file mode 100644 index 00000000..30bd2718 --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/submodule.py @@ -0,0 +1,265 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class BasicConv(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, bn=True, relu=True, **kwargs): + super(BasicConv, self).__init__() + + self.relu = relu + self.use_bn = bn + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.bn = nn.BatchNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_bn: + x = self.bn(x) + if self.relu: + x = nn.LeakyReLU()(x) # , inplace=True) + return x + + +class Conv2x(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, bn=True, + relu=True, keep_dispc=False): + super(Conv2x, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, + stride=stride, padding=padding) + else: + self.conv1 = BasicConv(in_channels, out_channels, deconv, is_3d, bn=True, relu=True, kernel_size=kernel, + stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = BasicConv(out_channels * 2, out_channels * mul, False, is_3d, bn, relu, kernel_size=3, + stride=1, padding=1) + else: + self.conv2 = BasicConv(out_channels, out_channels, False, is_3d, bn, relu, kernel_size=3, stride=1, + padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate( + x, + size=(rem.shape[-2], rem.shape[-1]), + mode='nearest') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + +class BasicConv_IN(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, IN=True, relu=True, **kwargs): + super(BasicConv_IN, self).__init__() + + self.relu = relu + self.use_in = IN + if is_3d: + if deconv: + self.conv = nn.ConvTranspose3d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv3d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm3d(out_channels) + else: + if deconv: + self.conv = nn.ConvTranspose2d(in_channels, out_channels, bias=False, **kwargs) + else: + self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) + self.IN = nn.InstanceNorm2d(out_channels) + + def forward(self, x): + x = self.conv(x) + if self.use_in: + x = self.IN(x) + if self.relu: + x = nn.LeakyReLU()(x) # , inplace=True) + return x + + +class Conv2x_IN(nn.Module): + + def __init__(self, in_channels, out_channels, deconv=False, is_3d=False, concat=True, keep_concat=True, IN=True, + relu=True, keep_dispc=False): + super(Conv2x_IN, self).__init__() + self.concat = concat + self.is_3d = is_3d + if deconv and is_3d: + kernel = (4, 4, 4) + elif deconv: + kernel = 4 + else: + kernel = 3 + + if deconv and is_3d and keep_dispc: + kernel = (1, 4, 4) + stride = (1, 2, 2) + padding = (0, 1, 1) + self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, + stride=stride, padding=padding) + else: + self.conv1 = BasicConv_IN(in_channels, out_channels, deconv, is_3d, IN=True, relu=True, kernel_size=kernel, + stride=2, padding=1) + + if self.concat: + mul = 2 if keep_concat else 1 + self.conv2 = BasicConv_IN(out_channels * 2, out_channels * mul, False, is_3d, IN, relu, kernel_size=3, + stride=1, padding=1) + else: + self.conv2 = BasicConv_IN(out_channels, out_channels, False, is_3d, IN, relu, kernel_size=3, stride=1, + padding=1) + + def forward(self, x, rem): + x = self.conv1(x) + if x.shape != rem.shape: + x = F.interpolate( + x, + size=(rem.shape[-2], rem.shape[-1]), + mode='nearest') + if self.concat: + x = torch.cat((x, rem), 1) + else: + x = x + rem + x = self.conv2(x) + return x + + +def groupwise_correlation(fea1, fea2, num_groups): + B, C, H, W = fea1.shape + assert C % num_groups == 0 + channels_per_group = C // num_groups + cost = (fea1 * fea2).view([B, num_groups, channels_per_group, H, W]).mean(dim=2) + assert cost.shape == (B, num_groups, H, W) + return cost + + +def build_gwc_volume(refimg_fea, targetimg_fea, maxdisp, num_groups): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, num_groups, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = groupwise_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i], + num_groups) + else: + volume[:, :, i, :, :] = groupwise_correlation(refimg_fea, targetimg_fea, num_groups) + volume = volume.contiguous() + return volume + + +def norm_correlation(fea1, fea2): + cost = torch.mean( + ((fea1 / (torch.norm(fea1, 2, 1, True) + 1e-05)) * (fea2 / (torch.norm(fea2, 2, 1, True) + 1e-05))), dim=1, + keepdim=True) + return cost + + +def build_norm_correlation_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = norm_correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i]) + else: + volume[:, :, i, :, :] = norm_correlation(refimg_fea, targetimg_fea) + volume = volume.contiguous() + return volume + + +def correlation(fea1, fea2): + cost = torch.sum((fea1 * fea2), dim=1, keepdim=True) + return cost + + +def build_correlation_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 1, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :, i, :, i:] = correlation(refimg_fea[:, :, :, i:], targetimg_fea[:, :, :, :-i]) + else: + volume[:, :, i, :, :] = correlation(refimg_fea, targetimg_fea) + volume = volume.contiguous() + return volume + + +def build_concat_volume(refimg_fea, targetimg_fea, maxdisp): + B, C, H, W = refimg_fea.shape + volume = refimg_fea.new_zeros([B, 2 * C, maxdisp, H, W]) + for i in range(maxdisp): + if i > 0: + volume[:, :C, i, :, :] = refimg_fea[:, :, :, :] + volume[:, C:, i, :, i:] = targetimg_fea[:, :, :, :-i] + else: + volume[:, :C, i, :, :] = refimg_fea + volume[:, C:, i, :, :] = targetimg_fea + volume = volume.contiguous() + return volume + + +def disparity_regression(x, maxdisp): + assert len(x.shape) == 4 + disp_values = torch.arange(0, maxdisp, dtype=x.dtype, device=x.device) + disp_values = disp_values.view(1, maxdisp, 1, 1) + return torch.sum(x * disp_values, 1, keepdim=True) + + +class FeatureAtt(nn.Module): + def __init__(self, cv_chan, feat_chan): + super(FeatureAtt, self).__init__() + + self.feat_att = nn.Sequential( + BasicConv(feat_chan, feat_chan // 2, kernel_size=1, stride=1, padding=0), + nn.Conv2d(feat_chan // 2, cv_chan, 1)) + + def forward(self, cv, feat): + ''' + ''' + feat_att = self.feat_att(feat).unsqueeze(2) + cv = torch.sigmoid(feat_att) * cv + return cv + + +def context_upsample(disp_low, up_weights): + ### + # cv (b,1,h,w) + # sp (b,9,4*h,4*w) + ### + b, c, h, w = disp_low.shape + + disp_unfold = F.unfold(disp_low.reshape(b, c, h, w), 3, 1, 1).reshape(b, -1, h, w) + disp_unfold = F.interpolate(disp_unfold, (h * 4, w * 4), mode='nearest').reshape(b, 9, h * 4, w * 4) + + disp = (disp_unfold * up_weights).sum(1) + + return disp diff --git a/stereo/modeling/models/igevpp_rt/trainer.py b/stereo/modeling/models/igevpp_rt/trainer.py new file mode 100644 index 00000000..ce9c8836 --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/trainer.py @@ -0,0 +1,13 @@ +# @Time : 2024/2/9 11:39 +# @Author : zhangchenming +from stereo.modeling.trainer_template import TrainerTemplate +from .igev_stereo import IGEVStereo + +__all__ = { + 'IGEVPP_rt': IGEVStereo, +} + +class Trainer(TrainerTemplate): + def __init__(self, args, cfgs, local_rank, global_rank, logger, tb_writer): + model = __all__[cfgs.MODEL.NAME](cfgs.MODEL) + super().__init__(args, cfgs, local_rank, global_rank, logger, tb_writer, model) diff --git a/stereo/modeling/models/igevpp_rt/update.py b/stereo/modeling/models/igevpp_rt/update.py new file mode 100644 index 00000000..583e5bdd --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/update.py @@ -0,0 +1,133 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class FlowHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=2): + super(FlowHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class DispHead(nn.Module): + def __init__(self, input_dim=128, hidden_dim=256, output_dim=1): + super(DispHead, self).__init__() + self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1) + self.conv2 = nn.Conv2d(hidden_dim, output_dim, 3, padding=1) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.conv2(self.relu(self.conv1(x))) + + +class ConvGRU(nn.Module): + def __init__(self, hidden_dim, input_dim, kernel_size=3): + super(ConvGRU, self).__init__() + self.convz = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2) + self.convr = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2) + self.convq = nn.Conv2d(hidden_dim+input_dim, hidden_dim, kernel_size, padding=kernel_size//2) + + def forward(self, h, c, x): + # x = torch.cat(x_list, dim=1) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz(hx) + c[0]) + r = torch.sigmoid(self.convr(hx) + c[1]) + q = torch.tanh(self.convq(torch.cat([r*h, x], dim=1)) + c[2]) + h = (1-z) * h + z * q + return h + + +class SepConvGRU(nn.Module): + def __init__(self, hidden_dim=128, input_dim=192 + 128): + super(SepConvGRU, self).__init__() + self.convz1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convr1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + self.convq1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2)) + + self.convz2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convr2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + self.convq2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0)) + + def forward(self, h, *x): + # horizontal + x = torch.cat(x, dim=1) + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz1(hx)) + r = torch.sigmoid(self.convr1(hx)) + q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + # vertical + hx = torch.cat([h, x], dim=1) + z = torch.sigmoid(self.convz2(hx)) + r = torch.sigmoid(self.convr2(hx)) + q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1))) + h = (1 - z) * h + z * q + + return h + + +class BasicMotionEncoder(nn.Module): + def __init__(self, args): + super(BasicMotionEncoder, self).__init__() + self.args = args + cor_planes = args.CORR_LEVELS * (2 * args.CORR_RADIUS + 1) * 8 + + self.convc1 = nn.Conv2d(cor_planes, 64, 1, padding=0) + self.convc2 = nn.Conv2d(64, 64, 3, padding=1) + + self.convd1 = nn.Conv2d(1, 64, 7, padding=3) + self.convd2 = nn.Conv2d(64, 64, 3, padding=1) + + self.conv = nn.Conv2d(64 + 64, 96 - 1, 3, padding=1) + + def forward(self, disp, corr): + cor = F.relu(self.convc1(corr)) + cor = F.relu(self.convc2(cor)) + disp_ = F.relu(self.convd1(disp)) + disp_ = F.relu(self.convd2(disp_)) + + cor_disp = torch.cat([cor, disp_], dim=1) + out = F.relu(self.conv(cor_disp)) + return torch.cat([out, disp], dim=1) + +def pool2x(x): + return F.avg_pool2d(x, 3, stride=2, padding=1) + + +def pool4x(x): + return F.avg_pool2d(x, 5, stride=4, padding=1) + + +# def interp(x, dest): +# interp_args = {'mode': 'bilinear', 'align_corners': True} +# return F.interpolate(x, dest.shape[2:], **interp_args) + + +class BasicUpdateBlock(nn.Module): + def __init__(self, args, hidden_dim=96): + super().__init__() + self.args = args + self.encoder = BasicMotionEncoder(args) + + input_dim = 96 + self.gru = ConvGRU(hidden_dim, input_dim) + self.disp_head = DispHead(hidden_dim, hidden_dim=128, output_dim=1) + + self.mask_feat_4 = nn.Sequential( + nn.Conv2d(hidden_dim, 32, 3, padding=1), + nn.ReLU(inplace=True)) + + def forward(self, net, inp, corr=None, disp=None): + + motion_features = self.encoder(disp, corr) + net = self.gru(net, inp, motion_features) + + delta_disp = self.disp_head(net) + mask_feat_4 = self.mask_feat_4(net) + return net, mask_feat_4, delta_disp diff --git a/stereo/modeling/models/igevpp_rt/utils.py b/stereo/modeling/models/igevpp_rt/utils.py new file mode 100644 index 00000000..94c1c4ca --- /dev/null +++ b/stereo/modeling/models/igevpp_rt/utils.py @@ -0,0 +1,212 @@ +import numpy as np +import torch +import torch.nn.functional as F +from scipy import interpolate + +def bilinear_grid_sample(im, + grid, + align_corners: bool = False): + """Given an input and a flow-field grid, computes the output using input + values and pixel locations from grid. Supported only bilinear interpolation + method to sample the input pixels. + + Args: + im (torch.Tensor): Input feature map, shape (N, C, H, W) + grid (torch.Tensor): Point coordinates, shape (N, Hg, Wg, 2) + align_corners (bool): If set to True, the extrema (-1 and 1) are + considered as referring to the center points of the input’s + corner pixels. If set to False, they are instead considered as + referring to the corner points of the input’s corner pixels, + making the sampling more resolution agnostic. + + Returns: + torch.Tensor: A tensor with sampled points, shape (N, C, Hg, Wg) + """ + n, c, h, w = im.shape + gn, gh, gw, _ = grid.shape + assert n == gn + + x = grid[:, :, :, 0] + y = grid[:, :, :, 1] + + if align_corners: + x = ((x + 1) / 2) * (w - 1) + y = ((y + 1) / 2) * (h - 1) + else: + x = ((x + 1) * w - 1) / 2 + y = ((y + 1) * h - 1) / 2 + + x = x.view(n, -1) + y = y.view(n, -1) + + x0 = torch.floor(x).long() + y0 = torch.floor(y).long() + x1 = x0 + 1 + y1 = y0 + 1 + + wa = ((x1 - x) * (y1 - y)).unsqueeze(1) + wb = ((x1 - x) * (y - y0)).unsqueeze(1) + wc = ((x - x0) * (y1 - y)).unsqueeze(1) + wd = ((x - x0) * (y - y0)).unsqueeze(1) + + # Apply default for grid_sample function zero padding + im_padded = F.pad(im, pad=[1, 1, 1, 1], mode='constant', value=0) + padded_h = h + 2 + padded_w = w + 2 + # save points positions after padding + x0, x1, y0, y1 = x0 + 1, x1 + 1, y0 + 1, y1 + 1 + + # Clip coordinates to padded image size + device = im.device + x0 = torch.where(x0 < 0, torch.tensor(0).to(device), x0) + x0 = torch.where(x0 > padded_w - 1, torch.tensor(padded_w - 1).to(device), x0) + x1 = torch.where(x1 < 0, torch.tensor(0).to(device), x1) + x1 = torch.where(x1 > padded_w - 1, torch.tensor(padded_w - 1).to(device), x1) + y0 = torch.where(y0 < 0, torch.tensor(0).to(device), y0) + y0 = torch.where(y0 > padded_h - 1, torch.tensor(padded_h - 1).to(device), y0) + y1 = torch.where(y1 < 0, torch.tensor(0).to(device), y1) + y1 = torch.where(y1 > padded_h - 1, torch.tensor(padded_h - 1).to(device), y1) + + im_padded = im_padded.view(n, c, -1) + + x0_y0 = (x0 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x0_y1 = (x0 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y0 = (x1 + y0 * padded_w).unsqueeze(1).expand(-1, c, -1) + x1_y1 = (x1 + y1 * padded_w).unsqueeze(1).expand(-1, c, -1) + + Ia = torch.gather(im_padded, 2, x0_y0) + Ib = torch.gather(im_padded, 2, x0_y1) + Ic = torch.gather(im_padded, 2, x1_y0) + Id = torch.gather(im_padded, 2, x1_y1) + + return (Ia * wa + Ib * wb + Ic * wc + Id * wd).reshape(n, c, gh, gw) + +class InputPadder: + """ Pads images such that dimensions are divisible by 8 """ + + def __init__(self, dims, mode='sintel', divis_by=8): + self.ht, self.wd = dims[-2:] + pad_ht = (((self.ht // divis_by) + 1) * divis_by - self.ht) % divis_by + pad_wd = (((self.wd // divis_by) + 1) * divis_by - self.wd) % divis_by + if mode == 'sintel': + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, pad_ht // 2, pad_ht - pad_ht // 2] + else: + self._pad = [pad_wd // 2, pad_wd - pad_wd // 2, 0, pad_ht] + + def pad(self, *inputs): + assert all((x.ndim == 4) for x in inputs) + return [F.pad(x, self._pad, mode='replicate') for x in inputs] + + def unpad(self, x): + assert x.ndim == 4 + ht, wd = x.shape[-2:] + c = [self._pad[2], ht - self._pad[3], self._pad[0], wd - self._pad[1]] + return x[..., c[0]:c[1], c[2]:c[3]] + + +def forward_interpolate(flow): + flow = flow.detach().cpu().numpy() + dx, dy = flow[0], flow[1] + + ht, wd = dx.shape + x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + + x1 = x0 + dx + y1 = y0 + dy + + x1 = x1.reshape(-1) + y1 = y1.reshape(-1) + dx = dx.reshape(-1) + dy = dy.reshape(-1) + + valid = (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + x1 = x1[valid] + y1 = y1[valid] + dx = dx[valid] + dy = dy[valid] + + flow_x = interpolate.griddata( + (x1, y1), dx, (x0, y0), method='nearest', fill_value=0) + + flow_y = interpolate.griddata( + (x1, y1), dy, (x0, y0), method='nearest', fill_value=0) + + flow = np.stack([flow_x, flow_y], axis=0) + return torch.from_numpy(flow).float() + + +def bilinear_sampler(img, coords, mode='bilinear', mask=False): + """ Wrapper for grid_sample, uses pixel coordinates """ + H, W = img.shape[-2:] + + # print("$$$55555", img.shape, coords.shape) + xgrid, ygrid = coords.split([1, 1], dim=-1) + xgrid = 2 * xgrid / (W - 1) - 1 + + # print("######88888", xgrid) + assert torch.unique(ygrid).numel() == 1 and H == 1 # This is a stereo problem + + grid = torch.cat([xgrid, ygrid], dim=-1) + # print("###37777", grid.shape) + img = F.grid_sample(img.contiguous(), grid.contiguous(), align_corners=True) + # img = bilinear_grid_sample(img, grid, align_corners=True) + if mask: + mask = (xgrid > -1) & (ygrid > -1) & (xgrid < 1) & (ygrid < 1) + return img, mask.float() + + return img + + +def coords_grid(batch, ht, wd): + coords = torch.meshgrid(torch.arange(ht), torch.arange(wd)) + coords = torch.stack(coords[::-1], dim=0).float() + return coords[None].repeat(batch, 1, 1, 1) + + +def upflow8(flow, mode='bilinear'): + new_size = (8 * flow.shape[2], 8 * flow.shape[3]) + return 8 * F.interpolate(flow, size=new_size, mode=mode, align_corners=True) + + +def gauss_blur(input, N=5, std=1): + B, D, H, W = input.shape + x, y = torch.meshgrid(torch.arange(N).float() - N // 2, torch.arange(N).float() - N // 2) + unnormalized_gaussian = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * std ** 2)) + weights = unnormalized_gaussian / unnormalized_gaussian.sum().clamp(min=1e-4) + weights = weights.view(1, 1, N, N).to(input) + output = F.conv2d(input.reshape(B * D, 1, H, W), weights, padding=N // 2) + return output.view(B, D, H, W) + + +class Map(dict): + """ + Example: + m = Map({'first_name': 'Eduardo'}, last_name='Pool', age=24, sports=['Soccer']) + """ + def __init__(self, *args, **kwargs): + super(Map, self).__init__(*args, **kwargs) + for arg in args: + if isinstance(arg, dict): + for k, v in arg.items(): + self[k] = v + + if kwargs: + for k, v in kwargs.items(): + self[k] = v + + def __getattr__(self, attr): + return self.get(attr) + + def __setattr__(self, key, value): + self.__setitem__(key, value) + + def __setitem__(self, key, value): + super(Map, self).__setitem__(key, value) + self.__dict__.update({key: value}) + + def __delattr__(self, item): + self.__delitem__(item) + + def __delitem__(self, key): + super(Map, self).__delitem__(key) + del self.__dict__[key] From 016be2ac6bb348a1ca02b8881910bde4a2987f60 Mon Sep 17 00:00:00 2001 From: reveLATONG <47379453+reveLATONG@users.noreply.github.com> Date: Wed, 30 Jul 2025 12:02:14 +0800 Subject: [PATCH 2/4] Update igevpp_stereo.py --- stereo/modeling/models/igevpp/igevpp_stereo.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/stereo/modeling/models/igevpp/igevpp_stereo.py b/stereo/modeling/models/igevpp/igevpp_stereo.py index 4f345fbe..c577d3c3 100644 --- a/stereo/modeling/models/igevpp/igevpp_stereo.py +++ b/stereo/modeling/models/igevpp/igevpp_stereo.py @@ -246,18 +246,24 @@ def get_loss(self, model_pred, input_data): disp_gt = disp_gt.unsqueeze(1) mag = torch.sum(disp_gt ** 2, dim=1).sqrt() valid = ((valid >= 0.5) & (mag < self.max_disp)).unsqueeze(1) + # valid = ((valid >= 0.5) & (mag < self.max_disp)) assert valid.shape == disp_gt.shape, [valid.shape, disp_gt.shape] assert not torch.isinf(disp_gt[valid.bool()]).any() disp_loss = 0.0 mag = torch.sum(disp_gt**2, dim=1).sqrt() - mask0 = ((valid >= 0.5) & (mag < max_disp0)).unsqueeze(1) - mask1 = ((valid >= 0.5) & (mag < max_disp1)).unsqueeze(1) - mask = ((valid >= 0.5) & (mag < max_disp)).unsqueeze(1) + mask0 = (valid >= 0.5) & (mag < max_disp0).unsqueeze(1) + mask1 = (valid >= 0.5) & (mag < max_disp1).unsqueeze(1) + mask = (valid >= 0.5) & (mag < max_disp).unsqueeze(1) disp_init_pred = model_pred['init_disp'] - disp_loss += 1.0 * F.smooth_l1_loss(disp_init_pred[0][mask0.bool()], disp_gt[mask0.bool()], reduction='mean') - disp_loss += 0.5 * F.smooth_l1_loss(disp_init_pred[1][mask1.bool()], disp_gt[mask1.bool()], reduction='mean') - disp_loss += 0.2 * F.smooth_l1_loss(disp_init_pred[2][mask.bool()], disp_gt[mask.bool()], reduction='mean') + disp_init_pred1 = [] + for disp_init in disp_init_pred: + disp_init = disp_init.unsqueeze(1) + disp_init_pred1.append(disp_init) + disp_loss += 1.0 * F.smooth_l1_loss(disp_init_pred1[0][mask0.bool()], disp_gt[mask0.bool()], reduction='mean') + disp_loss += 0.5 * F.smooth_l1_loss(disp_init_pred1[1][mask1.bool()], disp_gt[mask1.bool()], reduction='mean') + disp_loss += 0.2 * F.smooth_l1_loss(disp_init_pred1[2][mask.bool()], disp_gt[mask.bool()], reduction='mean') + # gru loss loss_gamma = 0.9 From 7f8c7f4c249e592797abdce0ff33278cea347248 Mon Sep 17 00:00:00 2001 From: reveLATONG <47379453+reveLATONG@users.noreply.github.com> Date: Wed, 30 Jul 2025 13:11:13 +0800 Subject: [PATCH 3/4] Update igevpp_stereo.py --- stereo/modeling/models/igevpp/igevpp_stereo.py | 1 + 1 file changed, 1 insertion(+) diff --git a/stereo/modeling/models/igevpp/igevpp_stereo.py b/stereo/modeling/models/igevpp/igevpp_stereo.py index c577d3c3..f5bd4fc0 100644 --- a/stereo/modeling/models/igevpp/igevpp_stereo.py +++ b/stereo/modeling/models/igevpp/igevpp_stereo.py @@ -81,6 +81,7 @@ class IGEVPPStereo(nn.Module): def __init__(self, args): super().__init__() self.args = args + self.max_disp = args.MAX_DISP context_dims = args.HIDDEN_DIMS From 7bfe9047c58c19bbfa08421a961cea4d1a8efa7d Mon Sep 17 00:00:00 2001 From: wuwentao Date: Tue, 23 Dec 2025 09:42:08 +0800 Subject: [PATCH 4/4] igevpp(not rt version) returned tensor instead of dict in test mode so error occured in running eval.py --- stereo/modeling/models/igevpp/igevpp_stereo.py | 6 +----- stereo/modeling/models/igevpp_rt/igev_stereo.py | 2 -- 2 files changed, 1 insertion(+), 7 deletions(-) diff --git a/stereo/modeling/models/igevpp/igevpp_stereo.py b/stereo/modeling/models/igevpp/igevpp_stereo.py index f5bd4fc0..45cc9564 100644 --- a/stereo/modeling/models/igevpp/igevpp_stereo.py +++ b/stereo/modeling/models/igevpp/igevpp_stereo.py @@ -147,10 +147,7 @@ def upsample_disp(self, disp, mask_feat_4, stem_2x): def forward(self, data): image1 = data['left'] image2 = data['right'] - """ Estimate disparity between pair of frames """ test_mode = not self.training - image1 = (2 * (image1 / 255.0) - 1.0).contiguous() - image2 = (2 * (image2 / 255.0) - 1.0).contiguous() features_left = self.feature(image1) features_right = self.feature(image2) @@ -221,7 +218,7 @@ def forward(self, data): iter_preds.append(disp_up) if test_mode: - return disp_up + return {'disp_pred': disp_up} xspx = self.spx_4(features_left[0]) xspx = self.spx_2(xspx, stem_2x) @@ -265,7 +262,6 @@ def get_loss(self, model_pred, input_data): disp_loss += 0.5 * F.smooth_l1_loss(disp_init_pred1[1][mask1.bool()], disp_gt[mask1.bool()], reduction='mean') disp_loss += 0.2 * F.smooth_l1_loss(disp_init_pred1[2][mask.bool()], disp_gt[mask.bool()], reduction='mean') - # gru loss loss_gamma = 0.9 disp_preds = model_pred['disp_preds'] diff --git a/stereo/modeling/models/igevpp_rt/igev_stereo.py b/stereo/modeling/models/igevpp_rt/igev_stereo.py index bc79f98c..12e6e897 100644 --- a/stereo/modeling/models/igevpp_rt/igev_stereo.py +++ b/stereo/modeling/models/igevpp_rt/igev_stereo.py @@ -141,8 +141,6 @@ def forward(self, data): image2 = data['right'] """ Estimate disparity between pair of frames """ test_mode = not self.training - image1 = (2 * (image1 / 255.0) - 1.0).contiguous() - image2 = (2 * (image2 / 255.0) - 1.0).contiguous() features_left = self.feature(image1) features_right = self.feature(image2) stem_2x = self.stem_2(image1)