Vincentqyw
fix: roma
8b973ee
raw
history blame
7.21 kB
from typing import Tuple
import numpy as np
import torch
def to_homogeneous(points):
"""Convert N-dimensional points to homogeneous coordinates.
Args:
points: torch.Tensor or numpy.ndarray with size (..., N).
Returns:
A torch.Tensor or numpy.ndarray with size (..., N+1).
"""
if isinstance(points, torch.Tensor):
pad = points.new_ones(points.shape[:-1] + (1,))
return torch.cat([points, pad], dim=-1)
elif isinstance(points, np.ndarray):
pad = np.ones((points.shape[:-1] + (1,)), dtype=points.dtype)
return np.concatenate([points, pad], axis=-1)
else:
raise ValueError
def from_homogeneous(points, eps=0.0):
"""Remove the homogeneous dimension of N-dimensional points.
Args:
points: torch.Tensor or numpy.ndarray with size (..., N+1).
Returns:
A torch.Tensor or numpy ndarray with size (..., N).
"""
return points[..., :-1] / (points[..., -1:] + eps)
def skew_symmetric(v):
"""Create a skew-symmetric matrix from a (batched) vector of size (..., 3)."""
z = torch.zeros_like(v[..., 0])
M = torch.stack(
[
z,
-v[..., 2],
v[..., 1],
v[..., 2],
z,
-v[..., 0],
-v[..., 1],
v[..., 0],
z,
],
dim=-1,
).reshape(v.shape[:-1] + (3, 3))
return M
def T_to_E(T):
"""Convert batched poses (..., 4, 4) to batched essential matrices."""
return skew_symmetric(T[..., :3, 3]) @ T[..., :3, :3]
def warp_points_torch(points, H, inverse=True):
"""
Warp a list of points with the INVERSE of the given homography.
The inverse is used to be coherent with tf.contrib.image.transform
Arguments:
points: batched list of N points, shape (B, N, 2).
homography: batched or not (shapes (B, 8) and (8,) respectively).
Returns: a Tensor of shape (B, N, 2) containing the new coordinates of the warped points.
"""
# H = np.expand_dims(homography, axis=0) if len(homography.shape) == 1 else homography
# Get the points to the homogeneous format
points = to_homogeneous(points)
# Apply the homography
out_shape = tuple(list(H.shape[:-1]) + [3, 3])
H_mat = torch.cat([H, torch.ones_like(H[..., :1])], axis=-1).reshape(out_shape)
if inverse:
H_mat = torch.inverse(H_mat)
warped_points = torch.einsum("...nj,...ji->...ni", points, H_mat.transpose(-2, -1))
warped_points = from_homogeneous(warped_points, eps=1e-5)
return warped_points
def seg_equation(segs):
# calculate list of start, end and midpoints points from both lists
start_points, end_points = to_homogeneous(segs[..., 0, :]), to_homogeneous(
segs[..., 1, :]
)
# Compute the line equations as ax + by + c = 0 , where x^2 + y^2 = 1
lines = torch.cross(start_points, end_points, dim=-1)
lines_norm = torch.sqrt(lines[..., 0] ** 2 + lines[..., 1] ** 2)[..., None]
assert torch.all(
lines_norm > 0
), "Error: trying to compute the equation of a line with a single point"
lines = lines / lines_norm
return lines
def is_inside_img(pts: torch.Tensor, img_shape: Tuple[int, int]):
h, w = img_shape
return (
(pts >= 0).all(dim=-1)
& (pts[..., 0] < w)
& (pts[..., 1] < h)
& (~torch.isinf(pts).any(dim=-1))
)
def shrink_segs_to_img(segs: torch.Tensor, img_shape: Tuple[int, int]) -> torch.Tensor:
"""
Shrink an array of segments to fit inside the image.
:param segs: The tensor of segments with shape (N, 2, 2)
:param img_shape: The image shape in format (H, W)
"""
EPS = 1e-4
device = segs.device
w, h = img_shape[1], img_shape[0]
# Project the segments to the reference image
segs = segs.clone()
eqs = seg_equation(segs)
x0, y0 = torch.tensor([1.0, 0, 0.0], device=device), torch.tensor(
[0.0, 1, 0], device=device
)
x0 = x0.repeat(eqs.shape[:-1] + (1,))
y0 = y0.repeat(eqs.shape[:-1] + (1,))
pt_x0s = torch.cross(eqs, x0, dim=-1)
pt_x0s = pt_x0s[..., :-1] / pt_x0s[..., None, -1]
pt_x0s_valid = is_inside_img(pt_x0s, img_shape)
pt_y0s = torch.cross(eqs, y0, dim=-1)
pt_y0s = pt_y0s[..., :-1] / pt_y0s[..., None, -1]
pt_y0s_valid = is_inside_img(pt_y0s, img_shape)
xW, yH = torch.tensor([1.0, 0, EPS - w], device=device), torch.tensor(
[0.0, 1, EPS - h], device=device
)
xW = xW.repeat(eqs.shape[:-1] + (1,))
yH = yH.repeat(eqs.shape[:-1] + (1,))
pt_xWs = torch.cross(eqs, xW, dim=-1)
pt_xWs = pt_xWs[..., :-1] / pt_xWs[..., None, -1]
pt_xWs_valid = is_inside_img(pt_xWs, img_shape)
pt_yHs = torch.cross(eqs, yH, dim=-1)
pt_yHs = pt_yHs[..., :-1] / pt_yHs[..., None, -1]
pt_yHs_valid = is_inside_img(pt_yHs, img_shape)
# If the X coordinate of the first endpoint is out
mask = (segs[..., 0, 0] < 0) & pt_x0s_valid
segs[mask, 0, :] = pt_x0s[mask]
mask = (segs[..., 0, 0] > (w - 1)) & pt_xWs_valid
segs[mask, 0, :] = pt_xWs[mask]
# If the X coordinate of the second endpoint is out
mask = (segs[..., 1, 0] < 0) & pt_x0s_valid
segs[mask, 1, :] = pt_x0s[mask]
mask = (segs[:, 1, 0] > (w - 1)) & pt_xWs_valid
segs[mask, 1, :] = pt_xWs[mask]
# If the Y coordinate of the first endpoint is out
mask = (segs[..., 0, 1] < 0) & pt_y0s_valid
segs[mask, 0, :] = pt_y0s[mask]
mask = (segs[..., 0, 1] > (h - 1)) & pt_yHs_valid
segs[mask, 0, :] = pt_yHs[mask]
# If the Y coordinate of the second endpoint is out
mask = (segs[..., 1, 1] < 0) & pt_y0s_valid
segs[mask, 1, :] = pt_y0s[mask]
mask = (segs[..., 1, 1] > (h - 1)) & pt_yHs_valid
segs[mask, 1, :] = pt_yHs[mask]
assert (
torch.all(segs >= 0)
and torch.all(segs[..., 0] < w)
and torch.all(segs[..., 1] < h)
)
return segs
def warp_lines_torch(
lines, H, inverse=True, dst_shape: Tuple[int, int] = None
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
:param lines: A tensor of shape (B, N, 2, 2) where B is the batch size, N the number of lines.
:param H: The homography used to convert the lines. batched or not (shapes (B, 8) and (8,) respectively).
:param inverse: Whether to apply H or the inverse of H
:param dst_shape:If provided, lines are trimmed to be inside the image
"""
device = lines.device
batch_size, n = lines.shape[:2]
lines = warp_points_torch(lines.reshape(batch_size, -1, 2), H, inverse).reshape(
lines.shape
)
if dst_shape is None:
return lines, torch.ones(lines.shape[:-2], dtype=torch.bool, device=device)
out_img = torch.any(
(lines < 0) | (lines >= torch.tensor(dst_shape[::-1], device=device)), -1
)
valid = ~out_img.all(-1)
any_out_of_img = out_img.any(-1)
lines_to_trim = valid & any_out_of_img
for b in range(batch_size):
lines_to_trim_mask_b = lines_to_trim[b]
lines_to_trim_b = lines[b][lines_to_trim_mask_b]
corrected_lines = shrink_segs_to_img(lines_to_trim_b, dst_shape)
lines[b][lines_to_trim_mask_b] = corrected_lines
return lines, valid