o oi, @sddlmZmZddlZddlmZddlmZm Z m Z m Z ddl m Z mZmZmZddlmZmZddlmZddlmZdd lmZmZdd ejd ed eejejffddZ  ddejdejdejdeejded ejf ddZdS))OptionalTupleN)qr)arange ones_likewherezeros) KORNIA_CHECKKORNIA_CHECK_IS_TENSORKORNIA_CHECK_SAME_SHAPEKORNIA_CHECK_SHAPE)convert_points_to_homogeneous normalize_points_with_intrinsics)transform_points)eye_like)_torch_linalg_svdvals_torch_svd_cast:0yE>pointsepsreturnc Cst|gdtj|ddd}||jdddjdd}|jd}tj|jdtj|jd }t|||}t |d|}t |tj |jd }|d d ||f|d d d f|d d ||f<|d d ||f|d d d f |d d d |f|d d ||f<t ||}||fS) a\Normalize points. Args: points : Tensor containing the points to be normalized with shape :math:`(B, N, D)`. eps : Small value to avoid division by zero error. Returns: Tuple containing the normalized points in the shape :math:`(B, N, D)` and the transformation matrix in the shape :math:`(B, D+1, D+1)`. )BNDT)dimkeepdim)rprdtypedeviceNr) r torchmeannormshapetensorfloat64r#sqrtrrint64r) rrx_meanscaleD_intD_float transformidxs points_normr3S/home/ubuntu/.local/lib/python3.10/site-packages/kornia/geometry/calibration/pnp.py_mean_isotropic_scale_normalizes  0D r5-C6? world_points img_points intrinsicsweightssvd_epsc#Cstjtjf}t|t|t|tt|tt|j|vt|j|vt|j|vt|gdt|gdt|gdt |dddddf|dddddft|j ddk|durkt||j dd\}}t |\}} t |} t | ddd f|krtd |d t|} t||} t | \} }t|}t|d|d f|j|jd }| |dddddddf<| |dddddddf<| d | dddf|ddddddd f<| d | dddf|ddddddd f<|dur,|j ||fkrtd|j d|d ddd|d|d}||}t|\}}}|d}||dd}td|}||ddddddf<t|| }t||ddddddf}t|ddddddf}t|}t|dk|d |}||ddddf}tj|dddddfddd}d|ddddf}||}t|ddddddf\}}td|}t ||} t|| }!tj!|!|ddddddfgd d}"|"S)a9 Attempt to solve the Perspective-n-Point (PnP) problem using Direct Linear Transform (DLT). Given a batch (where batch size is :math:`B`) of :math:`N` 3D points (where :math:`N \geq 6`) in the world space, a batch of :math:`N` corresponding 2D points in the image space and a batch of intrinsic matrices, this function tries to estimate a batch of world to camera transformation matrices. This implementation needs at least 6 points (i.e. :math:`N \geq 6`) to provide solutions. This function cannot be used if all the 3D world points (of any element of the batch) lie on a line or if all the 3D world points (of any element of the batch) lie on a plane. This function attempts to check for these conditions and throws an AssertionError if found. Do note that this check is sensitive to the value of the svd_eps parameter. Another bad condition occurs when the camera and the points lie on a twisted cubic. However, this function does not check for this condition. Args: world_points : A tensor with shape :math:`(B, N, 3)` representing the points in the world space. img_points : A tensor with shape :math:`(B, N, 2)` representing the points in the image space. intrinsics : A tensor with shape :math:`(B, 3, 3)` representing the intrinsic matrices. weights : A tensor with shape :math:`(B, N)` representing the weights for each point. If None, all points are considered to be equally important. svd_eps : A small float value to avoid numerical precision issues. Returns: A tensor with shape :math:`(B, 3, 4)` representing the estimated world to camera transformation matrices (also known as the extrinsic matrices). Example: >>> world_points = torch.tensor([[ ... [ 5. , -5. , 0. ], [ 0. , 0. , 1.5], ... [ 2.5, 3. , 6. ], [ 9. , -2. , 3. ], ... [-4. , 5. , 2. ], [-5. , 5. , 1. ], ... ]], dtype=torch.float64) >>> >>> img_points = torch.tensor([[ ... [1409.1504, -800.936 ], [ 407.0207, -182.1229], ... [ 392.7021, 177.9428], [1016.838 , -2.9416], ... [ -63.1116, 142.9204], [-219.3874, 99.666 ], ... ]], dtype=torch.float64) >>> >>> intrinsics = torch.tensor([[ ... [ 500., 0., 250.], ... [ 0., 500., 250.], ... [ 0., 0., 1.], ... ]], dtype=torch.float64) >>> >>> print(world_points.shape, img_points.shape, intrinsics.shape) torch.Size([1, 6, 3]) torch.Size([1, 6, 2]) torch.Size([1, 3, 3]) >>> >>> pred_world_to_cam = kornia.geometry.solve_pnp_dlt(world_points, img_points, intrinsics) >>> >>> print(pred_world_to_cam.shape) torch.Size([1, 3, 4]) >>> >>> pred_world_to_cam tensor([[[ 0.9392, -0.3432, -0.0130, 1.6734], [ 0.3390, 0.9324, -0.1254, -4.3634], [ 0.0552, 0.1134, 0.9920, 3.7785]]], dtype=torch.float64) )rr3)rr2)rr<r<NrrrrzQThe last singular value of one/more of the elements of the batch is smaller than z. This function cannot be used if all world_points (of any element of the batch) lie on a line or if all world_points (of any element of the batch) lie on a plane. r!.z&Weights should have shape (B, N). Got .).r)inputrrr )"r$float32r)r r isinstancefloatr"r r r'r5ranyAssertionErrorr rinverserr# unsqueezerepeatviewrreshaperbmmdetrrr& linalg_qrsigncat)#r7r8r9r:r;accepted_dtypesrrworld_points_normworld_transform_normsworld_points_norm_himg_points_invimg_points_normimg_transform_norminv_img_transform_normsystemweights_expanded_vsolution solution_4x4 intermediaterPonessign_fixnorm_col mul_factortemporthorightmask col_sign_fixrot_matpred_world_to_camr3r3r4 solve_pnp_dlt=st T.    00 "  " ""  (ro)r)Nr6)typingrrr$ torch.linalgrrQ kornia.corerrrrkornia.core.checkr r r r kornia.geometry.conversionsr rkornia.geometry.linalgr kornia.utilsrkornia.utils.helpersrrTensorrGr5ror3r3r3r4s2   &"