:j3ddlZddlmZddlmZddlmZmZmZm Z ddl Z ddl m Z m Z ddl mZddlmZdd lmZd d lmZd d lmZmZmZd d lmZd dlmZmZd dlmZm Z ddlmZ!ddl"m#Z#ddl$m%Z%ddl&m'Z'm(Z(ddgZ)de*de*de*dede jVfde jXf dZ-de*de*dede jVfde jXfdZ.de jVfd Z/Gd!d"e jVZ0Gd#d$e(Z1Gd%d&e(Z2Gd'd(e jVZ3d)e ejhejjfd*e*dede jVffd+Z6Gd,deZ7eed-e7jpfd.e jrf/dd0de jrddd1d2ee7d3e:d4ee*d5ee d6ee*deede jVfd7ede'fd8Z;y)9N) OrderedDict)partial)AnyCallableOptionalUnion)nnTensor)Conv2dNormActivation)ObjectDetection)_log_api_usage_once) mobilenet)register_modelWeights WeightsEnum)_COCO_CATEGORIES)_ovewrite_value_paramhandle_legacy_interface)mobilenet_v3_largeMobileNet_V3_Large_Weights)_utils)DefaultBoxGenerator)_validate_trainable_layers)SSDSSDScoringHead%SSDLite320_MobileNet_V3_Large_Weightsssdlite320_mobilenet_v3_large in_channels out_channels kernel_size norm_layer.returnc tjt|||||tjtj||dS)N)r#groupsr$activation_layerr)r Sequentialr ReLU6Conv2d)r!r"r#r$s y/home/mcse/projects/srt_converter/srt-converter-venv/lib/python3.12/site-packages/torchvision/models/detection/ssdlite.py_prediction_blockr-sE ==  #!XX   +|Q/  c tj}|dz}tjt||d||t||dd|||t||d||S)Nrr)r#r$r(r )r#strider'r$r()r r*r)r )r!r"r$ activationintermediate_channelss r, _extra_blockr30sqJ(A- == .A*gq   ! !(!'   ! a Hg   /+:OQRT^_ ` aZ  [1r. rVrWrXrYrZrr r[rJr^r_s@r,rKrK`sG2923792KN2\dehjljsjses\t22r.rKcVeZdZdeedeededejfffd ZxZ S)rMr!rEr$.c tj}t||D]%\}}|jt |d|zd|'t |t ||dy)Nr rb)rOr!rEr$bbox_regrfrgrPs r,rJzSSDLiteRegressionHead.__init__ls^==?!$[+!> U Hg OO-hG Q S T UX 1%r.rhr_s@r,rMrMks=&DI&DI&S[\_acajaj\jSk&&r.rMc eZdZ d dejdededejfdedef fd Zde d e e e ffd Z xZ S) SSDLiteFeatureExtractorMobileNetbackbonec4_posr$. width_mult min_depthc rt|t|||jr t dt j t j g|d|||jdt j ||jddg||dzd|_fd}t jt|dj|d|t|d|d|t|d|d|t|d|d|g}t|||_ y) Nz0backbone[c4_pos].use_res_connect should be Falserrc4tt|zSrH)maxrZ)drrrqs r,z;SSDLiteFeatureExtractorMobileNet.__init__..sc)SZ-@Ar.i)rIrJruse_res_connect ValueErrorr r)blockfeaturesrcr3r"rBextra) rOrorpr$rqrr get_depthrrPs `` r,rJz)SSDLiteFeatureExtractorMobileNet.__init__us D! F  + +OP P MM H8GV, Hhv.>.D.DQ.G H MM(6*004 Nx 7M N  B  Xb\66 # SYs^Ys^ZHYs^Ys^ZHYs^Ys^ZH    U r.rQr%c g}|jD]}||}|j||jD]}||}|j|tt |Dcgc]\}}t ||fc}}Scc}}wrH)r~rerr enumerater])rOrQoutputr}ivs r,rUz(SSDLiteFeatureExtractorMobileNet.forwards]] EaA MM!  ZZ EaA MM!  If4EFDAqSVQKFGGFs*B )g?)rVrWrXr r[rZrfloatrJr r\r]rUr^r_s@r,rnrntsq   ))S"))^,     B H HDf$5 Hr.rnrotrainable_layersc |j}dgt|Dcgc]\}}t|dds|c}}zt|dz gz}t|}d|cxkr|kst dt d|dk(r t|n|||z }|d|D](}|j D]}|j d*t||d|Scc}}w)Nr_is_cnFrzYtrainable_layers should be in the range [0, {num_stages}], instead got {trainable_layers})r~rgetattrlenr| parametersrequires_grad_rn) rorr$rb stage_indices num_stages freeze_before parameters r,_mobilenet_extractorrs   HC8)<\A8UZ@[1\\`cdl`mpq`q_rrM]#J  .J .tuu /tuu%5%:CM j[kNk@lM n} %, ,I  $ $U + ,, ,HmB6G TT]s CCc @eZdZedededddddiidd d d  ZeZy )rzShttps://download.pytorch.org/models/ssdlite320_mobilenet_v3_large_coco-a79551df.pthi}4)rrz]https://github.com/pytorch/vision/tree/main/references/detection#ssdlite320-mobilenetv3-largez COCO-val2017box_mapgL5@g-?gt*@zSThese weights were produced by following a similar training recipe as on the paper.) num_params categoriesmin_sizerecipe_metrics_ops _file_size_docs)url transformsmetaN)rVrWrXrr rCOCO_V1DEFAULTr.r,rrsG a"!*ut!  n  G$Gr. pretrainedpretrained_backbone)weightsweights_backboneT)rprogressrFrtrainable_backbone_layersr$rrrFrrkwargsc tj|}tj|}d|vrtjd|&d}t d|t |jd}n|d}t|duxs|du|dd}|du}|ttjdd }td||||d |}| t|t|||}d } tt!dD cgc]} d dgc} dd} t#j$|| } | j'} t | t | j(k7r-t+dt | dt | j(ddddgdgdd}i||}t-|| | |fdt/| | ||i|}|"|j1|j3|d|Scc} w)aSSDlite model architecture with input size 320x320 and a MobileNetV3 Large backbone, as described at `Searching for MobileNetV3 `__ and `MobileNetV2: Inverted Residuals and Linear Bottlenecks `__. .. betastatus:: detection module See :func:`~torchvision.models.detection.ssd300_vgg16` for more details. Example: >>> model = torchvision.models.detection.ssdlite320_mobilenet_v3_large(weights=SSDLite320_MobileNet_V3_Large_Weights.DEFAULT) >>> model.eval() >>> x = [torch.rand(3, 320, 320), torch.rand(3, 500, 400)] >>> predictions = model(x) Args: weights (:class:`~torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool, optional): If True, displays a progress bar of the download to stderr. Default is True. num_classes (int, optional): number of output classes of the model (including the background). weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained weights for the backbone. trainable_backbone_layers (int, optional): number of trainable (not frozen) layers starting from final block. Valid values are between 0 and 6, with 6 meaning all backbone layers are trainable. If ``None`` is passed (the default) this value is set to 6. norm_layer (callable, optional): Module specifying the normalization layer to use. **kwargs: parameters passed to the ``torchvision.models.detection.ssd.SSD`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.detection.SSDLite320_MobileNet_V3_Large_Weights :members: sizez?The size of the model is already fixed; ignoring the parameter.NrFr[gMbP?r6)epsmomentum)rrr$ reduced_tail)@rrr g?gffffff?) min_ratio max_ratioz4The length of the output channels from the backbone z? do not match the length of the anchor generator aspect ratios g?i,)?rr) score_thresh nms_threshdetections_per_imgtopk_candidates image_mean image_stdheadT)r check_hashr)rverifyrwarningswarnrrrrrr BatchNorm2drrBrrrange det_utilsretrieve_out_channelsnum_anchors_per_location aspect_ratiosr|rrDload_state_dictget_state_dict)rrrFrrr$r reduce_tailror_anchor_generatorr"rEdefaultsmodels r,r r sIp4::7CG1889IJ  WX+M;GLLYeLfHgh   :t;/t;=VXY[\! #d*KR^^F ! 8 YdhnHX#!H D*E!H+EqQF+EQT`de228TBL";;=K <C 0 > >??B3|CTBUVUVYZjZxZxVyUz {  !&$ H)X((F     {K L     E g44hSW4XY L?,Fs& G)rsD#11 ,2(77$CH!-6$,# $'69GOPSUWU^U^P^G_]]$c(3PRPYPY>BZ_a_l_l29ryy9  "))  22&N&-Hryy-H`UI))9+@+@@AUUbii(U.K, @HH I+-G-U-UV @D!%=W=e=e/359u ; <uu# u 9: u (} u#ryy.12uu u  ur.