a h: @sddlmZddlmZddlmZmZddlZddlmZddl m Z ddl m Z d d lmZmZmZd d lmZd d lmZmZmZd d lmZmZmZd dlmZmZmZm Z m!Z!ddlm"Z"ddl#m$Z$gdZ%Gddde"Z&Gdddej'Z(Gdddej'Z)Gdddej'Z*Gdddej+Z,ee-ee.e&dddZ/ed d!d"Z0Gd#d$d$eZ1Gd%d&d&eZ2Gd'd(d(eZ3ee-ee.e&dd)d*Z4eed+e1j5fd,e!j6fd-dd.dde!j6d/ee1e.ee-ee.ee!ee&d0d1d2Z7eed+e2j5fd,ej6fd-dd.ddej6d/ee2e.ee-ee.eeee&d0d3d4Z8eed+e3j5fd,ej6fd-dd.ddej6d/ee3e.ee-ee.eeee&d0d5d6Z9dS)7)Sequence)partial)AnyOptionalN)nn) functional)SemanticSegmentation)register_modelWeights WeightsEnum)_VOC_CATEGORIES)_ovewrite_value_paramhandle_legacy_interfaceIntermediateLayerGetter)mobilenet_v3_largeMobileNet_V3_Large_Weights MobileNetV3)ResNet resnet101ResNet101_Weightsresnet50ResNet50_Weights)_SimpleSegmentationModel)FCNHead) DeepLabV3DeepLabV3_ResNet50_WeightsDeepLabV3_ResNet101_Weights$DeepLabV3_MobileNet_V3_Large_Weightsdeeplabv3_mobilenet_v3_largedeeplabv3_resnet50deeplabv3_resnet101c@seZdZdZdS)ra Implements DeepLabV3 model from `"Rethinking Atrous Convolution for Semantic Image Segmentation" `_. Args: backbone (nn.Module): the network used to compute the features for the model. The backbone should return an OrderedDict[Tensor], with the key being "out" for the last feature map used, and "aux" if an auxiliary classifier is used. classifier (nn.Module): module that takes the "out" element returned from the backbone and returns a dense prediction. aux_classifier (nn.Module, optional): auxiliary classifier used during training N)__name__ __module__ __qualname____doc__r(r(W/var/www/auris/lib/python3.9/site-packages/torchvision/models/segmentation/deeplabv3.pyrsrcs.eZdZdeeeeddfdd ZZS) DeepLabHead $N) in_channels num_classes atrous_ratesreturnc sBtt||tjddddddtdttd|ddS)NrrF)paddingbias)super__init__ASPPrConv2d BatchNorm2dReLU)selfr/r0r1 __class__r(r)r72s zDeepLabHead.__init__)r+)r$r%r&intrr7 __classcell__r(r(r=r)r*1sr*cs(eZdZeeeddfdd ZZS)ASPPConvN)r/ out_channelsdilationr2cs6tj||d||ddt|tg}tj|dS)NrF)r4rCr5)rr9r:r;r6r7)r<r/rBrCmodulesr=r(r)r7=s zASPPConv.__init__)r$r%r&r?r7r@r(r(r=r)rA<srAcs:eZdZeeddfdd ZejejdddZZS) ASPPPoolingN)r/rBr2c s4ttdtj||dddt|tdS)NrFr5)r6r7rZAdaptiveAvgPool2dr9r:r;)r<r/rBr=r(r)r7Gs zASPPPooling.__init__xr2cCs2|jdd}|D] }||}qtj||dddS)NZbilinearF)sizemodeZ align_corners)shapeFZ interpolate)r<rHrJmodr(r(r)forwardOs zASPPPooling.forward) r$r%r&r?r7torchTensorrOr@r(r(r=r)rEFsrEcsBeZdZd eeeeddfdd ZejejdddZZ S) r8r3N)r/r1rBr2c stg}|ttj||dddt|tt|}|D]}|t |||qF|t ||t ||_ ttjt |j ||dddt|ttd|_dS)NrFrFg?)r6r7appendr Sequentialr9r:r;tuplerArEZ ModuleListconvslenZDropoutproject)r<r/r1rBrDZratesZrater=r(r)r7Ws  $ z ASPP.__init__rGcCs6g}|jD]}|||q tj|dd}||S)Nr)Zdim)rUrRrPcatrW)r<rHZ_resconvresr(r(r)rOms  z ASPP.forward)r3) r$r%r&r?rr7rPrQrOr@r(r(r=r)r8Vsr8)backboner0auxr2cCsHddi}|rd|d<t||d}|r.td|nd}td|}t|||S)NZlayer4outr\Zlayer3 return_layersii)rrr*r)r[r0r\r_aux_classifier classifierr(r(r)_deeplabv3_resnetus  rb)rrz These weights were trained on a subset of COCO, using only the 20 categories that are present in the Pascal VOC dataset. ) categoriesZmin_sizeZ_docsc @sDeZdZedeeddiedddddd id d d d ZeZdS)rzHhttps://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pthZ resize_sizeijzVhttps://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_resnet50COCO-val2017-VOC-labelsgP@皙W@ZmiouZ pixel_accgvWf@gGzd@Z num_paramsZrecipeZ_metricsZ_ops _file_sizeurlZ transformsmetaN r$r%r&r rr _COMMON_METACOCO_WITH_VOC_LABELS_V1DEFAULTr(r(r(r)rs$ rc @sDeZdZedeeddiedddddd id d d d ZeZdS)rzIhttps://download.pytorch.org/models/deeplabv3_resnet101_coco-586e9e4e.pthrdreijzQhttps://github.com/pytorch/vision/tree/main/references/segmentation#fcn_resnet101rfgP@rgrhgS+p@gm&m@rirkNrnr(r(r(r)rs$ rc @sDeZdZedeeddiedddddd id d d d ZeZdS)r zMhttps://download.pytorch.org/models/deeplabv3_mobilenet_v3_large-fc3c493d.pthrdreiPKz`https://github.com/pytorch/vision/tree/main/references/segmentation#deeplabv3_mobilenet_v3_largerfgfffff&N@gV@rhgCl$@gJ +&E@rirkNrnr(r(r(r)r s$ r c Cs|j}dgddt|Dt|dg}|d}||j}|d}||j}t|di}|rld|t|<t||d }|rt||nd} t||} t|| | S) NrcSs g|]\}}t|ddr|qS)Z_is_cnF)getattr).0ibr(r(r) z*_deeplabv3_mobilenetv3..rr]r\r^) features enumeraterVrBstrrrr*r) r[r0r\Z stage_indicesZout_posZ out_inplanesZaux_posZ aux_inplanesr_r`rar(r(r)_deeplabv3_mobilenetv3s&      r}Z pretrainedZpretrained_backbone)weightsweights_backboneT)r~progressr0aux_lossr)r~rr0rrkwargsr2cKst|}t|}|durDd}td|t|jd}td|d}n |durPd}t|gdd}t|||}|dur||j |dd |S) adConstructs a DeepLabV3 model with a ResNet-50 backbone. .. betastatus:: segmentation module Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. Args: weights (:class:`~torchvision.models.segmentation.DeepLabV3_ResNet50_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.segmentation.DeepLabV3_ResNet50_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) aux_loss (bool, optional): If True, it uses an auxiliary loss weights_backbone (:class:`~torchvision.models.ResNet50_Weights`, optional): The pretrained weights for the backbone **kwargs: unused .. autoclass:: torchvision.models.segmentation.DeepLabV3_ResNet50_Weights :members: Nr0rcrTFTTr~Zreplace_stride_with_dilationrZ check_hash) rverifyrrrVrmrrbload_state_dictget_state_dictr~rr0rrrr[modelr(r(r)r"s%   r"cKst|}t|}|durDd}td|t|jd}td|d}n |durPd}t|gdd}t|||}|dur||j |dd |S) aiConstructs a DeepLabV3 model with a ResNet-101 backbone. .. betastatus:: segmentation module Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. Args: weights (:class:`~torchvision.models.segmentation.DeepLabV3_ResNet101_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.segmentation.DeepLabV3_ResNet101_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) aux_loss (bool, optional): If True, it uses an auxiliary loss weights_backbone (:class:`~torchvision.models.ResNet101_Weights`, optional): The pretrained weights for the backbone **kwargs: unused .. autoclass:: torchvision.models.segmentation.DeepLabV3_ResNet101_Weights :members: Nr0rcrTrrrr) rrrrrVrmrrbrrrr(r(r)r#s%   r#cKst|}t|}|durDd}td|t|jd}td|d}n |durPd}t|dd}t|||}|dur||j |dd|S) akConstructs a DeepLabV3 model with a MobileNetV3-Large backbone. Reference: `Rethinking Atrous Convolution for Semantic Image Segmentation `__. Args: weights (:class:`~torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.segmentation.DeepLabV3_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) aux_loss (bool, optional): If True, it uses an auxiliary loss weights_backbone (:class:`~torchvision.models.MobileNet_V3_Large_Weights`, optional): The pretrained weights for the backbone **kwargs: unused .. autoclass:: torchvision.models.segmentation.DeepLabV3_MobileNet_V3_Large_Weights :members: Nr0rcrTr)r~Zdilatedr) r rrrrVrmrr}rrrr(r(r)r!Ts#    r!):collections.abcr functoolsrtypingrrrPrZtorch.nnrrMZtransforms._presetsr Z_apir r r Z_metar_utilsrrrZ mobilenetv3rrrZresnetrrrrrrZfcnr__all__rrSr*rArEModuler8r?boolrbrorrr r}rpZ IMAGENET1K_V1r"r#r!r(r(r(r)s               33