| 9 | |
| 10 | |
| 11 | class DepthNet(nn.Module): |
| 12 | def __init__(self, mode="regression"): |
| 13 | super(DepthNet, self).__init__() |
| 14 | |
| 15 | def forward(self, cost_reg, depth_values, num_depth, interval, prob_volume_init=None,stage=0): |
| 16 | |
| 17 | |
| 18 | |
| 19 | prob_volume = F.softmax(cost_reg, dim=2) # (b,2, ndepth, h, w) |
| 20 | depth_sub_plus = depth_regression(prob_volume, depth_values=depth_values.unsqueeze(1),axis=2) # (b, h, w) |
| 21 | |
| 22 | depth_sup_plus_small,depth_sup_plus_huge=depth_sub_plus.split([2,2],dim=1) |
| 23 | |
| 24 | |
| 25 | small_min,small_max=depth_sup_plus_small.min(1)[0],depth_sup_plus_small.max(1)[0] |
| 26 | huge_min,huge_max=depth_sup_plus_huge.min(1)[0],depth_sup_plus_huge.max(1)[0] |
| 27 | huge_min_d,huge_max_d=2*huge_min-huge_max,2*huge_max-huge_min |
| 28 | small_min_d,small_max_d=2*small_min-small_max,2*small_max-small_min |
| 29 | |
| 30 | coors=torch.stack( |
| 31 | [item.unsqueeze(0).expand_as(depth_sub_plus[:,0]) for item in torch.meshgrid(*[torch.arange(0, s) for s in depth_sub_plus[:,0].shape[-2:]])], |
| 32 | axis=-1).to(depth_sub_plus[:,0].device) |
| 33 | mask_00=((coors[:,:,:,0]%4==0)&(coors[:,:,:,1]%2==0)) |
| 34 | mask_01=((coors[:,:,:,0]%4==0)&(coors[:,:,:,1]%2==1)) |
| 35 | mask_10=((coors[:,:,:,0]%4==1)&(coors[:,:,:,1]%2==0)) |
| 36 | mask_11=((coors[:,:,:,0]%4==1)&(coors[:,:,:,1]%2==1)) |
| 37 | mask_20=((coors[:,:,:,0]%4==2)&(coors[:,:,:,1]%2==0)) |
| 38 | mask_21=((coors[:,:,:,0]%4==2)&(coors[:,:,:,1]%2==1)) |
| 39 | mask_30=((coors[:,:,:,0]%4==3)&(coors[:,:,:,1]%2==0)) |
| 40 | mask_31=((coors[:,:,:,0]%4==3)&(coors[:,:,:,1]%2==1)) |
| 41 | |
| 42 | small_stack=torch.stack((3*small_min-2*small_max,2*small_min-small_max,small_min,small_max,2*small_max-small_min,3*small_max-2*small_min),1) |
| 43 | small_stack_d=torch.stack((3*small_min_d-2*small_max_d,2*small_min_d-small_max_d,small_min_d,small_max_d,2*small_max_d-small_min_d,3*small_max_d-2*small_min_d),1) |
| 44 | huge_stack=torch.stack((3*huge_min-2*huge_max,2*huge_min-huge_max,huge_min,huge_max,2*huge_max-huge_min,3*huge_max-2*huge_min),1) |
| 45 | huge_stack_d=torch.stack((3*huge_min_d-2*huge_max_d,2*huge_min_d-huge_max_d,huge_min_d,huge_max_d,2*huge_max_d-huge_min_d,3*huge_max_d-2*huge_min_d),1) |
| 46 | |
| 47 | # depth=torch.zeros_like(depth_sub_plus[:,0]) |
| 48 | depth_values_c=torch.zeros_like(depth_sub_plus) |
| 49 | depth_values_c=torch.where(mask_00.unsqueeze(1),small_stack[:,:-2],depth_values_c) |
| 50 | depth_values_c=torch.where(mask_01.unsqueeze(1),small_stack[:,2:],depth_values_c) |
| 51 | depth_values_c=torch.where(mask_10.unsqueeze(1),huge_stack[:,2:],depth_values_c) |
| 52 | depth_values_c=torch.where(mask_11.unsqueeze(1),huge_stack[:,:-2],depth_values_c) |
| 53 | depth_values_c=torch.where(mask_20.unsqueeze(1),small_stack_d[:,:-2],depth_values_c) |
| 54 | depth_values_c=torch.where(mask_21.unsqueeze(1),small_stack_d[:,2:],depth_values_c) |
| 55 | depth_values_c=torch.where(mask_30.unsqueeze(1),huge_stack_d[:,2:],depth_values_c) |
| 56 | depth_values_c=torch.where(mask_31.unsqueeze(1),huge_stack_d[:,:-2],depth_values_c) |
| 57 | |
| 58 | |
| 59 | with torch.no_grad(): |
| 60 | # photometric confidence |
| 61 | temp_photometric_confidence=torch.sigmoid(interval/(depth_sub_plus.var(1,unbiased=False).sqrt()+1e-5)) |
| 62 | photometric_confidence=2*(temp_photometric_confidence-0.5) |
| 63 | |
| 64 | |
| 65 | return {"photometric_confidence": photometric_confidence, "prob_volume": prob_volume,"depth_sub_plus":depth_sub_plus,"depth_values_c":depth_values_c, |
| 66 | "depth_values": depth_values, "interval": interval} |
| 67 | def refine(self, cost_reg, depth_values, num_depth, interval,alpha=5): |
| 68 | prob_volume = F.softmax(cost_reg*alpha, dim=2) # (b,2, ndepth, h, w) |