o oi @spUddlmZddlZddlmZddlmZiZeeefed<ded<ded <d ed <Gd d d ej Z dS))DictN)nn)KORNIA_CHECK_SHAPEurlszQhttps://github.com/vbalnt/tfeat/raw/master/pretrained-models/tfeat-liberty.paramslibertyzShttps://github.com/vbalnt/tfeat/raw/master/pretrained-models/tfeat-notredame.params notredamezRhttps://github.com/vbalnt/tfeat/raw/master/pretrained-models/tfeat-yosemite.paramsyosemitecsFeZdZdZdZd deddffdd Zd ejdejfd d Z Z S) TFeataModule, which computes TFeat descriptors of given grayscale patches of 32x32. This is based on the original code from paper "Learning local feature descriptors with triplets and shallow convolutional neural networks". See :cite:`TFeat2016` for more details Args: pretrained: Download and set pretrained weights to the model. Returns: torch.Tensor: TFeat descriptor of the patches. Shape: - Input: :math:`(B, 1, 32, 32)` - Output: :math:`(B, 128)` Examples: >>> input = torch.rand(16, 1, 32, 32) >>> tfeat = TFeat() >>> descs = tfeat(input) # 16x128 F pretrainedreturnNc stttjdddtjddddttjdddtjdd d dt|_tt d d t|_ |rMt j j td t dd}|j|dd|dS)NF)affiner ) kernel_size)rstride@ircpu) map_locationT)strict)super__init__r SequentialInstanceNorm2dConv2dTanh MaxPool2dfeaturesLineardescrtorchhubload_state_dict_from_urlrdeviceload_state_dicteval)selfr pretrained_dict __class__H/home/ubuntu/.local/lib/python3.10/site-packages/kornia/feature/tfeat.pyr9s    zTFeat.__init__inputcCs8t|gd||}||dd}||}|S)N)B132r2r)rr viewsizer")r)r/xr-r-r.forwardJs   z TFeat.forward)F) __name__ __module__ __qualname____doc__ patch_sizeboolrr#Tensorr7 __classcell__r-r-r+r.r s r ) typingrr#rkornia.core.checkrrstr__annotations__Moduler r-r-r-r.s