Spaces:
Runtime error
Runtime error
File size: 3,911 Bytes
fc16538 |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
# TRI-VIDAR - Copyright 2022 Toyota Research Institute. All rights reserved.
import numpy as np
import torch
from vidar.geometry.camera import Camera
from vidar.utils.tensor import grid_sample
from vidar.utils.types import is_tensor
def warp_bins(rgb, cam, bins):
"""
Warp an image based on depth bins
Parameters
----------
rgb : torch.Tensor [B,?,H,W]
Input image for warping
cam : Camera
Input camera
bins : torch.Tensor
Depth bins for warping
Returns
-------
warped : torch.Tensor
Warped images for each depth bin
"""
ones = torch.ones((1, *cam.hw)).to(rgb.device)
volume = torch.stack([depth * ones for depth in bins], 1)
coords_volume = cam.coords_from_cost_volume(volume)
return grid_sample(
rgb.repeat(len(bins), 1, 1, 1), coords_volume[0].type(rgb.dtype),
padding_mode='zeros', mode='bilinear', align_corners=True)
def sample(grid, pred):
"""
Sample a grid based on predictions
Parameters
----------
grid : torch.Tensor
Grid to be sampled [B,?,H,W]
pred : torch.Tensor
Coordinate predictions [B,2,H,W]
Returns
-------
values : torch.Tensor
Sampled grid[B,?,H,W]
"""
n, _, h, w = grid.shape
coords = pred.permute(1, 2, 0).reshape(-1, 1, 1, 1).repeat(1, 1, 1, 2)
coords = 2 * coords / (n - 1) - 1
grid = grid.permute(2, 3, 0, 1).reshape(-1, 1, n, 1).repeat(1, 1, 1, 2)
values = grid_sample(grid, coords,
padding_mode='zeros', mode='bilinear', align_corners=True)
return values.reshape(h, w, 1, 1).permute(2, 3, 0, 1)
def compute_depth_bin(min_depth, max_depth, num_bins, i):
"""
Calculate a single SID depth bin
Parameters
----------
min_depth : Float
Minimum depth value
max_depth : Float
Maximum depth value
num_bins : Int
Number of depth bins
i : Int
Index of the depth bin in the interval
Returns
-------
bin : torch.Tensor
Corresponding depth bin
"""
return torch.exp(np.log(min_depth) + np.log(max_depth / min_depth) * i / (num_bins - 1)).\
clamp(min=min_depth, max=max_depth)
def uncompute_depth_bin(min_depth, max_depth, num_bins, depth):
"""
Recover the SID bin index from a depth value
Parameters
----------
min_depth : Float
Minimum depth value
max_depth : Float
Maximum depth value
num_bins : Int
Number of depth bins
depth : torch.Tensor
Depth value
Returns
-------
index : torch.Tensor
Index for the depth value in the SID interval
"""
return (num_bins - 1) * ((torch.log(depth) - np.log(min_depth)) /
np.log(max_depth / min_depth)).clamp(min=0, max=num_bins)
def compute_depth_bins(min_depth, max_depth, num_bins, mode):
"""
Compute depth bins for an interval
Parameters
----------
min_depth : Float
Minimum depth value
max_depth : Float
Maximum depth value
num_bins : Int
Number of depth bins
mode : String
Depth discretization mode
Returns
-------
bins : torch.Tensor
Discretized depth bins
"""
if is_tensor(min_depth):
min_depth = min_depth.detach().cpu()
if is_tensor(max_depth):
max_depth = max_depth.detach().cpu()
if mode == 'inverse':
depth_bins = 1. / np.linspace(
1. / max_depth, 1. / min_depth, num_bins)[::-1]
elif mode == 'linear':
depth_bins = np.linspace(
min_depth, max_depth, num_bins)
elif mode == 'sid':
depth_bins = np.array(
[np.exp(np.log(min_depth) + np.log(max_depth / min_depth) * i / (num_bins - 1))
for i in range(num_bins)])
else:
raise NotImplementedError
return torch.from_numpy(depth_bins).float()
|