import torch import torch.nn.functional as F import numpy as np import matplotlib.pyplot as plt import torchext import co class TimedModule(torch.nn.Module): def __init__(self, mod_name): super().__init__() self.mod_name = mod_name def tforward(self, *args, **kwargs): raise Exception('not implemented') def forward(self, *args, **kwargs): torch.cuda.synchronize() with co.gtimer.Ctx(self.mod_name): x = self.tforward(*args, **kwargs) torch.cuda.synchronize() return x class PosOutput(TimedModule): def __init__(self, channels_in, type, im_height, im_width, alpha=1, beta=0, gamma=1, offset=0): super().__init__(mod_name='PosOutput') self.im_width = im_width self.im_width = im_width if type == 'pos': self.layer = torch.nn.Sequential( torch.nn.Conv2d(channels_in, 1, kernel_size=3, padding=1), SigmoidAffine(alpha=alpha, beta=beta, gamma=gamma, offset=offset) ) elif type == 'pos_row': self.layer = torch.nn.Sequential( MultiLinear(im_height, channels_in, 1), SigmoidAffine(alpha=alpha, beta=beta, gamma=gamma, offset=offset) ) self.u_pos = None def tforward(self, x): if self.u_pos is None: self.u_pos = torch.arange(x.shape[3], dtype=torch.float32).view(1,1,1,-1) self.u_pos = self.u_pos.to(x.device) pos = self.layer(x) disp = self.u_pos - pos return disp class OutputLayerFactory(object): ''' Define type of output type options: linear: apply only conv channel, used for the edge decoder disp: estimate the disparity disp_row: independently estimate the disparity per row pos: estimate the absolute location pos_row: independently estimate the absolute location per row ''' def __init__(self, type='disp', params={}): self.type = type self.params = params def __call__(self, channels_in, imsize): if self.type == 'linear': return torch.nn.Conv2d(channels_in, 1, kernel_size=3, padding=1) elif self.type == 'disp': return torch.nn.Sequential( torch.nn.Conv2d(channels_in, 1, kernel_size=3, padding=1), SigmoidAffine(**self.params) ) elif self.type == 'disp_row': return torch.nn.Sequential( MultiLinear(imsize[0], channels_in, 1), SigmoidAffine(**self.params) ) elif self.type == 'pos' or self.type == 'pos_row': return PosOutput(channels_in, **self.params) else: raise Exception('unknown output layer type') class SigmoidAffine(TimedModule): def __init__(self, alpha=1, beta=0, gamma=1, offset=0): super().__init__(mod_name='SigmoidAffine') self.alpha = alpha self.beta = beta self.gamma = gamma self.offset = offset def tforward(self, x): return torch.sigmoid(x/self.gamma - self.offset) * self.alpha + self.beta class MultiLinear(TimedModule): def __init__(self, n, channels_in, channels_out): super().__init__(mod_name='MultiLinear') self.channels_out = channels_out self.mods = torch.nn.ModuleList() for idx in range(n): self.mods.append(torch.nn.Linear(channels_in, channels_out)) def tforward(self, x): x = x.permute(2,0,3,1) # BxCxHxW => HxBxWxC y = x.new_empty(*x.shape[:-1], self.channels_out) for hidx in range(x.shape[0]): y[hidx] = self.mods[hidx](x[hidx]) y = y.permute(1,3,0,2) # HxBxWxC => BxCxHxW return y class DispNetS(TimedModule): ''' Disparity Decoder based on DispNetS ''' def __init__(self, channels_in, imsizes, output_facs, output_ms=True, coordconv=False, weight_init=False, channel_multiplier=1): super(DispNetS, self).__init__(mod_name='DispNetS') self.output_ms = output_ms self.coordconv = coordconv conv_planes = channel_multiplier * np.array( [32, 64, 128, 256, 512, 512, 512] ) self.conv1 = self.downsample_conv(channels_in, conv_planes[0], kernel_size=7) self.conv2 = self.downsample_conv(conv_planes[0], conv_planes[1], kernel_size=5) self.conv3 = self.downsample_conv(conv_planes[1], conv_planes[2]) self.conv4 = self.downsample_conv(conv_planes[2], conv_planes[3]) self.conv5 = self.downsample_conv(conv_planes[3], conv_planes[4]) self.conv6 = self.downsample_conv(conv_planes[4], conv_planes[5]) self.conv7 = self.downsample_conv(conv_planes[5], conv_planes[6]) upconv_planes = channel_multiplier * np.array( [512, 512, 256, 128, 64, 32, 16] ) self.upconv7 = self.upconv(conv_planes[6], upconv_planes[0]) self.upconv6 = self.upconv(upconv_planes[0], upconv_planes[1]) self.upconv5 = self.upconv(upconv_planes[1], upconv_planes[2]) self.upconv4 = self.upconv(upconv_planes[2], upconv_planes[3]) self.upconv3 = self.upconv(upconv_planes[3], upconv_planes[4]) self.upconv2 = self.upconv(upconv_planes[4], upconv_planes[5]) self.upconv1 = self.upconv(upconv_planes[5], upconv_planes[6]) self.iconv7 = self.conv(upconv_planes[0] + conv_planes[5], upconv_planes[0]) self.iconv6 = self.conv(upconv_planes[1] + conv_planes[4], upconv_planes[1]) self.iconv5 = self.conv(upconv_planes[2] + conv_planes[3], upconv_planes[2]) self.iconv4 = self.conv(upconv_planes[3] + conv_planes[2], upconv_planes[3]) self.iconv3 = self.conv(1 + upconv_planes[4] + conv_planes[1], upconv_planes[4]) self.iconv2 = self.conv(1 + upconv_planes[5] + conv_planes[0], upconv_planes[5]) self.iconv1 = self.conv(1 + upconv_planes[6], upconv_planes[6]) if isinstance(output_facs, list): self.predict_disp4 = output_facs[3](upconv_planes[3], imsizes[3]) self.predict_disp3 = output_facs[2](upconv_planes[4], imsizes[2]) self.predict_disp2 = output_facs[1](upconv_planes[5], imsizes[1]) self.predict_disp1 = output_facs[0](upconv_planes[6], imsizes[0]) else: self.predict_disp4 = output_facs(upconv_planes[3], imsizes[3]) self.predict_disp3 = output_facs(upconv_planes[4], imsizes[2]) self.predict_disp2 = output_facs(upconv_planes[5], imsizes[1]) self.predict_disp1 = output_facs(upconv_planes[6], imsizes[0]) def init_weights(self): for m in self.modules(): if isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.ConvTranspose2d): torch.nn.init.xavier_uniform_(m.weight, gain=0.1) if m.bias is not None: torch.nn.init.zeros_(m.bias) def downsample_conv(self, in_planes, out_planes, kernel_size=3): if self.coordconv: conv = torchext.CoordConv2d(in_planes, out_planes, kernel_size=kernel_size, stride=2, padding=(kernel_size-1)//2) else: conv = torch.nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=2, padding=(kernel_size-1)//2) return torch.nn.Sequential( conv, torch.nn.ReLU(inplace=True), torch.nn.Conv2d(out_planes, out_planes, kernel_size=kernel_size, padding=(kernel_size-1)//2), torch.nn.ReLU(inplace=True) ) def conv(self, in_planes, out_planes): return torch.nn.Sequential( torch.nn.Conv2d(in_planes, out_planes, kernel_size=3, padding=1), torch.nn.ReLU(inplace=True) ) def upconv(self, in_planes, out_planes): return torch.nn.Sequential( torch.nn.ConvTranspose2d(in_planes, out_planes, kernel_size=3, stride=2, padding=1, output_padding=1), torch.nn.ReLU(inplace=True) ) def crop_like(self, input, ref): assert(input.size(2) >= ref.size(2) and input.size(3) >= ref.size(3)) return input[:, :, :ref.size(2), :ref.size(3)] def tforward(self, x): out_conv1 = self.conv1(x) out_conv2 = self.conv2(out_conv1) out_conv3 = self.conv3(out_conv2) out_conv4 = self.conv4(out_conv3) out_conv5 = self.conv5(out_conv4) out_conv6 = self.conv6(out_conv5) out_conv7 = self.conv7(out_conv6) out_upconv7 = self.crop_like(self.upconv7(out_conv7), out_conv6) concat7 = torch.cat((out_upconv7, out_conv6), 1) out_iconv7 = self.iconv7(concat7) out_upconv6 = self.crop_like(self.upconv6(out_iconv7), out_conv5) concat6 = torch.cat((out_upconv6, out_conv5), 1) out_iconv6 = self.iconv6(concat6) out_upconv5 = self.crop_like(self.upconv5(out_iconv6), out_conv4) concat5 = torch.cat((out_upconv5, out_conv4), 1) out_iconv5 = self.iconv5(concat5) out_upconv4 = self.crop_like(self.upconv4(out_iconv5), out_conv3) concat4 = torch.cat((out_upconv4, out_conv3), 1) out_iconv4 = self.iconv4(concat4) disp4 = self.predict_disp4(out_iconv4) out_upconv3 = self.crop_like(self.upconv3(out_iconv4), out_conv2) disp4_up = self.crop_like(torch.nn.functional.interpolate(disp4, scale_factor=2, mode='bilinear', align_corners=False), out_conv2) concat3 = torch.cat((out_upconv3, out_conv2, disp4_up), 1) out_iconv3 = self.iconv3(concat3) disp3 = self.predict_disp3(out_iconv3) out_upconv2 = self.crop_like(self.upconv2(out_iconv3), out_conv1) disp3_up = self.crop_like(torch.nn.functional.interpolate(disp3, scale_factor=2, mode='bilinear', align_corners=False), out_conv1) concat2 = torch.cat((out_upconv2, out_conv1, disp3_up), 1) out_iconv2 = self.iconv2(concat2) disp2 = self.predict_disp2(out_iconv2) out_upconv1 = self.crop_like(self.upconv1(out_iconv2), x) disp2_up = self.crop_like(torch.nn.functional.interpolate(disp2, scale_factor=2, mode='bilinear', align_corners=False), x) concat1 = torch.cat((out_upconv1, disp2_up), 1) out_iconv1 = self.iconv1(concat1) disp1 = self.predict_disp1(out_iconv1) if self.output_ms: return disp1, disp2, disp3, disp4 else: return disp1 class DispNetShallow(DispNetS): ''' Edge Decoder based on DispNetS with fewer layers ''' def __init__(self, channels_in, imsizes, output_facs, output_ms=True, coordconv=False, weight_init=False): super(DispNetShallow, self).__init__(channels_in, imsizes, output_facs, output_ms, coordconv, weight_init) self.mod_name = 'DispNetShallow' conv_planes = [32, 64, 128, 256, 512, 512, 512] upconv_planes = [512, 512, 256, 128, 64, 32, 16] self.iconv3 = self.conv(upconv_planes[4] + conv_planes[1], upconv_planes[4]) def tforward(self, x): out_conv1 = self.conv1(x) out_conv2 = self.conv2(out_conv1) out_conv3 = self.conv3(out_conv2) out_upconv3 = self.crop_like(self.upconv3(out_conv3), out_conv2) concat3 = torch.cat((out_upconv3, out_conv2), 1) out_iconv3 = self.iconv3(concat3) disp3 = self.predict_disp3(out_iconv3) out_upconv2 = self.crop_like(self.upconv2(out_iconv3), out_conv1) disp3_up = self.crop_like(torch.nn.functional.interpolate(disp3, scale_factor=2, mode='bilinear', align_corners=False), out_conv1) concat2 = torch.cat((out_upconv2, out_conv1, disp3_up), 1) out_iconv2 = self.iconv2(concat2) disp2 = self.predict_disp2(out_iconv2) out_upconv1 = self.crop_like(self.upconv1(out_iconv2), x) disp2_up = self.crop_like(torch.nn.functional.interpolate(disp2, scale_factor=2, mode='bilinear', align_corners=False), x) concat1 = torch.cat((out_upconv1, disp2_up), 1) out_iconv1 = self.iconv1(concat1) disp1 = self.predict_disp1(out_iconv1) if self.output_ms: return disp1, disp2, disp3 else: return disp1 class DispEdgeDecoders(TimedModule): ''' Disparity Decoder and Edge Decoder ''' def __init__(self, *args, max_disp=128, **kwargs): super(DispEdgeDecoders, self).__init__(mod_name='DispEdgeDecoders') output_facs = [OutputLayerFactory( type='disp', params={ 'alpha': max_disp/(2**s), 'beta': 0, 'gamma': 1, 'offset': 3}) for s in range(4)] self.disp_decoder = DispNetS(*args, output_facs=output_facs, **kwargs) output_facs = [OutputLayerFactory( type='linear' ) for s in range(4)] self.edge_decoder = DispNetShallow(*args, output_facs=output_facs, **kwargs) def tforward(self, x): disp = self.disp_decoder(x) edge = self.edge_decoder(x) return disp, edge class DispToDepth(TimedModule): def __init__(self, focal_length, baseline): super().__init__(mod_name='DispToDepth') self.baseline_focal_length = baseline * focal_length def tforward(self, disp): disp = torch.nn.functional.relu(disp) + 1e-12 depth = self.baseline_focal_length / disp return depth class PosToDepth(DispToDepth): def __init__(self, focal_length, baseline, im_height, im_width): super().__init__(focal_length, baseline) self.mod_name = 'PosToDepth' self.im_height = im_height self.im_width = im_width self.u_pos = torch.arange(im_width, dtype=torch.float32).view(1,1,1,-1) def tforward(self, pos): self.u_pos = self.u_pos.to(pos.device) disp = self.u_pos - pos return super().forward(disp) class RectifiedPatternSimilarityLoss(TimedModule): ''' Photometric Loss ''' def __init__(self, im_height, im_width, pattern, loss_type='census_sad', loss_eps=0.5): super().__init__(mod_name='RectifiedPatternSimilarityLoss') self.im_height = im_height self.im_width = im_width self.pattern = pattern.mean(dim=1, keepdim=True).contiguous() u, v = np.meshgrid(range(im_width), range(im_height)) uv0 = np.stack((u,v), axis=2).reshape(-1,1) uv0 = uv0.astype(np.float32).reshape(1,-1,2) self.uv0 = torch.from_numpy(uv0) self.loss_type = loss_type self.loss_eps = loss_eps def tforward(self, disp0, im, std=None): self.pattern = self.pattern.to(disp0.device) self.uv0 = self.uv0.to(disp0.device) uv0 = self.uv0.expand(disp0.shape[0], *self.uv0.shape[1:]) uv1 = torch.empty_like(uv0) uv1[...,0] = uv0[...,0] - disp0.contiguous().view(disp0.shape[0],-1) uv1[...,1] = uv0[...,1] uv1[..., 0] = 2 * (uv1[..., 0] / (self.im_width-1) - 0.5) uv1[..., 1] = 2 * (uv1[..., 1] / (self.im_height-1) - 0.5) uv1 = uv1.view(-1, self.im_height, self.im_width, 2).clone() pattern = self.pattern.expand(disp0.shape[0], *self.pattern.shape[1:]) pattern_proj = torch.nn.functional.grid_sample(pattern, uv1, padding_mode='border') mask = torch.ones_like(im) if std is not None: mask = mask*std diff = torchext.photometric_loss(pattern_proj.contiguous(), im.contiguous(), 9, self.loss_type, self.loss_eps) val = (mask*diff).sum() / mask.sum() return val, pattern_proj class DisparityLoss(TimedModule): ''' Disparity Loss ''' def __init__(self): super().__init__(mod_name='DisparityLoss') self.sobel = SobelFilter(norm=False) #if not edge_gt: self.b0=0.0503428816795 self.b1=1.07274045944 #else: # self.b0=0.0587115108967 # self.b1=1.51931190491 def tforward(self, disp, edge=None): self.sobel=self.sobel.to(disp.device) if edge is not None: grad = self.sobel(disp) grad = torch.sqrt(grad[:,0:1,...]**2 + grad[:,1:2,...]**2 + 1e-8) pdf = (1-edge)/self.b0 * torch.exp(-torch.abs(grad)/self.b0) + \ edge/self.b1 * torch.exp(-torch.abs(grad)/self.b1) val = torch.mean(-torch.log(pdf.clamp(min=1e-4))) else: # on qifeng's data we don't have ambient info # therefore we supress edge everywhere grad = self.sobel(disp) grad = torch.sqrt(grad[:,0:1,...]**2 + grad[:,1:2,...]**2 + 1e-8) grad= torch.clamp(grad, 0, 1.0) val = torch.mean(grad) return val class ProjectionBaseLoss(TimedModule): ''' Base module of the Geometric Loss ''' def __init__(self, K, Ki, im_height, im_width): super().__init__(mod_name='ProjectionBaseLoss') self.K = K.view(-1,3,3) self.im_height = im_height self.im_width = im_width u, v = np.meshgrid(range(im_width), range(im_height)) uv = np.stack((u,v,np.ones_like(u)), axis=2).reshape(-1,3) ray = uv @ Ki.numpy().T ray = ray.reshape(1,-1,3).astype(np.float32) self.ray = torch.from_numpy(ray) def transform(self, xyz, R=None, t=None): if t is not None: bs = xyz.shape[0] xyz = xyz - t.reshape(bs,1,3) if R is not None: xyz = torch.bmm(xyz, R) return xyz def unproject(self, depth, R=None, t=None): self.ray = self.ray.to(depth.device) bs = depth.shape[0] xyz = depth.reshape(bs,-1,1) * self.ray xyz = self.transform(xyz, R, t) return xyz def project(self, xyz, R, t): self.K = self.K.to(xyz.device) bs = xyz.shape[0] xyz = torch.bmm(xyz, R.transpose(1,2)) xyz = xyz + t.reshape(bs,1,3) Kt = self.K.transpose(1,2).expand(bs,-1,-1) uv = torch.bmm(xyz, Kt) d = uv[:,:,2:3] # avoid division by zero uv = uv[:,:,:2] / (torch.nn.functional.relu(d) + 1e-12) return uv, d def tforward(self, depth0, R0, t0, R1, t1): xyz = self.unproject(depth0, R0, t0) return self.project(xyz, R1, t1) class ProjectionDepthSimilarityLoss(ProjectionBaseLoss): ''' Geometric Loss ''' def __init__(self, *args, clamp=-1): super().__init__(*args) self.mod_name = 'ProjectionDepthSimilarityLoss' self.clamp = clamp def fwd(self, depth0, depth1, R0, t0, R1, t1): uv1, d1 = super().tforward(depth0, R0, t0, R1, t1) uv1[..., 0] = 2 * (uv1[..., 0] / (self.im_width-1) - 0.5) uv1[..., 1] = 2 * (uv1[..., 1] / (self.im_height-1) - 0.5) uv1 = uv1.view(-1, self.im_height, self.im_width, 2).clone() depth10 = torch.nn.functional.grid_sample(depth1, uv1, padding_mode='border') diff = torch.abs(d1.view(-1) - depth10.view(-1)) if self.clamp > 0: diff = torch.clamp(diff, 0, self.clamp) # return diff without clamping for debugging return diff.mean() def tforward(self, depth0, depth1, R0, t0, R1, t1): l0 = self.fwd(depth0, depth1, R0, t0, R1, t1) l1 = self.fwd(depth1, depth0, R1, t1, R0, t0) return l0+l1 class LCN(TimedModule): ''' Local Contract Normalization ''' def __init__(self, radius, epsilon): super().__init__(mod_name='LCN') self.box_conv = torch.nn.Sequential( torch.nn.ReflectionPad2d(radius), torch.nn.Conv2d(1, 1, kernel_size=2*radius+1, bias=False) ) self.box_conv[1].weight.requires_grad=False self.box_conv[1].weight.fill_(1.) self.epsilon = epsilon self.radius = radius def tforward(self, data): boxs = self.box_conv(data) avgs = boxs / (2*self.radius+1)**2 boxs_n2 = boxs**2 boxs_2n = self.box_conv(data**2) stds = torch.sqrt(boxs_2n / (2*self.radius+1)**2 - avgs**2 + 1e-6) stds = stds + self.epsilon return (data - avgs) / stds, stds class SobelFilter(TimedModule): ''' Sobel Filter ''' def __init__(self, norm=False): super(SobelFilter, self).__init__(mod_name='SobelFilter') kx = np.array([[-5, -4, 0, 4, 5], [-8, -10, 0, 10, 8], [-10, -20, 0, 20, 10], [-8, -10, 0, 10, 8], [-5, -4, 0, 4, 5]])/240.0 ky = kx.copy().transpose(1,0) self.conv_x=torch.nn.Conv2d(1, 1, kernel_size=5, stride=1, padding=0, bias=False) self.conv_x.weight=torch.nn.Parameter(torch.from_numpy(kx).float().unsqueeze(0).unsqueeze(0)) self.conv_y=torch.nn.Conv2d(1, 1, kernel_size=5, stride=1, padding=0, bias=False) self.conv_y.weight=torch.nn.Parameter(torch.from_numpy(ky).float().unsqueeze(0).unsqueeze(0)) self.norm=norm def tforward(self,x): x = F.pad(x, (2,2,2,2), "replicate") gx = self.conv_x(x) gy = self.conv_y(x) if self.norm: return torch.sqrt(gx**2 + gy**2 + 1e-8) else: return torch.cat((gx, gy), dim=1)