a d[g@sddlmZmZddlZddlmZmZddlmZm Z ddl m Z GdddZ Gdd d e Z d!ejeejd d d Zd"ejeejdddZejejejdddZejejdddZejejejdddZejejejejdddZd#ejejeejddd ZdS)$)IterableOptionalN)convert_points_from_homogeneousconvert_points_to_homogeneous)inverse_transformationtransform_points)_torch_inverse_castc@sNeZdZdZejejejejddddZeeeje dddZ eeje e d d d Z eeje e d d d Z eejdddZeejdddZeedddZeejdddZeejdddZeejdddZeejdddZeejdddZejddddZeejdd d!Zejddd"d!Zeejdd#d$Zejddd%d$Zeejdd&d'Zeejdd(d)Zeejdd*d+Zeejdd,d-Zddd.d/Zejdd0d1Z ddd2d3Z!ddd4d5Z"ejejd6d7d8Z#ejejd9d:d;Z$e%d@e&ej'e&ej(d=d>d?Z)dS)A PinholeCameraadClass that represents a Pinhole Camera model. Args: intrinsics: tensor with shape :math:`(B, 4, 4)` containing the full 4x4 camera calibration matrix. extrinsics: tensor with shape :math:`(B, 4, 4)` containing the full 4x4 rotation-translation matrix. height: tensor with shape :math:`(B)` containing the image height. width: tensor with shape :math:`(B)` containing the image width. .. note:: We assume that the class attributes are in batch form in order to take advantage of PyTorch parallelism to boost computing performance. N) intrinsics extrinsicsheightwidthreturncCs^|||||g||d||d||d||d||_||_||_||_dS)Nr r r r ) _check_valid_check_valid_params_check_valid_shaper r _intrinsics _extrinsics)selfr r r r rg/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/kornia/geometry/camera/pinhole.py__init__s    zPinholeCamera.__init__) data_iterrcCstdd|DstddS)Ncss|]}|jdVqdS)rN)shape).0datarrr +z-PinholeCamera._check_valid..zArguments shapes must matchT)all ValueError)rrrrr)szPinholeCamera._check_valid)r data_namercCs6t|jdvr2|jdddkr2td||jdS)N)r"r"zIArgument {} shape must be in the following shape Bx4x4 or BxNx4x4. Got {}Tlenrrformatrr rrrr/s z!PinholeCamera._check_valid_paramscCs$t|jdks td||jdS)Nz:Argument {} shape must be in the following shape B. Got {}Tr%r(rrrr8s  z PinholeCamera._check_valid_shapercCs||jdst|jS)ziThe full 4x4 intrinsics matrix. Returns: tensor of shape :math:`(B, 4, 4)`. r )rrAssertionErrorrrrrr @szPinholeCamera.intrinsicscCs||jdst|jS)ziThe full 4x4 extrinsics matrix. Returns: tensor of shape :math:`(B, 4, 4)`. r )rrr+r,rrrr KszPinholeCamera.extrinsicscCs |jjdS)zhReturn the batch size of the storage. Returns: scalar with the batch size. r)r rr,rrr batch_sizeVszPinholeCamera.batch_sizecCs |jdS)zoReturn the focal length in the x-direction. Returns: tensor of shape :math:`(B)`. .rrr r,rrrfx_szPinholeCamera.fxcCs |jdS)zoReturn the focal length in the y-direction. Returns: tensor of shape :math:`(B)`. .r)r)r/r,rrrfyhszPinholeCamera.fycCs |jdS)zsReturn the x-coordinate of the principal point. Returns: tensor of shape :math:`(B)`. .rr/r,rrrcxqszPinholeCamera.cxcCs |jdS)zsReturn the y-coordinate of the principal point. Returns: tensor of shape :math:`(B)`. .r)r4r/r,rrrcyzszPinholeCamera.cycCs |jdS)zvReturn the x-coordinate of the translation vector. Returns: tensor of shape :math:`(B)`. .rr r,rrrtxszPinholeCamera.txcCs||jd<|S)zDSet the x-coordinate of the translation vector with the given value.r8r:rvaluerrrr;s cCs |jdS)zvReturn the y-coordinate of the translation vector. Returns: tensor of shape :math:`(B)`. .r)r9r:r,rrrtyszPinholeCamera.tycCs||jd<|S)DSet the y-coordinate of the translation vector with the given value.r>r:r<rrrr?s cCs |jdS)zwReturns the z-coordinate of the translation vector. Returns: tensor of shape :math:`(B)`. .r4r9r:r,rrrtzszPinholeCamera.tzcCs||jd<|S)r@rAr:r<rrrrBs cCs|jdddddfS)zuReturn the 3x4 rotation-translation matrix. Returns: tensor of shape :math:`(B, 3, 4)`. .Nr!r"r:r,rrr rt_matrixszPinholeCamera.rt_matrixcCs|jdddddfS)zReturn the 3x3 camera matrix containing the intrinsics. Returns: tensor of shape :math:`(B, 3, 3)`. .Nr!r/r,rrr camera_matrixszPinholeCamera.camera_matrixcCs|jdddddfS)z}Return the 3x3 rotation matrix from the extrinsics. Returns: tensor of shape :math:`(B, 3, 3)`. .Nr!r:r,rrrrotation_matrixszPinholeCamera.rotation_matrixcCs|jdddddfS)z|Return the translation vector from the extrinsics. Returns: tensor of shape :math:`(B, 3, 1)`. .Nr!r9r:r,rrrtranslation_vectorsz PinholeCamera.translation_vectorcCs6|j}|j}|j}|j}t||||S)z2Return a deep copy of the current object instance.)r cloner r r r )rr r r r rrrrGs     zPinholeCamera.clonecCs |jS)zzReturn the inverse of the 4x4 instrisics matrix. Returns: tensor of shape :math:`(B, 4, 4)`. )r Zinverser,rrrintrinsics_inversesz PinholeCamera.intrinsics_inversecCsv|j}|d|9<|d|9<|d|9<|d|9<||j}||j}t||j||S)a.Scale the pinhole model. Args: scale_factor: a tensor with the scale factor. It has to be broadcastable with class members. The expected shape is :math:`(B)` or :math:`(1)`. Returns: the camera model with scaled parameters. r.r1r3r6)r rGr r r r )r scale_factorr r r rrrscales zPinholeCamera.scalecCsh|jd|9<|jd|9<|jd|9<|jd|9<|j|9_|j|9_|S)a7Scale the pinhole model in-place. Args: scale_factor: a tensor with the scale factor. It has to be broadcastable with class members. The expected shape is :math:`(B)` or :math:`(1)`. Returns: the camera model with scaled parameters. r.r1r3r6)r r r )rrIrrrscale_s zPinholeCamera.scale_)point_3drcCs|j|j}tt||S)aProject a 3d point in world coordinates onto the 2d camera plane. Args: point3d: tensor containing the 3d points to be projected to the camera plane. The shape of the tensor can be :math:`(*, 3)`. Returns: tensor of (u, v) cam coordinates with shape :math:`(*, 2)`. Example: >>> _ = torch.manual_seed(0) >>> X = torch.rand(1, 3) >>> K = torch.eye(4)[None] >>> E = torch.eye(4)[None] >>> h = torch.ones(1) >>> w = torch.ones(1) >>> pinhole = kornia.geometry.camera.PinholeCamera(K, E, h, w) >>> pinhole.project(X) tensor([[5.6088, 8.6827]]) )r r rr)rrLPrrrprojects zPinholeCamera.project)point_2ddepthcCs&|j|j}t|}t|t||S)aUnproject a 2d point in 3d. Transform coordinates in the pixel frame to the world frame. Args: point2d: tensor containing the 2d to be projected to world coordinates. The shape of the tensor can be :math:`(*, 2)`. depth: tensor containing the depth value of each 2d points. The tensor shape must be equal to point2d :math:`(*, 1)`. normalize: whether to normalize the pointcloud. This must be set to `True` when the depth is represented as the Euclidean ray length from the camera position. Returns: tensor of (x, y, z) world coordinates with shape :math:`(*, 3)`. Example: >>> _ = torch.manual_seed(0) >>> x = torch.rand(1, 2) >>> depth = torch.ones(1, 1) >>> K = torch.eye(4)[None] >>> E = torch.eye(4)[None] >>> h = torch.ones(1) >>> w = torch.ones(1) >>> pinhole = kornia.geometry.camera.PinholeCamera(K, E, h, w) >>> pinhole.unproject(x, depth) tensor([[0.4963, 0.7682, 1.0000]]) )r r rrr)rrOrPrMZP_invrrr unproject's zPinholeCamera.unprojectr)devicedtypec Cs tj| dd| | d} | d|7<| d|7<| d|7<| d|7<| dd7<| d d7<tjd| | d| d d }|d |7<|d |7<|d | 7<tj| | | d}|d|7<tj| | | d}|d|7<|| |||S)Nr"rRr.r1r3r6).r4r4?).r!r!r)r8r>rA.r)torchzeroseyerepeat)rr0r2r5r7r r r;r?rBr-rSrTr r Z height_tmpZ width_tmprrrfrom_parametersIs zPinholeCamera.from_parameters)r)NN)*__name__ __module__ __qualname____doc__rWTensorr staticmethodrboolrstrrrpropertyr r intr-r0r2r5r7r;setterr?rBrCrDrErFrGrHrJrKrNrQ classmethodrrSrTr[rrrrr st   "  r c@sZeZdZdZeeddddZeeddddZee d d d Z e ed d dZ dS)PinholeCamerasListalClass that represents a list of pinhole cameras. The class inherits from :class:`~kornia.PinholeCamera` meaning that it will keep the same class properties but with an extra dimension. .. note:: The underlying data tensor will be stacked in the first dimension. That's it, given a list of two camera instances, the intrinsics tensor will have a shape :math:`(B, N, 4, 4)` where :math:`B` is the batch size and :math:`N` is the numbers of cameras (in this case two). Args: pinholes_list: a python tuple or list containing a set of `PinholeCamera` instances. N) pinholes_listrcCs||dS)N)_initialize_parameters)rrirrrr~szPinholeCamerasList.__init__pinholesrcCst|ttfs tdt|gg}}gg}}|D]P}t|tsXtdt|||j||j ||j ||j q8t j |dd|_t j |dd|_ t j |dd|_t j |dd|_|S)z5Initialise the class attributes given a cameras list.z(pinhole must of type list or tuple. Got z8Argument pinhole must be from type PinholeCamera. Got {}r)Zdim) isinstancelisttuple TypeErrortyper r'appendr r r r rWstackrr)rrlr r r r pinholerrrrjs       z)PinholeCamerasList._initialize_parametersr*cCs"d}|jdurt|jjd}|S)z0Return the number of pinholes cameras per batch.r9Nr))r rer)r num_camerasrrrrvs zPinholeCamerasList.num_cameras)idxrcCsN|jd|f}|jd|f}|jdd|f}|jdd|f}t||||S)z)rlepsrcCst|jdkr|jddks&t|jtj|dddfddd\}}}}tjd|j|jd|}|ddd |jd dd}||dd d df<||dd dd f<||ddddf<||dddd f<|S) a,Function that returns the pinhole matrix from a pinhole model. .. note:: This method is going to be deprecated in version 0.2 in favour of :attr:`kornia.PinholeCamera.camera_matrix`. Args: pinholes: tensor of pinhole models. Returns: tensor of pinhole matrices. Shape: - Input: :math:`(N, 12)` - Output: :math:`(N, 4, 4)` Example: >>> rng = torch.manual_seed(0) >>> pinhole = torch.rand(1, 12) # Nx12 >>> pinhole_matrix(pinhole) # Nx4x4 tensor([[[4.9626e-01, 1.0000e-06, 8.8477e-02, 1.0000e-06], [1.0000e-06, 7.6822e-01, 1.3203e-01, 1.0000e-06], [1.0000e-06, 1.0000e-06, 1.0000e+00, 1.0000e-06], [1.0000e-06, 1.0000e-06, 1.0000e-06, 1.0000e+00]]]) r4r) .Nr"rmrRrr! r&rr+rWchunkrYrSrTviewrZ)rlrzr0r2r5r7krrrpinhole_matrixs $r)rurzrcCst|jdkr|jddks&t|jtj|dddfddd\}}}}tjd|j|jd}|ddd |jd dd}d |||dd d df<d |||ddddf<d ||||dd dd f<d ||||dddd f<|S) aReturn the inverted pinhole matrix from a pinhole model. .. note:: This method is going to be deprecated in version 0.2 in favour of :attr:`kornia.PinholeCamera.intrinsics_inverse()`. Args: pinholes: tensor with pinhole models. Returns: tensor of inverted pinhole matrices. Shape: - Input: :math:`(N, 12)` - Output: :math:`(N, 4, 4)` Example: >>> rng = torch.manual_seed(0) >>> pinhole = torch.rand(1, 12) # Nx12 >>> inverse_pinhole_matrix(pinhole) # Nx4x4 tensor([[[ 2.0151, 0.0000, -0.1783, 0.0000], [ 0.0000, 1.3017, -0.1719, 0.0000], [ 0.0000, 0.0000, 1.0000, 0.0000], [ 0.0000, 0.0000, 0.0000, 1.0000]]]) r4r)r{.Nr"rmrRrrUgr!r|)rurzr0r2r5r7rrrrinverse_pinhole_matrixs $r)rlrJrcCspt|jdkr|jddks&t|jt|jdkr>t|j|}|dddf|d|dddf<|S)aScale the pinhole matrix for each pinhole model. .. note:: This method is going to be deprecated in version 0.2 in favour of :attr:`kornia.PinholeCamera.scale()`. Args: pinholes: tensor with the pinhole model. scale: tensor of scales. Returns: tensor of scaled pinholes. Shape: - Input: :math:`(N, 12)` and :math:`(N, 1)` - Output: :math:`(N, 12)` Example: >>> rng = torch.manual_seed(0) >>> pinhole_i = torch.rand(1, 12) # Nx12 >>> scales = 2.0 * torch.ones(1) # N >>> scale_pinhole(pinhole_i, scales) # Nx12 tensor([[0.9925, 1.5364, 0.1770, 0.2641, 0.6148, 1.2682, 0.4901, 0.8964, 0.4556, 0.6323, 0.3489, 0.4017]]) r4r)r{.Nr9)r&rr+rGZ unsqueeze)rlrJZpinholes_scaledrrr scale_pinholes  &rrkcCs.t|jdkr|jddks&t|jtdS)a Compute extrinsic transformation matrices for pinholes. Args: pinholes: tensor of form [fx fy cx cy h w rx ry rz tx ty tz] of size (N, 12). Returns: tensor of extrinsic transformation matrices of size (N, 4, 4). r4r)r{N)r&rr+NotImplementedError)rlrrrget_optical_pose_base$s  r) pinhole_i pinhole_refrcCsxt|jdkr|jddks&t|j|j|jkrInput cam_coords_src has to be in the shape of BxHxWx3. Got {}r#Nr$z:Input dst_proj_src has to be in the shape of Bx4x4. Got {}rV).r)).r4r9rm)r&rrr'rrWrt) rrrzZ point_coordsZx_coordZy_coordZz_coordZu_coordZv_coordZpixel_coords_dstrrr cam2pixels     r)ry)ry)r)typingrrrWZkornia.geometry.conversionsrrZkornia.geometry.linalgrrZkornia.utils.helpersrr rhr`floatrrrrrrrrrrrs  f8,,&3"