
    3j^                       S r SSKJr  SSKJrJrJrJr  SSKrSSK	J
s  Jr  SSKJ
r
  SSKJr  SSKJr   " S S	\
R"                  5      r " S
 S\
R"                  5      r " S S\
R"                  5      r " S S\
R"                  5      rg)aK  XFeat: Accelerated Features for Lightweight Image Matching.

This module is adapted from the original XFeat implementation:
    https://github.com/verlab/accelerated_features

Original work copyright (c) 2024 Verlab / UFMG, licensed under the Apache License 2.0.
Modifications copyright (c) 2024 Kornia Team, licensed under the Apache License 2.0.

Reference:
    Guilherme Potje, Felipe Cadar, Andre Araujo, Renato Martins, Erickson R. Nascimento.
    "XFeat: Accelerated Features for Lightweight Image Matching", CVPR 2024.
    https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/
    )annotations)DictListOptionalTupleN)nn)KORNIA_CHECK)nms2dc                  f   ^  \ rS rSrSr     S               SU 4S jjjrSS jrSrU =r$ )	
BasicLayer,   a  Basic convolutional layer: Conv2d -> BatchNorm2d (no affine) -> ReLU.

Args:
    in_channels: number of input channels.
    out_channels: number of output channels.
    kernel_size: convolution kernel size. Default: ``3``.
    stride: convolution stride. Default: ``1``.
    padding: convolution padding. Default: ``1``.
    dilation: convolution dilation. Default: ``1``.
    bias: whether to use bias. Default: ``False``.
c                   > [         TU ]  5         [        R                  " [        R                  " XX5XFUS9[        R
                  " USS9[        R                  " SS95      U l        g )N)paddingstridedilationbiasFaffineTinplace)super__init__r   
SequentialConv2dBatchNorm2dReLUlayer)	selfin_channelsout_channelskernel_sizer   r   r   r   	__class__s	           N/home/wildlama/miniconda3/lib/python3.13/site-packages/kornia/feature/xfeat.pyr   BasicLayer.__init__9   sS     	]]II;PVpt NN<6GGD!

    c                $    U R                  U5      $ Nr   )r   xs     r#   forwardBasicLayer.forwardL   s    zz!}r%   r(   )      r-   r-   F)r   intr    r.   r!   r.   r   r.   r   r.   r   r.   r   boolreturnNone)r)   torch.Tensorr0   r2   )	__name__
__module____qualname____firstlineno____doc__r   r*   __static_attributes____classcell__r"   s   @r#   r   r   ,   s{    
  

 
 	

 
 
 
 
 

 
& r%   r   c                  H   ^  \ rS rSrSrSU 4S jjrSS	S jjrS
S jrSrU =r	$ )
XFeatModelP   ag  XFeat backbone: CNN feature extractor, keypoint and reliability heads.

Implements the architecture from
"XFeat: Accelerated Features for Lightweight Image Matching", CVPR 2024.

Input: float image tensor :math:`(B, C, H, W)` (grayscale or RGB, any channel count).
Output:

- ``feats``: dense descriptors :math:`(B, 64, H/8, W/8)`.
- ``keypoints``: keypoint logits :math:`(B, 65, H/8, W/8)`.
- ``heatmap``: reliability map :math:`(B, 1, H/8, W/8)`.

Note:
    Image normalisation (``InstanceNorm2d``) is wrapped in ``torch.no_grad()``
    following the original design; backpropagating through it is not supported.
