conversions.py 52 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597
  1. # LICENSE HEADER MANAGED BY add-license-header
  2. #
  3. # Copyright 2018 Kornia Team
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. #
  17. from __future__ import annotations
  18. from typing import Optional
  19. import torch
  20. import torch.nn.functional as F
  21. from kornia.constants import pi
  22. from kornia.core import Tensor, concatenate, cos, pad, sin, stack, tensor, where, zeros_like
  23. from kornia.core.check import KORNIA_CHECK, KORNIA_CHECK_SHAPE
  24. from kornia.utils import deprecated
  25. from kornia.utils.helpers import _torch_inverse_cast
  26. __all__ = [
  27. "ARKitQTVecs_to_ColmapQTVecs",
  28. "Rt_to_matrix4x4",
  29. "angle_axis_to_quaternion",
  30. "angle_axis_to_rotation_matrix",
  31. "angle_to_rotation_matrix",
  32. "axis_angle_to_quaternion",
  33. "axis_angle_to_rotation_matrix",
  34. "camtoworld_graphics_to_vision_4x4",
  35. "camtoworld_graphics_to_vision_Rt",
  36. "camtoworld_to_worldtocam_Rt",
  37. "camtoworld_vision_to_graphics_4x4",
  38. "camtoworld_vision_to_graphics_Rt",
  39. "cart2pol",
  40. "convert_affinematrix_to_homography",
  41. "convert_affinematrix_to_homography3d",
  42. "convert_points_from_homogeneous",
  43. "convert_points_to_homogeneous",
  44. "deg2rad",
  45. "denormalize_homography",
  46. "denormalize_pixel_coordinates",
  47. "denormalize_pixel_coordinates3d",
  48. "denormalize_points_with_intrinsics",
  49. "euler_from_quaternion",
  50. "matrix4x4_to_Rt",
  51. "normal_transform_pixel",
  52. "normal_transform_pixel3d",
  53. "normalize_homography",
  54. "normalize_homography3d",
  55. "normalize_pixel_coordinates",
  56. "normalize_pixel_coordinates3d",
  57. "normalize_points_with_intrinsics",
  58. "normalize_quaternion",
  59. "pol2cart",
  60. "quaternion_exp_to_log",
  61. "quaternion_from_euler",
  62. "quaternion_log_to_exp",
  63. "quaternion_to_angle_axis",
  64. "quaternion_to_axis_angle",
  65. "quaternion_to_rotation_matrix",
  66. "rad2deg",
  67. "rotation_matrix_to_angle_axis",
  68. "rotation_matrix_to_axis_angle",
  69. "rotation_matrix_to_quaternion",
  70. "vector_to_skew_symmetric_matrix",
  71. "worldtocam_to_camtoworld_Rt",
  72. ]
  73. def rad2deg(tensor: Tensor) -> Tensor:
  74. r"""Convert angles from radians to degrees.
  75. Args:
  76. tensor: Tensor of arbitrary shape.
  77. Returns:
  78. Tensor with same shape as input.
  79. Example:
  80. >>> input = tensor(3.1415926535)
  81. >>> rad2deg(input)
  82. tensor(180.)
  83. """
  84. if not isinstance(tensor, Tensor):
  85. raise TypeError(f"Input type is not a Tensor. Got {type(tensor)}")
  86. return 180.0 * tensor / pi.to(tensor.device).type(tensor.dtype)
  87. def deg2rad(tensor: Tensor) -> Tensor:
  88. r"""Convert angles from degrees to radians.
  89. Args:
  90. tensor: Tensor of arbitrary shape.
  91. Returns:
  92. tensor with same shape as input.
  93. Examples:
  94. >>> input = tensor(180.)
  95. >>> deg2rad(input)
  96. tensor(3.1416)
  97. """
  98. if not isinstance(tensor, Tensor):
  99. raise TypeError(f"Input type is not a Tensor. Got {type(tensor)}")
  100. return tensor * pi.to(tensor.device).type(tensor.dtype) / 180.0
  101. def pol2cart(rho: Tensor, phi: Tensor) -> tuple[Tensor, Tensor]:
  102. r"""Convert polar coordinates to cartesian coordinates.
  103. Args:
  104. rho: Tensor of arbitrary shape.
  105. phi: Tensor of same arbitrary shape.
  106. Returns:
  107. - x: Tensor with same shape as input.
  108. - y: Tensor with same shape as input.
  109. Example:
  110. >>> rho = torch.rand(1, 3, 3)
  111. >>> phi = torch.rand(1, 3, 3)
  112. >>> x, y = pol2cart(rho, phi)
  113. """
  114. if not (isinstance(rho, Tensor) & isinstance(phi, Tensor)):
  115. raise TypeError(f"Input type is not a Tensor. Got {type(rho)}, {type(phi)}")
  116. x = rho * cos(phi)
  117. y = rho * sin(phi)
  118. return x, y
  119. def cart2pol(x: Tensor, y: Tensor, eps: float = 1.0e-8) -> tuple[Tensor, Tensor]:
  120. """Convert cartesian coordinates to polar coordinates.
  121. Args:
  122. x: Tensor of arbitrary shape.
  123. y: Tensor of same arbitrary shape.
  124. eps: To avoid division by zero.
  125. Returns:
  126. - rho: Tensor with same shape as input.
  127. - phi: Tensor with same shape as input.
  128. Example:
  129. >>> x = torch.rand(1, 3, 3)
  130. >>> y = torch.rand(1, 3, 3)
  131. >>> rho, phi = cart2pol(x, y)
  132. """
  133. if not (isinstance(x, Tensor) & isinstance(y, Tensor)):
  134. raise TypeError(f"Input type is not a Tensor. Got {type(x)}, {type(y)}")
  135. rho = torch.sqrt(x**2 + y**2 + eps)
  136. phi = torch.atan2(y, x)
  137. return rho, phi
  138. def convert_points_from_homogeneous(points: Tensor, eps: float = 1e-8) -> Tensor:
  139. r"""Convert points from homogeneous to Euclidean space.
  140. Args:
  141. points: the points to be transformed of shape :math:`(B, N, D)`.
  142. eps: to avoid division by zero.
  143. Returns:
  144. the points in Euclidean space :math:`(B, N, D-1)`.
  145. Examples:
  146. >>> input = tensor([[0., 0., 1.]])
  147. >>> convert_points_from_homogeneous(input)
  148. tensor([[0., 0.]])
  149. """
  150. if not isinstance(points, Tensor):
  151. raise TypeError(f"Input type is not a Tensor. Got {type(points)}")
  152. if len(points.shape) < 2:
  153. raise ValueError(f"Input must be at least a 2D tensor. Got {points.shape}")
  154. # we check for points at max_val
  155. z_vec: Tensor = points[..., -1:]
  156. # set the results of division by zeror/near-zero to 1.0
  157. # follow the convention of opencv:
  158. # https://github.com/opencv/opencv/pull/14411/files
  159. mask: Tensor = torch.abs(z_vec) > eps
  160. scale = where(mask, 1.0 / (z_vec + eps), torch.ones_like(z_vec))
  161. return scale * points[..., :-1]
  162. def convert_points_to_homogeneous(points: Tensor) -> Tensor:
  163. r"""Convert points from Euclidean to homogeneous space.
  164. Args:
  165. points: the points to be transformed with shape :math:`(*, N, D)`.
  166. Returns:
  167. the points in homogeneous coordinates :math:`(*, N, D+1)`.
  168. Examples:
  169. >>> input = tensor([[0., 0.]])
  170. >>> convert_points_to_homogeneous(input)
  171. tensor([[0., 0., 1.]])
  172. """
  173. if not isinstance(points, Tensor):
  174. raise TypeError(f"Input type is not a Tensor. Got {type(points)}")
  175. if len(points.shape) < 2:
  176. raise ValueError(f"Input must be at least a 2D tensor. Got {points.shape}")
  177. return pad(points, [0, 1], "constant", 1.0)
  178. def _convert_affinematrix_to_homography_impl(A: Tensor) -> Tensor:
  179. H: Tensor = pad(A, [0, 0, 0, 1], "constant", value=0.0)
  180. H[..., -1, -1] += 1.0
  181. return H
  182. def convert_affinematrix_to_homography(A: Tensor) -> Tensor:
  183. r"""Convert batch of affine matrices.
  184. Args:
  185. A: the affine matrix with shape :math:`(B,2,3)`.
  186. Returns:
  187. the homography matrix with shape of :math:`(B,3,3)`.
  188. Examples:
  189. >>> A = tensor([[[1., 0., 0.],
  190. ... [0., 1., 0.]]])
  191. >>> convert_affinematrix_to_homography(A)
  192. tensor([[[1., 0., 0.],
  193. [0., 1., 0.],
  194. [0., 0., 1.]]])
  195. """
  196. if not isinstance(A, Tensor):
  197. raise TypeError(f"Input type is not a Tensor. Got {type(A)}")
  198. if not (len(A.shape) == 3 and A.shape[-2:] == (2, 3)):
  199. raise ValueError(f"Input matrix must be a Bx2x3 tensor. Got {A.shape}")
  200. return _convert_affinematrix_to_homography_impl(A)
  201. def convert_affinematrix_to_homography3d(A: Tensor) -> Tensor:
  202. r"""Convert batch of 3d affine matrices.
  203. Args:
  204. A: the affine matrix with shape :math:`(B,3,4)`.
  205. Returns:
  206. the homography matrix with shape of :math:`(B,4,4)`.
  207. Examples:
  208. >>> A = tensor([[[1., 0., 0., 0.],
  209. ... [0., 1., 0., 0.],
  210. ... [0., 0., 1., 0.]]])
  211. >>> convert_affinematrix_to_homography3d(A)
  212. tensor([[[1., 0., 0., 0.],
  213. [0., 1., 0., 0.],
  214. [0., 0., 1., 0.],
  215. [0., 0., 0., 1.]]])
  216. """
  217. if not isinstance(A, Tensor):
  218. raise TypeError(f"Input type is not a Tensor. Got {type(A)}")
  219. if not (len(A.shape) == 3 and A.shape[-2:] == (3, 4)):
  220. raise ValueError(f"Input matrix must be a Bx3x4 tensor. Got {A.shape}")
  221. return _convert_affinematrix_to_homography_impl(A)
  222. def axis_angle_to_rotation_matrix(axis_angle: Tensor) -> Tensor:
  223. r"""Convert 3d vector of axis-angle rotation to 3x3 rotation matrix.
  224. Args:
  225. axis_angle: tensor of 3d vector of axis-angle rotations in radians with shape :math:`(N, 3)`.
  226. Returns:
  227. tensor of rotation matrices of shape :math:`(N, 3, 3)`.
  228. Example:
  229. >>> input = tensor([[0., 0., 0.]])
  230. >>> axis_angle_to_rotation_matrix(input) # doctest: +ELLIPSIS
  231. tensor([[[1., ...0., 0.],
  232. [0., 1., ...0.],
  233. [...0., 0., 1.]]])
  234. >>> input = tensor([[1.5708, 0., 0.]])
  235. >>> axis_angle_to_rotation_matrix(input)
  236. tensor([[[ 1.0000e+00, 0.0000e+00, 0.0000e+00],
  237. [ 0.0000e+00, -3.6200e-06, -1.0000e+00],
  238. [ 0.0000e+00, 1.0000e+00, -3.6200e-06]]])
  239. """
  240. if not isinstance(axis_angle, Tensor):
  241. raise TypeError(f"Input type is not a Tensor. Got {type(axis_angle)}")
  242. if not axis_angle.shape[-1] == 3:
  243. raise ValueError(f"Input size must be a (*, 3) tensor. Got {axis_angle.shape}")
  244. def _compute_rotation_matrix(axis_angle: Tensor, theta2: Tensor, eps: float = 1e-6) -> Tensor:
  245. theta = torch.sqrt(theta2.clamp(min=1e-12)) # clamping to ensure no nan gradients
  246. wxyz = axis_angle / (theta.unsqueeze(-1) + eps) # (B, 3)
  247. wx, wy, wz = wxyz.unbind(dim=1) # (B,)
  248. cos_theta = torch.cos(theta)
  249. sin_theta = torch.sin(theta)
  250. one_minus_cos = 1.0 - cos_theta
  251. wxwy = wx * wy
  252. wxwz = wx * wz
  253. wywz = wy * wz
  254. r00 = cos_theta + wx * wx * one_minus_cos
  255. r01 = wxwy * one_minus_cos - wz * sin_theta
  256. r02 = wy * sin_theta + wxwz * one_minus_cos
  257. r10 = wz * sin_theta + wxwy * one_minus_cos
  258. r11 = cos_theta + wy * wy * one_minus_cos
  259. r12 = -wx * sin_theta + wywz * one_minus_cos
  260. r20 = -wy * sin_theta + wxwz * one_minus_cos
  261. r21 = wx * sin_theta + wywz * one_minus_cos
  262. r22 = cos_theta + wz * wz * one_minus_cos
  263. rot = torch.stack(
  264. [
  265. torch.stack([r00, r01, r02], dim=-1),
  266. torch.stack([r10, r11, r12], dim=-1),
  267. torch.stack([r20, r21, r22], dim=-1),
  268. ],
  269. dim=1,
  270. )
  271. return rot
  272. def _compute_rotation_matrix_taylor(axis_angle: Tensor) -> Tensor:
  273. rx, ry, rz = axis_angle.unbind(-1)
  274. k_one = torch.ones_like(rx)
  275. rot = torch.stack(
  276. [
  277. k_one,
  278. -rz,
  279. ry,
  280. rz,
  281. k_one,
  282. -rx,
  283. -ry,
  284. rx,
  285. k_one,
  286. ],
  287. dim=-1,
  288. ).view(-1, 3, 3)
  289. return rot
  290. theta2 = (axis_angle * axis_angle).sum(dim=-1)
  291. rot_normal = _compute_rotation_matrix(axis_angle, theta2) # (N,3,3)
  292. rot_taylor = _compute_rotation_matrix_taylor(axis_angle) # (N,3,3)
  293. mask = (theta2 > 1e-6).view(-1, 1, 1) # shape (N,1,1)
  294. rotation_matrix = torch.where(mask, rot_normal, rot_taylor)
  295. return rotation_matrix
  296. @deprecated(replace_with="axis_angle_to_rotation_matrix", version="0.7.0")
  297. def angle_axis_to_rotation_matrix(axis_angle: Tensor) -> Tensor: # noqa: D103
  298. return axis_angle_to_rotation_matrix(axis_angle)
  299. def rotation_matrix_to_axis_angle(rotation_matrix: Tensor) -> Tensor:
  300. r"""Convert 3x3 rotation matrix to Rodrigues vector in radians.
  301. Args:
  302. rotation_matrix: rotation matrix of shape :math:`(N, 3, 3)`.
  303. Returns:
  304. Rodrigues vector transformation of shape :math:`(N, 3)`.
  305. Example:
  306. >>> input = tensor([[1., 0., 0.],
  307. ... [0., 1., 0.],
  308. ... [0., 0., 1.]])
  309. >>> rotation_matrix_to_axis_angle(input)
  310. tensor([0., 0., 0.])
  311. >>> input = tensor([[1., 0., 0.],
  312. ... [0., 0., -1.],
  313. ... [0., 1., 0.]])
  314. >>> rotation_matrix_to_axis_angle(input)
  315. tensor([1.5708, 0.0000, 0.0000])
  316. """
  317. if not isinstance(rotation_matrix, Tensor):
  318. raise TypeError(f"Input type is not a Tensor. Got {type(rotation_matrix)}")
  319. if not rotation_matrix.shape[-2:] == (3, 3):
  320. raise ValueError(f"Input size must be a (*, 3, 3) tensor. Got {rotation_matrix.shape}")
  321. quaternion: Tensor = rotation_matrix_to_quaternion(rotation_matrix)
  322. return quaternion_to_axis_angle(quaternion)
  323. @deprecated(replace_with="rotation_matrix_to_axis_angle", version="0.7.0")
  324. def rotation_matrix_to_angle_axis(rotation_matrix: Tensor) -> Tensor: # noqa: D103
  325. return rotation_matrix_to_axis_angle(rotation_matrix)
  326. def rotation_matrix_to_quaternion(rotation_matrix: Tensor, eps: float = 1.0e-8) -> Tensor:
  327. r"""Convert 3x3 rotation matrix to 4d quaternion vector.
  328. The quaternion vector has components in (w, x, y, z) format.
  329. Args:
  330. rotation_matrix: the rotation matrix to convert with shape :math:`(*, 3, 3)`.
  331. eps: small value to avoid zero division.
  332. Return:
  333. the rotation in quaternion with shape :math:`(*, 4)`.
  334. Example:
  335. >>> input = tensor([[1., 0., 0.],
  336. ... [0., 1., 0.],
  337. ... [0., 0., 1.]])
  338. >>> rotation_matrix_to_quaternion(input, eps=torch.finfo(input.dtype).eps)
  339. tensor([1., 0., 0., 0.])
  340. """
  341. if not isinstance(rotation_matrix, Tensor):
  342. raise TypeError(f"Input type is not a Tensor. Got {type(rotation_matrix)}")
  343. if not rotation_matrix.shape[-2:] == (3, 3):
  344. raise ValueError(f"Input size must be a (*, 3, 3) tensor. Got {rotation_matrix.shape}")
  345. def safe_zero_division(numerator: Tensor, denominator: Tensor) -> Tensor:
  346. eps: float = torch.finfo(numerator.dtype).tiny
  347. return numerator / torch.clamp(denominator, min=eps)
  348. rotation_matrix_vec: Tensor = rotation_matrix.reshape(*rotation_matrix.shape[:-2], 9)
  349. m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.chunk(rotation_matrix_vec, chunks=9, dim=-1)
  350. trace: Tensor = m00 + m11 + m22
  351. def trace_positive_cond() -> Tensor:
  352. sq = torch.sqrt(trace + 1.0 + eps) * 2.0 # sq = 4 * qw.
  353. qw = 0.25 * sq
  354. qx = safe_zero_division(m21 - m12, sq)
  355. qy = safe_zero_division(m02 - m20, sq)
  356. qz = safe_zero_division(m10 - m01, sq)
  357. return concatenate((qw, qx, qy, qz), dim=-1)
  358. def cond_1() -> Tensor:
  359. sq = torch.sqrt(1.0 + m00 - m11 - m22 + eps) * 2.0 # sq = 4 * qx.
  360. qw = safe_zero_division(m21 - m12, sq)
  361. qx = 0.25 * sq
  362. qy = safe_zero_division(m01 + m10, sq)
  363. qz = safe_zero_division(m02 + m20, sq)
  364. return concatenate((qw, qx, qy, qz), dim=-1)
  365. def cond_2() -> Tensor:
  366. sq = torch.sqrt(1.0 + m11 - m00 - m22 + eps) * 2.0 # sq = 4 * qy.
  367. qw = safe_zero_division(m02 - m20, sq)
  368. qx = safe_zero_division(m01 + m10, sq)
  369. qy = 0.25 * sq
  370. qz = safe_zero_division(m12 + m21, sq)
  371. return concatenate((qw, qx, qy, qz), dim=-1)
  372. def cond_3() -> Tensor:
  373. sq = torch.sqrt(1.0 + m22 - m00 - m11 + eps) * 2.0 # sq = 4 * qz.
  374. qw = safe_zero_division(m10 - m01, sq)
  375. qx = safe_zero_division(m02 + m20, sq)
  376. qy = safe_zero_division(m12 + m21, sq)
  377. qz = 0.25 * sq
  378. return concatenate((qw, qx, qy, qz), dim=-1)
  379. where_2 = where(m11 > m22, cond_2(), cond_3())
  380. where_1 = where((m00 > m11) & (m00 > m22), cond_1(), where_2)
  381. quaternion: Tensor = where(trace > 0.0, trace_positive_cond(), where_1)
  382. return quaternion
  383. def normalize_quaternion(quaternion: Tensor, eps: float = 1.0e-12) -> Tensor:
  384. r"""Normalize a quaternion.
  385. The quaternion should be in (x, y, z, w) or (w, x, y, z) format.
  386. Args:
  387. quaternion: a tensor containing a quaternion to be normalized.
  388. The tensor can be of shape :math:`(*, 4)`.
  389. eps: small value to avoid division by zero.
  390. Return:
  391. the normalized quaternion of shape :math:`(*, 4)`.
  392. Example:
  393. >>> quaternion = tensor((1., 0., 1., 0.))
  394. >>> normalize_quaternion(quaternion)
  395. tensor([0.7071, 0.0000, 0.7071, 0.0000])
  396. """
  397. if not isinstance(quaternion, Tensor):
  398. raise TypeError(f"Input type is not a Tensor. Got {type(quaternion)}")
  399. if not quaternion.shape[-1] == 4:
  400. raise ValueError(f"Input must be a tensor of shape (*, 4). Got {quaternion.shape}")
  401. return F.normalize(quaternion, p=2.0, dim=-1, eps=eps)
  402. # based on:
  403. # https://github.com/matthew-brett/transforms3d/blob/8965c48401d9e8e66b6a8c37c65f2fc200a076fa/transforms3d/quaternions.py#L101
  404. # https://github.com/tensorflow/graphics/blob/master/tensorflow_graphics/geometry/transformation/rotation_matrix_3d.py#L247
  405. def quaternion_to_rotation_matrix(quaternion: Tensor) -> Tensor:
  406. r"""Convert a quaternion to a rotation matrix.
  407. The quaternion should be in (w, x, y, z) format.
  408. Args:
  409. quaternion: a tensor containing a quaternion to be converted.
  410. The tensor can be of shape :math:`(*, 4)`.
  411. Return:
  412. the rotation matrix of shape :math:`(*, 3, 3)`.
  413. Example:
  414. >>> quaternion = tensor((0., 0., 0., 1.))
  415. >>> quaternion_to_rotation_matrix(quaternion)
  416. tensor([[-1., 0., 0.],
  417. [ 0., -1., 0.],
  418. [ 0., 0., 1.]])
  419. """
  420. if not isinstance(quaternion, Tensor):
  421. raise TypeError(f"Input type is not a Tensor. Got {type(quaternion)}")
  422. if not quaternion.shape[-1] == 4:
  423. raise ValueError(f"Input must be a tensor of shape (*, 4). Got {quaternion.shape}")
  424. # normalize the input quaternion
  425. quaternion_norm: Tensor = normalize_quaternion(quaternion)
  426. # unpack the normalized quaternion components
  427. w = quaternion_norm[..., 0]
  428. x = quaternion_norm[..., 1]
  429. y = quaternion_norm[..., 2]
  430. z = quaternion_norm[..., 3]
  431. # compute the actual conversion
  432. tx: Tensor = 2.0 * x
  433. ty: Tensor = 2.0 * y
  434. tz: Tensor = 2.0 * z
  435. twx: Tensor = tx * w
  436. twy: Tensor = ty * w
  437. twz: Tensor = tz * w
  438. txx: Tensor = tx * x
  439. txy: Tensor = ty * x
  440. txz: Tensor = tz * x
  441. tyy: Tensor = ty * y
  442. tyz: Tensor = tz * y
  443. tzz: Tensor = tz * z
  444. one: Tensor = tensor(1.0)
  445. matrix_flat: Tensor = stack(
  446. (
  447. one - (tyy + tzz),
  448. txy - twz,
  449. txz + twy,
  450. txy + twz,
  451. one - (txx + tzz),
  452. tyz - twx,
  453. txz - twy,
  454. tyz + twx,
  455. one - (txx + tyy),
  456. ),
  457. dim=-1,
  458. )
  459. # this slightly awkward construction of the output shape is to satisfy torchscript
  460. output_shape = [*list(quaternion.shape[:-1]), 3, 3]
  461. matrix = matrix_flat.reshape(output_shape)
  462. return matrix
  463. def quaternion_to_axis_angle(quaternion: Tensor) -> Tensor:
  464. """Convert quaternion vector to axis angle of rotation in radians.
  465. The quaternion should be in (w, x, y, z) format.
  466. Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
  467. Args:
  468. quaternion: tensor with quaternions.
  469. Return:
  470. tensor with axis angle of rotation.
  471. Shape:
  472. - Input: :math:`(*, 4)` where `*` means, any number of dimensions
  473. - Output: :math:`(*, 3)`
  474. Example:
  475. >>> quaternion = tensor((1., 0., 0., 0.))
  476. >>> quaternion_to_axis_angle(quaternion)
  477. tensor([0., 0., 0.])
  478. """
  479. if not torch.is_tensor(quaternion):
  480. raise TypeError(f"Input type is not a Tensor. Got {type(quaternion)}")
  481. if not quaternion.shape[-1] == 4:
  482. raise ValueError(f"Input must be a tensor of shape Nx4 or 4. Got {quaternion.shape}")
  483. # unpack input and compute conversion
  484. q1: Tensor = tensor([])
  485. q2: Tensor = tensor([])
  486. q3: Tensor = tensor([])
  487. cos_theta: Tensor = tensor([])
  488. cos_theta = quaternion[..., 0]
  489. q1 = quaternion[..., 1]
  490. q2 = quaternion[..., 2]
  491. q3 = quaternion[..., 3]
  492. sin_squared_theta: Tensor = q1 * q1 + q2 * q2 + q3 * q3
  493. sin_theta: Tensor = torch.sqrt(sin_squared_theta)
  494. two_theta: Tensor = 2.0 * where(
  495. cos_theta < 0.0, torch.atan2(-sin_theta, -cos_theta), torch.atan2(sin_theta, cos_theta)
  496. )
  497. k_pos: Tensor = two_theta / sin_theta
  498. k_neg: Tensor = 2.0 * torch.ones_like(sin_theta)
  499. k: Tensor = where(sin_squared_theta > 0.0, k_pos, k_neg)
  500. axis_angle: Tensor = torch.zeros_like(quaternion)[..., :3]
  501. axis_angle[..., 0] += q1 * k
  502. axis_angle[..., 1] += q2 * k
  503. axis_angle[..., 2] += q3 * k
  504. return axis_angle
  505. @deprecated(replace_with="quaternion_to_axis_angle", version="0.7.0")
  506. def quaternion_to_angle_axis(quaternion: Tensor) -> Tensor: # noqa: D103
  507. return quaternion_to_axis_angle(quaternion)
  508. def quaternion_log_to_exp(quaternion: Tensor, eps: float = 1.0e-8) -> Tensor:
  509. r"""Apply exponential map to log quaternion.
  510. The quaternion should be in (w, x, y, z) format.
  511. Args:
  512. quaternion: a tensor containing a quaternion to be converted.
  513. The tensor can be of shape :math:`(*, 3)`.
  514. eps: a small number for clamping.
  515. Return:
  516. the quaternion exponential map of shape :math:`(*, 4)`.
  517. Example:
  518. >>> quaternion = tensor((0., 0., 0.))
  519. >>> quaternion_log_to_exp(quaternion, eps=torch.finfo(quaternion.dtype).eps)
  520. tensor([1., 0., 0., 0.])
  521. """
  522. if not isinstance(quaternion, Tensor):
  523. raise TypeError(f"Input type is not a Tensor. Got {type(quaternion)}")
  524. if not quaternion.shape[-1] == 3:
  525. raise ValueError(f"Input must be a tensor of shape (*, 3). Got {quaternion.shape}")
  526. # compute quaternion norm
  527. norm_q: Tensor = torch.norm(quaternion, p=2, dim=-1, keepdim=True).clamp(min=eps)
  528. # compute scalar and vector
  529. quaternion_vector: Tensor = quaternion * sin(norm_q) / norm_q
  530. quaternion_scalar: Tensor = cos(norm_q)
  531. # compose quaternion and return
  532. quaternion_exp: Tensor = tensor([])
  533. quaternion_exp = concatenate((quaternion_scalar, quaternion_vector), dim=-1)
  534. return quaternion_exp
  535. def quaternion_exp_to_log(quaternion: Tensor, eps: float = 1.0e-8) -> Tensor:
  536. r"""Apply the log map to a quaternion.
  537. The quaternion should be in (w, x, y, z) format.
  538. Args:
  539. quaternion: a tensor containing a quaternion to be converted.
  540. The tensor can be of shape :math:`(*, 4)`.
  541. eps: a small number for clamping.
  542. Return:
  543. the quaternion log map of shape :math:`(*, 3)`.
  544. Example:
  545. >>> quaternion = tensor((1., 0., 0., 0.))
  546. >>> quaternion_exp_to_log(quaternion, eps=torch.finfo(quaternion.dtype).eps)
  547. tensor([0., 0., 0.])
  548. """
  549. if not isinstance(quaternion, Tensor):
  550. raise TypeError(f"Input type is not a Tensor. Got {type(quaternion)}")
  551. if not quaternion.shape[-1] == 4:
  552. raise ValueError(f"Input must be a tensor of shape (*, 4). Got {quaternion.shape}")
  553. # unpack quaternion vector and scalar
  554. quaternion_vector: Tensor = tensor([])
  555. quaternion_scalar: Tensor = tensor([])
  556. quaternion_scalar = quaternion[..., 0:1]
  557. quaternion_vector = quaternion[..., 1:4]
  558. # compute quaternion norm
  559. norm_q: Tensor = torch.norm(quaternion_vector, p=2, dim=-1, keepdim=True).clamp(min=eps)
  560. # apply log map
  561. quaternion_log: Tensor = quaternion_vector * torch.acos(torch.clamp(quaternion_scalar, min=-1.0, max=1.0)) / norm_q
  562. return quaternion_log
  563. # based on:
  564. # https://github.com/facebookresearch/QuaterNet/blob/master/common/quaternion.py#L138
  565. def axis_angle_to_quaternion(axis_angle: Tensor) -> Tensor:
  566. r"""Convert an axis angle to a quaternion.
  567. The quaternion vector has components in (w, x, y, z) format.
  568. Adapted from ceres C++ library: ceres-solver/include/ceres/rotation.h
  569. Args:
  570. axis_angle: tensor with axis angle in radians.
  571. Return:
  572. tensor with quaternion.
  573. Shape:
  574. - Input: :math:`(*, 3)` where `*` means, any number of dimensions
  575. - Output: :math:`(*, 4)`
  576. Example:
  577. >>> axis_angle = tensor((0., 1., 0.))
  578. >>> axis_angle_to_quaternion(axis_angle)
  579. tensor([0.8776, 0.0000, 0.4794, 0.0000])
  580. """
  581. if not torch.is_tensor(axis_angle):
  582. raise TypeError(f"Input type is not a Tensor. Got {type(axis_angle)}")
  583. if not axis_angle.shape[-1] == 3:
  584. raise ValueError(f"Input must be a tensor of shape Nx3 or 3. Got {axis_angle.shape}")
  585. # unpack input and compute conversion
  586. a0: Tensor = axis_angle[..., 0:1]
  587. a1: Tensor = axis_angle[..., 1:2]
  588. a2: Tensor = axis_angle[..., 2:3]
  589. theta_squared: Tensor = a0 * a0 + a1 * a1 + a2 * a2
  590. theta: Tensor = torch.sqrt(theta_squared)
  591. half_theta: Tensor = theta * 0.5
  592. mask: Tensor = theta_squared > 0.0
  593. ones: Tensor = torch.ones_like(half_theta)
  594. k_neg: Tensor = 0.5 * ones
  595. k_pos: Tensor = sin(half_theta) / theta
  596. k: Tensor = where(mask, k_pos, k_neg)
  597. w: Tensor = where(mask, cos(half_theta), ones)
  598. quaternion: Tensor = torch.zeros(size=(*axis_angle.shape[:-1], 4), dtype=axis_angle.dtype, device=axis_angle.device)
  599. quaternion[..., 1:2] = a0 * k
  600. quaternion[..., 2:3] = a1 * k
  601. quaternion[..., 3:4] = a2 * k
  602. quaternion[..., 0:1] = w
  603. return quaternion
  604. @deprecated(replace_with="axis_angle_to_quaternion", version="0.7.0")
  605. def angle_axis_to_quaternion(axis_angle: Tensor) -> Tensor: # noqa: D103
  606. return axis_angle_to_quaternion(axis_angle)
  607. # inspired by: https://stackoverflow.com/questions/56207448/efficient-quaternions-to-euler-transformation
  608. def euler_from_quaternion(w: Tensor, x: Tensor, y: Tensor, z: Tensor) -> tuple[Tensor, Tensor, Tensor]:
  609. """Convert a quaternion coefficients to Euler angles.
  610. Returned angles are in radians in XYZ convention.
  611. Args:
  612. w: quaternion :math:`q_w` coefficient.
  613. x: quaternion :math:`q_x` coefficient.
  614. y: quaternion :math:`q_y` coefficient.
  615. z: quaternion :math:`q_z` coefficient.
  616. Return:
  617. A tuple with euler angles`roll`, `pitch`, `yaw`.
  618. """
  619. KORNIA_CHECK(w.shape == x.shape)
  620. KORNIA_CHECK(x.shape == y.shape)
  621. KORNIA_CHECK(y.shape == z.shape)
  622. yy = y * y
  623. sinr_cosp = 2.0 * (w * x + y * z)
  624. cosr_cosp = 1.0 - 2.0 * (x * x + yy)
  625. roll = sinr_cosp.atan2(cosr_cosp)
  626. sinp = 2.0 * (w * y - z * x)
  627. sinp = sinp.clamp(min=-1.0, max=1.0)
  628. pitch = sinp.asin()
  629. siny_cosp = 2.0 * (w * z + x * y)
  630. cosy_cosp = 1.0 - 2.0 * (yy + z * z)
  631. yaw = siny_cosp.atan2(cosy_cosp)
  632. return roll, pitch, yaw
  633. def quaternion_from_euler(roll: Tensor, pitch: Tensor, yaw: Tensor) -> tuple[Tensor, Tensor, Tensor, Tensor]:
  634. """Convert Euler angles to quaternion coefficients.
  635. Euler angles are assumed to be in radians in XYZ convention.
  636. Args:
  637. roll: the roll euler angle.
  638. pitch: the pitch euler angle.
  639. yaw: the yaw euler angle.
  640. Return:
  641. A tuple with quaternion coefficients in order of `wxyz`.
  642. """
  643. KORNIA_CHECK(roll.shape == pitch.shape)
  644. KORNIA_CHECK(pitch.shape == yaw.shape)
  645. roll_half = roll * 0.5
  646. pitch_half = pitch * 0.5
  647. yaw_half = yaw * 0.5
  648. cy = yaw_half.cos()
  649. sy = yaw_half.sin()
  650. cp = pitch_half.cos()
  651. sp = pitch_half.sin()
  652. cr = roll_half.cos()
  653. sr = roll_half.sin()
  654. qw = cy * cp * cr + sy * sp * sr
  655. qx = cy * cp * sr - sy * sp * cr
  656. qy = sy * cp * sr + cy * sp * cr
  657. qz = sy * cp * cr - cy * sp * sr
  658. return qw, qx, qy, qz
  659. # based on:
  660. # https://github.com/ClementPinard/SfmLearner-Pytorch/blob/master/inverse_warp.py#L65-L71
  661. def normalize_pixel_coordinates(pixel_coordinates: Tensor, height: int, width: int, eps: float = 1e-8) -> Tensor:
  662. r"""Normalize pixel coordinates between -1 and 1.
  663. Normalized, -1 if on extreme left, 1 if on extreme right (x = w-1).
  664. Args:
  665. pixel_coordinates: the grid with pixel coordinates. Shape can be :math:`(*, 2)`.
  666. width: the maximum width in the x-axis.
  667. height: the maximum height in the y-axis.
  668. eps: safe division by zero.
  669. Return:
  670. the normalized pixel coordinates with shape :math:`(*, 2)`.
  671. Examples:
  672. >>> coords = tensor([[50., 100.]])
  673. >>> normalize_pixel_coordinates(coords, 100, 50)
  674. tensor([[1.0408, 1.0202]])
  675. """
  676. if pixel_coordinates.shape[-1] != 2:
  677. raise ValueError(f"Input pixel_coordinates must be of shape (*, 2). Got {pixel_coordinates.shape}")
  678. # compute normalization factor
  679. hw: Tensor = stack(
  680. [
  681. tensor(width, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype),
  682. tensor(height, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype),
  683. ]
  684. )
  685. factor: Tensor = tensor(2.0, device=pixel_coordinates.device, dtype=pixel_coordinates.dtype) / (hw - 1).clamp(eps)
  686. return factor * pixel_coordinates - 1
  687. def denormalize_pixel_coordinates(pixel_coordinates: Tensor, height: int, width: int, eps: float = 1e-8) -> Tensor:
  688. r"""Denormalize pixel coordinates.
  689. The input is assumed to be -1 if on extreme left, 1 if on extreme right (x = w-1).
  690. Args:
  691. pixel_coordinates: the normalized grid coordinates. Shape can be :math:`(*, 2)`.
  692. width: the maximum width in the x-axis.
  693. height: the maximum height in the y-axis.
  694. eps: safe division by zero.
  695. Return:
  696. the denormalized pixel coordinates with shape :math:`(*, 2)`.
  697. Examples:
  698. >>> coords = tensor([[-1., -1.]])
  699. >>> denormalize_pixel_coordinates(coords, 100, 50)
  700. tensor([[0., 0.]])
  701. """
  702. if pixel_coordinates.shape[-1] != 2:
  703. raise ValueError(f"Input pixel_coordinates must be of shape (*, 2). Got {pixel_coordinates.shape}")
  704. # compute normalization factor
  705. hw: Tensor = stack([tensor(width), tensor(height)]).to(pixel_coordinates.device).to(pixel_coordinates.dtype)
  706. factor: Tensor = tensor(2.0) / (hw - 1).clamp(eps)
  707. return tensor(1.0) / factor * (pixel_coordinates + 1)
  708. def normalize_pixel_coordinates3d(
  709. pixel_coordinates: Tensor, depth: int, height: int, width: int, eps: float = 1e-8
  710. ) -> Tensor:
  711. r"""Normalize pixel coordinates between -1 and 1.
  712. Normalized, -1 if on extreme left, 1 if on extreme right (x = w-1).
  713. Args:
  714. pixel_coordinates: the grid with pixel coordinates. Shape can be :math:`(*, 3)`.
  715. depth: the maximum depth in the z-axis.
  716. height: the maximum height in the y-axis.
  717. width: the maximum width in the x-axis.
  718. eps: safe division by zero.
  719. Return:
  720. the normalized pixel coordinates.
  721. """
  722. if pixel_coordinates.shape[-1] != 3:
  723. raise ValueError(f"Input pixel_coordinates must be of shape (*, 3). Got {pixel_coordinates.shape}")
  724. # compute normalization factor
  725. dhw: Tensor = (
  726. stack([tensor(depth), tensor(width), tensor(height)]).to(pixel_coordinates.device).to(pixel_coordinates.dtype)
  727. )
  728. factor: Tensor = tensor(2.0) / (dhw - 1).clamp(eps)
  729. return factor * pixel_coordinates - 1
  730. def denormalize_pixel_coordinates3d(
  731. pixel_coordinates: Tensor, depth: int, height: int, width: int, eps: float = 1e-8
  732. ) -> Tensor:
  733. r"""Denormalize pixel coordinates.
  734. The input is assumed to be -1 if on extreme left, 1 if on extreme right (x = w-1).
  735. Args:
  736. pixel_coordinates: the normalized grid coordinates. Shape can be :math:`(*, 3)`.
  737. depth: the maximum depth in the x-axis.
  738. height: the maximum height in the y-axis.
  739. width: the maximum width in the x-axis.
  740. eps: safe division by zero.
  741. Return:
  742. the denormalized pixel coordinates.
  743. """
  744. if pixel_coordinates.shape[-1] != 3:
  745. raise ValueError(f"Input pixel_coordinates must be of shape (*, 3). Got {pixel_coordinates.shape}")
  746. # compute normalization factor
  747. dhw: Tensor = (
  748. stack([tensor(depth), tensor(width), tensor(height)]).to(pixel_coordinates.device).to(pixel_coordinates.dtype)
  749. )
  750. factor: Tensor = tensor(2.0) / (dhw - 1).clamp(eps)
  751. return tensor(1.0) / factor * (pixel_coordinates + 1)
  752. def angle_to_rotation_matrix(angle: Tensor) -> Tensor:
  753. r"""Create a rotation matrix out of angles in degrees.
  754. Args:
  755. angle: tensor of angles in degrees, any shape :math:`(*)`.
  756. Returns:
  757. tensor of rotation matrices with shape :math:`(*, 2, 2)`.
  758. Example:
  759. >>> input = torch.rand(1, 3) # Nx3
  760. >>> output = angle_to_rotation_matrix(input) # Nx3x2x2
  761. """
  762. ang_rad = deg2rad(angle)
  763. cos_a: Tensor = cos(ang_rad)
  764. sin_a: Tensor = sin(ang_rad)
  765. return stack([cos_a, sin_a, -sin_a, cos_a], dim=-1).view(*angle.shape, 2, 2)
  766. def normalize_homography(
  767. dst_pix_trans_src_pix: Tensor, dsize_src: tuple[int, int], dsize_dst: tuple[int, int]
  768. ) -> Tensor:
  769. r"""Normalize a given homography in pixels to [-1, 1].
  770. Args:
  771. dst_pix_trans_src_pix: homography/ies from source to destination to be
  772. normalized. :math:`(B, 3, 3)`
  773. dsize_src: size of the source image (height, width).
  774. dsize_dst: size of the destination image (height, width).
  775. Returns:
  776. the normalized homography of shape :math:`(B, 3, 3)`.
  777. """
  778. if not isinstance(dst_pix_trans_src_pix, Tensor):
  779. raise TypeError(f"Input type is not a Tensor. Got {type(dst_pix_trans_src_pix)}")
  780. if not (len(dst_pix_trans_src_pix.shape) == 3 or dst_pix_trans_src_pix.shape[-2:] == (3, 3)):
  781. raise ValueError(f"Input dst_pix_trans_src_pix must be a Bx3x3 tensor. Got {dst_pix_trans_src_pix.shape}")
  782. # source and destination sizes
  783. src_h, src_w = dsize_src
  784. dst_h, dst_w = dsize_dst
  785. # compute the transformation pixel/norm for src/dst
  786. src_norm_trans_src_pix: Tensor = normal_transform_pixel(src_h, src_w).to(dst_pix_trans_src_pix)
  787. src_pix_trans_src_norm = _torch_inverse_cast(src_norm_trans_src_pix)
  788. dst_norm_trans_dst_pix: Tensor = normal_transform_pixel(dst_h, dst_w).to(dst_pix_trans_src_pix)
  789. # compute chain transformations
  790. dst_norm_trans_src_norm: Tensor = dst_norm_trans_dst_pix @ (dst_pix_trans_src_pix @ src_pix_trans_src_norm)
  791. return dst_norm_trans_src_norm
  792. def normal_transform_pixel(
  793. height: int,
  794. width: int,
  795. eps: float = 1e-14,
  796. device: Optional[torch.device] = None,
  797. dtype: Optional[torch.dtype] = None,
  798. ) -> Tensor:
  799. r"""Compute the normalization matrix from image size in pixels to [-1, 1].
  800. Args:
  801. height: image height.
  802. width: image width.
  803. eps: epsilon to prevent divide-by-zero errors
  804. device: device to place the result on.
  805. dtype: dtype of the result.
  806. Returns:
  807. normalized transform with shape :math:`(1, 3, 3)`.
  808. """
  809. tr_mat = tensor([[1.0, 0.0, -1.0], [0.0, 1.0, -1.0], [0.0, 0.0, 1.0]], device=device, dtype=dtype) # 3x3
  810. # prevent divide by zero bugs
  811. width_denom: float = eps if width == 1 else width - 1.0
  812. height_denom: float = eps if height == 1 else height - 1.0
  813. tr_mat[0, 0] = tr_mat[0, 0] * 2.0 / width_denom
  814. tr_mat[1, 1] = tr_mat[1, 1] * 2.0 / height_denom
  815. return tr_mat.unsqueeze(0) # 1x3x3
  816. def normal_transform_pixel3d(
  817. depth: int,
  818. height: int,
  819. width: int,
  820. eps: float = 1e-14,
  821. device: Optional[torch.device] = None,
  822. dtype: Optional[torch.dtype] = None,
  823. ) -> Tensor:
  824. r"""Compute the normalization matrix from image size in pixels to [-1, 1].
  825. Args:
  826. depth: image depth.
  827. height: image height.
  828. width: image width.
  829. eps: epsilon to prevent divide-by-zero errors
  830. device: device to place the result on.
  831. dtype: dtype of the result.
  832. Returns:
  833. normalized transform with shape :math:`(1, 4, 4)`.
  834. """
  835. tr_mat = tensor(
  836. [[1.0, 0.0, 0.0, -1.0], [0.0, 1.0, 0.0, -1.0], [0.0, 0.0, 1.0, -1.0], [0.0, 0.0, 0.0, 1.0]],
  837. device=device,
  838. dtype=dtype,
  839. ) # 4x4
  840. # prevent divide by zero bugs
  841. width_denom: float = eps if width == 1 else width - 1.0
  842. height_denom: float = eps if height == 1 else height - 1.0
  843. depth_denom: float = eps if depth == 1 else depth - 1.0
  844. tr_mat[0, 0] = tr_mat[0, 0] * 2.0 / width_denom
  845. tr_mat[1, 1] = tr_mat[1, 1] * 2.0 / height_denom
  846. tr_mat[2, 2] = tr_mat[2, 2] * 2.0 / depth_denom
  847. return tr_mat.unsqueeze(0) # 1x4x4
  848. def denormalize_homography(
  849. dst_pix_trans_src_pix: Tensor, dsize_src: tuple[int, int], dsize_dst: tuple[int, int]
  850. ) -> Tensor:
  851. r"""De-normalize a given homography in pixels from [-1, 1] to actual height and width.
  852. Args:
  853. dst_pix_trans_src_pix: homography/ies from source to destination to be
  854. denormalized. :math:`(B, 3, 3)`
  855. dsize_src: size of the source image (height, width).
  856. dsize_dst: size of the destination image (height, width).
  857. Returns:
  858. the denormalized homography of shape :math:`(B, 3, 3)`.
  859. """
  860. if not isinstance(dst_pix_trans_src_pix, Tensor):
  861. raise TypeError(f"Input type is not a Tensor. Got {type(dst_pix_trans_src_pix)}")
  862. if not (len(dst_pix_trans_src_pix.shape) == 3 or dst_pix_trans_src_pix.shape[-2:] == (3, 3)):
  863. raise ValueError(f"Input dst_pix_trans_src_pix must be a Bx3x3 tensor. Got {dst_pix_trans_src_pix.shape}")
  864. # source and destination sizes
  865. src_h, src_w = dsize_src
  866. dst_h, dst_w = dsize_dst
  867. # compute the transformation pixel/norm for src/dst
  868. src_norm_trans_src_pix: Tensor = normal_transform_pixel(src_h, src_w).to(dst_pix_trans_src_pix)
  869. dst_norm_trans_dst_pix: Tensor = normal_transform_pixel(dst_h, dst_w).to(dst_pix_trans_src_pix)
  870. dst_denorm_trans_dst_pix = _torch_inverse_cast(dst_norm_trans_dst_pix)
  871. # compute chain transformations
  872. dst_norm_trans_src_norm: Tensor = dst_denorm_trans_dst_pix @ (dst_pix_trans_src_pix @ src_norm_trans_src_pix)
  873. return dst_norm_trans_src_norm
  874. def normalize_homography3d(
  875. dst_pix_trans_src_pix: Tensor, dsize_src: tuple[int, int, int], dsize_dst: tuple[int, int, int]
  876. ) -> Tensor:
  877. r"""Normalize a given homography in pixels to [-1, 1].
  878. Args:
  879. dst_pix_trans_src_pix: homography/ies from source to destination to be
  880. normalized. :math:`(B, 4, 4)`
  881. dsize_src: size of the source image (depth, height, width).
  882. dsize_dst: size of the destination image (depth, height, width).
  883. Returns:
  884. the normalized homography.
  885. Shape:
  886. Output: :math:`(B, 4, 4)`
  887. """
  888. if not isinstance(dst_pix_trans_src_pix, Tensor):
  889. raise TypeError(f"Input type is not a Tensor. Got {type(dst_pix_trans_src_pix)}")
  890. if not (len(dst_pix_trans_src_pix.shape) == 3 or dst_pix_trans_src_pix.shape[-2:] == (4, 4)):
  891. raise ValueError(f"Input dst_pix_trans_src_pix must be a Bx3x3 tensor. Got {dst_pix_trans_src_pix.shape}")
  892. # source and destination sizes
  893. src_d, src_h, src_w = dsize_src
  894. dst_d, dst_h, dst_w = dsize_dst
  895. # compute the transformation pixel/norm for src/dst
  896. src_norm_trans_src_pix: Tensor = normal_transform_pixel3d(src_d, src_h, src_w).to(dst_pix_trans_src_pix)
  897. src_pix_trans_src_norm = _torch_inverse_cast(src_norm_trans_src_pix)
  898. dst_norm_trans_dst_pix: Tensor = normal_transform_pixel3d(dst_d, dst_h, dst_w).to(dst_pix_trans_src_pix)
  899. # compute chain transformations
  900. dst_norm_trans_src_norm: Tensor = dst_norm_trans_dst_pix @ (dst_pix_trans_src_pix @ src_pix_trans_src_norm)
  901. return dst_norm_trans_src_norm
  902. def normalize_points_with_intrinsics(point_2d: Tensor, camera_matrix: Tensor) -> Tensor:
  903. """Normalize points with intrinsics. Useful for conversion of keypoints to be used with essential matrix.
  904. Args:
  905. point_2d: tensor containing the 2d points in the image pixel coordinates. The shape of the tensor can be
  906. :math:`(*, 2)`.
  907. camera_matrix: tensor containing the intrinsics camera matrix. The tensor shape must be :math:`(*, 3, 3)`.
  908. Returns:
  909. tensor of (u, v) cam coordinates with shape :math:`(*, 2)`.
  910. Example:
  911. >>> _ = torch.manual_seed(0)
  912. >>> X = torch.rand(1, 2)
  913. >>> K = torch.eye(3)[None]
  914. >>> normalize_points_with_intrinsics(X, K)
  915. tensor([[0.4963, 0.7682]])
  916. """
  917. KORNIA_CHECK_SHAPE(point_2d, ["*", "2"])
  918. KORNIA_CHECK_SHAPE(camera_matrix, ["*", "3", "3"])
  919. # projection eq. K_inv * [u v 1]'
  920. # x = (u - cx) * Z / fx
  921. # y = (v - cy) * Z / fy
  922. # unpack coordinates
  923. cxcy = camera_matrix[..., :2, 2]
  924. fxfy = camera_matrix[..., :2, :2].diagonal(dim1=-2, dim2=-1)
  925. if len(cxcy.shape) < len(point_2d.shape): # broadcast intrinsics:
  926. cxcy, fxfy = cxcy.unsqueeze(-2), fxfy.unsqueeze(-2)
  927. xy = (point_2d - cxcy) / fxfy
  928. return xy
  929. def denormalize_points_with_intrinsics(point_2d_norm: Tensor, camera_matrix: Tensor) -> Tensor:
  930. """Normalize points with intrinsics. Useful for conversion of keypoints to be used with essential matrix.
  931. Args:
  932. point_2d_norm: tensor containing the 2d points in the image pixel coordinates. The shape of the tensor can be
  933. :math:`(*, 2)`.
  934. camera_matrix: tensor containing the intrinsics camera matrix. The tensor shape must be :math:`(*, 3, 3)`.
  935. Returns:
  936. tensor of (u, v) cam coordinates with shape :math:`(*, 2)`.
  937. Example:
  938. >>> _ = torch.manual_seed(0)
  939. >>> X = torch.rand(1, 2)
  940. >>> K = torch.eye(3)[None]
  941. >>> denormalize_points_with_intrinsics(X, K)
  942. tensor([[0.4963, 0.7682]])
  943. """
  944. KORNIA_CHECK_SHAPE(point_2d_norm, ["*", "2"])
  945. KORNIA_CHECK_SHAPE(camera_matrix, ["*", "3", "3"])
  946. # projection eq. [u, v, w]' = K * [x y z 1]'
  947. # u = fx * X + cx
  948. # v = fy * Y + cy
  949. # unpack coordinates
  950. x_coord: Tensor = point_2d_norm[..., 0]
  951. y_coord: Tensor = point_2d_norm[..., 1]
  952. # unpack intrinsics
  953. fx: Tensor = camera_matrix[..., 0, 0]
  954. fy: Tensor = camera_matrix[..., 1, 1]
  955. cx: Tensor = camera_matrix[..., 0, 2]
  956. cy: Tensor = camera_matrix[..., 1, 2]
  957. if len(cx.shape) < len(x_coord.shape): # broadcast intrinsics
  958. cx, cy, fx, fy = cx.unsqueeze(-1), cy.unsqueeze(-1), fx.unsqueeze(-1), fy.unsqueeze(-1)
  959. # apply intrinsics ans return
  960. u_coord: Tensor = x_coord * fx + cx
  961. v_coord: Tensor = y_coord * fy + cy
  962. return stack([u_coord, v_coord], dim=-1)
  963. def Rt_to_matrix4x4(R: Tensor, t: Tensor) -> Tensor:
  964. r"""Combine 3x3 rotation matrix R and 1x3 translation vector t into 4x4 extrinsics.
  965. Args:
  966. R: Rotation matrix, :math:`(B, 3, 3).`
  967. t: Translation matrix :math:`(B, 3, 1)`.
  968. Returns:
  969. the extrinsics :math:`(B, 4, 4)`.
  970. Example:
  971. >>> R, t = torch.eye(3)[None], torch.ones(3).reshape(1, 3, 1)
  972. >>> Rt_to_matrix4x4(R, t)
  973. tensor([[[1., 0., 0., 1.],
  974. [0., 1., 0., 1.],
  975. [0., 0., 1., 1.],
  976. [0., 0., 0., 1.]]])
  977. """
  978. KORNIA_CHECK_SHAPE(R, ["B", "3", "3"])
  979. KORNIA_CHECK_SHAPE(t, ["B", "3", "1"])
  980. Rt = concatenate([R, t], dim=2)
  981. return convert_affinematrix_to_homography3d(Rt)
  982. def matrix4x4_to_Rt(extrinsics: Tensor) -> tuple[Tensor, Tensor]:
  983. r"""Convert 4x4 extrinsics into 3x3 rotation matrix R and 1x3 translation vector ts.
  984. Args:
  985. extrinsics: pose matrix :math:`(B, 4, 4)`.
  986. Returns:
  987. R: Rotation matrix, :math:`(B, 3, 3).`
  988. t: Translation matrix :math:`(B, 3, 1)`.
  989. Example:
  990. >>> ext = torch.eye(4)[None]
  991. >>> matrix4x4_to_Rt(ext)
  992. (tensor([[[1., 0., 0.],
  993. [0., 1., 0.],
  994. [0., 0., 1.]]]), tensor([[[0.],
  995. [0.],
  996. [0.]]]))
  997. """
  998. KORNIA_CHECK_SHAPE(extrinsics, ["B", "4", "4"])
  999. R, t = extrinsics[:, :3, :3], extrinsics[:, :3, 3:]
  1000. return R, t
  1001. def camtoworld_graphics_to_vision_4x4(extrinsics_graphics: Tensor) -> Tensor:
  1002. r"""Convert graphics coordinate frame (e.g. OpenGL) to vision coordinate frame (e.g. OpenCV.).
  1003. I.e. flips y and z axis. Graphics convention: [+x, +y, +z] == [right, up, backwards].
  1004. Vision convention: [+x, +y, +z] == [right, down, forwards].
  1005. Args:
  1006. extrinsics_graphics: pose matrix :math:`(B, 4, 4)`.
  1007. Returns:
  1008. extrinsics: pose matrix :math:`(B, 4, 4)`.
  1009. Example:
  1010. >>> ext = torch.eye(4)[None]
  1011. >>> camtoworld_graphics_to_vision_4x4(ext)
  1012. tensor([[[ 1., 0., 0., 0.],
  1013. [ 0., -1., 0., 0.],
  1014. [ 0., 0., -1., 0.],
  1015. [ 0., 0., 0., 1.]]])
  1016. """
  1017. KORNIA_CHECK_SHAPE(extrinsics_graphics, ["B", "4", "4"])
  1018. invert_yz = tensor(
  1019. [[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1.0]]],
  1020. dtype=extrinsics_graphics.dtype,
  1021. device=extrinsics_graphics.device,
  1022. )
  1023. return extrinsics_graphics @ invert_yz
  1024. def camtoworld_graphics_to_vision_Rt(R: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
  1025. r"""Convert graphics coordinate frame (e.g. OpenGL) to vision coordinate frame (e.g. OpenCV.).
  1026. I.e. flips y and z axis. Graphics convention: [+x, +y, +z] == [right, up, backwards].
  1027. Vision convention: [+x, +y, +z] == [right, down, forwards].
  1028. Args:
  1029. R: Rotation matrix, :math:`(B, 3, 3).`
  1030. t: Translation matrix :math:`(B, 3, 1)`.
  1031. Returns:
  1032. R: Rotation matrix, :math:`(B, 3, 3).`
  1033. t: Translation matrix :math:`(B, 3, 1)`.
  1034. Example:
  1035. >>> R, t = torch.eye(3)[None], torch.ones(3).reshape(1, 3, 1)
  1036. >>> camtoworld_graphics_to_vision_Rt(R, t)
  1037. (tensor([[[ 1., 0., 0.],
  1038. [ 0., -1., 0.],
  1039. [ 0., 0., -1.]]]), tensor([[[1.],
  1040. [1.],
  1041. [1.]]]))
  1042. """
  1043. KORNIA_CHECK_SHAPE(R, ["B", "3", "3"])
  1044. KORNIA_CHECK_SHAPE(t, ["B", "3", "1"])
  1045. mat4x4 = camtoworld_graphics_to_vision_4x4(Rt_to_matrix4x4(R, t))
  1046. return matrix4x4_to_Rt(mat4x4)
  1047. def camtoworld_vision_to_graphics_4x4(extrinsics_vision: Tensor) -> Tensor:
  1048. r"""Convert vision coordinate frame (e.g. OpenCV) to graphics coordinate frame (e.g. OpenGK.).
  1049. I.e. flips y and z axis Graphics convention: [+x, +y, +z] == [right, up, backwards].
  1050. Vision convention: [+x, +y, +z] == [right, down, forwards].
  1051. Args:
  1052. extrinsics_vision: pose matrix :math:`(B, 4, 4)`.
  1053. Returns:
  1054. extrinsics: pose matrix :math:`(B, 4, 4)`.
  1055. Example:
  1056. >>> ext = torch.eye(4)[None]
  1057. >>> camtoworld_vision_to_graphics_4x4(ext)
  1058. tensor([[[ 1., 0., 0., 0.],
  1059. [ 0., -1., 0., 0.],
  1060. [ 0., 0., -1., 0.],
  1061. [ 0., 0., 0., 1.]]])
  1062. """
  1063. KORNIA_CHECK_SHAPE(extrinsics_vision, ["B", "4", "4"])
  1064. invert_yz = tensor(
  1065. [[[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1.0]]],
  1066. dtype=extrinsics_vision.dtype,
  1067. device=extrinsics_vision.device,
  1068. )
  1069. return extrinsics_vision @ invert_yz
  1070. def camtoworld_vision_to_graphics_Rt(R: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
  1071. r"""Convert graphics coordinate frame (e.g. OpenGL) to vision coordinate frame (e.g. OpenCV.).
  1072. I.e. flips y and z axis. Graphics convention: [+x, +y, +z] == [right, up, backwards].
  1073. Vision convention: [+x, +y, +z] == [right, down, forwards]
  1074. Args:
  1075. R: Rotation matrix, :math:`(B, 3, 3).`
  1076. t: Translation matrix :math:`(B, 3, 1)`.
  1077. Returns:
  1078. R: Rotation matrix, :math:`(B, 3, 3).`
  1079. t: Translation matrix :math:`(B, 3, 1)`.
  1080. Example:
  1081. >>> R, t = torch.eye(3)[None], torch.ones(3).reshape(1, 3, 1)
  1082. >>> camtoworld_vision_to_graphics_Rt(R, t)
  1083. (tensor([[[ 1., 0., 0.],
  1084. [ 0., -1., 0.],
  1085. [ 0., 0., -1.]]]), tensor([[[1.],
  1086. [1.],
  1087. [1.]]]))
  1088. """
  1089. KORNIA_CHECK_SHAPE(R, ["B", "3", "3"])
  1090. KORNIA_CHECK_SHAPE(t, ["B", "3", "1"])
  1091. mat4x4 = camtoworld_vision_to_graphics_4x4(Rt_to_matrix4x4(R, t))
  1092. return matrix4x4_to_Rt(mat4x4)
  1093. def camtoworld_to_worldtocam_Rt(R: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
  1094. r"""Convert camtoworld to worldtocam frame used in Colmap.
  1095. See
  1096. long-url: https://colmap.github.io/format.html#output-format
  1097. Args:
  1098. R: Rotation matrix, :math:`(B, 3, 3).`
  1099. t: Translation matrix :math:`(B, 3, 1)`.
  1100. Returns:
  1101. Rinv: Rotation matrix, :math:`(B, 3, 3).`
  1102. tinv: Translation matrix :math:`(B, 3, 1)`.
  1103. Example:
  1104. >>> R, t = torch.eye(3)[None], torch.ones(3).reshape(1, 3, 1)
  1105. >>> camtoworld_to_worldtocam_Rt(R, t)
  1106. (tensor([[[1., 0., 0.],
  1107. [0., 1., 0.],
  1108. [0., 0., 1.]]]), tensor([[[-1.],
  1109. [-1.],
  1110. [-1.]]]))
  1111. """
  1112. KORNIA_CHECK_SHAPE(R, ["B", "3", "3"])
  1113. KORNIA_CHECK_SHAPE(t, ["B", "3", "1"])
  1114. R_inv = R.transpose(1, 2)
  1115. new_t: Tensor = -R_inv @ t
  1116. return (R_inv, new_t)
  1117. def worldtocam_to_camtoworld_Rt(R: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
  1118. r"""Convert worldtocam frame used in Colmap to camtoworld.
  1119. Args:
  1120. R: Rotation matrix, :math:`(B, 3, 3).`
  1121. t: Translation matrix :math:`(B, 3, 1)`.
  1122. Returns:
  1123. Rinv: Rotation matrix, :math:`(B, 3, 3).`
  1124. tinv: Translation matrix :math:`(B, 3, 1)`.
  1125. Example:
  1126. >>> R, t = torch.eye(3)[None], torch.ones(3).reshape(1, 3, 1)
  1127. >>> worldtocam_to_camtoworld_Rt(R, t)
  1128. (tensor([[[1., 0., 0.],
  1129. [0., 1., 0.],
  1130. [0., 0., 1.]]]), tensor([[[-1.],
  1131. [-1.],
  1132. [-1.]]]))
  1133. """
  1134. KORNIA_CHECK_SHAPE(R, ["B", "3", "3"])
  1135. KORNIA_CHECK_SHAPE(t, ["B", "3", "1"])
  1136. R_inv = R.transpose(1, 2)
  1137. new_t: Tensor = -R_inv @ t
  1138. return (R_inv, new_t)
  1139. def ARKitQTVecs_to_ColmapQTVecs(qvec: Tensor, tvec: Tensor) -> tuple[Tensor, Tensor]:
  1140. r"""Convert output of Apple ARKit screen pose to the camera-to-world transformation, expected by Colmap.
  1141. Both poses in quaternion representation.
  1142. Args:
  1143. qvec: ARKit rotation quaternion :math:`(B, 4)`, [w, x, y, z] format.
  1144. tvec: translation vector :math:`(B, 3, 1)`, [x, y, z]
  1145. Returns:
  1146. qvec: Colmap rotation quaternion :math:`(B, 4)`, [w, x, y, z] format.
  1147. tvec: translation vector :math:`(B, 3, 1)`, [x, y, z]
  1148. Example:
  1149. >>> q, t = tensor([0, 1, 0, 1.])[None], torch.ones(3).reshape(1, 3, 1)
  1150. >>> ARKitQTVecs_to_ColmapQTVecs(q, t)
  1151. (tensor([[0.7071, 0.0000, 0.7071, 0.0000]]), tensor([[[-1.0000],
  1152. [-1.0000],
  1153. [ 1.0000]]]))
  1154. """
  1155. # ToDo: integrate QuaterniaonAPI
  1156. Rcg = quaternion_to_rotation_matrix(qvec)
  1157. Rcv, Tcv = camtoworld_graphics_to_vision_Rt(Rcg, tvec)
  1158. R_colmap, t_colmap = camtoworld_to_worldtocam_Rt(Rcv, Tcv)
  1159. t_colmap = t_colmap.reshape(-1, 3, 1)
  1160. q_colmap = rotation_matrix_to_quaternion(R_colmap.contiguous())
  1161. return q_colmap, t_colmap
  1162. def vector_to_skew_symmetric_matrix(vec: Tensor) -> Tensor:
  1163. r"""Convert a vector to a skew symmetric matrix.
  1164. A vector :math:`(v1, v2, v3)` has a corresponding skew-symmetric matrix, which is of the form:
  1165. .. math::
  1166. \begin{bmatrix} 0 & -v3 & v2 \\
  1167. v3 & 0 & -v1 \\
  1168. -v2 & v1 & 0\end{bmatrix}
  1169. Args:
  1170. vec: tensor of shape :math:`(B, 3)`.
  1171. Returns:
  1172. tensor of shape :math:`(B, 3, 3)`.
  1173. Example:
  1174. >>> vec = torch.tensor([1.0, 2.0, 3.0])
  1175. >>> vector_to_skew_symmetric_matrix(vec)
  1176. tensor([[ 0., -3., 2.],
  1177. [ 3., 0., -1.],
  1178. [-2., 1., 0.]])
  1179. """
  1180. # KORNIA_CHECK_SHAPE(vec, ["B", "3"])
  1181. if vec.shape[-1] != 3 or len(vec.shape) > 2:
  1182. raise ValueError(f"Input vector must be of shape (B, 3) or (3,). Got {vec.shape}")
  1183. v1, v2, v3 = vec[..., 0], vec[..., 1], vec[..., 2]
  1184. zeros = zeros_like(v1)
  1185. skew_symmetric_matrix = stack(
  1186. [stack([zeros, -v3, v2], dim=-1), stack([v3, zeros, -v1], dim=-1), stack([-v2, v1, zeros], dim=-1)], dim=-2
  1187. )
  1188. return skew_symmetric_matrix