3jG SSKJr SSKJrJrJrJrJrJrJ r SSK r SSK J r J r SSKJr SSKJr SSKJr S\S \4S \S \44S jrS\S \4S \S \44S jrS \ \ R,\\ R,\ R,44S \ R.4SjrS \ \ R,\\ R,\ R,44S \ R.4SjrS:S \ R,S\ R,S\S \ R,4SjjrS:S \ R,S\ R,S\S \ R,4SjjrS \ R,S \ R,4SjrS \ R,S \ R,4SjrS \ R,S\\ R>S S4Sjr SS.S\ R,S\\!S 4S\\ R,S \ R,4Sjjr"S;S\ \\!S 4\ R.4S\\#S 4S S4Sjjr$S \ R,S\!S\!S \4S jr%SS \ \ R,\4S4\#S \4S5jjr0S \ \ R,\4S \4S6jr1S \ R,S \ R,4S7jr2"S8S95r3g)?)wraps)AnyCallableDictListOptionalTupleUnionN)BetaUniform)_extract_device_dtype)Boxes) Keypointsf.returnc x^[T5S[RS[S[S[4U4Sjj5nU$)zValidate the 2D input of the wrapped function. Args: f: a function that takes the first argument as torch.Tensor. Returns: the wrapped function after input is validated. inputargskwargsrc4>[R"U5(d[S[U535e[ UR SS9 [ U[R[R[R[R/S9 T"U/UQ70UD6$)N&Input type is not a torch.Tensor. Got BCHW)required_shapesaccepted_dtypes) torch is_tensor TypeErrortype_validate_shapeshape_validate_input_dtypebfloat16float16float32float64)rrrrs [/home/wildlama/miniconda3/lib/python3.13/site-packages/kornia/augmentation/utils/helpers.pywrapper _validate_input..wrapper(suu%%DT%[MRS S Y?eennemmUZUbUbdidqdq5rs((((rrTensorrrr)s` r(_validate_inputr/sB 1X)u||)C)3)3)) Nr+c x^[T5S[RS[S[S[4U4Sjj5nU$)zValidate the 3D input of the wrapped function. Args: f: a function that takes the first argument as torch.Tensor. Returns: the wrapped function after input is validated. rrrrc`>[R"U5(d[S[U535e[ UR 5nUS:wa[ SUS35e[U[R[R[R[R/S9 T"U/UQ70UD6$)Nrz"Expect input of 5 dimensions, got z insteadr) rrrr lenr"AssertionErrorr#r$r%r&r')rrr input_shapers r(r)"_validate_input3d..wrapper@su%%DT%[MRS S%++& !  #Ek]RZ![\ \eennemmUZUbUbdidqdq5rs((((r+r,r.s` r(_validate_input3dr75sB 1X )u|| )C )3 )3 ) ) Nr+rc[U[5(a[US5nUR$[U5nUR$z\Infer input shape. Input may be either (torch.Tensor,) or (torch.Tensor, transform_matrix) r) isinstancetuple_transform_inputr"r tensor_vars r(_infer_batch_shaper?OsE %%eAh/    &e,   r+c[U[5(a[US5nUR$[U5nUR$r9)r:r;_transform_input3dr"r=s r(_infer_batch_shape3drB[sE %'a1    (.   r+reference_shape match_channelc[U5S:aUSOSn[U5S:aUSOSn[UR5S:XaURS5n[UR5S:Xa4X0RS:XaURS5OURS5nU(a(U(a!URSU:Xd [S 5eU$U(aUc [S 5eU$) Reshape an input torch.Tensor to have the same dimensions as the reference_shape. Arguments: input: torch.Tensor to be transformed reference_shape: shape used as reference match_channel: if True, C_{src} == C_{ref}. otherwise, no constrain. C =1 by default NrNThe C dimension of torch.Tensor did not match with the reference torch.Tensor.5The reference torch.Tensor do not have a C dimension!r3r" unsqueeze ValueErrorrrCrDBCs r(_transform_input_by_shaperUgs #?3q8dA"?3q8dA 5;;1" 5;;1'(;;r?&:"PQ@R{{2!#mn n$ L 19PQQ Lr+c[U5S:aUSOSn[U5S:aUSOSn[UR5S:XaURS5n[UR5S:Xa#X0RS:XaURS5nU(a(U(a!URSU:Xd [S 5eU$U(aUc [S 5eU$) rFr2NrGrHrIrrKrMrNrOrRs r(_transform_input3d_by_shaperXs #?3q8dA"?3q8dA 5;;1" 5;;1kk"o!5"{{2!#mn n$ L 19PQQ Lr+ct[R"U5(d[S[U535e[ UR 5S;a[ SUR 35e[ UR 5S:XaURS5n[ UR 5S:XaURS5nU$)zReshape an input torch.Tensor to be (*, C, H, W). Accept either (H, W), (C, H, W) or (*, C, H, W). Args: input: torch.Tensor Returns: torch.Tensor r)rKrIrGzNInput size must have a shape of either (H, W), (C, H, W) or (*, C, H, W). Got rKrrIrrrr r3r"rQrPrs r(r<r<s ??5 ! !@e NOO 5;;y(ijojujuivwxx 5;;1" 5;;1" Lr+ct[R"U5(d[S[U535e[ UR 5S;a[ SUR 35e[ UR 5S:XaURS5n[ UR 5S:XaURS5nU$)zReshape an input torch.Tensor to be (*, C, D, H, W). Accept either (D, H, W), (C, D, H, W) or (*, C, D, H, W). Args: input: torch.Tensor Returns: torch.Tensor r)rIrGr2zWInput size must have a shape of either (D, H, W), (C, D, H, W) or (*, C, D, H, W). Got rIrrGrZr[s r(rArAs ??5 ! !@e NOO 5;;y(efkfqfqer s   5;;1" 5;;1" Lr+rcZURU;a[SUSUR35eg)zCheck if the dtype of the input torch.Tensor is in the range of accepted_dtypes. Args: input: torch.Tensor accepted_dtypes: List. e.g. [torch.float32, torch.float64] zExpected input of . Got N)dtyper)rrs r(r#r#s3 {{/),_,=VEKK=QRR*r+)rCoutputr"c\Un[[UR5[U5- 5H{nSnUb1URSUSs=:XaS:waO O[U5S:aSnURUS:wa[SUSURU35eUR U5nM} U$)aCollapse the broadcasted batch dimensions an input torch.Tensor to be the specified shape. Args: output: torch.Tensor shape: List/tuple of int reference_shape: torch.Tensor representation of shape to control which dimensions are collapsed. Returns: torch.Tensor rrLrKz Dimension z# of input is expected to be 1, got )ranger3r"r4squeeze)r`r"rC out_tensordimidxs r(_transform_output_shapergsJS))*SZ78  &:+;+;A+>/RSBT+YXY+Y^abg^hkl^lC   C A % :cU2UV`VfVfgjVkUl!mn n'', 9 r+rcSnUHn[U5[U5:XdMSn O U(d[SWSUS35eg)zCheck if the dtype of the input torch.Tensor is in the range of accepted_dtypes. Args: shape: torch.Tensor shape required_shapes: List. e.g. ["BCHW", "BCDHW"] FTzExpected input shape in r^.N)r3r)r"rpassedrequired_shapes r(r!r!sRF) u:^, ,F * 2>2B&qQRR r+ channel_indexnumberc&URUU:H$)a2Validate if an input has the right shape. e.g. to check if an input is channel first. If channel first, the second channel of an RGB input shall be fixed to 3. To verify using: _validate_input_shape(input, 1, 3) Args: input: torch.Tensor channel_index: int number: int Returns: bool )r")rrlrms r(_validate_input_shaperos ;;} % //r+dist same_on_batchcD[U[5(a[R"U5nU(a^[R"S/USSQ75nUR U5nUR "US/S/[ UR5S- -Q76$UR U5$)zSample from a uniform reparameterized sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. rLNr)r:r;rSizersamplerepeatr3r")r"rprq rsample_sizerts r(_adapted_rsamplingrw!s% 5!zz1/uQRy/2 ,,|,~~eAhH!GMM0BQ0F)GHH << r+c([U[5(a[R"U5nU(aPUR [R"S/USSQ755R "US/S/[ U5S- -Q76$UR U5$)zSample from a uniform sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. rLNr)r:r;rrssamplerur3)r"rprqs r(_adapted_samplingrz5s}% 5!{{5::qo59o67>>uQxa1#QTUZQ[^_Q_J`aa ;;u r+lowhighc$[[U[R5(aUOS[U[R5(aUOS/5upE[R"XUS9n[R"X$US9n[ XSS9n[ XU5$)a{Sample from a uniform sampling function that accepts 'same_on_batch'. If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default, same_on_batch is set to False. By default, sampling happens on the default device and dtype. If low/high is a torch.Tensor, sampling will happen in the same device/dtype as low/high torch.Tensor. Ndevicer_F validate_args)r r:rr- as_tensorr rw)r"r{r|rqrr_rps r(_adapted_uniformrGs}*3 --4DRWR^R^A_A_eijMF //#E :C ??4e `_ defined by it's 4 vertices coordinates (A, B, C, D). Coordinates must be in ``x, y`` order. The height and width of a box is defined as ``width = xmax - xmin + 1`` and ``height = ymax - ymin + 1``. Examples of `quadrilaterals `_ are rectangles, rhombus and trapezoids. rGrKNz,Only BxNx4x2 torch.Tensor is supported. Got ri)rzExpect `Boxes` type. Got ) r:rr-r3r"rsrr from_tensorr )rrs r(preprocess_boxesrsF%&&EKK A%%++ab/UZZA=O*O!Mekk]Z[\] ]!!%3 eU # #6tE{m1EFF Lr+cj[U[R5(af[UR5S:Xa(URSS[R "S/5:Xd[ SURS35e[US5n[U[5(a[ S[U5S35eU$) zPreprocess input keypoints.rIrLNrKz*Only BxNx2 torch.Tensor is supported. Got riFzExpect `Keypoints` type. Got ) r:rr-r3r"rsrrr r[s r(preprocess_keypointsrs%&&EKK A%%++ab/UZZ_*L!KEKK=XYZ[ [%'%##:4;-qIJJ Lr+cU$)zPreprocess input class tags.r[s r(preprocess_classesrs  Lr+c<^\rSrSrSrS\S\SS4U4SjjrSrU=r$) MultiprocessWrapperizYWhen used as a base class, makes the class work with the 'spawn' multiprocessing context.rrrNc >[SU55nUR5VVs0sH7up4U[U[R5(aUR 5OU_M9 nnn[ TU]"U0UD6 gs snnf)Nc3# UH6n[U[R5(aUR5OUv M8 g7f)N)r:rr-r).0args r( /MultiprocessWrapper.__init__..s-[VZsJsELL$A$ASYY[sJVZs>A)r;rr:rr-rsuper__init__)selfrrkeyval __class__s r(rMultiprocessWrapper.__init__sl[VZ[[]c]i]i]kl]kQYQT#jell&C&Csyy{L]kl $)&)ms>A9r) __name__ __module__ __qualname____firstlineno____doc__rr__static_attributes__ __classcell__)rs@r(rrs%c*c*S*T**r+r)T)r)F)NrF) vertices_plus)4 functoolsrtypingrrrrrr r rtorch.distributionsr r kornia.core.utilsr kornia.geometry.boxesrkornia.geometry.keypointsrr/r7r-rsr?rBboolrUrXr<rAr_r#intrgstrr!ro distributions Distributionrwrzfloatrrrrrrrrrrr+r(rs$DDD -3'/xS)hsCx.@0#s(+c0B4 eELL% ell8R2S$ST Y^YcYc  ellE%,, :T4U&U V [`[e[e OS <<*/,,GK \\@OS <<*/,,GK \\>ELLU\\2ellu||6SSU[[@QSVZS`d LL!&sCxFNu||F\ \\6 S5sCx%**!<= SPUVY[^V^P_ Spt S 00c030SW0*%* sCx%**, -    * *D> \\ .%* sCx%**, -    * *D> \\ , : sCx%**, -: uell" #: u||# $: :  \\ :: : sCx%**, -: UELL !: UELL !: :  \\ :._U\\_%c3hc8R2S_[^_cg_ $sCx. T#s(^ "15!  cNd38n-   #s(^ >)E%,,"56)c)X])X ellI&= > 9 ellu|| **r+