File size: 1,536 Bytes
bd86ed9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import torch
import torch.nn.functional as F
import copy
def update_sample(bin_edges, target_bin_left, target_bin_right, depth_r, pred_label, depth_num, min_depth, max_depth, uncertainty_range):
with torch.no_grad():
b, _, h, w = bin_edges.shape
mode = 'direct'
if mode == 'direct':
depth_range = uncertainty_range
depth_start_update = torch.clamp_min(depth_r - 0.5 * depth_range, min_depth)
else:
depth_range = uncertainty_range + (target_bin_right - target_bin_left).abs()
depth_start_update = torch.clamp_min(target_bin_left - 0.5 * uncertainty_range, min_depth)
interval = depth_range / depth_num
interval = interval.repeat(1, depth_num, 1, 1)
interval = torch.cat([torch.ones([b, 1, h, w], device=bin_edges.device) * depth_start_update, interval], 1)
bin_edges = torch.cumsum(interval, 1).clamp(min_depth, max_depth)
curr_depth = 0.5 * (bin_edges[:, :-1] + bin_edges[:, 1:])
return bin_edges.detach(), curr_depth.detach()
def get_label(gt_depth_img, bin_edges, depth_num):
with torch.no_grad():
gt_label = torch.zeros(gt_depth_img.size(), dtype=torch.int64, device=gt_depth_img.device)
for i in range(depth_num):
bin_mask = torch.ge(gt_depth_img, bin_edges[:, i])
bin_mask = torch.logical_and(bin_mask,
torch.lt(gt_depth_img, bin_edges[:, i + 1]))
gt_label[bin_mask] = i
return gt_label
|