o )j@snddlZddlmZdgZdZdededefddZd d d d d ddddedededededede fddZ ddde fddZ ddde fddZ ddde fddZ ddde fddZdd d!dd"d#ed$ed%ede fd&d'ZGd(d)d)ejZd d!dd*ded%ede fd+d,Zd!dd-d%ede fd.d/Zddde fd0d1Zddde fd2d3Zddde fd4d5ZdS)6Ntorchz%https://dl.fbaipublicfiles.com/dinov2 arch_name patch_sizereturncCs"|dddd}d||S)N_Zdinov2_)replace)rrZcompact_arch_namer l/var/www/html/Deteccion_Ine/venv/lib/python3.10/site-packages/modelscope/models/cv/anydoor/dinov2/hubconf.py_make_dinov2_model_namesr vit_largeig?ZmlpT)rimg_sizer init_values ffn_layer block_chunks pretrainedrrrrrc KsPddlm}t||} t|||||d} | jdi||j|di| } | S)N)vision_transformer)rrrrrr )Z dinov2.modelsrr dictupdate__dict__) rrrrrrrkwargsZvitsrZ vit_kwargsmodelr r r _make_dinov2_models r)rcKtdd|d|S)zP DINOv2 ViT-S/14 model (optionally) pretrained on the LVD-142M dataset. vit_smallrrNr rrrr r r dinov2_vits142 r!cKr)zC DINOv2 ViT-B/14 model pretrained on the LVD-142M dataset. vit_baserNr rr r r r dinov2_vitb14:r"r$cKr)zP DINOv2 ViT-L/14 model (optionally) pretrained on the LVD-142M dataset. r rNr rr r r r dinov2_vitl14Br"r%cKtddd|d|S)zP DINOv2 ViT-g/14 model (optionally) pretrained on the LVD-142M dataset. vit_giant2 swiglufusedrrrNr rr r r r dinov2_vitg14Jsr*ir model_name embed_dimlayersrr,r-r.c Ks|dvs Jd|td||d}|r>|dkrt|nd}td|d|d|d }tjj|d d }|j|d d |S)N)rrUnsupported number of layers: rirr/Z_linearz _head.pthcpu)Z map_locationF)strict)nnZLinearstr_DINOV2_BASE_URLrZhubZload_state_dict_from_urlZload_state_dict) r,r-r.rr linear_headZ layers_strurlZ state_dictr r r _make_dinov2_linear_headUsr8cs<eZdZdddejdejdeffddZdd ZZS) _LinearClassifierWrapperr)r.backboner6r.cs t||_||_||_dS)N)super__init__r:r6r.)selfr:r6r. __class__r r r<ls  z!_LinearClassifierWrapper.__init__c Cs|jdkr$|j|}|dd}|dd}t||dg}nG|jdkrc|jj|ddd}t|ddd|ddd|ddd|d dd|d dddg}nJd |j||S) NrZx_norm_clstokenrZx_norm_patchtokensrT)nZreturn_class_tokenFr/) r.r:Zforward_featuresZsqueezercatmeanZget_intermediate_layersr6)r=xZ cls_tokenZ patch_tokensZ linear_inputr r r forwardvs   0& z _LinearClassifierWrapper.forward) __name__ __module__ __qualname__r3Moduleintr<rF __classcell__r r r>r r9js r9rr.rc KsHtd||d|}|j}|j}t||}t||||d}t|||dS)Nrr+)r:r6r.r )rr-rr r8r9) rr.rrr:r-rr,r6r r r _make_dinov2_linear_classifiers" rN)r.rcKstdd||d|S)z Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-S/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k. rrMNr rN)r.rrr r r dinov2_vits14_lcs rPcKr)z Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-B/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k. r#rNr rOr r r r dinov2_vitb14_lc rQcKr)z Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-L/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k. r rNr rOr r r r dinov2_vitl14_lcrRrScKr&)z Linear classifier (1 or 4 layers) on top of a DINOv2 ViT-g/14 backbone (optionally) pretrained on the LVD-142M dataset and trained on ImageNet-1k. r'r(r)Nr rOr r r r dinov2_vitg14_lcsrT)rZtorch.nnr3 dependenciesr5r4rKr floatboolrr!r$r%r*r8rJr9rNrPrQrSrTr r r r sv