Vincentqyw
fix: roma
c74a070
import time
import numpy as np
import torch
import torch.nn.functional as F
def rnd_sample(inputs, n_sample):
cur_size = inputs[0].shape[0]
rnd_idx = torch.randperm(cur_size)[0:n_sample]
outputs = [i[rnd_idx] for i in inputs]
return outputs
def _grid_positions(h, w, bs):
x_rng = torch.arange(0, w.int())
y_rng = torch.arange(0, h.int())
xv, yv = torch.meshgrid(x_rng, y_rng, indexing="xy")
return (
torch.reshape(torch.stack((yv, xv), axis=-1), (1, -1, 2))
.repeat(bs, 1, 1)
.float()
)
def getK(ori_img_size, cur_feat_size, K):
# WARNING: cur_feat_size's order is [h, w]
r = ori_img_size / cur_feat_size[[1, 0]]
r_K0 = torch.stack(
[K[:, 0] / r[:, 0][..., None], K[:, 1] / r[:, 1][..., None], K[:, 2]], axis=1
)
return r_K0
def gather_nd(params, indices):
"""The same as tf.gather_nd but batched gather is not supported yet.
indices is an k-dimensional integer tensor, best thought of as a (k-1)-dimensional tensor of indices into params, where each element defines a slice of params:
output[\\(i_0, ..., i_{k-2}\\)] = params[indices[\\(i_0, ..., i_{k-2}\\)]]
Args:
params (Tensor): "n" dimensions. shape: [x_0, x_1, x_2, ..., x_{n-1}]
indices (Tensor): "k" dimensions. shape: [y_0,y_2,...,y_{k-2}, m]. m <= n.
Returns: gathered Tensor.
shape [y_0,y_2,...y_{k-2}] + params.shape[m:]
"""
orig_shape = list(indices.shape)
num_samples = np.prod(orig_shape[:-1])
m = orig_shape[-1]
n = len(params.shape)
if m <= n:
out_shape = orig_shape[:-1] + list(params.shape)[m:]
else:
raise ValueError(
f"the last dimension of indices must less or equal to the rank of params. Got indices:{indices.shape}, params:{params.shape}. {m} > {n}"
)
indices = indices.reshape((num_samples, m)).transpose(0, 1).tolist()
output = params[indices] # (num_samples, ...)
return output.reshape(out_shape).contiguous()
# input: pos [kpt_n, 2]; inputs [H, W, 128] / [H, W]
# output: [kpt_n, 128] / [kpt_n]
def interpolate(pos, inputs, nd=True):
h = inputs.shape[0]
w = inputs.shape[1]
i = pos[:, 0]
j = pos[:, 1]
i_top_left = torch.clamp(torch.floor(i).int(), 0, h - 1)
j_top_left = torch.clamp(torch.floor(j).int(), 0, w - 1)
i_top_right = torch.clamp(torch.floor(i).int(), 0, h - 1)
j_top_right = torch.clamp(torch.ceil(j).int(), 0, w - 1)
i_bottom_left = torch.clamp(torch.ceil(i).int(), 0, h - 1)
j_bottom_left = torch.clamp(torch.floor(j).int(), 0, w - 1)
i_bottom_right = torch.clamp(torch.ceil(i).int(), 0, h - 1)
j_bottom_right = torch.clamp(torch.ceil(j).int(), 0, w - 1)
dist_i_top_left = i - i_top_left.float()
dist_j_top_left = j - j_top_left.float()
w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
w_top_right = (1 - dist_i_top_left) * dist_j_top_left
w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
w_bottom_right = dist_i_top_left * dist_j_top_left
if nd:
w_top_left = w_top_left[..., None]
w_top_right = w_top_right[..., None]
w_bottom_left = w_bottom_left[..., None]
w_bottom_right = w_bottom_right[..., None]
interpolated_val = (
w_top_left * gather_nd(inputs, torch.stack([i_top_left, j_top_left], axis=-1))
+ w_top_right
* gather_nd(inputs, torch.stack([i_top_right, j_top_right], axis=-1))
+ w_bottom_left
* gather_nd(inputs, torch.stack([i_bottom_left, j_bottom_left], axis=-1))
+ w_bottom_right
* gather_nd(inputs, torch.stack([i_bottom_right, j_bottom_right], axis=-1))
)
return interpolated_val
def validate_and_interpolate(
pos, inputs, validate_corner=True, validate_val=None, nd=False
):
if nd:
h, w, c = inputs.shape
else:
h, w = inputs.shape
ids = torch.arange(0, pos.shape[0])
i = pos[:, 0]
j = pos[:, 1]
i_top_left = torch.floor(i).int()
j_top_left = torch.floor(j).int()
i_top_right = torch.floor(i).int()
j_top_right = torch.ceil(j).int()
i_bottom_left = torch.ceil(i).int()
j_bottom_left = torch.floor(j).int()
i_bottom_right = torch.ceil(i).int()
j_bottom_right = torch.ceil(j).int()
if validate_corner:
# Valid corner
valid_top_left = torch.logical_and(i_top_left >= 0, j_top_left >= 0)
valid_top_right = torch.logical_and(i_top_right >= 0, j_top_right < w)
valid_bottom_left = torch.logical_and(i_bottom_left < h, j_bottom_left >= 0)
valid_bottom_right = torch.logical_and(i_bottom_right < h, j_bottom_right < w)
valid_corner = torch.logical_and(
torch.logical_and(valid_top_left, valid_top_right),
torch.logical_and(valid_bottom_left, valid_bottom_right),
)
i_top_left = i_top_left[valid_corner]
j_top_left = j_top_left[valid_corner]
i_top_right = i_top_right[valid_corner]
j_top_right = j_top_right[valid_corner]
i_bottom_left = i_bottom_left[valid_corner]
j_bottom_left = j_bottom_left[valid_corner]
i_bottom_right = i_bottom_right[valid_corner]
j_bottom_right = j_bottom_right[valid_corner]
ids = ids[valid_corner]
if validate_val is not None:
# Valid depth
valid_depth = torch.logical_and(
torch.logical_and(
gather_nd(inputs, torch.stack([i_top_left, j_top_left], axis=-1)) > 0,
gather_nd(inputs, torch.stack([i_top_right, j_top_right], axis=-1)) > 0,
),
torch.logical_and(
gather_nd(inputs, torch.stack([i_bottom_left, j_bottom_left], axis=-1))
> 0,
gather_nd(
inputs, torch.stack([i_bottom_right, j_bottom_right], axis=-1)
)
> 0,
),
)
i_top_left = i_top_left[valid_depth]
j_top_left = j_top_left[valid_depth]
i_top_right = i_top_right[valid_depth]
j_top_right = j_top_right[valid_depth]
i_bottom_left = i_bottom_left[valid_depth]
j_bottom_left = j_bottom_left[valid_depth]
i_bottom_right = i_bottom_right[valid_depth]
j_bottom_right = j_bottom_right[valid_depth]
ids = ids[valid_depth]
# Interpolation
i = i[ids]
j = j[ids]
dist_i_top_left = i - i_top_left.float()
dist_j_top_left = j - j_top_left.float()
w_top_left = (1 - dist_i_top_left) * (1 - dist_j_top_left)
w_top_right = (1 - dist_i_top_left) * dist_j_top_left
w_bottom_left = dist_i_top_left * (1 - dist_j_top_left)
w_bottom_right = dist_i_top_left * dist_j_top_left
if nd:
w_top_left = w_top_left[..., None]
w_top_right = w_top_right[..., None]
w_bottom_left = w_bottom_left[..., None]
w_bottom_right = w_bottom_right[..., None]
interpolated_val = (
w_top_left * gather_nd(inputs, torch.stack([i_top_left, j_top_left], axis=-1))
+ w_top_right
* gather_nd(inputs, torch.stack([i_top_right, j_top_right], axis=-1))
+ w_bottom_left
* gather_nd(inputs, torch.stack([i_bottom_left, j_bottom_left], axis=-1))
+ w_bottom_right
* gather_nd(inputs, torch.stack([i_bottom_right, j_bottom_right], axis=-1))
)
pos = torch.stack([i, j], axis=1)
return [interpolated_val, pos, ids]
# pos0: [2, 230400, 2]
# depth0: [2, 480, 480]
def getWarp(pos0, rel_pose, depth0, K0, depth1, K1, bs):
def swap_axis(data):
return torch.stack([data[:, 1], data[:, 0]], axis=-1)
all_pos0 = []
all_pos1 = []
all_ids = []
for i in range(bs):
z0, new_pos0, ids = validate_and_interpolate(pos0[i], depth0[i], validate_val=0)
uv0_homo = torch.cat(
[
swap_axis(new_pos0),
torch.ones((new_pos0.shape[0], 1)).to(new_pos0.device),
],
axis=-1,
)
xy0_homo = torch.matmul(torch.linalg.inv(K0[i]), uv0_homo.t())
xyz0_homo = torch.cat(
[
torch.unsqueeze(z0, 0) * xy0_homo,
torch.ones((1, new_pos0.shape[0])).to(z0.device),
],
axis=0,
)
xyz1 = torch.matmul(rel_pose[i], xyz0_homo)
xy1_homo = xyz1 / torch.unsqueeze(xyz1[-1, :], axis=0)
uv1 = torch.matmul(K1[i], xy1_homo).t()[:, 0:2]
new_pos1 = swap_axis(uv1)
annotated_depth, new_pos1, new_ids = validate_and_interpolate(
new_pos1, depth1[i], validate_val=0
)
ids = ids[new_ids]
new_pos0 = new_pos0[new_ids]
estimated_depth = xyz1.t()[new_ids][:, -1]
inlier_mask = torch.abs(estimated_depth - annotated_depth) < 0.05
all_ids.append(ids[inlier_mask])
all_pos0.append(new_pos0[inlier_mask])
all_pos1.append(new_pos1[inlier_mask])
# all_pos0 & all_pose1: [inlier_num, 2] * batch_size
return all_pos0, all_pos1, all_ids
# pos0: [2, 230400, 2]
# depth0: [2, 480, 480]
def getWarpNoValidate(pos0, rel_pose, depth0, K0, depth1, K1, bs):
def swap_axis(data):
return torch.stack([data[:, 1], data[:, 0]], axis=-1)
all_pos0 = []
all_pos1 = []
all_ids = []
for i in range(bs):
z0, new_pos0, ids = validate_and_interpolate(pos0[i], depth0[i], validate_val=0)
uv0_homo = torch.cat(
[
swap_axis(new_pos0),
torch.ones((new_pos0.shape[0], 1)).to(new_pos0.device),
],
axis=-1,
)
xy0_homo = torch.matmul(torch.linalg.inv(K0[i]), uv0_homo.t())
xyz0_homo = torch.cat(
[
torch.unsqueeze(z0, 0) * xy0_homo,
torch.ones((1, new_pos0.shape[0])).to(z0.device),
],
axis=0,
)
xyz1 = torch.matmul(rel_pose[i], xyz0_homo)
xy1_homo = xyz1 / torch.unsqueeze(xyz1[-1, :], axis=0)
uv1 = torch.matmul(K1[i], xy1_homo).t()[:, 0:2]
new_pos1 = swap_axis(uv1)
_, new_pos1, new_ids = validate_and_interpolate(
new_pos1, depth1[i], validate_val=0
)
ids = ids[new_ids]
new_pos0 = new_pos0[new_ids]
all_ids.append(ids)
all_pos0.append(new_pos0)
all_pos1.append(new_pos1)
# all_pos0 & all_pose1: [inlier_num, 2] * batch_size
return all_pos0, all_pos1, all_ids
# pos0: [2, 230400, 2]
# depth0: [2, 480, 480]
def getWarpNoValidate2(pos0, rel_pose, depth0, K0, depth1, K1):
def swap_axis(data):
return torch.stack([data[:, 1], data[:, 0]], axis=-1)
z0 = interpolate(pos0, depth0, nd=False)
uv0_homo = torch.cat(
[swap_axis(pos0), torch.ones((pos0.shape[0], 1)).to(pos0.device)], axis=-1
)
xy0_homo = torch.matmul(torch.linalg.inv(K0), uv0_homo.t())
xyz0_homo = torch.cat(
[
torch.unsqueeze(z0, 0) * xy0_homo,
torch.ones((1, pos0.shape[0])).to(z0.device),
],
axis=0,
)
xyz1 = torch.matmul(rel_pose, xyz0_homo)
xy1_homo = xyz1 / torch.unsqueeze(xyz1[-1, :], axis=0)
uv1 = torch.matmul(K1, xy1_homo).t()[:, 0:2]
new_pos1 = swap_axis(uv1)
return new_pos1
def get_dist_mat(feat1, feat2, dist_type):
eps = 1e-6
cos_dist_mat = torch.matmul(feat1, feat2.t())
if dist_type == "cosine_dist":
dist_mat = torch.clamp(cos_dist_mat, -1, 1)
elif dist_type == "euclidean_dist":
dist_mat = torch.sqrt(torch.clamp(2 - 2 * cos_dist_mat, min=eps))
elif dist_type == "euclidean_dist_no_norm":
norm1 = torch.sum(feat1 * feat1, axis=-1, keepdims=True)
norm2 = torch.sum(feat2 * feat2, axis=-1, keepdims=True)
dist_mat = torch.sqrt(
torch.clamp(norm1 - 2 * cos_dist_mat + norm2.t(), min=0.0) + eps
)
else:
raise NotImplementedError()
return dist_mat