mmagic.models.editors.edvr.edvr_net

Module Contents

Classes

EDVRNet	EDVR network structure for video super-resolution.
ModulatedDCNPack	Modulated Deformable Convolutional Pack.
PCDAlignment	Alignment module using Pyramid, Cascading and Deformable convolution
TSAFusion	Temporal Spatial Attention (TSA) fusion module. It is used in EDVRNet.

class mmagic.models.editors.edvr.edvr_net.EDVRNet(in_channels, out_channels, mid_channels=64, num_frames=5, deform_groups=8, num_blocks_extraction=5, num_blocks_reconstruction=10, center_frame_idx=2, with_tsa=True, init_cfg=None)

Bases: mmengine.model.BaseModule

EDVR network structure for video super-resolution.

Now only support X4 upsampling factor. Paper: EDVR: Video Restoration with Enhanced Deformable Convolutional Networks.

Parameters

• in_channels (int) – Channel number of inputs.

• out_channels (int) – Channel number of outputs.

• mid_channels (int) – Channel number of intermediate features. Default: 64.

• num_frames (int) – Number of input frames. Default: 5.

• deform_groups (int) – Deformable groups. Defaults: 8.

• num_blocks_extraction (int) – Number of blocks for feature extraction. Default: 5.

• num_blocks_reconstruction (int) – Number of blocks for reconstruction. Default: 10.

• center_frame_idx (int) – The index of center frame. Frame counting from 0. Default: 2.

• with_tsa (bool) – Whether to use TSA module. Default: True.

• init_cfg (dict, optional) – Initialization config dict. Default: None.

forward(x)

Forward function for EDVRNet.

Parameters

x (Tensor) – Input tensor with shape (n, t, c, h, w).

Returns

SR center frame with shape (n, c, h, w).

Return type

Tensor

init_weights()

Init weights for models.

class mmagic.models.editors.edvr.edvr_net.ModulatedDCNPack(*args, **kwargs)

Bases: mmcv.ops.ModulatedDeformConv2d

Modulated Deformable Convolutional Pack.

Different from the official DCN, which generates offsets and masks from the preceding features, this ModulatedDCNPack takes another different feature to generate masks and offsets.

Parameters

• in_channels (int) – Same as nn.Conv2d.

• out_channels (int) – Same as nn.Conv2d.

• kernel_size (int or tuple[int]) – Same as nn.Conv2d.

• stride (int or tuple[int]) – Same as nn.Conv2d.

• padding (int or tuple[int]) – Same as nn.Conv2d.

• dilation (int or tuple[int]) – Same as nn.Conv2d.

• groups (int) – Same as nn.Conv2d.

• bias (bool or str) – If specified as auto, it will be decided by the norm_cfg. Bias will be set as True if norm_cfg is None, otherwise False.

init_offset()

Init constant offset.

forward(x, extra_feat)

Forward function.

class mmagic.models.editors.edvr.edvr_net.PCDAlignment(mid_channels=64, deform_groups=8, act_cfg=dict(type='LeakyReLU', negative_slope=0.1))

Bases: mmengine.model.BaseModule

Alignment module using Pyramid, Cascading and Deformable convolution (PCD). It is used in EDVRNet.

Parameters

• mid_channels (int) – Number of the channels of middle features. Default: 64.

• deform_groups (int) – Deformable groups. Defaults: 8.

• act_cfg (dict) – Activation function config for ConvModule. Default: LeakyReLU with negative_slope=0.1.

forward(neighbor_feats, ref_feats)

Forward function for PCDAlignment.

Align neighboring frames to the reference frame in the feature level.

Parameters

• neighbor_feats (list[Tensor]) – List of neighboring features. It contains three pyramid levels (L1, L2, L3), each with shape (n, c, h, w).

• ref_feats (list[Tensor]) – List of reference features. It contains three pyramid levels (L1, L2, L3), each with shape (n, c, h, w).

Returns

Aligned features.

Return type

Tensor

class mmagic.models.editors.edvr.edvr_net.TSAFusion(mid_channels=64, num_frames=5, center_frame_idx=2, act_cfg=dict(type='LeakyReLU', negative_slope=0.1))

Bases: mmengine.model.BaseModule

Temporal Spatial Attention (TSA) fusion module. It is used in EDVRNet.

Parameters

• mid_channels (int) – Number of the channels of middle features. Default: 64.

• num_frames (int) – Number of frames. Default: 5.

• center_frame_idx (int) – The index of center frame. Default: 2.

• act_cfg (dict) – Activation function config for ConvModule. Default: LeakyReLU with negative_slope=0.1.

forward(aligned_feat)

Forward function for TSAFusion.

Parameters

aligned_feat (Tensor) – Aligned features with shape (n, t, c, h, w).

Returns

Features after TSA with the shape (n, c, h, w).

Return type

Tensor