o €o™iíã@shUddlmZddlZddlmZddlmZiZeeefed<ded<ded <Gd d„dej ƒZ dS)é)ÚDictN)Únn)ÚKORNIA_CHECK_SHAPEÚurlszVhttps://github.com/yuruntian/SOSNet/raw/master/sosnet-weights/sosnet_32x32_liberty.pthÚlibzYhttps://github.com/yuruntian/SOSNet/raw/master/sosnet-weights/sosnet_32x32_hpatches_a.pthÚhp_acsLeZdZdZdZddeddf‡fdd„ Zdd ejde dejfdd „Z ‡ZS)ÚSOSNetaò128-dimensional SOSNet model definition for 32x32 patches. This is based on the original code from paper "SOSNet:Second Order Similarity Regularization for Local Descriptor Learning". Args: pretrained: Download and set pretrained weights to the model. Shape: - Input: :math:`(B, 1, 32, 32)` - Output: :math:`(B, 128)` Examples: >>> input = torch.rand(8, 1, 32, 32) >>> sosnet = SOSNet() >>> descs = sosnet(input) # 8x128 é FÚ pretrainedÚreturnNcsrtƒ ¡t tjdddtjddddddtjdddt ¡tjddddddtjdddt ¡tjddddddd tjdddt ¡tjddddddtjdddt ¡tjdd ddddd tjd ddt ¡tjd d ddddtjd ddt ¡t d¡tjd d ddd tjd dd¡|_ t tj ddddd¡|_|r³tj jtdt d¡d}|j|dd| ¡dS)NéF)Úaffiner é)Úkernel_sizeÚpaddingÚbiasé@é)rÚstriderré€gš™™™™™¹?é)rrégp@gà?g)ÚalphaÚbetaÚkrÚcpu)Úmap_locationT)Ústrict)ÚsuperÚ__init__rÚ SequentialÚInstanceNorm2dÚConv2dÚBatchNorm2dÚReLUÚDropoutÚlayersÚLocalResponseNormÚ desc_normÚtorchÚhubÚload_state_dict_from_urlrÚdeviceÚload_state_dictÚeval)Úselfr Úpretrained_dict©Ú __class__©úI/home/ubuntu/.local/lib/python3.10/site-packages/kornia/feature/sosnet.pyr4s< êzSOSNet.__init__ç»½×Ùß|Û=ÚinputÚepscCs8t|gd¢ƒ| | |¡|¡}| | d¡d¡}|S)N)ÚBÚ1Ú32r:réÿÿÿÿ)rr(r&ÚviewÚsize)r/r6r7Údescrr3r3r4ÚforwardUszSOSNet.forward)F)r5)Ú__name__Ú __module__Ú__qualname__Ú__doc__Ú patch_sizeÚboolrr)ÚTensorÚfloatr?Ú __classcell__r3r3r1r4rs $!r)Útypingrr)rÚkornia.core.checkrrÚstrÚ__annotations__ÚModulerr3r3r3r4Ús