c                T  > [         TU ]  5         [        R                  " S5      U l        [        R
                  " [        R                  " SSS9[        R                  " SSSSSS95      U l        [        R
                  " [        SSSS9[        SSSS9[        SSSS9[        SSSS95      U l
        [        R
                  " [        SSSS9[        SSSS95      U l        [        R
                  " [        SS	SS9[        S	S	SS9[        S	S	SSS
95      U l        [        R
                  " [        S	S	SS9[        S	S	SS9[        S	S	SS95      U l        [        R
                  " [        S	SSS9[        SSSS9[        SSSS9[        SS	SSS
95      U l        [        R
                  " [        S	S	SS9[        S	S	SS9[        R                  " S	S	SSS
95      U l        [        R
                  " [        S	S	SSS
9[        S	S	SSS
9[        R                  " S	SS5      [        R                   " 5       5      U l        [        R
                  " [        S	S	SSS
9[        S	S	SSS
9[        S	S	SSS
9[        R                  " S	SS5      5      U l        [        R
                  " [        R&                  " SS5      [        R(                  " SSS9[        R*                  " SS9[        R&                  " SS5      [        R(                  " SSS9[        R*                  " SS9[        R&                  " SS5      [        R(                  " SSS9[        R*                  " SS9[        R&                  " SS5      [        R(                  " SSS9[        R*                  " SS9[        R&                  " SS	5      5      U l        g )Nr-      )r      r   )r   r         @   )r      A   i   Fr   Tr   )r   r   r   InstanceNorm2dnormr   	AvgPool2dr   skip1r   block1block2block3block4block5block_fusionSigmoidheatmap_headkeypoint_headLinearBatchNorm1dr   fine_matcher)r   r"   s    r#   r   XFeatModel.__init__b   s   %%a(	]]LL1%IIaQq!4


 mmq!A&q!A&q!A&q"Q'	
 mmr2a(r2a(

 mmr2a(r2a(r2q!,
 mmr2a(r2a(r2a(
 mmr3q)sC*sC*sB1-	
 MMr2a(r2a(IIb"a+
 MMr2q!,r2q!,IIb!QJJL	
  ]]r2q!,r2q!,r2q!,IIb"a 	
 MMIIc3NN3u-GGD!IIc3NN3u-GGD!IIc3NN3u-GGD!IIc3NN3u-GGD!IIc2
r%   c                    UR                   u  p4pVUR                  SX"5      R                  SX"5      R                  X4XR-  Xb-  US-  5      nUR                  SSSSS5      R                  USXR-  Xb-  5      $ )zHUnfold spatial dims by window size ``ws`` and concatenate into channels.rB   r,   r   r-   r?   )shapeunfoldreshapepermute)r   r)   wsBCHWs          r#   	_unfold2dXFeatModel._unfold2d   st    WW
aHHQ&&q"199!RTVWRWXyyAq!Q'//2qwHHr%   c                   [         R                  " 5          UR                  SSS9nU R                  U5      nSSS5        U R	                  U5      nU R                  X R                  U5      -   5      nU R                  U5      nU R                  U5      nU R                  U5      n[        R                  " XTR                  S   UR                  S   4SSS	9n[        R                  " XdR                  S   UR                  S   4SSS	9nU R                  XE-   U-   5      nU R                  U5      nU R                  U R!                  US
S95      n	XyU4$ ! , (       d  f       GN#= f)a  Run the XFeat backbone.

Args:
    x: image tensor of shape :math:`(B, C, H, W)`.

Returns:
    Tuple of:
    - feats: dense descriptors :math:`(B, 64, H/8, W/8)`.
    - keypoints: keypoint logits :math:`(B, 65, H/8, W/8)`.
    - heatmap: reliability map :math:`(B, 1, H/8, W/8)`.
r-   T)dimkeepdimNrX   bilinearFmodealign_cornersrA   )r]   )torchno_gradmeanrG   rJ   rK   rI   rL   rM   rN   FinterpolaterY   rO   rQ   rR   rb   )
r   r)   x1x2x3x4x5featsheatmap	keypointss
             r#   r*   XFeatModel.forward   s&    ]]_1d+A		!A  [[^[[jjm+,[[_[[_[[_]]2bhhrl;*\ab]]2bhhrl;*\ab!!"'B,/##E*&&t~~aA~'>?	((! _s   "E
E#)rJ   rK   rL   rM   rN   rO   rU   rQ   rR   rG   rI   )r0   r1   )rB   )r)   r2   r]   r.   r0   r2   )r)   r2   r0   /Tuple[torch.Tensor, torch.Tensor, torch.Tensor])
r3   r4   r5   r6   r7   r   rb   r*   r8   r9   r:   s   @r#   r<   r<   P   s    "J
XI) )r%   r<   c                  H   ^  \ rS rSrSrSSU 4S jjjrS	S jrS
S jrSrU =r	$ )InterpolateSparse2d   a  Bilinearly or bicubically sample a dense feature map at sparse 2-D positions.

Args:
    mode: interpolation mode for :func:`torch.nn.functional.grid_sample`.
        Default: ``'bicubic'``.
    align_corners: passed to ``grid_sample``. Default: ``False``.

Shape:
    - Input ``x``: :math:`(B, C, H, W)`.
    - Input ``pos``: :math:`(B, N, 2)` integer or float (x, y) coordinates.
    - Output: :math:`(B, N, C)`.
c                :   > [         TU ]  5         Xl        X l        g r'   )r   r   rj   rk   )r   rj   rk   r"   s      r#   r   InterpolateSparse2d.__init__   s    	*r%   c                x    SU[         R                  " US-
  US-
  /UR                  UR                  S9-  -  S-
  $ )aB  Normalise pixel coordinates to the ``[-1, 1]`` range expected by ``grid_sample``.

Uses the formula from the original XFeat implementation:
``2 * x / [W-1, H-1] - 1``.  Note that ``grid_sample`` is always
called with ``align_corners=False``; this asymmetry matches the
convention the pretrained weights were trained with.
g       @r-   )devicedtype      ?)rl   tensorr   r   )r   r)   r`   ra   s       r#   normgridInterpolateSparse2d.normgrid   s;     a%,,Aq1u~ahhaggVVWZ]]]r%   c                    U R                  X#U5      R                  S5      R                  UR                  5      n[        R
                  " XU R                  SS9nUR                  SSSS5      R                  S5      $ )a*  Sample ``x`` at positions ``pos``.

Args:
    x: feature map :math:`(B, C, H, W)`.
    pos: sampling positions :math:`(B, N, 2)` in pixel coordinates.
    H: height used for coordinate normalisation.
    W: width used for coordinate normalisation.

Returns:
    Sampled features :math:`(B, N, C)`.
rg   Fri   r   rB   r,   r-   )	r   	unsqueezetor   ro   grid_samplerj   r\   squeeze)r   r)   posr`   ra   grids         r#   r*   InterpolateSparse2d.forward   sf     }}SQ'11"588A MM!		GyyAq!$,,R00r%   )rk   rj   )bicubicF)rj   strrk   r/   r0   r1   )r)   r2   r`   r.   ra   r.   r0   r2   )
r)   r2   r   r2   r`   r.   ra   r.   r0   r2   )
r3   r4   r5   r6   r7   r   r   r*   r8   r9   r:   s   @r#   r|   r|      s"    + +
^1 1r%   r|   c                  v  ^  \ rS rSr% SrSrS\S'   SSU 4S jjjr\SSS jj5       r	SS jr
\SSS	 jj5       r\SSS
 jj5       r\SS j5       r\S S!S jj5       r S"       S#S jjr S$       S%S jjrS&S'S jjr S(         S)S jjr S*           S+S jjr\R,                  " 5         S,       S-S jj5       r\R,                  " 5         S.       S/S jj5       r\R,                  " 5         S0         S1S jj5       r\R,                  " 5        S2       S3S jj5       rSrU =r$ )4XFeati  a  XFeat sparse and semi-dense local feature extractor and matcher.

Wraps :class:`XFeatModel` with NMS keypoint detection, descriptor
interpolation, and mutual nearest-neighbour matching helpers.

Reference:
    "XFeat: Accelerated Features for Lightweight Image Matching", CVPR 2024.
    https://www.verlab.dcc.ufmg.br/descriptors/xfeat_cvpr24/

.. image:: _static/img/XFeat.png

Args:
    top_k: maximum number of keypoints to keep per image. Default: ``4096``.
    detection_threshold: minimum keypoint score. Default: ``0.05``.

Example:
    >>> model = XFeat()
    >>> img = torch.rand(1, 3, 256, 256)
    >>> out = model.detectAndCompute(img)
    >>> out[0]['keypoints'].shape
    torch.Size([..., 2])
zHhttps://github.com/verlab/accelerated_features/raw/main/weights/xfeat.ptr   weights_urlc                   > [         TU ]  5         Xl        X l        [	        5       U l        [        S5      U l        [        S5      U l        [        S5      U l	        g )Nr   nearestrh   )
r   r   top_kdetection_thresholdr<   netr|   interpolator_nearest	_bilinear)r   r   r   r"   s      r#   r   XFeat.__init__  sG    
#6 </	:+I6,Z8r%   c                    U " XS9n[         R                  R                  U R                  SS9nUR                  R                  U5        UR                  5         U$ )a*  Instantiate XFeat with pretrained weights downloaded from the official release.

Args:
    top_k: maximum number of keypoints to keep. Default: ``4096``.
    detection_threshold: minimum keypoint score. Default: ``0.05``.

Returns:
    XFeat model with pretrained weights loaded, set to eval mode.
)r   r   zxfeat.pt)	file_name)rl   hubload_state_dict_from_urlr   r   load_state_dicteval)clsr   r   model
state_dicts        r#   from_pretrainedXFeat.from_pretrained'  sL     %IYY77S]7^
		!!*-

r%   c                   [        UR                  5       S:H  S5        UR                  SS u  p#[        SUS-  S-  5      [        SUS-  S-  5      pTX$-  X5-  pv[        R
                  " UR                  5       XE4SSS9nXU4$ )	zjResize to the largest multiple of 32 not exceeding the input size (min 32); return image and scale ratios.r?   z'Input must be a 4-D tensor (B, C, H, W)rg   N    rh   Fri   )r	   re   rY   maxro   rp   float)r   r)   r`   ra   _H_Wrhrws           r#   _preprocess_tensorXFeat._preprocess_tensor<  s~    QUUW\#LMwwrs|R!r'R(#b17b.*ABBMM!'')bXJeTbyr%   c                
   [         R                  " X-  S5      SS2SS24   nUR                  u  p4pVUR                  SSSS5      R	                  X5USS5      nUR                  SSSSS5      R	                  USUS-  US-  5      $ )	z@Convert 65-channel keypoint logits to a full-resolution heatmap.r-   NrC   r   rB   r,   rA   r?   )ro   softmaxrY   r\   r[   )kptssoftmax_tempscoresr^   _r`   ra   rw   s           r#   _get_kpts_heatmapXFeat._get_kpts_heatmapE  s     4.21crc6:\\
a..Aq!,44Q1aCq!Q1-55aAE1q5IIr%   c           	        U R                   S   n[        XU4SS9X:  -  nU Vs/ s H(  oUR                  5       SSS24   R                  S5      PM*     nn[	        S U 5       5      n[
        R                  " X7S	4S[
        R                  U R                  S
9n[        U5       H  n	Xi   XS[        Xi   5      24'   M     U$ s  snf )zSNon-maximum suppression on a heatmap; returns (B, N, 2) integer keypoint positions.r   T)	mask_only.r-   NrX   c              3  8   #    U  H  n[        U5      v   M     g 7fr'   )len).0ps     r#   	<genexpr>XFeat._nms.<locals>.<genexpr>S  s     2kc!ffks   rB   )r   r   )rY   r
   nonzeroflipr   rl   fulllongr   ranger   )
r)   	thresholdr!   r^   r   kpos_batchedpad_valpos_outbs
             r#   _nms
XFeat._nmsM  s     GGAJA[1TBamT>ABcyy{37+004cB2k22**a!_b

188TqA0;G,[^,,,-  Cs   /Cc                    [         R                  " [         R                  " XS9[         R                  " XS9SS9u  p4[         R                  " US   US   /S5      R	                  SS5      $ )zDCreate a grid of (x, y) pixel coordinates of shape :math:`(H*W, 2)`.r   ijindexing.NrX   rB   )rl   meshgridarangecatr[   )hwr   yr)   s        r#   
_create_xyXFeat._create_xyY  sV     ~~ell1<ell1>\gklyy!I,)5r:BB2qIIr%   c                   U R                   u  p#n[        R                  " XR                  X#U-  5      -  S5      R                  X#U5      n [        R                  " [        R
                  " X@R                  S9[        R
                  " X0R                  S9SS9u  pVXTS-  -
  nXcS-  -
  n[        R                  " US   U -  S   US   U -  S   /S5      nUR                  X#U-  S5      R                  S5      $ )	z6Compute soft-argmax offsets from an (N, H, W) heatmap.rX   r   xyr   rB   Nr   r-   )	rY   rl   r   viewr   r   r   r   sum)heatmapstempNr`   ra   r)   r   coordss           r#   _subpix_softmax2dXFeat._subpix_softmax2d_  s     ..a==aQ(?!?DII!PQR~~LL??3U\\!OO5\gk
 aLaLQtWx/;ag>PR[=\]_ab{{1!eQ'++A..r%   c                Z   XR                  5       -  nX!R                  5       -  nUR                  SS9u  pgUR                  SS9u  ph[        R                  " [	        U5      UR
                  S9n	X   U	:H  n
US:  a&  UR                  SS9u  pX:  nXU-     n	XzU-     nX4$ X   n	Xz   nX4$ )zDMutual nearest-neighbour matching on L2-normalised descriptor pairs.r-   re   r   r   )tr   rl   r   r   r   )r   feats1feats2
min_cossimcossimcossim_tr   match12match21idx0mutual
cossim_maxgoodidx1s                 r#   
_match_mnnXFeat._match_mnnl  s     ((*$HHJ&ZZAZ&
\\a\(
||CL@!T)>"JJ1J-MJ*D&DD=)D z <D?Dzr%   c                   UR                   S   n[        R                  " XR                  SSS5      5      n[        R                  " USS9n[        R                  " UR                  SSS5      SS9n[        R
                  " UR                   S   UR                  S9n/ n	[        U5       H[  n
Xz   Xj      U:H  nUS:  a(  XZ   R                  SS9u  pX:  nXU-     nXj   X-     nOX   nXj   U   nU	R                  UU45        M]     U	$ )z@Batched MNN matching; returns a list of ``(idx0, idx1)`` tuples.r   rB   r-   rX   r   r   )
rY   rl   bmmr\   argmaxr   r   r   r   append)r   r   r   r   r^   r   r   r   r   batched_matchesr   r   r   r   r   idx0_bidx1_bs                    r#   _batch_matchXFeat._batch_match  s    LLO6>>!Q#:;,,v2.,,v~~aA6B?||GMM!,W^^DqAZ
+t3FA~ &	! 4
!.tm, FM2 F+""FF#34  r%   c           	        US:  a  SnU R                  U5      u  pnU R                  U5      u  pVnUR                  u  ppU R                  XUR                  5      S-  R                  USS5      nUR                  SSSS5      R                  USU	5      nUR                  SSSS5      R                  US5      n[        R                  " U[        UR                  S   U5      SS9u  p[        R                  " USUS	   R                  SSU	5      5      n[        R                  " USUS	   R                  SSS5      5      nU[        R                  " XC/UR                  S
9R                  SS5      -  nUU4$ )zIExtract coarse descriptors from an FPN feature map, top-k by reliability.r-   i rA   rX   r   rB   r,   )r   re   r   r   )r   r   rY   r   r   expandr\   r[   rl   topkmingatherr   r   )r   r)   r   rh1rw1M1_K1H1r^   r_   _H1_W1xy1r   	top_k_idxrv   mkptss                    r#   _extract_denseXFeat._extract_dense  sM   19E--a0hhqkcs3a7??2rJZZ1a#++Ar15ZZ1a#++Ar2zz"BHHQK(?RHRIi$8$?$?B$JKS!Yy%9%@%@R%KLcZEJJ1bQQe|r%   c                2   [         R                  " XSSS9n[         R                  " XSSS9nU R                  U[        US-  5      5      u  pxU R                  U[        US-  5      5      u  p[        R
                  " Xs-  X-  /SS9n[        R                  " UR                  SS	 UR                  S
9SU-  -  n[        R                  " U	R                  SS	 U	R                  S
9SU-  -  nU[        R
                  " X/SS9[        R
                  " X/SS94$ )z/Extract dense features at two scales and merge.Frh   )scale_factorrk   rj   g?g?r-   r   NrB   r   )	ro   rp   r  r.   rl   r   onesrY   r   )r   r)   r   s1s2rq   rr   mkpts_1feats_1mkpts_2feats_2r  sc1sc2s                 r#   _extract_dualscaleXFeat._extract_dualscale  s     ]]1UT]]1UT..r3ut|3DE..r3ut|3DE		7<6A>jjr*7>>Ba"fMjjr*7>>Ba"fMeii
2EIIw>PVW4XXXr%   c                   X4   u  pgUS   U   U   nUS   U   U   n	US   U   U   R                  5       n
US   U   U   nUS   U   U   nU R                  R                  [        R                  " X/SS95      n[
        R                  " US-  SS9R                  SS9S   nU R                  UR                  SSS5      5      nXUS	S	2S	4   -  -   n
X:  n[        R                  " X   X   /SS9$ )
zORefine coarse match positions using the fine-matcher MLP with subpixel softmax.descriptorsrx   scalesrX   r   r,   r   rA   N)
cloner   rU   rl   r   ro   r   r   r   r   )r   d0d1matches	batch_idx	fine_confr   r   r   r   mkpts_0r  sc0offsetsconf	mask_goods                   r#   _refine_matchesXFeat._refine_matches  s    '
M"9-d3M"9-d3[/),T288:[/),T2l9%d+((''		62B(KLyy1"-11b19!<((b!Q)?@c!T'l22$	yy',g.@ArJJr%   c                   Uc  U R                   nUc  U R                  nU R                  U5      u  pnUR                  u  pgpU R	                  U5      u  pn[
        R                  " U
SS9n
U R                  U5      nU R                  XSS9nU R                  XX5      U R                  XX5      -  R                  S5      nSU[        R                  " US:H  SS9'   [        X/R                  S   5      n[        R                  " UUSSSS9u  nn[        R                   " US   SU5      n[        R                   " US	   SU5      n[        R"                  " US
   US
   /SS9nU R%                  XXS9n[
        R                  " USS9nU[        R&                  " XT/UR(                  S9R+                  SSS5      -  nUS:  n[-        U5       Vs/ s H$  nUU   UU      UU   UU      UU   UU      S.PM&     sn$ s  snf )a  Detect sparse keypoints and compute descriptors.

Args:
    x: image tensor of shape :math:`(B, C, H, W)`.
    top_k: number of keypoints to keep (overrides ``self.top_k``).
    detection_threshold: minimum score (overrides ``self.detection_threshold``).

Returns:
    List of length ``B``. Each element is a dict with:

    - ``'keypoints'``: :math:`(N, 2)` keypoints in (x, y) pixel coordinates.
    - ``'scores'``: :math:`(N,)` reliability scores.
    - ``'descriptors'``: :math:`(N, 64)` L2-normalised descriptors.
r-   r      )r   r!   rX   T)r   re   largestsorted).r   ).r-   r   )r`   ra   r   r   )rx   r   r   )r   r   r   rY   r   ro   	normalizer   r   r   r   r   rl   allr  r  r  r   r   r   r   r   r   )r   r)   r   r   r  r  r^   r   r  r  r  K1r
  K1hr  r   r   idxsmkpts_xmkpts_yrv   validr   s                          r#   detectAndComputeXFeat.detectAndCompute  s    * =JJE&"&":":--a0cXXa[
[[#$$R(		#!	L--C5rRU8[[ddegh13uyy""-.||B'(zz&A2tDQ,,uV}b$7,,uV}b$7		79-wy/ABK!!"s!:Er*cZEJJ1aQSTT
 1X
 	 #1XeAh/ )E!H-$Qxa1
 
 	
 
s   +G:c                    Uc  U R                   nU(       a  U R                  X5      u  pEnO?U R                  X5      u  pF[        R                  " UR
                  SS UR                  S9nXFUS.$ )a  Detect keypoints and compute dense coarse descriptors.

Args:
    x: image tensor of shape :math:`(B, C, H, W)`.
    top_k: number of features to keep (overrides ``self.top_k``).
    multiscale: use dual-scale (0.6x and 1.3x) extraction. Default: ``True``.

Returns:
    Dict with:

    - ``'keypoints'``: :math:`(B, K, 2)` coarse keypoints.
    - ``'descriptors'``: :math:`(B, K, 64)` coarse descriptors.
    - ``'scales'``: :math:`(B, K)` extraction scale per keypoint.
NrB   r   )rx   r   r!  )r   r  r  rl   r  rY   r   )r   r)   r   
multiscaler  scrv   s          r#   detectAndComputeDenseXFeat.detectAndComputeDense  se    * =JJE#66q@Eu..q8LEEKKOELLAB"BGGr%   c                    Uc  U R                   nU R                  XS9S   nU R                  X#S9S   nU R                  US   US   US9u  pxUS   U   US   U   4$ )a  Detect, describe and mutually match keypoints from two images.

Args:
    img1: first image tensor of shape :math:`(1, C, H, W)`.
    img2: second image tensor of shape :math:`(1, C, H, W)`.
    top_k: number of top keypoints to use.
    min_cossim: minimum cosine similarity threshold. Use ``-1`` to disable.

Returns:
    Tuple ``(mkpts0, mkpts1)`` of matched keypoints, each :math:`(N, 2)`.
r   r   r   )r   rx   )r   r;  r   )	r   img1img2r   r   out1out2r   r   s	            r#   match_xfeatXFeat.match_xfeat&  s    & =JJE$$T$7:$$T$7:__T-%8$}:MZd_e
K &[(9$(???r%   c           
     L   Uc  U R                   nU R                  XS9nU R                  X#S9nU R                  US   US   5      nUR                  S   n[	        U5       Vs/ s H  oR                  XEXhS9PM     n	nUS:  a  U	$ U	S   SS2SS24   U	S   SS2SS24   4$ s  snf )a|  Extract coarse features, match pairs and refine matches (XFeat*).

Args:
    im_set1: batch of images :math:`(B, C, H, W)`.
    im_set2: batch of images :math:`(B, C, H, W)`.
    top_k: number of top features to use.

Returns:
    If ``B > 1``: list of :math:`(N, 4)` tensors with ``(x1, y1, x2, y2)`` matches.
    If ``B == 1``: tuple of :math:`(N, 2)` matched keypoint tensors.
NrC  r   r   )r%  r&  r-   rB   )r   r@  r   rY   r   r-  )
r   im_set1im_set2r   rF  rG  	idxs_listr^   r   r%  s
             r#   match_xfeat_starXFeat.match_xfeat_star@  s    $ =JJE))')?))')?%%d=&94;NO	MM!]bcd]ef]eXY''I'S]efq5Nqz!RaR% '!*QU"333 gs   "B!)r   r   r   r   r   r   )i   皙?)r   r.   r   r   r0   r1   )r   r.   r   r   r0   r   )r)   r2   r0   z!Tuple[torch.Tensor, float, float])r   )r   r2   r   r   r0   r2   )rP  r0  )r)   r2   r   r   r!   r.   r0   r2   )r   r.   r   r.   r   ztorch.devicer0   r2   )g      @)r   r2   r   r   r0   r2   )g=
ףp=?)r   r2   r   r2   r   r   r0   !Tuple[torch.Tensor, torch.Tensor])rX   )r   r2   r   r2   r   r   r0   'List[Tuple[torch.Tensor, torch.Tensor]])i@  )r)   r2   r   r.   r0   rQ  )g333333?g?)
r)   r2   r   r.   r  r   r  r   r0   rz   )g      ?)r#  Dict[str, torch.Tensor]r$  rS  r%  rR  r&  r.   r'  r   r0   r2   )NN)r)   r2   r   Optional[int]r   zOptional[float]r0   zList[Dict[str, torch.Tensor]])NT)r)   r2   r   rT  r>  r/   r0   rS  )NrX   )
rD  r2   rE  r2   r   rT  r   r   r0   rQ  r'   )rK  r2   rL  r2   r   rT  r0   z6List[torch.Tensor] | Tuple[torch.Tensor, torch.Tensor])r3   r4   r5   r6   r7   r   __annotations__r   classmethodr   r   staticmethodr   r   r   r   r   r   r  r  r-  rl   inference_moder;  r@  rH  rN  r8   r9   r:   s   @r#   r   r     s   . bKa9 9  ( J J 	 	 J J
 
/ 
/ OS",8FK	** OQ",8FK	0." ILYY&)Y/4Y@EY	8Y&  K#K $K 9	K
 K K 
K4   $/3	7
7
 7
 -	7

 
'7
 7
r   $	HH H 	H
 
!H H: 
  $@@ @ 	@
 @ 
+@ @2 
  $	44 4 	4
 
@4 4r%   r   )r7   
__future__r   typingr   r   r   r   rl   torch.nn.functionalr   
functionalro   kornia.core.checkr	   kornia.geometry.subpixr
   Moduler   r<   r|   r    r%   r#   <module>ra     sn   $ # . .     * (! !H@) @)F.1")) .1bW4BII W4r%   