U (dk@sddlmZmZddlZddlmZddlmmZddlm Z ddl m Z ddl m Z ddlmZddlmZdd lmZd d lmZmZd d lmZd dlmZmZmZdZGdddejZGdddejZGdddejZ GdddejZ!GdddejZ"ddZ#GdddejZ$GdddejZ%Gd d!d!ejZ&Gd"d#d#ejZ'Gd$d%d%ejZ(Gd&d'd'ejZ)d(d)iZ*Gd*d+d+eZ+Gd,d-d-eZ,dd.d/d0d1Z-ed2e+j.fd3dd4d/ee+e)d5d6d7Z/ed2e,j0fd3dd4d/ee,e)d5d8d9Z1dS):)ListOptionalN)Tensor) BatchNorm2d)InstanceNorm2d)Conv2dNormActivation) OpticalFlow)_log_api_usage_once)Weights WeightsEnum)handle_legacy_interface) grid_samplemake_coords_grid upsample_flow)RAFT raft_large raft_smallRaft_Large_WeightsRaft_Small_Weightscs.eZdZdZddfdd ZddZZS) ResidualBlockzGsr>cs<eZdZdZedejdfdd ZddZdd Z Z S) FeatureEncoderzThe feature encoder, used both as the actual feature encoder, and as the context encoder. It must downsample its input by 8. @rC`blocklayersrcs2tt|dkr(tdt|td|d|dddd|_|j||d|d |d d |_|j||d |d|dd |_|j||d|d|dd |_ t j |d|d d d |_ | D]n}t|t j rt jj|jd ddqt|t jt jfr|jdk rt j|jd |jdk rt j|jdqdS)Nz0The expected number of layers is 5, instead got rrr Trr)r first_strider?)rZfan_outr*)modeZ nonlinearity)r"r#len ValueErrorr convnormrelu_make_2_blockslayer1layer2layer3r&Conv2dconvmodules isinstanceinitZkaiming_normal_ZweightrrZ constant_r)r,rHrIrmr/r1r2r#ws.      zFeatureEncoder.__init__cCs,|||||d}||||dd}t||S)N)rrr)r&Z Sequential)r,rHr-r.rrLZblock1Zblock2r1r1r2rQszFeatureEncoder._make_2_blockscCs6||}||}||}||}||}|Sr3)rPrRrSrTrVr,r5r1r1r2r7s      zFeatureEncoder.forward) r9r:r;r<rr&rr#rQr7r=r1r1r/r2rAqsrAcs2eZdZdZddddfdd Zdd ZZS) MotionEncoderzThe motion encoder, part of the update block. Takes the current predicted flow and the correlation features as input and returns an encoded version of these. rFrErCrE) corr_layers flow_layersr.cstt|dkr(tdt|t|dkrFtdt|t||dddd|_t|dkrt|d|dddd|_n t|_td|ddd d|_ t|d|dddd|_ t|d |d |dddd|_ ||_ dS) Nr z5The expected number of flow_layers is 2, instead got rr z8The number of corr_layers should be 1 or 2, instead got rrrrrrK) r"r#rNrOr convcorr1 convcorr2r&r' convflow1 convflow2rVr.)r,in_channels_corrr`rar.r/r1r2r#s$     zMotionEncoder.__init__cCsZ||}||}|}||}||}tj||gdd}||}tj||gddSNrdim)rerfrgrhtorchcatrV)r,flow corr_featurescorrZ flow_origZ corr_flowr1r1r2r7s     zMotionEncoder.forwardr8r1r1r/r2r\sr\cs(eZdZdZfddZddZZS)ConvGRUzConvolutional Gru unit.csVttj|||||d|_tj|||||d|_tj|||||d|_dS)Nrpadding)r"r#r&rUconvzconvrconvqr, input_size hidden_sizerrtr/r1r2r#s zConvGRU.__init__cCsltj||gdd}t||}t||}t|tj|||gdd}d||||}|Srj)rmrnZsigmoidrurvtanhrw)r,hr5hxzrqr1r1r2r7s "zConvGRU.forwardr8r1r1r/r2rrs rrcCs|Sr3r1)r|_r1r1r2_pass_through_hsrcs0eZdZdZdddfdd ZddZZS) RecurrentBlockzRecurrent block, part of the update block. Takes the current hidden state and the concatenation of (motion encoder output, context) as input. Returns an updated hidden state. )rrJ)rJr)rr )r rrscstt|t|kr6tdt|dt|t|dkrTtdt|t|||d|dd|_t|dkrt|||d|dd|_nt|_||_dS) NzSkernel_size should have the same length as padding, instead got len(kernel_size) = z and len(padding) = rbz.kernel_size should either 1 or 2, instead got rryrzrrtr r) r"r#rNrOrrconvgru1convgru2rrzrxr/r1r2r#s,    zRecurrentBlock.__init__cCs|||}|||}|Sr3)rr)r,r|r5r1r1r2r7s  zRecurrentBlock.forwardr8r1r1r/r2rsrcs(eZdZdZfddZddZZS)FlowHeadzFlow head, part of the update block. Takes the hidden state of the recurrent unit as input, and outputs the predicted "delta flow". csDttj||ddd|_tj|dddd|_tjdd|_dS)Nrrrtr Tr!)r"r#r&rUconv1conv2r)r*)r,r-rzr/r1r2r# s zFlowHead.__init__cCs||||Sr3)rr*rr[r1r1r2r7szFlowHead.forwardr8r1r1r/r2rs rcs(eZdZdZfddZddZZS) UpdateBlockzThe update block which contains the motion encoder, the recurrent block, and the flow head. It must expose a ``hidden_state_size`` attribute which is the hidden state size of its recurrent block. cs(t||_||_||_|j|_dSr3)r"r#motion_encoderrecurrent_block flow_headrzhidden_state_size)r,rrrr/r1r2r#s  zUpdateBlock.__init__cCs<|||}tj||gdd}|||}||}||fSrj)rrmrnrr)r, hidden_statecontextrproZmotion_featuresr5 delta_flowr1r1r2r7%s    zUpdateBlock.forwardr8r1r1r/r2rs rcs.eZdZdZddfdd ZddZZS) MaskPredictorzMask predictor to be used when upsampling the predicted flow. It takes the hidden state of the recurrent unit as input and outputs the mask. This is not used in the raft-small model. ?) multipliercs:tt||ddd|_tj|dddd|_||_dS)Nrrci@rrr)r"r#rconvrelur&rUrVr)r,r-rzrr/r1r2r#5s zMaskPredictor.__init__cCs||}||}|j|Sr3)rrVrr[r1r1r2r7Bs  zMaskPredictor.forwardr8r1r1r/r2r.s rcsHeZdZdZdddeedfddZddZdd Zd d ZZ S) CorrBlockaThe correlation block. Creates a correlation pyramid with ``num_levels`` levels from the outputs of the feature encoder, and then indexes from this pyramid to create correlation features. The "indexing" of a given centroid pixel x' is done by concatenating its surrounding neighbors that are within a ``radius``, according to the infinity norm (see paper section 3.2). Note: typo in the paper, it should be infinity norm, not 1-norm. r? num_levelsradiuscs>t||_||_tdg|_|d|dd|_dS)Nrr r)r"r#rrrmtensor corr_pyramidr.)r,rrr/r1r2r#Rs  zCorrBlock.__init__c Cs|j|jkr&td|jd|jd|||}|j\}}}}}}|||||||}|g|_t|jdD] }tj|ddd}|j |qrdS)aLBuild the correlation pyramid from two feature maps. The correlation volume is first computed as the dot product of each pair (pixel_in_fmap1, pixel_in_fmap2) The last 2 dimensions of the correlation volume are then pooled num_levels times at different resolutions to build the correlation pyramid. z;Input feature maps should have the same shape, instead got z (fmap1.shape) != z (fmap2.shape)rr )rrN) shaperO_compute_corr_volumereshaperrangerFZ avg_pool2dappend) r,fmap1fmap2 corr_volume batch_sizer|w num_channelsrr1r1r2 build_pyramid_s  zCorrBlock.build_pyramidcCs2d|jd}t|j |j|}t|j |j|}tjtj||dddd|j}|d||d}|j\}}}} | dddd ||| ddd}g} |j D]:} ||} t | | d d d ||| d} | | |d}qtj| dd dddd}||j|| f}|j|kr.td |d |j|S)z9Return correlation features by indexing from the pyramid.r rZij)ZindexingrdrkrrTZbilinear)Z align_cornersrMz6Output shape of index pyramid is incorrect. Should be z, got )rrmZlinspacestackZmeshgridtodeviceviewrZpermuterrrrrn contiguousr.rO)r,centroids_coordsZneighborhood_side_lenZdiZdjdeltarrr|rZindexed_pyramidrZsampling_coordsZindexed_corr_volumerpZexpected_output_shaper1r1r2 index_pyramidts2"$    zCorrBlock.index_pyramidcCsn|j\}}}}|||||}|||||}t|dd|}||||d||}|tt|S)Nrr )rrrmmatmulZ transposesqrtr)r,rrrrr|rrqr1r1r2rs zCorrBlock._compute_corr_volume) r9r:r;r<intr#rrrr=r1r1r/r2rHs   rcs2eZdZddfdd Zd edddZZS) rN)mask_predictorcsHtt|||_||_||_||_||_t|jdsDt ddS)a RAFT model from `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `_. args: feature_encoder (nn.Module): The feature encoder. It must downsample the input by 8. Its input is the concatenation of ``image1`` and ``image2``. context_encoder (nn.Module): The context encoder. It must downsample the input by 8. Its input is ``image1``. As in the original implementation, its output will be split into 2 parts: - one part will be used as the actual "context", passed to the recurrent unit of the ``update_block`` - one part will be used to initialize the hidden state of the of the recurrent unit of the ``update_block`` These 2 parts are split according to the ``hidden_state_size`` of the ``update_block``, so the output of the ``context_encoder`` must be strictly greater than ``hidden_state_size``. corr_block (nn.Module): The correlation block, which creates a correlation pyramid from the output of the ``feature_encoder``, and then indexes from this pyramid to create correlation features. It must expose 2 methods: - a ``build_pyramid`` method that takes ``feature_map_1`` and ``feature_map_2`` as input (these are the output of the ``feature_encoder``). - a ``index_pyramid`` method that takes the coordinates of the centroid pixels as input, and returns the correlation features. See paper section 3.2. It must expose an ``out_channels`` attribute. update_block (nn.Module): The update block, which contains the motion encoder, the recurrent unit, and the flow head. It takes as input the hidden state of its recurrent unit, the context, the correlation features, and the current predicted flow. It outputs an updated hidden state, and the ``delta_flow`` prediction (see paper appendix A). It must expose a ``hidden_state_size`` attribute. mask_predictor (nn.Module, optional): Predicts the mask that will be used to upsample the predicted flow. The output channel must be 8 * 8 * 9 - see paper section 3.3, and Appendix B. If ``None`` (default), the flow is upsampled using interpolation. rzIThe update_block parameter should expose a 'hidden_state_size' attribute.N) r"r#r feature_encodercontext_encoder corr_block update_blockrhasattrrO)r,rrrrrr/r1r2r#s$  z RAFT.__init__ )num_flow_updatesc Cs.|j\}}}}||f|jddkrHtd|d|d|jdd|ddksv|ddkrvtd|d|d |tj||gdd }tj|d dd \} } | jdd|d|dfkrtd |j| | ||} | jdd|d|dfkrtd|j j } | jd| } | dkrDtd| jdd| dtj | | | gdd \}}t |}t |}t||d|d| j}t||d|d| j}g}t|D]t}|}|jj|d}||}| ||||\}}||}|jdkrdn||}t|||d}||q|S)Nz6input images should have the same shape, instead got (z, z) != rz9input image H and W should be divisible by 8, insted got z (h) and z (w)rkr )chunksrlz2The feature encoder should downsample H and W by 8z2The context encoder should downsample H and W by 8rzThe context encoder outputs zA channels, but it should have at strictly more than hidden_state=z channels)r)roup_mask)rrOrrmrnchunkrrrrrsplitr{rr*rrrrdetachrrrr)r,Zimage1Zimage2rrrr|rZfmapsrrZ context_outrout_channels_contextrrZcoords0Zcoords1Zflow_predictionsrprorrZupsampled_flowr1r1r2r7sF$      z RAFT.forward)r)r9r:r;r#rr7r=r1r1r/r2rs1rZmin_size)rErEc @seZdZdZedeeddddiddidd d d d d dZedeeddddiddiddd d dd dZedeeddddiddiddd dZ edeeddddiddiddd dZ edeeddd d!d"iid#d dZ ed$eeddd d!d%iid&d dZ e Z d'S)(rThe metrics reported here are as follows. ``epe`` is the "end-point-error" and indicates how far (in pixels) the predicted flow is from its true value. This is averaged over all pixels of all images. ``per_image_epe`` is similar, but the average is different: the epe is first computed on each image independently, and then averaged over all images. This corresponds to "Fl-epe" (sometimes written "F1-epe") in the original paper, and it's only used on Kitti. ``fl-all`` is also a Kitti-specific metric, defined by the author of the dataset and used for the Kitti leaderboard. It corresponds to the average of pixels whose epe is either <3px, or <5% of flow's 2-norm. zBhttps://download.pytorch.org/models/raft_large_C_T_V1-22a6c225.pthi@9P$https://github.com/princeton-vl/RAFTepeg?߾?g{P@gu@gޓZs1@Z per_image_epefl_allzSintel-Train-CleanpasszSintel-Train-Finalpassz Kitti-TrainThese weights were ported from the original paper. They are trained on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D`.Z num_paramsZrecipeZ_metricsZ_docsurlZ transformsmetazBhttps://download.pytorch.org/models/raft_large_C_T_V2-1bb1363a.pthChttps://github.com/pytorch/vision/tree/main/references/optical_flowgH}?g&S@g_L @gea0@These weights were trained from scratch on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D`.zGhttps://download.pytorch.org/models/raft_large_C_T_SKHT_V1-0b8c9e55.pthg ףp= ?gq= ףp @)zSintel-Test-CleanpasszSintel-Test-Finalpassa0 These weights were ported from the original paper. They are trained on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D` and fine-tuned on Sintel. The Sintel fine-tuning step is a combination of :class:`~torchvision.datasets.Sintel`, :class:`~torchvision.datasets.KittiFlow`, :class:`~torchvision.datasets.HD1K`, and :class:`~torchvision.datasets.FlyingThings3D` (clean pass). zGhttps://download.pytorch.org/models/raft_large_C_T_SKHT_V2-ff5fadd5.pthgv?gK7@a/ These weights were trained from scratch. They are pre-trained on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D` and then fine-tuned on Sintel. The Sintel fine-tuning step is a combination of :class:`~torchvision.datasets.Sintel`, :class:`~torchvision.datasets.KittiFlow`, :class:`~torchvision.datasets.HD1K`, and :class:`~torchvision.datasets.FlyingThings3D` (clean pass). zIhttps://download.pytorch.org/models/raft_large_C_T_SKHT_K_V1-4a6a5039.pthz Kitti-Testrgffffff@a These weights were ported from the original paper. They are pre-trained on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D`, fine-tuned on Sintel, and then fine-tuned on :class:`~torchvision.datasets.KittiFlow`. The Sintel fine-tuning step was described above. zIhttps://download.pytorch.org/models/raft_large_C_T_SKHT_K_V2-b5c70766.pthg(\@a These weights were trained from scratch. They are pre-trained on :class:`~torchvision.datasets.FlyingChairs` + :class:`~torchvision.datasets.FlyingThings3D`, fine-tuned on Sintel, and then fine-tuned on :class:`~torchvision.datasets.KittiFlow`. The Sintel fine-tuning step was described above. N)r9r:r;r<r r _COMMON_METAC_T_V1C_T_V2Z C_T_SKHT_V1 C_T_SKHT_V2Z C_T_SKHT_K_V1Z C_T_SKHT_K_V2DEFAULTr1r1r1r2rs rc @sxeZdZdZedeeddddiddidd d d d d dZedeeddddiddiddd d dd dZeZ dS)rrzBhttps://download.pytorch.org/models/raft_small_C_T_V1-ad48884c.pthirrgQ@gZd; @go@g3G9@rrrrrzBhttps://download.pytorch.org/models/raft_small_C_T_V2-01064c6d.pthrgHPs?gC @g1%d@g"lxz<9@rN) r9r:r;r<r r rrrrr1r1r1r2rs8 rF)weightsprogresscKs|ddpt|||d}|ddp2t|||d}|ddpJt|| d}|dd}|dkrt|j| | | d}|d| }t|j|| ||d }t| |d }t|||d }|d d}|dkr|rt| d dd}t f|||||d|}|dk r | |j |d|S)NrrGrrrr)rir`rar.rdr)r-rz)rrrrrFr)r-rzr)rrrrr)r) poprArr\r.rrrrrZload_state_dictZget_state_dict)rrfeature_encoder_layersfeature_encoder_blockfeature_encoder_norm_layercontext_encoder_layerscontext_encoder_blockcontext_encoder_norm_layercorr_block_num_levelscorr_block_radiusmotion_encoder_corr_layersmotion_encoder_flow_layersmotion_encoder_out_channels!recurrent_block_hidden_state_sizerecurrent_block_kernel_sizerecurrent_block_paddingflow_head_hidden_sizeuse_mask_predictorkwargsrrrrrrrrrmodelr1r1r2_rafts^      rZ pretrained)rT)rreturncKs>t|}tf||dttdttdddddddddd d |S) aRAFT model from `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `_. Please see the example below for a tutorial on how to use this model. Args: weights(:class:`~torchvision.models.optical_flow.Raft_Large_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.optical_flow.Raft_Large_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.optical_flow.RAFT`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.optical_flow.Raft_Large_Weights :members: rBr?r]r_rErrrFTrrrrrrrrrrrrrrrrrr)rverifyrrrrrrrr1r1r2r!s. rcKs>t|}tf||dttdtddddddd d d d d d|S)aRAFT "small" model from `RAFT: Recurrent All Pairs Field Transforms for Optical Flow `__. Please see the example below for a tutorial on how to use this model. Args: weights(:class:`~torchvision.models.optical_flow.Raft_Small_Weights`, optional): The pretrained weights to use. See :class:`~torchvision.models.optical_flow.Raft_Small_Weights` below for more details, and possible values. By default, no pre-trained weights are used. progress (bool): If True, displays a progress bar of the download to stderr. Default is True. **kwargs: parameters passed to the ``torchvision.models.optical_flow.RAFT`` base class. Please refer to the `source code `_ for more details about this class. .. autoclass:: torchvision.models.optical_flow.Raft_Small_Weights :members: ) rrCrDrE)rrrCrDNr?r)rD)rCrRrD)r)rrEFr)rrrr>rrr1r1r2rXs. r)2typingrrrmZtorch.nnr&Ztorch.nn.functionalZ functionalrrZtorch.nn.modules.batchnormrZtorch.nn.modules.instancenormrZtorchvision.opsrZtransforms._presetsr utilsr Z_apir r _utilsrrrr__all__Modulerr>rAr\rrrrrrrrrrrrrrrrrr1r1r1r2sL         ,*1+#Tg7 S6